import math
from typing import Optional

from ecologits.tracers.utils import llm_impacts, ImpactsOutput
from fastapi import FastAPI
from pydantic import BaseModel


app = FastAPI()


class EstimationInputs(BaseModel):
    provider: str
    model_name: str
    output_token_count: int
    request_latency: Optional[float] = math.inf
    electricity_mix_zone: Optional[str] = "WOR"


@app.post("/estimations")
def create_estimation(inputs: EstimationInputs) -> ImpactsOutput:
    return llm_impacts(**inputs.model_dump())