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

feat: Save baml progress

Browse source
This commit is contained in:
Igor Ilic 2025-09-08 21:16:24 +02:00
parent f7eb482ba5
commit befb8ac237
12 changed files with 221 additions and 240 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
cognee/infrastructure/llm/LLMGateway.py
View file

@ -26,7 +26,7 @@ class LLMGateway:
)
return acreate_structured_output(
content=text_input,
text_input=text_input,
system_prompt=system_prompt,
response_model=response_model,
)
@ -142,19 +142,19 @@ class LLMGateway:
@staticmethod
def extract_summary(content: str, response_model: Type[BaseModel]) -> Coroutine:
llm_config = get_llm_config()
if llm_config.structured_output_framework.upper() == "BAML":
from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction import (
extract_summary,
)
# llm_config = get_llm_config()
# if llm_config.structured_output_framework.upper() == "BAML":
# from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction import (
# extract_summary,
# )
#
# return extract_summary(content=content, response_model=response_model)
# else:
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.extraction import (
extract_summary,
)
return extract_summary(content=content, response_model=response_model)
else:
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.extraction import (
extract_summary,
)
return extract_summary(content=content, response_model=response_model)
return extract_summary(content=content, response_model=response_model)
@staticmethod
def extract_event_graph(content: str, response_model: Type[BaseModel]) -> Coroutine:

43
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/async_client.py
View file

