jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Add JSON Mode option to OpenAIModelComponent (#2386)

Browse source 
* feat: Add JSON Mode option to OpenAIModelComponent

* ♻️ (OpenAIModel.py): add type ignore comment

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-26 15:02:32 -07:00 • committed by GitHub
commit a81f686c39
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 11 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

13
src/backend/base/langflow/components/models/OpenAIModel.py
View file

@ -34,6 +34,12 @@ class OpenAIModelComponent(LCModelComponent):
info="The maximum number of tokens to generate. Set to 0 for unlimited tokens.", info="The maximum number of tokens to generate. Set to 0 for unlimited tokens.",
), ),
DictInput(name="model_kwargs", display_name="Model Kwargs", advanced=True), DictInput(name="model_kwargs", display_name="Model Kwargs", advanced=True),
BoolInput(
name="json_mode",
display_name="JSON Mode",
advanced=True,
info="If True, it will output JSON regardless of passing a schema.",
),
DictInput( DictInput(
name="output_schema", name="output_schema",
is_list=True, is_list=True,
@ -84,7 +90,7 @@ class OpenAIModelComponent(LCModelComponent):
max_tokens = self.max_tokens max_tokens = self.max_tokens
model_kwargs = self.model_kwargs or {} 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) or self.json_mode
seed = self.seed seed = self.seed
model_kwargs["seed"] = seed model_kwargs["seed"] = seed
@ -101,7 +107,10 @@ class OpenAIModelComponent(LCModelComponent):
temperature=temperature or 0.1, temperature=temperature or 0.1,
) )
if json_mode: if json_mode:
output = output.with_structured_output(schema=output_schema_dict, method="json_mode") # type: ignore if output_schema_dict:
output = output.with_structured_output(schema=output_schema_dict, method="json_mode") # type: ignore
else:
output = output.bind(response_format={"type": "json_object"}) # type: ignore
return output return output