from typing import Optional

from langchain_community.chat_models import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from pydantic import BaseModel, ConfigDict


class OpenAIModelConfig(BaseModel):
    model_config = ConfigDict()
    model_config["protected_namespaces"] = ()
    prompt_template: str
    model_kwargs: dict = {}


class OpenAIModel:
    def __init__(self, config: OpenAIModelConfig):
        self.config = config
        self._model = None

    @property
    def model(self):
        return ChatOpenAI(**self.config.model_kwargs)

    @property
    def prompt(self) -> Optional[PromptTemplate]:
        if self.config.prompt_template:
            return PromptTemplate(
                input_variables=["context", "question"], template=self.config.prompt_template
            )