@ -81,19 +81,15 @@ class BamlAsyncClient:
async def AcreateStructuredOutput(
self,
content: str,
text_input: str,
system_prompt: str,
user_prompt: str,
baml_options: BamlCallOptions = {},
) -> types.DynamicOutputModel:
) -> types.ResponseModel:
# Check if on_tick is provided
if "on_tick" in baml_options:
# Use streaming internally when on_tick is provided
stream = self.stream.AcreateStructuredOutput(
content=content,
system_prompt=system_prompt,
user_prompt=user_prompt,
baml_options=baml_options,
text_input=text_input, system_prompt=system_prompt, baml_options=baml_options
)
return await stream.get_final_response()
else:
@ -101,14 +97,12 @@ class BamlAsyncClient:
result = await self.__options.merge_options(baml_options).call_function_async(
function_name="AcreateStructuredOutput",
args={
"content": content,
"text_input": text_input,
"system_prompt": system_prompt,
"user_prompt": user_prompt,
},
)
return typing.cast(
types.DynamicOutputModel,
result.cast_to(types, types, stream_types, False, __runtime__),
types.ResponseModel, result.cast_to(types, types, stream_types, False, __runtime__)
)
async def ExtractCategories(
@ -267,27 +261,24 @@ class BamlStreamClient:
def AcreateStructuredOutput(
self,
content: str,
text_input: str,
system_prompt: str,
user_prompt: str,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlStream[stream_types.DynamicOutputModel, types.DynamicOutputModel]:
) -> baml_py.BamlStream[stream_types.ResponseModel, types.ResponseModel]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream(
function_name="AcreateStructuredOutput",
args={
"content": content,
"text_input": text_input,
"system_prompt": system_prompt,
"user_prompt": user_prompt,
},
)
return baml_py.BamlStream[stream_types.DynamicOutputModel, types.DynamicOutputModel](
return baml_py.BamlStream[stream_types.ResponseModel, types.ResponseModel](
result,
lambda x: typing.cast(
stream_types.DynamicOutputModel,
x.cast_to(types, types, stream_types, True, __runtime__),
stream_types.ResponseModel, x.cast_to(types, types, stream_types, True, __runtime__)
),
lambda x: typing.cast(
types.DynamicOutputModel, x.cast_to(types, types, stream_types, False, __runtime__)
types.ResponseModel, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
)
@ -444,17 +435,15 @@ class BamlHttpRequestClient:
async def AcreateStructuredOutput(
self,
content: str,
text_input: str,
system_prompt: str,
user_prompt: str,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(
function_name="AcreateStructuredOutput",
args={
"content": content,
"text_input": text_input,
"system_prompt": system_prompt,
"user_prompt": user_prompt,
},
mode="request",
)
@ -563,17 +552,15 @@ class BamlHttpStreamRequestClient:
async def AcreateStructuredOutput(
self,
content: str,
text_input: str,
system_prompt: str,
user_prompt: str,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(
function_name="AcreateStructuredOutput",
args={
"content": content,
"text_input": text_input,
"system_prompt": system_prompt,
"user_prompt": user_prompt,
},
mode="stream",
)

2
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/inlinedbaml.py
View file

File diff suppressed because one or more lines are too long

8
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/parser.py
View file

@ -27,11 +27,11 @@ class LlmResponseParser:
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.DynamicOutputModel:
) -> types.ResponseModel:
result = self.__options.merge_options(baml_options).parse_response(
function_name="AcreateStructuredOutput", llm_response=llm_response, mode="request"
)
return typing.cast(types.DynamicOutputModel, result)
return typing.cast(types.ResponseModel, result)
def ExtractCategories(
self,
@ -94,11 +94,11 @@ class LlmStreamParser:
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.DynamicOutputModel:
) -> stream_types.ResponseModel:
result = self.__options.merge_options(baml_options).parse_response(
function_name="AcreateStructuredOutput", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.DynamicOutputModel, result)
return typing.cast(stream_types.ResponseModel, result)
def ExtractCategories(
self,

16
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/stream_types.py
View file

@ -27,7 +27,7 @@ class StreamState(BaseModel, typing.Generic[StreamStateValueT]):
# #########################################################################
# Generated classes (19)
# Generated classes (18)
# #########################################################################
@ -46,20 +46,10 @@ class DefaultContentPrediction(BaseModel):
label: typing.Optional["ContentLabel"] = None
class DynamicInputModel(BaseModel):
model_config = ConfigDict(extra="allow")
test: typing.Optional[str] = None
class DynamicKnowledgeGraph(BaseModel):
model_config = ConfigDict(extra="allow")
class DynamicOutputModel(BaseModel):
model_config = ConfigDict(extra="allow")
test: typing.Optional[str] = None
class Edge(BaseModel):
# doc string for edge
# doc string for source_node_id
@ -102,6 +92,10 @@ class ProceduralContent(BaseModel):
subclass: typing.List[str]
class ResponseModel(BaseModel):
model_config = ConfigDict(extra="allow")
class SummarizedClass(BaseModel):
name: typing.Optional[str] = None
description: typing.Optional[str] = None

43
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/sync_client.py
View file

@ -94,18 +94,14 @@ class BamlSyncClient:
def AcreateStructuredOutput(
self,
content: str,
text_input: str,
system_prompt: str,
user_prompt: str,
baml_options: BamlCallOptions = {},
) -> types.DynamicOutputModel:
) -> types.ResponseModel:
# Check if on_tick is provided
if "on_tick" in baml_options:
stream = self.stream.AcreateStructuredOutput(
content=content,
system_prompt=system_prompt,
user_prompt=user_prompt,
baml_options=baml_options,
text_input=text_input, system_prompt=system_prompt, baml_options=baml_options
)
return stream.get_final_response()
else:
@ -113,14 +109,12 @@ class BamlSyncClient:
result = self.__options.merge_options(baml_options).call_function_sync(
function_name="AcreateStructuredOutput",
args={
"content": content,
"text_input": text_input,
"system_prompt": system_prompt,
"user_prompt": user_prompt,
},
)
return typing.cast(
types.DynamicOutputModel,
result.cast_to(types, types, stream_types, False, __runtime__),
types.ResponseModel, result.cast_to(types, types, stream_types, False, __runtime__)
)
def ExtractCategories(
@ -274,27 +268,24 @@ class BamlStreamClient:
def AcreateStructuredOutput(
self,
content: str,
text_input: str,
system_prompt: str,
user_prompt: str,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlSyncStream[stream_types.DynamicOutputModel, types.DynamicOutputModel]:
) -> baml_py.BamlSyncStream[stream_types.ResponseModel, types.ResponseModel]:
ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
function_name="AcreateStructuredOutput",
args={
"content": content,
"text_input": text_input,
"system_prompt": system_prompt,
"user_prompt": user_prompt,
},
)
return baml_py.BamlSyncStream[stream_types.DynamicOutputModel, types.DynamicOutputModel](
return baml_py.BamlSyncStream[stream_types.ResponseModel, types.ResponseModel](
result,
lambda x: typing.cast(
stream_types.DynamicOutputModel,
x.cast_to(types, types, stream_types, True, __runtime__),
stream_types.ResponseModel, x.cast_to(types, types, stream_types, True, __runtime__)
),
lambda x: typing.cast(
types.DynamicOutputModel, x.cast_to(types, types, stream_types, False, __runtime__)
types.ResponseModel, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
)
@ -455,17 +446,15 @@ class BamlHttpRequestClient:
def AcreateStructuredOutput(
self,
content: str,
text_input: str,
system_prompt: str,
user_prompt: str,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(
function_name="AcreateStructuredOutput",
args={
"content": content,
"text_input": text_input,
"system_prompt": system_prompt,
"user_prompt": user_prompt,
},
mode="request",
)
@ -574,17 +563,15 @@ class BamlHttpStreamRequestClient:
def AcreateStructuredOutput(
self,
content: str,
text_input: str,
system_prompt: str,
user_prompt: str,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(
function_name="AcreateStructuredOutput",
args={
"content": content,
"text_input": text_input,
"system_prompt": system_prompt,
"user_prompt": user_prompt,
},
mode="stream",
)

170
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_builder.py
View file

@ -27,9 +27,7 @@ class TypeBuilder(type_builder.TypeBuilder):
"AudioContent",
"ContentLabel",
"DefaultContentPrediction",
"DynamicInputModel",
"DynamicKnowledgeGraph",
"DynamicOutputModel",
"Edge",
"ImageContent",
"KnowledgeGraph",
@ -37,6 +35,7 @@ class TypeBuilder(type_builder.TypeBuilder):
"MultimediaContent",
"Node",
"ProceduralContent",
"ResponseModel",
"SummarizedClass",
"SummarizedCode",
"SummarizedContent",
@ -54,7 +53,7 @@ class TypeBuilder(type_builder.TypeBuilder):
# #########################################################################
# #########################################################################
# Generated classes 19
# Generated classes 18
# #########################################################################
@property
@ -69,18 +68,10 @@ class TypeBuilder(type_builder.TypeBuilder):
def DefaultContentPrediction(self) -> "DefaultContentPredictionViewer":
return DefaultContentPredictionViewer(self)
@property
def DynamicInputModel(self) -> "DynamicInputModelBuilder":
return DynamicInputModelBuilder(self)
@property
def DynamicKnowledgeGraph(self) -> "DynamicKnowledgeGraphBuilder":
return DynamicKnowledgeGraphBuilder(self)
@property
def DynamicOutputModel(self) -> "DynamicOutputModelBuilder":
return DynamicOutputModelBuilder(self)
@property
def Edge(self) -> "EdgeViewer":
return EdgeViewer(self)
@ -109,6 +100,10 @@ class TypeBuilder(type_builder.TypeBuilder):
def ProceduralContent(self) -> "ProceduralContentViewer":
return ProceduralContentViewer(self)
@property
def ResponseModel(self) -> "ResponseModelBuilder":
return ResponseModelBuilder(self)
@property
def SummarizedClass(self) -> "SummarizedClassViewer":
return SummarizedClassViewer(self)
@ -140,7 +135,7 @@ class TypeBuilder(type_builder.TypeBuilder):
# #########################################################################
# Generated classes 19
# Generated classes 18
# #########################################################################
@ -279,59 +274,6 @@ class DefaultContentPredictionProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("label"))
class DynamicInputModelAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("DynamicInputModel")
self._properties: typing.Set[str] = set(
[
"test",
]
)
self._props = DynamicInputModelProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "DynamicInputModelProperties":
return self._props
class DynamicInputModelBuilder(DynamicInputModelAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
def add_property(self, name: str, type: baml_py.FieldType) -> baml_py.ClassPropertyBuilder:
if name in self._properties:
raise ValueError(f"Property {name} already exists.")
return self._bldr.property(name).type(type)
def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]:
return self._bldr.list_properties()
def remove_property(self, name: str) -> None:
self._bldr.remove_property(name)
def reset(self) -> None:
self._bldr.reset()
class DynamicInputModelProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
def __getattr__(self, name: str) -> baml_py.ClassPropertyBuilder:
if name not in self.__properties:
raise AttributeError(f"Property {name} not found.")
return self.__bldr.property(name)
@property
def test(self) -> baml_py.ClassPropertyBuilder:
return self.__bldr.property("test")
class DynamicKnowledgeGraphAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
@ -377,59 +319,6 @@ class DynamicKnowledgeGraphProperties:
return self.__bldr.property(name)
class DynamicOutputModelAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("DynamicOutputModel")
self._properties: typing.Set[str] = set(
[
"test",
]
)
self._props = DynamicOutputModelProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "DynamicOutputModelProperties":
return self._props
class DynamicOutputModelBuilder(DynamicOutputModelAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
def add_property(self, name: str, type: baml_py.FieldType) -> baml_py.ClassPropertyBuilder:
if name in self._properties:
raise ValueError(f"Property {name} already exists.")
return self._bldr.property(name).type(type)
def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]:
return self._bldr.list_properties()
def remove_property(self, name: str) -> None:
self._bldr.remove_property(name)
def reset(self) -> None:
self._bldr.reset()
class DynamicOutputModelProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
def __getattr__(self, name: str) -> baml_py.ClassPropertyBuilder:
if name not in self.__properties:
raise AttributeError(f"Property {name} not found.")
return self.__bldr.property(name)
@property
def test(self) -> baml_py.ClassPropertyBuilder:
return self.__bldr.property("test")
class EdgeAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
@ -773,6 +662,51 @@ class ProceduralContentProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
class ResponseModelAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("ResponseModel")
self._properties: typing.Set[str] = set([])
self._props = ResponseModelProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "ResponseModelProperties":
return self._props
class ResponseModelBuilder(ResponseModelAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
def add_property(self, name: str, type: baml_py.FieldType) -> baml_py.ClassPropertyBuilder:
if name in self._properties:
raise ValueError(f"Property {name} already exists.")
return self._bldr.property(name).type(type)
def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]:
return self._bldr.list_properties()
def remove_property(self, name: str) -> None:
self._bldr.remove_property(name)
def reset(self) -> None:
self._bldr.reset()
class ResponseModelProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
def __getattr__(self, name: str) -> baml_py.ClassPropertyBuilder:
if name not in self.__properties:
raise AttributeError(f"Property {name} not found.")
return self.__bldr.property(name)
class SummarizedClassAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)

6
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_map.py
View file

@ -21,12 +21,8 @@ type_map = {
"stream_types.ContentLabel": stream_types.ContentLabel,
"types.DefaultContentPrediction": types.DefaultContentPrediction,
"stream_types.DefaultContentPrediction": stream_types.DefaultContentPrediction,
"types.DynamicInputModel": types.DynamicInputModel,
"stream_types.DynamicInputModel": stream_types.DynamicInputModel,
"types.DynamicKnowledgeGraph": types.DynamicKnowledgeGraph,
"stream_types.DynamicKnowledgeGraph": stream_types.DynamicKnowledgeGraph,
"types.DynamicOutputModel": types.DynamicOutputModel,
"stream_types.DynamicOutputModel": stream_types.DynamicOutputModel,
"types.Edge": types.Edge,
"stream_types.Edge": stream_types.Edge,
"types.ImageContent": types.ImageContent,
@ -41,6 +37,8 @@ type_map = {
"stream_types.Node": stream_types.Node,
"types.ProceduralContent": types.ProceduralContent,
"stream_types.ProceduralContent": stream_types.ProceduralContent,
"types.ResponseModel": types.ResponseModel,
"stream_types.ResponseModel": stream_types.ResponseModel,
"types.SummarizedClass": types.SummarizedClass,
"stream_types.SummarizedClass": stream_types.SummarizedClass,
"types.SummarizedCode": types.SummarizedCode,

16
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/types.py
View file

@ -48,7 +48,7 @@ def all_succeeded(checks: typing.Dict[CheckName, Check]) -> bool:
# #########################################################################
# #########################################################################
# Generated classes (19)
# Generated classes (18)
# #########################################################################
@ -75,20 +75,10 @@ class DefaultContentPrediction(BaseModel):
label: "ContentLabel"
class DynamicInputModel(BaseModel):
model_config = ConfigDict(extra="allow")
test: str
class DynamicKnowledgeGraph(BaseModel):
model_config = ConfigDict(extra="allow")
class DynamicOutputModel(BaseModel):
model_config = ConfigDict(extra="allow")
test: str
class Edge(BaseModel):
# doc string for edge
# doc string for source_node_id
@ -131,6 +121,10 @@ class ProceduralContent(BaseModel):
subclass: typing.List[str]
class ResponseModel(BaseModel):
model_config = ConfigDict(extra="allow")
class SummarizedClass(BaseModel):
name: str
description: str

16
cognee/infrastructure/llm/structured_output_framework/baml/baml_src/acreate_structured_output.baml
View file

@ -1,25 +1,17 @@
class DynamicInputModel {
test string
@@dynamic
}
class DynamicOutputModel {
test string
class ResponseModel {
@@dynamic
}
function AcreateStructuredOutput(
content: string,
text_input: string,
system_prompt: string,
user_prompt: string,
) -> DynamicOutputModel {
) -> ResponseModel {
client OpenAI
prompt #"
{{ system_prompt }}
{{ ctx.output_format }}
{{ _.role('user') }}
{{ user_prompt }}
{{ content }}
{{ text_input }}
"#
}

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

@ -1,18 +1,113 @@
import os
import asyncio
from typing import Type
from pydantic import BaseModel
from cognee.shared.logging_utils import get_logger
from cognee.shared.data_models import SummarizedCode
from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.async_client import b
from cognee.infrastructure.llm.config import get_llm_config
from typing import List, Dict, Union, Optional, Literal
from enum import Enum
from baml_py import Image, Audio, Video, Pdf
from datetime import datetime
from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.type_builder import (
TypeBuilder,
)
from cognee.infrastructure.llm.structured_output_framework.baml.baml_client import b
from pydantic import BaseModel
logger = get_logger("extract_summary_baml")
def create_dynamic_baml_type(pydantic_model):
tb = TypeBuilder()
# if pydantic_model == str:
# b.ResponseModel.add_property("text", tb.string())
# return tb
#
# def map_type(field_type, field_info):
# # Handle Optional/Union types
# if getattr(field_type, "__origin__", None) == Union:
# # Extract types from Union
# types = field_type.__args__
# # Handle Optional (Union with NoneType)
# if type(None) in types:
# inner_type = next(t for t in types if t != type(None))
# return map_type(inner_type, field_info).optional()
# # Handle regular Union
# mapped_types = [map_type(t, field_info) for t in types]
# return tb.union(*mapped_types)
#
# # Handle Lists
# if getattr(field_type, "__origin__", None) == list:
# inner_type = field_type.__args__[0]
# return map_type(inner_type, field_info).list()
#
# # Handle Maps/Dictionaries
# if getattr(field_type, "__origin__", None) == dict:
# key_type, value_type = field_type.__args__
# # BAML only supports string or enum keys in maps
# if key_type not in [str, Enum]:
# raise ValueError("Map keys must be strings or enums in BAML")
# return tb.map(map_type(key_type, field_info), map_type(value_type, field_info))
#
# # Handle Literal types
# if getattr(field_type, "__origin__", None) == Literal:
# literal_values = field_type.__args__
# return tb.union(*[tb.literal(val) for val in literal_values])
#
# # Handle Enums
# if isinstance(field_type, type) and issubclass(field_type, Enum):
# enum_type = tb.add_enum(field_type.__name__)
# for member in field_type:
# enum_type.add_value(member.name)
# return enum_type.type()
#
# # Handle primitive and special types
# type_mapping = {
# str: tb.string(),
# int: tb.int(),
# float: tb.float(),
# bool: tb.bool(),
# Image: tb.image(),
# Audio: tb.audio(),
# Video: tb.video(),
# Pdf: tb.pdf(),
# # datetime is not natively supported in BAML, map to string
# datetime: tb.string(),
# }
#
# # Handle nested BaseModel classes
# if isinstance(field_type, type) and issubclass(field_type, BaseModel):
# nested_tb = create_dynamic_baml_type(field_type)
# # Get the last created class from the nested TypeBuilder
# return nested_tb.get_last_class().type()
#
# if field_type in type_mapping:
# return type_mapping[field_type]
#
# raise ValueError(f"Unsupported type: {field_type}")
#
# fields = pydantic_model.model_fields
#
# # Add fields
# for field_name, field_info in fields.items():
# field_type = field_info.annotation
# baml_type = map_type(field_type, field_info)
#
# # Add property with type
# prop = b.ResponseModel.add_property(field_name, baml_type)
#
# # Add description if available
# if field_info.description:
# prop.description(field_info.description)
return tb
async def acreate_structured_output(
content: str, system_prompt: str, user_prompt: str, response_model: Type[BaseModel]
text_input: str, system_prompt: str, response_model: Type[BaseModel]
):
"""
Extract summary using BAML framework.
@ -26,12 +121,12 @@ async def acreate_structured_output(
"""
config = get_llm_config()
# Use BAML's SummarizeContent function
type_builder = create_dynamic_baml_type(response_model)
result = await b.AcreateStructuredOutput(
content=content,
text_input=text_input,
system_prompt=system_prompt,
user_prompt=user_prompt,
baml_options={"client_registry": config.baml_registry},
baml_options={"client_registry": config.baml_registry, "tb": type_builder},
)
return result

4
cognee/infrastructure/llm/utils.py
View file

@ -3,6 +3,7 @@ import litellm
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.get_llm_client import (
get_llm_client,
)
from cognee.infrastructure.llm.LLMGateway import LLMGateway
from cognee.shared.logging_utils import get_logger
logger = get_logger()
@ -76,8 +77,7 @@ async def test_llm_connection():
the connection attempt and re-raise the exception for further handling.
"""
try:
llm_adapter = get_llm_client()
await llm_adapter.acreate_structured_output(
await LLMGateway.acreate_structured_output(
text_input="test",
system_prompt='Respond to me with the following string: "test"',
response_model=str,