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

added files

Browse source
This commit is contained in:
vasilije 2025-07-24 14:07:16 +02:00
parent 6a056ebd01
commit 0dfd2872ea
17 changed files with 1268 additions and 592 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/__init__.py
View file

@ -13,9 +13,9 @@
__version__ = "0.201.0" __version__ = "0.201.0"
try: try:
from baml_py.safe_import import EnsureBamlPyImport from baml_py.safe_import import EnsureBamlPyImport
except ImportError: except ImportError:
raise ImportError(f"""Update to baml-py required. raise ImportError(f"""Update to baml-py required.
Version of baml_client generator (see generators.baml): {__version__} Version of baml_client generator (see generators.baml): {__version__}
Please upgrade baml-py to version "{__version__}". Please upgrade baml-py to version "{__version__}".
@ -31,16 +31,15 @@ https://boundaryml.com/discord
with EnsureBamlPyImport(__version__) as e: with EnsureBamlPyImport(__version__) as e:
e.raise_if_incompatible_version(__version__) e.raise_if_incompatible_version(__version__)
from . import types from . import types
from . import tracing from . import tracing
from . import stream_types from . import stream_types
from . import config from . import config
from .config import reset_baml_env_vars from .config import reset_baml_env_vars
from .sync_client import b
from .sync_client import b
# FOR LEGACY COMPATIBILITY, expose "partial_types" as an alias for "stream_types" # FOR LEGACY COMPATIBILITY, expose "partial_types" as an alias for "stream_types"
@ -48,11 +47,11 @@ with EnsureBamlPyImport(__version__) as e:
partial_types = stream_types partial_types = stream_types
__all__ = [ __all__ = [
"b", "b",
"stream_types", "stream_types",
"partial_types", "partial_types",
"tracing", "tracing",
"types", "types",
"reset_baml_env_vars", "reset_baml_env_vars",
"config", "config",
] ]

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

@ -36,10 +36,13 @@ class BamlAsyncClient:
self.__llm_response_parser = LlmResponseParser(options) self.__llm_response_parser = LlmResponseParser(options)
self.__llm_stream_parser = LlmStreamParser(options) self.__llm_stream_parser = LlmStreamParser(options)
def with_options(self, def with_options(
self,
tb: typing.Optional[type_builder.TypeBuilder] = None, tb: typing.Optional[type_builder.TypeBuilder] = None,
client_registry: typing.Optional[baml_py.baml_py.ClientRegistry] = None, client_registry: typing.Optional[baml_py.baml_py.ClientRegistry] = None,
collector: typing.Optional[typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]] = None, collector: typing.Optional[
typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]
] = None,
env: typing.Optional[typing.Dict[str, typing.Optional[str]]] = None, env: typing.Optional[typing.Dict[str, typing.Optional[str]]] = None,
) -> "BamlAsyncClient": ) -> "BamlAsyncClient":
options: BamlCallOptions = {} options: BamlCallOptions = {}
@ -55,60 +58,124 @@ class BamlAsyncClient:
@property @property
def stream(self): def stream(self):
return self.__stream_client return self.__stream_client
@property @property
def request(self): def request(self):
return self.__http_request return self.__http_request
@property @property
def stream_request(self): def stream_request(self):
return self.__http_stream_request return self.__http_stream_request
@property @property
def parse(self): def parse(self):
return self.__llm_response_parser return self.__llm_response_parser
@property @property
def parse_stream(self): def parse_stream(self):
return self.__llm_stream_parser return self.__llm_stream_parser
async def ExtractCategories(self, content: str, async def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.DefaultContentPrediction: ) -> types.DefaultContentPrediction:
result = await self.__options.merge_options(baml_options).call_function_async(function_name="ExtractCategories", args={ result = await self.__options.merge_options(baml_options).call_function_async(
"content": content, function_name="ExtractCategories",
}) args={
return typing.cast(types.DefaultContentPrediction, result.cast_to(types, types, stream_types, False, __runtime__)) "content": content,
async def ExtractContentGraphGeneric(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None, },
)
return typing.cast(
types.DefaultContentPrediction,
result.cast_to(types, types, stream_types, False, __runtime__),
)
async def ExtractContentGraphGeneric(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.KnowledgeGraph: ) -> types.KnowledgeGraph:
result = await self.__options.merge_options(baml_options).call_function_async(function_name="ExtractContentGraphGeneric", args={ result = await self.__options.merge_options(baml_options).call_function_async(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractContentGraphGeneric",
}) args={
return typing.cast(types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)) "content": content,
async def ExtractDynamicContentGraph(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None, "mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)
)
async def ExtractDynamicContentGraph(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.DynamicKnowledgeGraph: ) -> types.DynamicKnowledgeGraph:
result = await self.__options.merge_options(baml_options).call_function_async(function_name="ExtractDynamicContentGraph", args={ result = await self.__options.merge_options(baml_options).call_function_async(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractDynamicContentGraph",
}) args={
return typing.cast(types.DynamicKnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)) "content": content,
async def SummarizeCode(self, content: str, "mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.DynamicKnowledgeGraph,
result.cast_to(types, types, stream_types, False, __runtime__),
)
async def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.SummarizedCode: ) -> types.SummarizedCode:
result = await self.__options.merge_options(baml_options).call_function_async(function_name="SummarizeCode", args={ result = await self.__options.merge_options(baml_options).call_function_async(
"content": content, function_name="SummarizeCode",
}) args={
return typing.cast(types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__)) "content": content,
async def SummarizeContent(self, content: str, },
)
return typing.cast(
types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__)
)
async def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.SummarizedContent: ) -> types.SummarizedContent:
result = await self.__options.merge_options(baml_options).call_function_async(function_name="SummarizeContent", args={ result = await self.__options.merge_options(baml_options).call_function_async(
"content": content, function_name="SummarizeContent",
}) args={
return typing.cast(types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__)) "content": content,
},
)
return typing.cast(
types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__)
)
class BamlStreamClient: class BamlStreamClient:
@ -117,65 +184,147 @@ class BamlStreamClient:
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
def ExtractCategories(self, content: str, def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.BamlStream[stream_types.DefaultContentPrediction, types.DefaultContentPrediction]: ) -> baml_py.BamlStream[stream_types.DefaultContentPrediction, types.DefaultContentPrediction]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="ExtractCategories", args={ ctx, result = self.__options.merge_options(baml_options).create_async_stream(
"content": content, function_name="ExtractCategories",
}) args={
return baml_py.BamlStream[stream_types.DefaultContentPrediction, types.DefaultContentPrediction]( "content": content,
result, },
lambda x: typing.cast(stream_types.DefaultContentPrediction, x.cast_to(types, types, stream_types, True, __runtime__)),
lambda x: typing.cast(types.DefaultContentPrediction, x.cast_to(types, types, stream_types, False, __runtime__)),
ctx,
) )
def ExtractContentGraphGeneric(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None, return baml_py.BamlStream[
stream_types.DefaultContentPrediction, types.DefaultContentPrediction
](
result,
lambda x: typing.cast(
stream_types.DefaultContentPrediction,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.DefaultContentPrediction,
x.cast_to(types, types, stream_types, False, __runtime__),
),
ctx,
)
def ExtractContentGraphGeneric(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.BamlStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]: ) -> baml_py.BamlStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="ExtractContentGraphGeneric", args={ ctx, result = self.__options.merge_options(baml_options).create_async_stream(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractContentGraphGeneric",
}) args={
return baml_py.BamlStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]( "content": content,
result, "mode": mode,
lambda x: typing.cast(stream_types.KnowledgeGraph, x.cast_to(types, types, stream_types, True, __runtime__)), "custom_prompt_content": custom_prompt_content,
lambda x: typing.cast(types.KnowledgeGraph, x.cast_to(types, types, stream_types, False, __runtime__)), },
ctx,
) )
def ExtractDynamicContentGraph(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None, return baml_py.BamlStream[stream_types.KnowledgeGraph, types.KnowledgeGraph](
result,
lambda x: typing.cast(
stream_types.KnowledgeGraph,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.KnowledgeGraph, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
)
def ExtractDynamicContentGraph(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.BamlStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph]: ) -> baml_py.BamlStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="ExtractDynamicContentGraph", args={ ctx, result = self.__options.merge_options(baml_options).create_async_stream(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractDynamicContentGraph",
}) args={
return baml_py.BamlStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph]( "content": content,
result, "mode": mode,
lambda x: typing.cast(stream_types.DynamicKnowledgeGraph, x.cast_to(types, types, stream_types, True, __runtime__)), "custom_prompt_content": custom_prompt_content,
lambda x: typing.cast(types.DynamicKnowledgeGraph, x.cast_to(types, types, stream_types, False, __runtime__)), },
ctx,
) )
def SummarizeCode(self, content: str, return baml_py.BamlStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph](
result,
lambda x: typing.cast(
stream_types.DynamicKnowledgeGraph,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.DynamicKnowledgeGraph,
x.cast_to(types, types, stream_types, False, __runtime__),
),
ctx,
)
def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.BamlStream[stream_types.SummarizedCode, types.SummarizedCode]: ) -> baml_py.BamlStream[stream_types.SummarizedCode, types.SummarizedCode]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="SummarizeCode", args={ ctx, result = self.__options.merge_options(baml_options).create_async_stream(
"content": content, function_name="SummarizeCode",
}) args={
return baml_py.BamlStream[stream_types.SummarizedCode, types.SummarizedCode]( "content": content,
result, },
lambda x: typing.cast(stream_types.SummarizedCode, x.cast_to(types, types, stream_types, True, __runtime__)),
lambda x: typing.cast(types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)),
ctx,
) )
def SummarizeContent(self, content: str, return baml_py.BamlStream[stream_types.SummarizedCode, types.SummarizedCode](
result,
lambda x: typing.cast(
stream_types.SummarizedCode,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
)
def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.BamlStream[stream_types.SummarizedContent, types.SummarizedContent]: ) -> baml_py.BamlStream[stream_types.SummarizedContent, types.SummarizedContent]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="SummarizeContent", args={ ctx, result = self.__options.merge_options(baml_options).create_async_stream(
"content": content, function_name="SummarizeContent",
}) args={
"content": content,
},
)
return baml_py.BamlStream[stream_types.SummarizedContent, types.SummarizedContent]( return baml_py.BamlStream[stream_types.SummarizedContent, types.SummarizedContent](
result, result,
lambda x: typing.cast(stream_types.SummarizedContent, x.cast_to(types, types, stream_types, True, __runtime__)), lambda x: typing.cast(
lambda x: typing.cast(types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)), stream_types.SummarizedContent,
ctx, x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
) )
@ -185,40 +334,98 @@ class BamlHttpRequestClient:
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
async def ExtractCategories(self, content: str, async def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractCategories", args={ result = await self.__options.merge_options(baml_options).create_http_request_async(
"content": content, function_name="ExtractCategories",
}, mode="request") args={
"content": content,
},
mode="request",
)
return result return result
async def ExtractContentGraphGeneric(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None,
async def ExtractContentGraphGeneric(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractContentGraphGeneric", args={ result = await self.__options.merge_options(baml_options).create_http_request_async(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractContentGraphGeneric",
}, mode="request") args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="request",
)
return result return result
async def ExtractDynamicContentGraph(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None,
async def ExtractDynamicContentGraph(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractDynamicContentGraph", args={ result = await self.__options.merge_options(baml_options).create_http_request_async(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractDynamicContentGraph",
}, mode="request") args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="request",
)
return result return result
async def SummarizeCode(self, content: str,
async def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="SummarizeCode", args={ result = await self.__options.merge_options(baml_options).create_http_request_async(
"content": content, function_name="SummarizeCode",
}, mode="request") args={
"content": content,
},
mode="request",
)
return result return result
async def SummarizeContent(self, content: str,
async def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="SummarizeContent", args={ result = await self.__options.merge_options(baml_options).create_http_request_async(
"content": content, function_name="SummarizeContent",
}, mode="request") args={
"content": content,
},
mode="request",
)
return result return result
@ -228,40 +435,98 @@ class BamlHttpStreamRequestClient:
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
async def ExtractCategories(self, content: str, async def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractCategories", args={ result = await self.__options.merge_options(baml_options).create_http_request_async(
"content": content, function_name="ExtractCategories",
}, mode="stream") args={
"content": content,
},
mode="stream",
)
return result return result
async def ExtractContentGraphGeneric(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None,
async def ExtractContentGraphGeneric(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractContentGraphGeneric", args={ result = await self.__options.merge_options(baml_options).create_http_request_async(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractContentGraphGeneric",
}, mode="stream") args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="stream",
)
return result return result
async def ExtractDynamicContentGraph(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None,
async def ExtractDynamicContentGraph(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractDynamicContentGraph", args={ result = await self.__options.merge_options(baml_options).create_http_request_async(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractDynamicContentGraph",
}, mode="stream") args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="stream",
)
return result return result
async def SummarizeCode(self, content: str,
async def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="SummarizeCode", args={ result = await self.__options.merge_options(baml_options).create_http_request_async(
"content": content, function_name="SummarizeCode",
}, mode="stream") args={
"content": content,
},
mode="stream",
)
return result return result
async def SummarizeContent(self, content: str,
async def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="SummarizeContent", args={ result = await self.__options.merge_options(baml_options).create_http_request_async(
"content": content, function_name="SummarizeContent",
}, mode="stream") args={
"content": content,
},
mode="stream",
)
return result return result

