from fastapi import FastAPI
from pydantic import BaseModel
from langchain_openai import AzureChatOpenAI
import os

app = FastAPI(title="Translation API")
os.environ['AZURE_OPENAI_API_KEY'] = os.getenv("AZURE_OPENAI_API_KEY")
os.environ['OPENAI_API_VERSION'] = os.getenv("OPENAI_API_VERSION")
os.environ['AZURE_OPENAI_ENDPOINT'] = os.getenv("AZURE_OPENAI_ENDPOINT")

llm = AzureChatOpenAI(
    azure_deployment="gpt-4o",
    temperature=0,
    max_tokens=None,
    timeout=None,
    max_retries=2,
)

class TranslateRequest(BaseModel):
    text: str

class TranslateResponse(BaseModel):
    translated_text: str


@app.post("/chat", response_model=TranslateResponse)
def translate_text(request: TranslateRequest):
    messages = [
        (
            "system",
            "You are a helpful assistant that answers questions about the user's query.",
        ),
        ("human", request.text),
    ]

    ai_msg = llm.invoke(messages)

    return TranslateResponse(
        translated_text=ai_msg.content
    )