from __future__ import annotations
import os
from pydantic import BaseModel
from core.ports import ChatModel


class ModalChatModel(ChatModel):
    """ChatModel backed by a Modal-hosted vLLM endpoint (OpenAI-compatible API)."""

    def __init__(self) -> None:
        from langchain_openai import ChatOpenAI  # deferred: must follow load_dotenv()
        url = os.environ["MODAL_INFERENCE_URL"]
        key = os.environ["MODAL_API_KEY"]
        model_id = os.environ["MODAL_MODEL_ID"]
        self._llm = ChatOpenAI(
            base_url=url.rstrip("/") + "/v1",
            api_key=key,
            model=model_id,
            max_tokens=2048,
            # temperature=0: deterministic structured calls. The model default (0.6, seen in the
            # vLLM generation_config) makes the model ramble to max_tokens → truncated JSON →
            # "Could not parse response content as the length limit was reached". This is the
            # primary fix for that truncation class.
            temperature=0,
        )

    def structured(self, messages: list[dict], schema: type[BaseModel]) -> BaseModel:
        # method="json_schema" is already the langchain_openai default; stated explicitly to
        # document that we rely on vLLM's structured-outputs (guided JSON) backend, not tool
        # calling. (strict=True is intentionally deferred until verified against the live vLLM
        # endpoint — it would gate every structured call in the app.)
        return self._llm.with_structured_output(schema, method="json_schema").invoke(messages)

    def complete(self, messages: list[dict]) -> str:
        return self._llm.invoke(messages).content