12
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/globals.py
View file

@ -19,17 +19,19 @@ from .inlinedbaml import get_baml_files
from typing import Dict from typing import Dict
DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME = BamlRuntime.from_files( DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME = BamlRuntime.from_files(
"baml_src", "baml_src", get_baml_files(), os.environ.copy()
get_baml_files(),
os.environ.copy()
) )
DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX = BamlCtxManager(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME) DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX = BamlCtxManager(
DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME
)
def reset_baml_env_vars(env_vars: Dict[str, str]): def reset_baml_env_vars(env_vars: Dict[str, str]):
warnings.warn( warnings.warn(
"reset_baml_env_vars is deprecated and should be removed. Environment variables are now lazily loaded on each function call", "reset_baml_env_vars is deprecated and should be removed. Environment variables are now lazily loaded on each function call",
DeprecationWarning, DeprecationWarning,
stacklevel=2 stacklevel=2,
) )
__all__ = [] __all__ = []

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

File diff suppressed because one or more lines are too long

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

@ -16,6 +16,7 @@ import typing_extensions
from . import stream_types, types from . import stream_types, types
from .runtime import DoNotUseDirectlyCallManager, BamlCallOptions from .runtime import DoNotUseDirectlyCallManager, BamlCallOptions
class LlmResponseParser: class LlmResponseParser:
__options: DoNotUseDirectlyCallManager __options: DoNotUseDirectlyCallManager
@ -23,37 +24,56 @@ class LlmResponseParser:
self.__options = options self.__options = options
def ExtractCategories( def ExtractCategories(
self, llm_response: str, baml_options: BamlCallOptions = {}, self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.DefaultContentPrediction: ) -> types.DefaultContentPrediction:
result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractCategories", llm_response=llm_response, mode="request") result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractCategories", llm_response=llm_response, mode="request"
)
return typing.cast(types.DefaultContentPrediction, result) return typing.cast(types.DefaultContentPrediction, result)
def ExtractContentGraphGeneric( def ExtractContentGraphGeneric(
self, llm_response: str, baml_options: BamlCallOptions = {}, self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.KnowledgeGraph: ) -> types.KnowledgeGraph:
result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractContentGraphGeneric", llm_response=llm_response, mode="request") result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractContentGraphGeneric", llm_response=llm_response, mode="request"
)
return typing.cast(types.KnowledgeGraph, result) return typing.cast(types.KnowledgeGraph, result)
def ExtractDynamicContentGraph( def ExtractDynamicContentGraph(
self, llm_response: str, baml_options: BamlCallOptions = {}, self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.DynamicKnowledgeGraph: ) -> types.DynamicKnowledgeGraph:
result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractDynamicContentGraph", llm_response=llm_response, mode="request") result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractDynamicContentGraph", llm_response=llm_response, mode="request"
)
return typing.cast(types.DynamicKnowledgeGraph, result) return typing.cast(types.DynamicKnowledgeGraph, result)
def SummarizeCode( def SummarizeCode(
self, llm_response: str, baml_options: BamlCallOptions = {}, self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.SummarizedCode: ) -> types.SummarizedCode:
result = self.__options.merge_options(baml_options).parse_response(function_name="SummarizeCode", llm_response=llm_response, mode="request") result = self.__options.merge_options(baml_options).parse_response(
function_name="SummarizeCode", llm_response=llm_response, mode="request"
)
return typing.cast(types.SummarizedCode, result) return typing.cast(types.SummarizedCode, result)
def SummarizeContent( def SummarizeContent(
self, llm_response: str, baml_options: BamlCallOptions = {}, self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.SummarizedContent: ) -> types.SummarizedContent:
result = self.__options.merge_options(baml_options).parse_response(function_name="SummarizeContent", llm_response=llm_response, mode="request") result = self.__options.merge_options(baml_options).parse_response(
function_name="SummarizeContent", llm_response=llm_response, mode="request"
)
return typing.cast(types.SummarizedContent, result) return typing.cast(types.SummarizedContent, result)
class LlmStreamParser: class LlmStreamParser:
__options: DoNotUseDirectlyCallManager __options: DoNotUseDirectlyCallManager
@ -61,33 +81,51 @@ class LlmStreamParser:
self.__options = options self.__options = options
def ExtractCategories( def ExtractCategories(
self, llm_response: str, baml_options: BamlCallOptions = {}, self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.DefaultContentPrediction: ) -> stream_types.DefaultContentPrediction:
result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractCategories", llm_response=llm_response, mode="stream") result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractCategories", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.DefaultContentPrediction, result) return typing.cast(stream_types.DefaultContentPrediction, result)
def ExtractContentGraphGeneric( def ExtractContentGraphGeneric(
self, llm_response: str, baml_options: BamlCallOptions = {}, self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.KnowledgeGraph: ) -> stream_types.KnowledgeGraph:
result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractContentGraphGeneric", llm_response=llm_response, mode="stream") result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractContentGraphGeneric", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.KnowledgeGraph, result) return typing.cast(stream_types.KnowledgeGraph, result)
def ExtractDynamicContentGraph( def ExtractDynamicContentGraph(
self, llm_response: str, baml_options: BamlCallOptions = {}, self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.DynamicKnowledgeGraph: ) -> stream_types.DynamicKnowledgeGraph:
result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractDynamicContentGraph", llm_response=llm_response, mode="stream") result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractDynamicContentGraph", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.DynamicKnowledgeGraph, result) return typing.cast(stream_types.DynamicKnowledgeGraph, result)
def SummarizeCode( def SummarizeCode(
self, llm_response: str, baml_options: BamlCallOptions = {}, self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.SummarizedCode: ) -> stream_types.SummarizedCode:
result = self.__options.merge_options(baml_options).parse_response(function_name="SummarizeCode", llm_response=llm_response, mode="stream") result = self.__options.merge_options(baml_options).parse_response(
function_name="SummarizeCode", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.SummarizedCode, result) return typing.cast(stream_types.SummarizedCode, result)
def SummarizeContent( def SummarizeContent(
self, llm_response: str, baml_options: BamlCallOptions = {}, self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.SummarizedContent: ) -> stream_types.SummarizedContent:
result = self.__options.merge_options(baml_options).parse_response(function_name="SummarizeContent", llm_response=llm_response, mode="stream") result = self.__options.merge_options(baml_options).parse_response(
function_name="SummarizeContent", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.SummarizedContent, result) return typing.cast(stream_types.SummarizedContent, result)

24
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/runtime.py
View file

@ -17,7 +17,10 @@ import typing_extensions
import baml_py import baml_py
from . import types, stream_types, type_builder from . import types, stream_types, type_builder
from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME as __runtime__, DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX as __ctx__manager__ from .globals import (
DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME as __runtime__,
DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX as __ctx__manager__,
)
class BamlCallOptions(typing.TypedDict, total=False): class BamlCallOptions(typing.TypedDict, total=False):
@ -48,9 +51,6 @@ class _ResolvedBamlOptions:
self.env_vars = env_vars self.env_vars = env_vars
class DoNotUseDirectlyCallManager: class DoNotUseDirectlyCallManager:
def __init__(self, baml_options: BamlCallOptions): def __init__(self, baml_options: BamlCallOptions):
self.__baml_options = baml_options self.__baml_options = baml_options
@ -74,7 +74,9 @@ class DoNotUseDirectlyCallManager:
collectors_as_list = ( collectors_as_list = (
collector collector
if isinstance(collector, list) if isinstance(collector, list)
else [collector] if collector is not None else [] else [collector]
if collector is not None
else []
) )
env_vars = os.environ.copy() env_vars = os.environ.copy()
for k, v in self.__baml_options.get("env", {}).items(): for k, v in self.__baml_options.get("env", {}).items():
@ -164,7 +166,9 @@ class DoNotUseDirectlyCallManager:
*, *,
function_name: str, function_name: str,
args: typing.Dict[str, typing.Any], args: typing.Dict[str, typing.Any],
) -> typing.Tuple[baml_py.baml_py.RuntimeContextManager, baml_py.baml_py.SyncFunctionResultStream]: ) -> typing.Tuple[
baml_py.baml_py.RuntimeContextManager, baml_py.baml_py.SyncFunctionResultStream
]:
resolved_options = self.__resolve() resolved_options = self.__resolve()
ctx = __ctx__manager__.get() ctx = __ctx__manager__.get()
result = __runtime__.stream_function_sync( result = __runtime__.stream_function_sync(
@ -232,7 +236,13 @@ class DoNotUseDirectlyCallManager:
mode == "stream", mode == "stream",
) )
def parse_response(self, *, function_name: str, llm_response: str, mode: typing_extensions.Literal["stream", "request"]) -> typing.Any: def parse_response(
self,
*,
function_name: str,
llm_response: str,
mode: typing_extensions.Literal["stream", "request"],
) -> typing.Any:
resolved_options = self.__resolve() resolved_options = self.__resolve()
return __runtime__.parse_llm_response( return __runtime__.parse_llm_response(
function_name, function_name,

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

@ -18,28 +18,37 @@ import baml_py
from . import types from . import types
StreamStateValueT = typing.TypeVar('StreamStateValueT') StreamStateValueT = typing.TypeVar("StreamStateValueT")
class StreamState(BaseModel, typing.Generic[StreamStateValueT]): class StreamState(BaseModel, typing.Generic[StreamStateValueT]):
value: StreamStateValueT value: StreamStateValueT
state: typing_extensions.Literal["Pending", "Incomplete", "Complete"] state: typing_extensions.Literal["Pending", "Incomplete", "Complete"]
# ######################################################################### # #########################################################################
# Generated classes (17) # Generated classes (17)
# ######################################################################### # #########################################################################
class AudioContent(BaseModel): class AudioContent(BaseModel):
type: typing.Optional[str] = None type: typing.Optional[str] = None
subclass: typing.List[str] subclass: typing.List[str]
class ContentLabel(BaseModel): class ContentLabel(BaseModel):
content_type: typing.Optional[typing.Union[str, str, str, str, str, str, str]] = None content_type: typing.Optional[typing.Union[str, str, str, str, str, str, str]] = None
type: typing.Optional[str] = None type: typing.Optional[str] = None
subclass: typing.List[str] subclass: typing.List[str]
class DefaultContentPrediction(BaseModel): class DefaultContentPrediction(BaseModel):
label: typing.Optional["ContentLabel"] = None label: typing.Optional["ContentLabel"] = None
class DynamicKnowledgeGraph(BaseModel): class DynamicKnowledgeGraph(BaseModel):
model_config = ConfigDict(extra='allow') model_config = ConfigDict(extra="allow")
class Edge(BaseModel): class Edge(BaseModel):
# doc string for edge # doc string for edge
@ -49,39 +58,47 @@ class Edge(BaseModel):
target_node_id: typing.Optional[str] = None target_node_id: typing.Optional[str] = None
relationship_name: typing.Optional[str] = None relationship_name: typing.Optional[str] = None
class ImageContent(BaseModel): class ImageContent(BaseModel):
type: typing.Optional[str] = None type: typing.Optional[str] = None
subclass: typing.List[str] subclass: typing.List[str]
class KnowledgeGraph(BaseModel): class KnowledgeGraph(BaseModel):
nodes: typing.List["types.Node"] nodes: typing.List["types.Node"]
edges: typing.List["Edge"] edges: typing.List["Edge"]
class Model3DContent(BaseModel): class Model3DContent(BaseModel):
type: typing.Optional[str] = None type: typing.Optional[str] = None
subclass: typing.List[str] subclass: typing.List[str]
class MultimediaContent(BaseModel): class MultimediaContent(BaseModel):
type: typing.Optional[str] = None type: typing.Optional[str] = None
subclass: typing.List[str] subclass: typing.List[str]
class Node(BaseModel): class Node(BaseModel):
model_config = ConfigDict(extra='allow') model_config = ConfigDict(extra="allow")
id: typing.Optional[str] = None id: typing.Optional[str] = None
name: typing.Optional[str] = None name: typing.Optional[str] = None
type: typing.Optional[str] = None type: typing.Optional[str] = None
description: typing.Optional[str] = None description: typing.Optional[str] = None
class ProceduralContent(BaseModel): class ProceduralContent(BaseModel):
type: typing.Optional[str] = None type: typing.Optional[str] = None
subclass: typing.List[str] subclass: typing.List[str]
class SummarizedClass(BaseModel): class SummarizedClass(BaseModel):
name: typing.Optional[str] = None name: typing.Optional[str] = None
description: typing.Optional[str] = None description: typing.Optional[str] = None
methods: typing.Optional[typing.List["SummarizedFunction"]] = None methods: typing.Optional[typing.List["SummarizedFunction"]] = None
decorators: typing.Optional[typing.List[str]] = None decorators: typing.Optional[typing.List[str]] = None
class SummarizedCode(BaseModel): class SummarizedCode(BaseModel):
high_level_summary: typing.Optional[str] = None high_level_summary: typing.Optional[str] = None
key_features: typing.List[str] key_features: typing.List[str]
@ -91,10 +108,12 @@ class SummarizedCode(BaseModel):
functions: typing.List["SummarizedFunction"] functions: typing.List["SummarizedFunction"]
workflow_description: typing.Optional[str] = None workflow_description: typing.Optional[str] = None
class SummarizedContent(BaseModel): class SummarizedContent(BaseModel):
summary: typing.Optional[str] = None summary: typing.Optional[str] = None
description: typing.Optional[str] = None description: typing.Optional[str] = None
class SummarizedFunction(BaseModel): class SummarizedFunction(BaseModel):
name: typing.Optional[str] = None name: typing.Optional[str] = None
description: typing.Optional[str] = None description: typing.Optional[str] = None
@ -102,14 +121,17 @@ class SummarizedFunction(BaseModel):
outputs: typing.Optional[typing.List[str]] = None outputs: typing.Optional[typing.List[str]] = None
decorators: typing.Optional[typing.List[str]] = None decorators: typing.Optional[typing.List[str]] = None
class TextContent(BaseModel): class TextContent(BaseModel):
type: typing.Optional[str] = None type: typing.Optional[str] = None
subclass: typing.List[str] subclass: typing.List[str]
class VideoContent(BaseModel): class VideoContent(BaseModel):
type: typing.Optional[str] = None type: typing.Optional[str] = None
subclass: typing.List[str] subclass: typing.List[str]
# ######################################################################### # #########################################################################
# Generated type aliases (0) # Generated type aliases (0)
# ######################################################################### # #########################################################################

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

@ -19,6 +19,7 @@ from .parser import LlmResponseParser, LlmStreamParser
from .runtime import DoNotUseDirectlyCallManager, BamlCallOptions from .runtime import DoNotUseDirectlyCallManager, BamlCallOptions
from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME as __runtime__ from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME as __runtime__
class BamlSyncClient: class BamlSyncClient:
__options: DoNotUseDirectlyCallManager __options: DoNotUseDirectlyCallManager
__stream_client: "BamlStreamClient" __stream_client: "BamlStreamClient"
@ -48,10 +49,13 @@ class BamlSyncClient:
self.__llm_response_parser = LlmResponseParser(self.__options) self.__llm_response_parser = LlmResponseParser(self.__options)
self.__llm_stream_parser = LlmStreamParser(self.__options) self.__llm_stream_parser = LlmStreamParser(self.__options)
def with_options(self, def with_options(
self,
tb: typing.Optional[type_builder.TypeBuilder] = None, tb: typing.Optional[type_builder.TypeBuilder] = None,
client_registry: typing.Optional[baml_py.baml_py.ClientRegistry] = None, client_registry: typing.Optional[baml_py.baml_py.ClientRegistry] = None,
collector: typing.Optional[typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]] = None, collector: typing.Optional[
typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]
] = None,
env: typing.Optional[typing.Dict[str, typing.Optional[str]]] = None, env: typing.Optional[typing.Dict[str, typing.Optional[str]]] = None,
) -> "BamlSyncClient": ) -> "BamlSyncClient":
options: BamlCallOptions = {} options: BamlCallOptions = {}
@ -67,60 +71,124 @@ class BamlSyncClient:
@property @property
def stream(self): def stream(self):
return self.__stream_client return self.__stream_client
@property @property
def request(self): def request(self):
return self.__http_request return self.__http_request
@property @property
def stream_request(self): def stream_request(self):
return self.__http_stream_request return self.__http_stream_request
@property @property
def parse(self): def parse(self):
return self.__llm_response_parser return self.__llm_response_parser
@property @property
def parse_stream(self): def parse_stream(self):
return self.__llm_stream_parser return self.__llm_stream_parser
def ExtractCategories(self, content: str, def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.DefaultContentPrediction: ) -> types.DefaultContentPrediction:
result = self.__options.merge_options(baml_options).call_function_sync(function_name="ExtractCategories", args={ result = self.__options.merge_options(baml_options).call_function_sync(
"content": content, function_name="ExtractCategories",
}) args={
return typing.cast(types.DefaultContentPrediction, result.cast_to(types, types, stream_types, False, __runtime__)) "content": content,
def ExtractContentGraphGeneric(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None, },
)
return typing.cast(
types.DefaultContentPrediction,
result.cast_to(types, types, stream_types, False, __runtime__),
)
def ExtractContentGraphGeneric(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.KnowledgeGraph: ) -> types.KnowledgeGraph:
result = self.__options.merge_options(baml_options).call_function_sync(function_name="ExtractContentGraphGeneric", args={ result = self.__options.merge_options(baml_options).call_function_sync(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractContentGraphGeneric",
}) args={
return typing.cast(types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)) "content": content,
def ExtractDynamicContentGraph(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None, "mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)
)
def ExtractDynamicContentGraph(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.DynamicKnowledgeGraph: ) -> types.DynamicKnowledgeGraph:
result = self.__options.merge_options(baml_options).call_function_sync(function_name="ExtractDynamicContentGraph", args={ result = self.__options.merge_options(baml_options).call_function_sync(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractDynamicContentGraph",
}) args={
return typing.cast(types.DynamicKnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)) "content": content,
def SummarizeCode(self, content: str, "mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.DynamicKnowledgeGraph,
result.cast_to(types, types, stream_types, False, __runtime__),
)
def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.SummarizedCode: ) -> types.SummarizedCode:
result = self.__options.merge_options(baml_options).call_function_sync(function_name="SummarizeCode", args={ result = self.__options.merge_options(baml_options).call_function_sync(
"content": content, function_name="SummarizeCode",
}) args={
return typing.cast(types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__)) "content": content,
def SummarizeContent(self, content: str, },
)
return typing.cast(
types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__)
)
def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.SummarizedContent: ) -> types.SummarizedContent:
result = self.__options.merge_options(baml_options).call_function_sync(function_name="SummarizeContent", args={ result = self.__options.merge_options(baml_options).call_function_sync(
"content": content, function_name="SummarizeContent",
}) args={
return typing.cast(types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__)) "content": content,
},
)
return typing.cast(
types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__)
)
class BamlStreamClient: class BamlStreamClient:
@ -129,65 +197,151 @@ class BamlStreamClient:
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
def ExtractCategories(self, content: str, def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.BamlSyncStream[stream_types.DefaultContentPrediction, types.DefaultContentPrediction]: ) -> baml_py.BamlSyncStream[
ctx, result = self.__options.merge_options(baml_options).create_sync_stream(function_name="ExtractCategories", args={ stream_types.DefaultContentPrediction, types.DefaultContentPrediction
"content": content, ]:
}) ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
return baml_py.BamlSyncStream[stream_types.DefaultContentPrediction, types.DefaultContentPrediction]( function_name="ExtractCategories",
result, args={
lambda x: typing.cast(stream_types.DefaultContentPrediction, x.cast_to(types, types, stream_types, True, __runtime__)), "content": content,
lambda x: typing.cast(types.DefaultContentPrediction, x.cast_to(types, types, stream_types, False, __runtime__)), },
ctx,
) )
def ExtractContentGraphGeneric(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None, return baml_py.BamlSyncStream[
stream_types.DefaultContentPrediction, types.DefaultContentPrediction
](
result,
lambda x: typing.cast(
stream_types.DefaultContentPrediction,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.DefaultContentPrediction,
x.cast_to(types, types, stream_types, False, __runtime__),
),
ctx,
)
def ExtractContentGraphGeneric(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]: ) -> baml_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_sync_stream(function_name="ExtractContentGraphGeneric", args={ ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractContentGraphGeneric",
}) args={
return baml_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]( "content": content,
result, "mode": mode,
lambda x: typing.cast(stream_types.KnowledgeGraph, x.cast_to(types, types, stream_types, True, __runtime__)), "custom_prompt_content": custom_prompt_content,
lambda x: typing.cast(types.KnowledgeGraph, x.cast_to(types, types, stream_types, False, __runtime__)), },
ctx,
) )
def ExtractDynamicContentGraph(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None, return baml_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph](
result,
lambda x: typing.cast(
stream_types.KnowledgeGraph,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.KnowledgeGraph, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
)
def ExtractDynamicContentGraph(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.BamlSyncStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph]: ) -> baml_py.BamlSyncStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_sync_stream(function_name="ExtractDynamicContentGraph", args={ ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractDynamicContentGraph",
}) args={
return baml_py.BamlSyncStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph]( "content": content,
result, "mode": mode,
lambda x: typing.cast(stream_types.DynamicKnowledgeGraph, x.cast_to(types, types, stream_types, True, __runtime__)), "custom_prompt_content": custom_prompt_content,
lambda x: typing.cast(types.DynamicKnowledgeGraph, x.cast_to(types, types, stream_types, False, __runtime__)), },
ctx,
) )
def SummarizeCode(self, content: str, return baml_py.BamlSyncStream[
stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph
](
result,
lambda x: typing.cast(
stream_types.DynamicKnowledgeGraph,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.DynamicKnowledgeGraph,
x.cast_to(types, types, stream_types, False, __runtime__),
),
ctx,
)
def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode]: ) -> baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode]:
ctx, result = self.__options.merge_options(baml_options).create_sync_stream(function_name="SummarizeCode", args={ ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
"content": content, function_name="SummarizeCode",
}) args={
return baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode]( "content": content,
result, },
lambda x: typing.cast(stream_types.SummarizedCode, x.cast_to(types, types, stream_types, True, __runtime__)),
lambda x: typing.cast(types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)),
ctx,
) )
def SummarizeContent(self, content: str, return baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode](
result,
lambda x: typing.cast(
stream_types.SummarizedCode,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
)
def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent]: ) -> baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent]:
ctx, result = self.__options.merge_options(baml_options).create_sync_stream(function_name="SummarizeContent", args={ ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
"content": content, function_name="SummarizeContent",
}) args={
"content": content,
},
)
return baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent]( return baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent](
result, result,
lambda x: typing.cast(stream_types.SummarizedContent, x.cast_to(types, types, stream_types, True, __runtime__)), lambda x: typing.cast(
lambda x: typing.cast(types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)), stream_types.SummarizedContent,
ctx, x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
) )
@ -197,40 +351,98 @@ class BamlHttpRequestClient:
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
def ExtractCategories(self, content: str, def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="ExtractCategories", args={ result = self.__options.merge_options(baml_options).create_http_request_sync(
"content": content, function_name="ExtractCategories",
}, mode="request") args={
"content": content,
},
mode="request",
)
return result return result
def ExtractContentGraphGeneric(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None,
def ExtractContentGraphGeneric(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="ExtractContentGraphGeneric", args={ result = self.__options.merge_options(baml_options).create_http_request_sync(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractContentGraphGeneric",
}, mode="request") args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="request",
)
return result return result
def ExtractDynamicContentGraph(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None,
def ExtractDynamicContentGraph(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="ExtractDynamicContentGraph", args={ result = self.__options.merge_options(baml_options).create_http_request_sync(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractDynamicContentGraph",
}, mode="request") args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="request",
)
return result return result
def SummarizeCode(self, content: str,
def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="SummarizeCode", args={ result = self.__options.merge_options(baml_options).create_http_request_sync(
"content": content, function_name="SummarizeCode",
}, mode="request") args={
"content": content,
},
mode="request",
)
return result return result
def SummarizeContent(self, content: str,
def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="SummarizeContent", args={ result = self.__options.merge_options(baml_options).create_http_request_sync(
"content": content, function_name="SummarizeContent",
}, mode="request") args={
"content": content,
},
mode="request",
)
return result return result
@ -240,40 +452,98 @@ class BamlHttpStreamRequestClient:
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
def ExtractCategories(self, content: str, def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="ExtractCategories", args={ result = self.__options.merge_options(baml_options).create_http_request_sync(
"content": content, function_name="ExtractCategories",
}, mode="stream") args={
"content": content,
},
mode="stream",
)
return result return result
def ExtractContentGraphGeneric(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None,
def ExtractContentGraphGeneric(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="ExtractContentGraphGeneric", args={ result = self.__options.merge_options(baml_options).create_http_request_sync(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractContentGraphGeneric",
}, mode="stream") args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="stream",
)
return result return result
def ExtractDynamicContentGraph(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None,
def ExtractDynamicContentGraph(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="ExtractDynamicContentGraph", args={ result = self.__options.merge_options(baml_options).create_http_request_sync(
"content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, function_name="ExtractDynamicContentGraph",
}, mode="stream") args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="stream",
)
return result return result
def SummarizeCode(self, content: str,
def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="SummarizeCode", args={ result = self.__options.merge_options(baml_options).create_http_request_sync(
"content": content, function_name="SummarizeCode",
}, mode="stream") args={
"content": content,
},
mode="stream",
)
return result return result
def SummarizeContent(self, content: str,
def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest: ) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="SummarizeContent", args={ result = self.__options.merge_options(baml_options).create_http_request_sync(
"content": content, function_name="SummarizeContent",
}, mode="stream") args={
"content": content,
},
mode="stream",
)
return result return result

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

@ -14,9 +14,13 @@ from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX
trace = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.trace_fn trace = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.trace_fn
set_tags = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.upsert_tags set_tags = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.upsert_tags
def flush(): def flush():
DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.flush() DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.flush()
on_log_event = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.on_log_event on_log_event = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.on_log_event
__all__ = ['trace', 'set_tags', "flush", "on_log_event"] __all__ = ["trace", "set_tags", "flush", "on_log_event"]

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

@ -15,19 +15,39 @@ from baml_py import type_builder
from baml_py import baml_py from baml_py import baml_py
from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME
class TypeBuilder(type_builder.TypeBuilder): class TypeBuilder(type_builder.TypeBuilder):
def __init__(self): def __init__(self):
super().__init__(classes=set( super().__init__(
["AudioContent","ContentLabel","DefaultContentPrediction","DynamicKnowledgeGraph","Edge","ImageContent","KnowledgeGraph","Model3DContent","MultimediaContent","Node","ProceduralContent","SummarizedClass","SummarizedCode","SummarizedContent","SummarizedFunction","TextContent","VideoContent",] classes=set(
), enums=set( [
[] "AudioContent",
), runtime=DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME) "ContentLabel",
"DefaultContentPrediction",
"DynamicKnowledgeGraph",
"Edge",
"ImageContent",
"KnowledgeGraph",
"Model3DContent",
"MultimediaContent",
"Node",
"ProceduralContent",
"SummarizedClass",
"SummarizedCode",
"SummarizedContent",
"SummarizedFunction",
"TextContent",
"VideoContent",
]
),
enums=set([]),
runtime=DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME,
)
# ######################################################################### # #########################################################################
# Generated enums 0 # Generated enums 0
# ######################################################################### # #########################################################################
# ######################################################################### # #########################################################################
# Generated classes 17 # Generated classes 17
# ######################################################################### # #########################################################################
@ -101,7 +121,6 @@ class TypeBuilder(type_builder.TypeBuilder):
return VideoContentViewer(self) return VideoContentViewer(self)
# ######################################################################### # #########################################################################
# Generated enums 0 # Generated enums 0
# ######################################################################### # #########################################################################
@ -111,11 +130,17 @@ class TypeBuilder(type_builder.TypeBuilder):
# Generated classes 17 # Generated classes 17
# ######################################################################### # #########################################################################
class AudioContentAst: class AudioContentAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("AudioContent") self._bldr = _tb.class_("AudioContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ]) self._properties: typing.Set[str] = set(
[
"type",
"subclass",
]
)
self._props = AudioContentProperties(self._bldr, self._properties) self._props = AudioContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -130,18 +155,17 @@ class AudioContentViewer(AudioContentAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class AudioContentProperties: class AudioContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def type(self) -> type_builder.ClassPropertyViewer: def type(self) -> type_builder.ClassPropertyViewer:
@ -152,13 +176,17 @@ class AudioContentProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass")) return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
class ContentLabelAst: class ContentLabelAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("ContentLabel") self._bldr = _tb.class_("ContentLabel")
self._properties: typing.Set[str] = set([ "content_type", "type", "subclass", ]) self._properties: typing.Set[str] = set(
[
"content_type",
"type",
"subclass",
]
)
self._props = ContentLabelProperties(self._bldr, self._properties) self._props = ContentLabelProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -173,18 +201,17 @@ class ContentLabelViewer(ContentLabelAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class ContentLabelProperties: class ContentLabelProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def content_type(self) -> type_builder.ClassPropertyViewer: def content_type(self) -> type_builder.ClassPropertyViewer:
@ -199,13 +226,15 @@ class ContentLabelProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass")) return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
class DefaultContentPredictionAst: class DefaultContentPredictionAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("DefaultContentPrediction") self._bldr = _tb.class_("DefaultContentPrediction")
self._properties: typing.Set[str] = set([ "label", ]) self._properties: typing.Set[str] = set(
[
"label",
]
)
self._props = DefaultContentPredictionProperties(self._bldr, self._properties) self._props = DefaultContentPredictionProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -220,31 +249,28 @@ class DefaultContentPredictionViewer(DefaultContentPredictionAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class DefaultContentPredictionProperties: class DefaultContentPredictionProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def label(self) -> type_builder.ClassPropertyViewer: def label(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("label")) return type_builder.ClassPropertyViewer(self.__bldr.property("label"))
class DynamicKnowledgeGraphAst: class DynamicKnowledgeGraphAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("DynamicKnowledgeGraph") self._bldr = _tb.class_("DynamicKnowledgeGraph")
self._properties: typing.Set[str] = set([ ]) self._properties: typing.Set[str] = set([])
self._props = DynamicKnowledgeGraphProperties(self._bldr, self._properties) self._props = DynamicKnowledgeGraphProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -259,7 +285,6 @@ class DynamicKnowledgeGraphBuilder(DynamicKnowledgeGraphAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def add_property(self, name: str, type: baml_py.FieldType) -> baml_py.ClassPropertyBuilder: def add_property(self, name: str, type: baml_py.FieldType) -> baml_py.ClassPropertyBuilder:
if name in self._properties: if name in self._properties:
raise ValueError(f"Property {name} already exists.") raise ValueError(f"Property {name} already exists.")
@ -269,13 +294,10 @@ class DynamicKnowledgeGraphBuilder(DynamicKnowledgeGraphAst):
return [(name, self._bldr.property(name)) for name in self._properties] return [(name, self._bldr.property(name)) for name in self._properties]
class DynamicKnowledgeGraphProperties: class DynamicKnowledgeGraphProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
def __getattr__(self, name: str) -> baml_py.ClassPropertyBuilder: def __getattr__(self, name: str) -> baml_py.ClassPropertyBuilder:
if name not in self.__properties: if name not in self.__properties:
@ -283,14 +305,17 @@ class DynamicKnowledgeGraphProperties:
return self.__bldr.property(name) return self.__bldr.property(name)
class EdgeAst: class EdgeAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("Edge") self._bldr = _tb.class_("Edge")
self._properties: typing.Set[str] = set([ "source_node_id", "target_node_id", "relationship_name", ]) self._properties: typing.Set[str] = set(
[
"source_node_id",
"target_node_id",
"relationship_name",
]
)
self._props = EdgeProperties(self._bldr, self._properties) self._props = EdgeProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -305,18 +330,17 @@ class EdgeViewer(EdgeAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class EdgeProperties: class EdgeProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def source_node_id(self) -> type_builder.ClassPropertyViewer: def source_node_id(self) -> type_builder.ClassPropertyViewer:
@ -331,13 +355,16 @@ class EdgeProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("relationship_name")) return type_builder.ClassPropertyViewer(self.__bldr.property("relationship_name"))
class ImageContentAst: class ImageContentAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("ImageContent") self._bldr = _tb.class_("ImageContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ]) self._properties: typing.Set[str] = set(
[
"type",
"subclass",
]
)
self._props = ImageContentProperties(self._bldr, self._properties) self._props = ImageContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -352,18 +379,17 @@ class ImageContentViewer(ImageContentAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class ImageContentProperties: class ImageContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def type(self) -> type_builder.ClassPropertyViewer: def type(self) -> type_builder.ClassPropertyViewer:
@ -374,13 +400,16 @@ class ImageContentProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass")) return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
class KnowledgeGraphAst: class KnowledgeGraphAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("KnowledgeGraph") self._bldr = _tb.class_("KnowledgeGraph")
self._properties: typing.Set[str] = set([ "nodes", "edges", ]) self._properties: typing.Set[str] = set(
[
"nodes",
"edges",
]
)
self._props = KnowledgeGraphProperties(self._bldr, self._properties) self._props = KnowledgeGraphProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -395,18 +424,17 @@ class KnowledgeGraphViewer(KnowledgeGraphAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class KnowledgeGraphProperties: class KnowledgeGraphProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def nodes(self) -> type_builder.ClassPropertyViewer: def nodes(self) -> type_builder.ClassPropertyViewer:
@ -417,13 +445,16 @@ class KnowledgeGraphProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("edges")) return type_builder.ClassPropertyViewer(self.__bldr.property("edges"))
class Model3DContentAst: class Model3DContentAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("Model3DContent") self._bldr = _tb.class_("Model3DContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ]) self._properties: typing.Set[str] = set(
[
"type",
"subclass",
]
)
self._props = Model3DContentProperties(self._bldr, self._properties) self._props = Model3DContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -438,18 +469,17 @@ class Model3DContentViewer(Model3DContentAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class Model3DContentProperties: class Model3DContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def type(self) -> type_builder.ClassPropertyViewer: def type(self) -> type_builder.ClassPropertyViewer:
@ -460,13 +490,16 @@ class Model3DContentProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass")) return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
class MultimediaContentAst: class MultimediaContentAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("MultimediaContent") self._bldr = _tb.class_("MultimediaContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ]) self._properties: typing.Set[str] = set(
[
"type",
"subclass",
]
)
self._props = MultimediaContentProperties(self._bldr, self._properties) self._props = MultimediaContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -481,18 +514,17 @@ class MultimediaContentViewer(MultimediaContentAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class MultimediaContentProperties: class MultimediaContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def type(self) -> type_builder.ClassPropertyViewer: def type(self) -> type_builder.ClassPropertyViewer:
@ -503,13 +535,18 @@ class MultimediaContentProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass")) return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
class NodeAst: class NodeAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("Node") self._bldr = _tb.class_("Node")
self._properties: typing.Set[str] = set([ "id", "name", "type", "description", ]) self._properties: typing.Set[str] = set(
[
"id",
"name",
"type",
"description",
]
)
self._props = NodeProperties(self._bldr, self._properties) self._props = NodeProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -524,7 +561,6 @@ class NodeBuilder(NodeAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def add_property(self, name: str, type: baml_py.FieldType) -> baml_py.ClassPropertyBuilder: def add_property(self, name: str, type: baml_py.FieldType) -> baml_py.ClassPropertyBuilder:
if name in self._properties: if name in self._properties:
raise ValueError(f"Property {name} already exists.") raise ValueError(f"Property {name} already exists.")
@ -534,20 +570,16 @@ class NodeBuilder(NodeAst):
return [(name, self._bldr.property(name)) for name in self._properties] return [(name, self._bldr.property(name)) for name in self._properties]
class NodeProperties: class NodeProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
def __getattr__(self, name: str) -> baml_py.ClassPropertyBuilder: def __getattr__(self, name: str) -> baml_py.ClassPropertyBuilder:
if name not in self.__properties: if name not in self.__properties:
raise AttributeError(f"Property {name} not found.") raise AttributeError(f"Property {name} not found.")
return self.__bldr.property(name) return self.__bldr.property(name)
@property @property
def id(self) -> baml_py.ClassPropertyBuilder: def id(self) -> baml_py.ClassPropertyBuilder:
return self.__bldr.property("id") return self.__bldr.property("id")
@ -565,13 +597,16 @@ class NodeProperties:
return self.__bldr.property("description") return self.__bldr.property("description")
class ProceduralContentAst: class ProceduralContentAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("ProceduralContent") self._bldr = _tb.class_("ProceduralContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ]) self._properties: typing.Set[str] = set(
[
"type",
"subclass",
]
)
self._props = ProceduralContentProperties(self._bldr, self._properties) self._props = ProceduralContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -586,18 +621,17 @@ class ProceduralContentViewer(ProceduralContentAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class ProceduralContentProperties: class ProceduralContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def type(self) -> type_builder.ClassPropertyViewer: def type(self) -> type_builder.ClassPropertyViewer:
@ -608,13 +642,18 @@ class ProceduralContentProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass")) return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
class SummarizedClassAst: class SummarizedClassAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("SummarizedClass") self._bldr = _tb.class_("SummarizedClass")
self._properties: typing.Set[str] = set([ "name", "description", "methods", "decorators", ]) self._properties: typing.Set[str] = set(
[
"name",
"description",
"methods",
"decorators",
]
)
self._props = SummarizedClassProperties(self._bldr, self._properties) self._props = SummarizedClassProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -629,18 +668,17 @@ class SummarizedClassViewer(SummarizedClassAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class SummarizedClassProperties: class SummarizedClassProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def name(self) -> type_builder.ClassPropertyViewer: def name(self) -> type_builder.ClassPropertyViewer:
@ -659,13 +697,21 @@ class SummarizedClassProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("decorators")) return type_builder.ClassPropertyViewer(self.__bldr.property("decorators"))
class SummarizedCodeAst: class SummarizedCodeAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("SummarizedCode") self._bldr = _tb.class_("SummarizedCode")
self._properties: typing.Set[str] = set([ "high_level_summary", "key_features", "imports", "constants", "classes", "functions", "workflow_description", ]) self._properties: typing.Set[str] = set(
[
"high_level_summary",
"key_features",
"imports",
"constants",
"classes",
"functions",
"workflow_description",
]
)
self._props = SummarizedCodeProperties(self._bldr, self._properties) self._props = SummarizedCodeProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -680,18 +726,17 @@ class SummarizedCodeViewer(SummarizedCodeAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class SummarizedCodeProperties: class SummarizedCodeProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def high_level_summary(self) -> type_builder.ClassPropertyViewer: def high_level_summary(self) -> type_builder.ClassPropertyViewer:
@ -722,13 +767,16 @@ class SummarizedCodeProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("workflow_description")) return type_builder.ClassPropertyViewer(self.__bldr.property("workflow_description"))
class SummarizedContentAst: class SummarizedContentAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("SummarizedContent") self._bldr = _tb.class_("SummarizedContent")
self._properties: typing.Set[str] = set([ "summary", "description", ]) self._properties: typing.Set[str] = set(
[
"summary",
"description",
]
)
self._props = SummarizedContentProperties(self._bldr, self._properties) self._props = SummarizedContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -743,18 +791,17 @@ class SummarizedContentViewer(SummarizedContentAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class SummarizedContentProperties: class SummarizedContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def summary(self) -> type_builder.ClassPropertyViewer: def summary(self) -> type_builder.ClassPropertyViewer:
@ -765,13 +812,19 @@ class SummarizedContentProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("description")) return type_builder.ClassPropertyViewer(self.__bldr.property("description"))
class SummarizedFunctionAst: class SummarizedFunctionAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("SummarizedFunction") self._bldr = _tb.class_("SummarizedFunction")
self._properties: typing.Set[str] = set([ "name", "description", "inputs", "outputs", "decorators", ]) self._properties: typing.Set[str] = set(
[
"name",
"description",
"inputs",
"outputs",
"decorators",
]
)
self._props = SummarizedFunctionProperties(self._bldr, self._properties) self._props = SummarizedFunctionProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -786,18 +839,17 @@ class SummarizedFunctionViewer(SummarizedFunctionAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class SummarizedFunctionProperties: class SummarizedFunctionProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def name(self) -> type_builder.ClassPropertyViewer: def name(self) -> type_builder.ClassPropertyViewer:
@ -820,13 +872,16 @@ class SummarizedFunctionProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("decorators")) return type_builder.ClassPropertyViewer(self.__bldr.property("decorators"))
class TextContentAst: class TextContentAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("TextContent") self._bldr = _tb.class_("TextContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ]) self._properties: typing.Set[str] = set(
[
"type",
"subclass",
]
)
self._props = TextContentProperties(self._bldr, self._properties) self._props = TextContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -841,18 +896,17 @@ class TextContentViewer(TextContentAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class TextContentProperties: class TextContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def type(self) -> type_builder.ClassPropertyViewer: def type(self) -> type_builder.ClassPropertyViewer:
@ -863,13 +917,16 @@ class TextContentProperties:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass")) return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
class VideoContentAst: class VideoContentAst:
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute) _tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("VideoContent") self._bldr = _tb.class_("VideoContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ]) self._properties: typing.Set[str] = set(
[
"type",
"subclass",
]
)
self._props = VideoContentProperties(self._bldr, self._properties) self._props = VideoContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType: def type(self) -> baml_py.FieldType:
@ -884,18 +941,17 @@ class VideoContentViewer(VideoContentAst):
def __init__(self, tb: type_builder.TypeBuilder): def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb) super().__init__(tb)
def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]: def list_properties(self) -> typing.List[typing.Tuple[str, type_builder.ClassPropertyViewer]]:
return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties] return [
(name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
for name in self._properties
]
class VideoContentProperties: class VideoContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property @property
def type(self) -> type_builder.ClassPropertyViewer: def type(self) -> type_builder.ClassPropertyViewer:
@ -904,6 +960,3 @@ class VideoContentProperties:
@property @property
def subclass(self) -> type_builder.ClassPropertyViewer: def subclass(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass")) return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))

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

@ -15,57 +15,38 @@ from . import stream_types
type_map = { type_map = {
"types.AudioContent": types.AudioContent, "types.AudioContent": types.AudioContent,
"stream_types.AudioContent": stream_types.AudioContent, "stream_types.AudioContent": stream_types.AudioContent,
"types.ContentLabel": types.ContentLabel, "types.ContentLabel": types.ContentLabel,
"stream_types.ContentLabel": stream_types.ContentLabel, "stream_types.ContentLabel": stream_types.ContentLabel,
"types.DefaultContentPrediction": types.DefaultContentPrediction, "types.DefaultContentPrediction": types.DefaultContentPrediction,
"stream_types.DefaultContentPrediction": stream_types.DefaultContentPrediction, "stream_types.DefaultContentPrediction": stream_types.DefaultContentPrediction,
"types.DynamicKnowledgeGraph": types.DynamicKnowledgeGraph, "types.DynamicKnowledgeGraph": types.DynamicKnowledgeGraph,
"stream_types.DynamicKnowledgeGraph": stream_types.DynamicKnowledgeGraph, "stream_types.DynamicKnowledgeGraph": stream_types.DynamicKnowledgeGraph,
"types.Edge": types.Edge, "types.Edge": types.Edge,
"stream_types.Edge": stream_types.Edge, "stream_types.Edge": stream_types.Edge,
"types.ImageContent": types.ImageContent, "types.ImageContent": types.ImageContent,
"stream_types.ImageContent": stream_types.ImageContent, "stream_types.ImageContent": stream_types.ImageContent,
"types.KnowledgeGraph": types.KnowledgeGraph, "types.KnowledgeGraph": types.KnowledgeGraph,
"stream_types.KnowledgeGraph": stream_types.KnowledgeGraph, "stream_types.KnowledgeGraph": stream_types.KnowledgeGraph,
"types.Model3DContent": types.Model3DContent, "types.Model3DContent": types.Model3DContent,
"stream_types.Model3DContent": stream_types.Model3DContent, "stream_types.Model3DContent": stream_types.Model3DContent,
"types.MultimediaContent": types.MultimediaContent, "types.MultimediaContent": types.MultimediaContent,
"stream_types.MultimediaContent": stream_types.MultimediaContent, "stream_types.MultimediaContent": stream_types.MultimediaContent,
"types.Node": types.Node, "types.Node": types.Node,
"stream_types.Node": stream_types.Node, "stream_types.Node": stream_types.Node,
"types.ProceduralContent": types.ProceduralContent, "types.ProceduralContent": types.ProceduralContent,
"stream_types.ProceduralContent": stream_types.ProceduralContent, "stream_types.ProceduralContent": stream_types.ProceduralContent,
"types.SummarizedClass": types.SummarizedClass, "types.SummarizedClass": types.SummarizedClass,
"stream_types.SummarizedClass": stream_types.SummarizedClass, "stream_types.SummarizedClass": stream_types.SummarizedClass,
"types.SummarizedCode": types.SummarizedCode, "types.SummarizedCode": types.SummarizedCode,
"stream_types.SummarizedCode": stream_types.SummarizedCode, "stream_types.SummarizedCode": stream_types.SummarizedCode,
"types.SummarizedContent": types.SummarizedContent, "types.SummarizedContent": types.SummarizedContent,
"stream_types.SummarizedContent": stream_types.SummarizedContent, "stream_types.SummarizedContent": stream_types.SummarizedContent,
"types.SummarizedFunction": types.SummarizedFunction, "types.SummarizedFunction": types.SummarizedFunction,
"stream_types.SummarizedFunction": stream_types.SummarizedFunction, "stream_types.SummarizedFunction": stream_types.SummarizedFunction,
"types.TextContent": types.TextContent, "types.TextContent": types.TextContent,
"stream_types.TextContent": stream_types.TextContent, "stream_types.TextContent": stream_types.TextContent,
"types.VideoContent": types.VideoContent, "types.VideoContent": types.VideoContent,
"stream_types.VideoContent": stream_types.VideoContent, "stream_types.VideoContent": stream_types.VideoContent,
} }

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

@ -20,22 +20,29 @@ from pydantic import BaseModel, ConfigDict
import baml_py import baml_py
CheckT = typing_extensions.TypeVar('CheckT') CheckT = typing_extensions.TypeVar("CheckT")
CheckName = typing_extensions.TypeVar('CheckName', bound=str) CheckName = typing_extensions.TypeVar("CheckName", bound=str)
class Check(BaseModel): class Check(BaseModel):
name: str name: str
expression: str expression: str
status: str status: str
class Checked(BaseModel, typing.Generic[CheckT, CheckName]): class Checked(BaseModel, typing.Generic[CheckT, CheckName]):
value: CheckT value: CheckT
checks: typing.Dict[CheckName, Check] checks: typing.Dict[CheckName, Check]
def get_checks(checks: typing.Dict[CheckName, Check]) -> typing.List[Check]: def get_checks(checks: typing.Dict[CheckName, Check]) -> typing.List[Check]:
return list(checks.values()) return list(checks.values())
def all_succeeded(checks: typing.Dict[CheckName, Check]) -> bool: def all_succeeded(checks: typing.Dict[CheckName, Check]) -> bool:
return all(check.status == "succeeded" for check in get_checks(checks)) return all(check.status == "succeeded" for check in get_checks(checks))
# ######################################################################### # #########################################################################
# Generated enums (0) # Generated enums (0)
# ######################################################################### # #########################################################################
@ -44,20 +51,33 @@ def all_succeeded(checks: typing.Dict[CheckName, Check]) -> bool:
# Generated classes (17) # Generated classes (17)
# ######################################################################### # #########################################################################
class AudioContent(BaseModel): class AudioContent(BaseModel):
type: str type: str
subclass: typing.List[str] subclass: typing.List[str]
class ContentLabel(BaseModel): class ContentLabel(BaseModel):
content_type: typing.Union[typing_extensions.Literal['text'], typing_extensions.Literal['audio'], typing_extensions.Literal['image'], typing_extensions.Literal['video'], typing_extensions.Literal['multimedia'], typing_extensions.Literal['3d_model'], typing_extensions.Literal['procedural']] content_type: typing.Union[
typing_extensions.Literal["text"],
typing_extensions.Literal["audio"],
typing_extensions.Literal["image"],
typing_extensions.Literal["video"],
typing_extensions.Literal["multimedia"],
typing_extensions.Literal["3d_model"],
typing_extensions.Literal["procedural"],
]
type: str type: str
subclass: typing.List[str] subclass: typing.List[str]
class DefaultContentPrediction(BaseModel): class DefaultContentPrediction(BaseModel):
label: "ContentLabel" label: "ContentLabel"
class DynamicKnowledgeGraph(BaseModel): class DynamicKnowledgeGraph(BaseModel):
model_config = ConfigDict(extra='allow') model_config = ConfigDict(extra="allow")
class Edge(BaseModel): class Edge(BaseModel):
# doc string for edge # doc string for edge
@ -67,39 +87,47 @@ class Edge(BaseModel):
target_node_id: str target_node_id: str
relationship_name: str relationship_name: str
class ImageContent(BaseModel): class ImageContent(BaseModel):
type: str type: str
subclass: typing.List[str] subclass: typing.List[str]
class KnowledgeGraph(BaseModel): class KnowledgeGraph(BaseModel):
nodes: typing.List["Node"] nodes: typing.List["Node"]
edges: typing.List["Edge"] edges: typing.List["Edge"]
class Model3DContent(BaseModel): class Model3DContent(BaseModel):
type: str type: str
subclass: typing.List[str] subclass: typing.List[str]
class MultimediaContent(BaseModel): class MultimediaContent(BaseModel):
type: str type: str
subclass: typing.List[str] subclass: typing.List[str]
class Node(BaseModel): class Node(BaseModel):
model_config = ConfigDict(extra='allow') model_config = ConfigDict(extra="allow")
id: str id: str
name: str name: str
type: str type: str
description: str description: str
class ProceduralContent(BaseModel): class ProceduralContent(BaseModel):
type: str type: str
subclass: typing.List[str] subclass: typing.List[str]
class SummarizedClass(BaseModel): class SummarizedClass(BaseModel):
name: str name: str
description: str description: str
methods: typing.Optional[typing.List["SummarizedFunction"]] = None methods: typing.Optional[typing.List["SummarizedFunction"]] = None
decorators: typing.Optional[typing.List[str]] = None decorators: typing.Optional[typing.List[str]] = None
class SummarizedCode(BaseModel): class SummarizedCode(BaseModel):
high_level_summary: str high_level_summary: str
key_features: typing.List[str] key_features: typing.List[str]
@ -109,10 +137,12 @@ class SummarizedCode(BaseModel):
functions: typing.List["SummarizedFunction"] functions: typing.List["SummarizedFunction"]
workflow_description: typing.Optional[str] = None workflow_description: typing.Optional[str] = None
class SummarizedContent(BaseModel): class SummarizedContent(BaseModel):
summary: str summary: str
description: str description: str
class SummarizedFunction(BaseModel): class SummarizedFunction(BaseModel):
name: str name: str
description: str description: str
@ -120,14 +150,17 @@ class SummarizedFunction(BaseModel):
outputs: typing.Optional[typing.List[str]] = None outputs: typing.Optional[typing.List[str]] = None
decorators: typing.Optional[typing.List[str]] = None decorators: typing.Optional[typing.List[str]] = None
class TextContent(BaseModel): class TextContent(BaseModel):
type: str type: str
subclass: typing.List[str] subclass: typing.List[str]
class VideoContent(BaseModel): class VideoContent(BaseModel):
type: str type: str
subclass: typing.List[str] subclass: typing.List[str]
# ######################################################################### # #########################################################################
# Generated type aliases (0) # Generated type aliases (0)
# ######################################################################### # #########################################################################

5
cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_categories.py
View file

@ -2,13 +2,14 @@ import os
from typing import Type from typing import Type
from pydantic import BaseModel from pydantic import BaseModel
from cognee.infrastructure.llm.structured_output_framework.baml_src.config import get_llm_config from cognee.infrastructure.llm.structured_output_framework.baml_src.config import get_llm_config
config = get_llm_config()
from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.async_client import b from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.async_client import b
from cognee.shared.data_models import SummarizedCode from cognee.shared.data_models import SummarizedCode
from cognee.shared.logging_utils import get_logger from cognee.shared.logging_utils import get_logger
from baml_py import ClientRegistry from baml_py import ClientRegistry
config = get_llm_config()
logger = get_logger("extract_summary_baml") logger = get_logger("extract_summary_baml")

28
cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_summary.py
View file

@ -1,13 +1,15 @@
import os import os
from typing import Type from typing import Type
from pydantic import BaseModel from pydantic import BaseModel
from cognee.infrastructure.llm.structured_output_framework.baml_src.config import get_llm_config from baml_py import ClientRegistry
config = get_llm_config() 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.structured_output_framework.baml.baml_client.async_client import b
from cognee.infrastructure.llm.structured_output_framework.baml_src.config import get_llm_config from cognee.infrastructure.llm.structured_output_framework.baml_src.config import get_llm_config
from cognee.shared.data_models import SummarizedCode
from cognee.shared.logging_utils import get_logger config = get_llm_config()
from baml_py import ClientRegistry
logger = get_logger("extract_summary_baml") logger = get_logger("extract_summary_baml")
@ -45,9 +47,10 @@ async def extract_summary(content: str, response_model: Type[BaseModel]):
options={ options={
"model": config.llm_model, "model": config.llm_model,
"temperature": config.llm_temperature, "temperature": config.llm_temperature,
"api_key": config.llm_api_key "api_key": config.llm_api_key,
}) },
baml_registry.set_primary('def') )
baml_registry.set_primary("def")
# Use BAML's SummarizeContent function # Use BAML's SummarizeContent function
summary_result = await b.SummarizeContent( summary_result = await b.SummarizeContent(
@ -97,12 +100,11 @@ async def extract_code_summary(content: str):
options={ options={
"model": config.llm_model, "model": config.llm_model,
"temperature": config.llm_temperature, "temperature": config.llm_temperature,
"api_key": config.llm_api_key "api_key": config.llm_api_key,
}) },
baml_registry.set_primary('def')
result = await b.SummarizeCode(
content, baml_options={"client_registry": baml_registry}
) )
baml_registry.set_primary("def")
result = await b.SummarizeCode(content, baml_options={"client_registry": baml_registry})
except Exception as e: except Exception as e:
logger.error( logger.error(
"Failed to extract code summary with BAML, falling back to mock summary", exc_info=e "Failed to extract code summary with BAML, falling back to mock summary", exc_info=e

7
cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/knowledge_graph/extract_content_graph.py
View file

@ -1,12 +1,11 @@
import os from baml_py import ClientRegistry
from typing import Type from typing import Type
from pydantic import BaseModel from pydantic import BaseModel
from cognee.infrastructure.llm.structured_output_framework.baml_src.config import get_llm_config from cognee.infrastructure.llm.structured_output_framework.baml_src.config import get_llm_config
from cognee.shared.logging_utils import get_logger, setup_logging
from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.async_client import b
config = get_llm_config() config = get_llm_config()
from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.async_client import b
from cognee.shared.logging_utils import get_logger, setup_logging
from baml_py import ClientRegistry
async def extract_content_graph( async def extract_content_graph(

3
cognee/tasks/graph/extract_graph_from_data.py
View file

@ -27,9 +27,6 @@ else:
extract_content_graph, extract_content_graph,
) )
from cognee.tasks.storage.add_data_points import add_data_points
async def integrate_chunk_graphs( async def integrate_chunk_graphs(
data_chunks: list[DocumentChunk], data_chunks: list[DocumentChunk],

2
pyproject.toml
View file

@ -56,7 +56,7 @@ dependencies = [
"dlt[sqlalchemy]>=1.9.0,<2", "dlt[sqlalchemy]>=1.9.0,<2",
"sentry-sdk[fastapi]>=2.9.0,<3", "sentry-sdk[fastapi]>=2.9.0,<3",
"structlog>=25.2.0,<26", "structlog>=25.2.0,<26",
"baml-py (>=0.201.0,<0.202.0)" "baml-py (>=0.201.0,<0.202.0)",
"pympler>=1.1,<2.0.0", "pympler>=1.1,<2.0.0",
"onnxruntime>=1.0.0,<2.0.0", "onnxruntime>=1.0.0,<2.0.0",
"pylance>=0.22.0,<1.0.0", "pylance>=0.22.0,<1.0.0",