kabylake's picture
commit
7bd11ed
raw
history blame contribute delete
812 Bytes
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
)