refactor: Update OpenAIModelComponent to include seed in model_kwargs

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-20 18:18:04 -03:00
commit ac627e831f

View file

@ -82,25 +82,26 @@ class OpenAIModelComponent(LCModelComponent):
temperature = self.temperature temperature = self.temperature
model_name: str = self.model_name model_name: str = self.model_name
max_tokens = self.max_tokens max_tokens = self.max_tokens
model_kwargs = self.model_kwargs model_kwargs = self.model_kwargs or {}
openai_api_base = self.openai_api_base or "https://api.openai.com/v1" openai_api_base = self.openai_api_base or "https://api.openai.com/v1"
json_mode = bool(output_schema_dict) json_mode = bool(output_schema_dict)
seed = self.seed seed = self.seed
model_kwargs["seed"] = seed
if openai_api_key: if openai_api_key:
api_key = SecretStr(openai_api_key) api_key = SecretStr(openai_api_key)
else: else:
api_key = None api_key = None
output = ChatOpenAI( output = ChatOpenAI(
max_tokens=max_tokens or None, max_tokens=max_tokens or None,
model_kwargs=model_kwargs or {}, model_kwargs=model_kwargs,
model=model_name, model=model_name,
base_url=openai_api_base, base_url=openai_api_base,
api_key=api_key, api_key=api_key,
temperature=temperature or 0.1, temperature=temperature or 0.1,
seed=seed,
) )
if json_mode: if json_mode:
output = output.with_structured_output(schema=output_schema_dict, method="json_mode") output = output.with_structured_output(schema=output_schema_dict, method="json_mode") # type: ignore
return output return output