from __future__ import annotations

from typing import Any

from pydantic import BaseModel


class LogFeaturesRequest(BaseModel):
    """
    Request body for POST /ai/log_features.
    """

    trace_id: str
    speedup_ratio: dict[str, float | int | None] | None = None
    original_runtime: float | int | None = None
    optimized_runtime: dict[str, float | int | None] | None = None
    is_correct: dict[str, bool] | None = None
    optimized_line_profiler_results: dict[str, str] | None = None
    metadata: dict[str, Any] | None = None
    optimizations_post: dict[str, str] | None = None


class LogFeaturesErrorResponse(BaseModel):
    """
    Error response from POST /ai/log_features.
    """

    error: str
