from typing import Optional, List, Any
from pydantic import Field
from langchain_core.language_models.llms import LLM
from langchain_core.callbacks.manager import CallbackManagerForLLMRun

from trim.trim_client import TrimClient


class TrimLLM(LLM):
    trim_client: TrimClient = Field(default_factory=TrimClient)

    def __init__(self, **kwargs: Any):
        super().__init__(**kwargs)

    @property
    def _llm_type(self) -> str:
        return "router-llm"

    def _call(self,
              prompt: str,
              stop: Optional[List[str]] = None,
              run_manager: Optional[CallbackManagerForLLMRun] = None,
              **kwargs) -> str:
        completion = self.trim_client.responses.create(
            model=None,
            input=prompt,
            **kwargs
        )
        return completion.output_text
