gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

docs: Add comments to baml structured outputs creator

Browse source
This commit is contained in:
Igor Ilic 2025-09-09 14:44:24 +02:00
parent 8ab89e8364
commit 11aefadb81
1 changed files with 4 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/acreate_structured_output.py
View file

@ -152,8 +152,9 @@ async def acreate_structured_output(
BaseModel: The summarized content in the specified format BaseModel: The summarized content in the specified format
""" """
config = get_llm_config() config = get_llm_config()
tb = TypeBuilder()
# Dynamically create BAML response model
tb = TypeBuilder()
type_builder = create_dynamic_baml_type(tb, tb.ResponseModel, response_model) type_builder = create_dynamic_baml_type(tb, tb.ResponseModel, response_model)
result = await b.AcreateStructuredOutput( result = await b.AcreateStructuredOutput(
@ -162,8 +163,9 @@ async def acreate_structured_output(
baml_options={"client_registry": config.baml_registry, "tb": type_builder}, baml_options={"client_registry": config.baml_registry, "tb": type_builder},
) )
# Transform BAML response to proper pydantic reponse model
if response_model is str: if response_model is str:
return result return str(result)
return response_model.model_validate(result.dict()) return response_model.model_validate(result.dict())