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

added infra improvements

Browse source
This commit is contained in:
vasilije 2025-07-13 12:36:29 +02:00
parent c0e557ada9
commit 6e305d4395
17 changed files with 1318 additions and 1491 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
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,15 +31,16 @@ 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"
@ -47,11 +48,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",
] ]

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

@ -36,13 +36,10 @@ 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( def with_options(self,
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[ collector: typing.Optional[typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]] = None,
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 = {}
@ -58,151 +55,60 @@ 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 ExtractContentGraph( async def ExtractCategories(self, content: str,
self, baml_options: BamlCallOptions = {},
content: str, ) -> types.DefaultContentPrediction:
mode: typing.Optional[ result = await self.__options.merge_options(baml_options).call_function_async(function_name="ExtractCategories", args={
typing.Union[ "content": content,
typing_extensions.Literal["simple"], })
typing_extensions.Literal["base"], return typing.cast(types.DefaultContentPrediction, result.cast_to(types, types, stream_types, False, __runtime__))
typing_extensions.Literal["guided"], 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,
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( result = await self.__options.merge_options(baml_options).call_function_async(function_name="ExtractContentGraphGeneric", args={
function_name="ExtractContentGraph", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ })
"content": content, return typing.cast(types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__))
"mode": mode, 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,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.KnowledgeGraph, 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.DynamicKnowledgeGraph:
result = await self.__options.merge_options(baml_options).call_function_async( result = await self.__options.merge_options(baml_options).call_function_async(function_name="ExtractDynamicContentGraph", args={
function_name="ExtractContentGraphGeneric", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ })
"content": content, return typing.cast(types.DynamicKnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__))
"mode": mode, async def SummarizeCode(self, content: str,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)
)
async def ExtractContentGraphWithAnthropic(
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 = {},
) -> types.KnowledgeGraph:
result = await self.__options.merge_options(baml_options).call_function_async(
function_name="ExtractContentGraphWithAnthropic",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)
)
async def ExtractContentGraphWithEnvPrompt(
self,
content: str,
prompt_override: typing.Optional[str] = None,
baml_options: BamlCallOptions = {},
) -> types.KnowledgeGraph:
result = await self.__options.merge_options(baml_options).call_function_async(
function_name="ExtractContentGraphWithEnvPrompt",
args={
"content": content,
"prompt_override": prompt_override,
},
)
return typing.cast(
types.KnowledgeGraph, 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( result = await self.__options.merge_options(baml_options).call_function_async(function_name="SummarizeCode", args={
function_name="SummarizeCode", "content": content,
args={ })
"content": content, return typing.cast(types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__))
}, 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( result = await self.__options.merge_options(baml_options).call_function_async(function_name="SummarizeContent", args={
function_name="SummarizeContent", "content": content,
args={ })
"content": content, return typing.cast(types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__))
},
)
return typing.cast(
types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__)
)
class BamlStreamClient: class BamlStreamClient:
@ -211,182 +117,67 @@ class BamlStreamClient:
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
def ExtractContentGraph( def ExtractCategories(self, content: str,
self, baml_options: BamlCallOptions = {},
content: str, ) -> baml_py.BamlStream[stream_types.DefaultContentPrediction, types.DefaultContentPrediction]:
mode: typing.Optional[ ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="ExtractCategories", args={
typing.Union[ "content": content,
typing_extensions.Literal["simple"], })
typing_extensions.Literal["base"], return baml_py.BamlStream[stream_types.DefaultContentPrediction, types.DefaultContentPrediction](
typing_extensions.Literal["guided"], result,
typing_extensions.Literal["strict"], lambda x: typing.cast(stream_types.DefaultContentPrediction, x.cast_to(types, types, stream_types, True, __runtime__)),
typing_extensions.Literal["custom"], lambda x: typing.cast(types.DefaultContentPrediction, x.cast_to(types, types, stream_types, False, __runtime__)),
] ctx,
] = 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.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( ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="ExtractContentGraphGeneric", args={
function_name="ExtractContentGraph", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ })
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return baml_py.BamlStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]( return baml_py.BamlStream[stream_types.KnowledgeGraph, types.KnowledgeGraph](
result, result,
lambda x: typing.cast( lambda x: typing.cast(stream_types.KnowledgeGraph, x.cast_to(types, types, stream_types, True, __runtime__)),
stream_types.KnowledgeGraph, lambda x: typing.cast(types.KnowledgeGraph, x.cast_to(types, types, stream_types, False, __runtime__)),
x.cast_to(types, types, stream_types, True, __runtime__), ctx,
),
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,
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.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream( ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="ExtractDynamicContentGraph", args={
function_name="ExtractContentGraphGeneric", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ })
"content": content, return baml_py.BamlStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph](
"mode": mode, result,
"custom_prompt_content": custom_prompt_content, 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,
) )
return baml_py.BamlStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]( def SummarizeCode(self, content: str,
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 ExtractContentGraphWithAnthropic(
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_py.BamlStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream(
function_name="ExtractContentGraphWithAnthropic",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
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 ExtractContentGraphWithEnvPrompt(
self,
content: str,
prompt_override: typing.Optional[str] = None,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream(
function_name="ExtractContentGraphWithEnvPrompt",
args={
"content": content,
"prompt_override": prompt_override,
},
)
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 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( ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="SummarizeCode", args={
function_name="SummarizeCode", "content": content,
args={ })
"content": content,
},
)
return baml_py.BamlStream[stream_types.SummarizedCode, types.SummarizedCode]( return baml_py.BamlStream[stream_types.SummarizedCode, types.SummarizedCode](
result, result,
lambda x: typing.cast( lambda x: typing.cast(stream_types.SummarizedCode, x.cast_to(types, types, stream_types, True, __runtime__)),
stream_types.SummarizedCode, lambda x: typing.cast(types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)),
x.cast_to(types, types, stream_types, True, __runtime__), ctx,
),
lambda x: typing.cast(
types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
) )
def SummarizeContent(self, content: str,
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( ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="SummarizeContent", args={
function_name="SummarizeContent", "content": content,
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( lambda x: typing.cast(stream_types.SummarizedContent, x.cast_to(types, types, stream_types, True, __runtime__)),
stream_types.SummarizedContent, lambda x: typing.cast(types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)),
x.cast_to(types, types, stream_types, True, __runtime__), ctx,
),
lambda x: typing.cast(
types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
) )
class BamlHttpRequestClient: class BamlHttpRequestClient:
__options: DoNotUseDirectlyCallManager __options: DoNotUseDirectlyCallManager
@ -394,128 +185,42 @@ class BamlHttpRequestClient:
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
async def ExtractContentGraph( async def ExtractCategories(self, content: str,
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( result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractCategories", args={
function_name="ExtractContentGraph", "content": content,
args={ }, mode="request")
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_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( result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractContentGraphGeneric", args={
function_name="ExtractContentGraphGeneric", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ }, mode="request")
"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 ExtractContentGraphWithAnthropic(
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( result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractDynamicContentGraph", args={
function_name="ExtractContentGraphWithAnthropic", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ }, mode="request")
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="request",
)
return result return result
async def SummarizeCode(self, content: str,
async def ExtractContentGraphWithEnvPrompt(
self,
content: str,
prompt_override: 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( result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="SummarizeCode", args={
function_name="ExtractContentGraphWithEnvPrompt", "content": content,
args={ }, mode="request")
"content": content,
"prompt_override": prompt_override,
},
mode="request",
)
return result return result
async def SummarizeContent(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( result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="SummarizeContent", args={
function_name="SummarizeCode", "content": content,
args={ }, mode="request")
"content": content,
},
mode="request",
)
return result return result
async def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(
function_name="SummarizeContent",
args={
"content": content,
},
mode="request",
)
return result
class BamlHttpStreamRequestClient: class BamlHttpStreamRequestClient:
__options: DoNotUseDirectlyCallManager __options: DoNotUseDirectlyCallManager
@ -523,127 +228,41 @@ class BamlHttpStreamRequestClient:
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
async def ExtractContentGraph( async def ExtractCategories(self, content: str,
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( result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractCategories", args={
function_name="ExtractContentGraph", "content": content,
args={ }, mode="stream")
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_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( result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractContentGraphGeneric", args={
function_name="ExtractContentGraphGeneric", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ }, mode="stream")
"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 ExtractContentGraphWithAnthropic(
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( result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractDynamicContentGraph", args={
function_name="ExtractContentGraphWithAnthropic", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ }, mode="stream")
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="stream",
)
return result return result
async def SummarizeCode(self, content: str,
async def ExtractContentGraphWithEnvPrompt(
self,
content: str,
prompt_override: 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( result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="SummarizeCode", args={
function_name="ExtractContentGraphWithEnvPrompt", "content": content,
args={ }, mode="stream")
"content": content,
"prompt_override": prompt_override,
},
mode="stream",
)
return result return result
async def SummarizeContent(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( result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="SummarizeContent", args={
function_name="SummarizeCode", "content": content,
args={ }, mode="stream")
"content": content,
},
mode="stream",
)
return result return result
async def SummarizeContent( b = BamlAsyncClient(DoNotUseDirectlyCallManager({}))
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(
function_name="SummarizeContent",
args={
"content": content,
},
mode="stream",
)
return result
b = BamlAsyncClient(DoNotUseDirectlyCallManager({}))

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

@ -19,19 +19,17 @@ 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", get_baml_files(), os.environ.copy() "baml_src",
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_CTX = BamlCtxManager(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME)
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__ = []

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

File diff suppressed because one or more lines are too long

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

@ -16,77 +16,43 @@ 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
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
def ExtractContentGraph( def ExtractCategories(
self, self, llm_response: str, baml_options: BamlCallOptions = {},
llm_response: str, ) -> types.DefaultContentPrediction:
baml_options: BamlCallOptions = {}, result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractCategories", llm_response=llm_response, mode="request")
) -> types.KnowledgeGraph: return typing.cast(types.DefaultContentPrediction, result)
result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractContentGraph", llm_response=llm_response, mode="request"
)
return typing.cast(types.KnowledgeGraph, result)
def ExtractContentGraphGeneric( def ExtractContentGraphGeneric(
self, self, llm_response: str, baml_options: BamlCallOptions = {},
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.KnowledgeGraph: ) -> types.KnowledgeGraph:
result = self.__options.merge_options(baml_options).parse_response( result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractContentGraphGeneric", llm_response=llm_response, mode="request")
function_name="ExtractContentGraphGeneric", llm_response=llm_response, mode="request"
)
return typing.cast(types.KnowledgeGraph, result) return typing.cast(types.KnowledgeGraph, result)
def ExtractContentGraphWithAnthropic( def ExtractDynamicContentGraph(
self, self, llm_response: str, baml_options: BamlCallOptions = {},
llm_response: str, ) -> types.DynamicKnowledgeGraph:
baml_options: BamlCallOptions = {}, result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractDynamicContentGraph", llm_response=llm_response, mode="request")
) -> types.KnowledgeGraph: return typing.cast(types.DynamicKnowledgeGraph, result)
result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractContentGraphWithAnthropic",
llm_response=llm_response,
mode="request",
)
return typing.cast(types.KnowledgeGraph, result)
def ExtractContentGraphWithEnvPrompt(
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.KnowledgeGraph:
result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractContentGraphWithEnvPrompt",
llm_response=llm_response,
mode="request",
)
return typing.cast(types.KnowledgeGraph, result)
def SummarizeCode( def SummarizeCode(
self, self, llm_response: str, baml_options: BamlCallOptions = {},
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.SummarizedCode: ) -> types.SummarizedCode:
result = self.__options.merge_options(baml_options).parse_response( result = self.__options.merge_options(baml_options).parse_response(function_name="SummarizeCode", llm_response=llm_response, mode="request")
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, self, llm_response: str, baml_options: BamlCallOptions = {},
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.SummarizedContent: ) -> types.SummarizedContent:
result = self.__options.merge_options(baml_options).parse_response( result = self.__options.merge_options(baml_options).parse_response(function_name="SummarizeContent", llm_response=llm_response, mode="request")
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
@ -94,66 +60,34 @@ class LlmStreamParser:
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
def ExtractContentGraph( def ExtractCategories(
self, self, llm_response: str, baml_options: BamlCallOptions = {},
llm_response: str, ) -> stream_types.DefaultContentPrediction:
baml_options: BamlCallOptions = {}, result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractCategories", llm_response=llm_response, mode="stream")
) -> stream_types.KnowledgeGraph: return typing.cast(stream_types.DefaultContentPrediction, result)
result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractContentGraph", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.KnowledgeGraph, result)
def ExtractContentGraphGeneric( def ExtractContentGraphGeneric(
self, self, llm_response: str, baml_options: BamlCallOptions = {},
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.KnowledgeGraph: ) -> stream_types.KnowledgeGraph:
result = self.__options.merge_options(baml_options).parse_response( result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractContentGraphGeneric", llm_response=llm_response, mode="stream")
function_name="ExtractContentGraphGeneric", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.KnowledgeGraph, result) return typing.cast(stream_types.KnowledgeGraph, result)
def ExtractContentGraphWithAnthropic( def ExtractDynamicContentGraph(
self, self, llm_response: str, baml_options: BamlCallOptions = {},
llm_response: str, ) -> stream_types.DynamicKnowledgeGraph:
baml_options: BamlCallOptions = {}, result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractDynamicContentGraph", llm_response=llm_response, mode="stream")
) -> stream_types.KnowledgeGraph: return typing.cast(stream_types.DynamicKnowledgeGraph, result)
result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractContentGraphWithAnthropic",
llm_response=llm_response,
mode="stream",
)
return typing.cast(stream_types.KnowledgeGraph, result)
def ExtractContentGraphWithEnvPrompt(
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.KnowledgeGraph:
result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractContentGraphWithEnvPrompt",
llm_response=llm_response,
mode="stream",
)
return typing.cast(stream_types.KnowledgeGraph, result)
def SummarizeCode( def SummarizeCode(
self, self, llm_response: str, baml_options: BamlCallOptions = {},
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.SummarizedCode: ) -> stream_types.SummarizedCode:
result = self.__options.merge_options(baml_options).parse_response( result = self.__options.merge_options(baml_options).parse_response(function_name="SummarizeCode", llm_response=llm_response, mode="stream")
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, self, llm_response: str, baml_options: BamlCallOptions = {},
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.SummarizedContent: ) -> stream_types.SummarizedContent:
result = self.__options.merge_options(baml_options).parse_response( result = self.__options.merge_options(baml_options).parse_response(function_name="SummarizeContent", llm_response=llm_response, mode="stream")
function_name="SummarizeContent", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.SummarizedContent, result) return typing.cast(stream_types.SummarizedContent, result)

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

@ -17,10 +17,7 @@ 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 ( 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__
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):
@ -51,6 +48,9 @@ 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,9 +74,7 @@ class DoNotUseDirectlyCallManager:
collectors_as_list = ( collectors_as_list = (
collector collector
if isinstance(collector, list) if isinstance(collector, list)
else [collector] else [collector] if collector is not None else []
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():
@ -166,9 +164,7 @@ class DoNotUseDirectlyCallManager:
*, *,
function_name: str, function_name: str,
args: typing.Dict[str, typing.Any], args: typing.Dict[str, typing.Any],
) -> typing.Tuple[ ) -> typing.Tuple[baml_py.baml_py.RuntimeContextManager, baml_py.baml_py.SyncFunctionResultStream]:
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(
@ -176,7 +172,7 @@ class DoNotUseDirectlyCallManager:
args, args,
# this is always None, we set this later! # this is always None, we set this later!
# on_event # on_event
None, None,
# ctx # ctx
ctx, ctx,
# tb # tb
@ -236,13 +232,7 @@ class DoNotUseDirectlyCallManager:
mode == "stream", mode == "stream",
) )
def parse_response( def parse_response(self, *, function_name: str, llm_response: str, mode: typing_extensions.Literal["stream", "request"]) -> typing.Any:
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,
@ -263,4 +253,4 @@ class DoNotUseDirectlyCallManager:
resolved_options.client_registry, resolved_options.client_registry,
# env_vars # env_vars
resolved_options.env_vars, resolved_options.env_vars,
) )

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

@ -18,18 +18,28 @@ 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 (7) # Generated classes (17)
# ######################################################################### # #########################################################################
class AudioContent(BaseModel):
type: typing.Optional[str] = None
subclass: typing.List[str]
class ContentLabel(BaseModel):
content_type: typing.Optional[typing.Union[str, str, str, str, str, str, str]] = None
type: typing.Optional[str] = None
subclass: typing.List[str]
class DefaultContentPrediction(BaseModel):
label: typing.Optional["ContentLabel"] = None
class DynamicKnowledgeGraph(BaseModel):
model_config = ConfigDict(extra='allow')
class Edge(BaseModel): class Edge(BaseModel):
# doc string for edge # doc string for edge
@ -39,19 +49,32 @@ 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):
type: typing.Optional[str] = None
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):
type: typing.Optional[str] = None
subclass: typing.List[str]
class MultimediaContent(BaseModel):
type: typing.Optional[str] = None
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):
type: typing.Optional[str] = None
subclass: typing.List[str]
class SummarizedClass(BaseModel): class SummarizedClass(BaseModel):
name: typing.Optional[str] = None name: typing.Optional[str] = None
@ -59,7 +82,6 @@ class SummarizedClass(BaseModel):
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]
@ -69,12 +91,10 @@ 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
@ -82,6 +102,13 @@ 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):
type: typing.Optional[str] = None
subclass: typing.List[str]
class VideoContent(BaseModel):
type: typing.Optional[str] = None
subclass: typing.List[str]
# ######################################################################### # #########################################################################
# Generated type aliases (0) # Generated type aliases (0)

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

@ -19,7 +19,6 @@ 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"
@ -49,13 +48,10 @@ 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( def with_options(self,
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[ collector: typing.Optional[typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]] = None,
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 = {}
@ -71,151 +67,60 @@ 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 ExtractContentGraph( def ExtractCategories(self, content: str,
self, baml_options: BamlCallOptions = {},
content: str, ) -> types.DefaultContentPrediction:
mode: typing.Optional[ result = self.__options.merge_options(baml_options).call_function_sync(function_name="ExtractCategories", args={
typing.Union[ "content": content,
typing_extensions.Literal["simple"], })
typing_extensions.Literal["base"], return typing.cast(types.DefaultContentPrediction, result.cast_to(types, types, stream_types, False, __runtime__))
typing_extensions.Literal["guided"], 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,
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( result = self.__options.merge_options(baml_options).call_function_sync(function_name="ExtractContentGraphGeneric", args={
function_name="ExtractContentGraph", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ })
"content": content, return typing.cast(types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__))
"mode": mode, 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,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.KnowledgeGraph, 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.DynamicKnowledgeGraph:
result = self.__options.merge_options(baml_options).call_function_sync( result = self.__options.merge_options(baml_options).call_function_sync(function_name="ExtractDynamicContentGraph", args={
function_name="ExtractContentGraphGeneric", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ })
"content": content, return typing.cast(types.DynamicKnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__))
"mode": mode, def SummarizeCode(self, content: str,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)
)
def ExtractContentGraphWithAnthropic(
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 = {},
) -> types.KnowledgeGraph:
result = self.__options.merge_options(baml_options).call_function_sync(
function_name="ExtractContentGraphWithAnthropic",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)
)
def ExtractContentGraphWithEnvPrompt(
self,
content: str,
prompt_override: typing.Optional[str] = None,
baml_options: BamlCallOptions = {},
) -> types.KnowledgeGraph:
result = self.__options.merge_options(baml_options).call_function_sync(
function_name="ExtractContentGraphWithEnvPrompt",
args={
"content": content,
"prompt_override": prompt_override,
},
)
return typing.cast(
types.KnowledgeGraph, 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( result = self.__options.merge_options(baml_options).call_function_sync(function_name="SummarizeCode", args={
function_name="SummarizeCode", "content": content,
args={ })
"content": content, return typing.cast(types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__))
}, 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( result = self.__options.merge_options(baml_options).call_function_sync(function_name="SummarizeContent", args={
function_name="SummarizeContent", "content": content,
args={ })
"content": content, return typing.cast(types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__))
},
)
return typing.cast(
types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__)
)
class BamlStreamClient: class BamlStreamClient:
@ -224,182 +129,67 @@ class BamlStreamClient:
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
def ExtractContentGraph( def ExtractCategories(self, content: str,
self, baml_options: BamlCallOptions = {},
content: str, ) -> baml_py.BamlSyncStream[stream_types.DefaultContentPrediction, types.DefaultContentPrediction]:
mode: typing.Optional[ ctx, result = self.__options.merge_options(baml_options).create_sync_stream(function_name="ExtractCategories", args={
typing.Union[ "content": content,
typing_extensions.Literal["simple"], })
typing_extensions.Literal["base"], return baml_py.BamlSyncStream[stream_types.DefaultContentPrediction, types.DefaultContentPrediction](
typing_extensions.Literal["guided"], result,
typing_extensions.Literal["strict"], lambda x: typing.cast(stream_types.DefaultContentPrediction, x.cast_to(types, types, stream_types, True, __runtime__)),
typing_extensions.Literal["custom"], lambda x: typing.cast(types.DefaultContentPrediction, x.cast_to(types, types, stream_types, False, __runtime__)),
] ctx,
] = 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.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( ctx, result = self.__options.merge_options(baml_options).create_sync_stream(function_name="ExtractContentGraphGeneric", args={
function_name="ExtractContentGraph", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ })
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return baml_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]( return baml_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph](
result, result,
lambda x: typing.cast( lambda x: typing.cast(stream_types.KnowledgeGraph, x.cast_to(types, types, stream_types, True, __runtime__)),
stream_types.KnowledgeGraph, lambda x: typing.cast(types.KnowledgeGraph, x.cast_to(types, types, stream_types, False, __runtime__)),
x.cast_to(types, types, stream_types, True, __runtime__), ctx,
),
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,
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.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_sync_stream( ctx, result = self.__options.merge_options(baml_options).create_sync_stream(function_name="ExtractDynamicContentGraph", args={
function_name="ExtractContentGraphGeneric", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ })
"content": content, return baml_py.BamlSyncStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph](
"mode": mode, result,
"custom_prompt_content": custom_prompt_content, 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,
) )
return baml_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]( def SummarizeCode(self, content: str,
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 ExtractContentGraphWithAnthropic(
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_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
function_name="ExtractContentGraphWithAnthropic",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
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 ExtractContentGraphWithEnvPrompt(
self,
content: str,
prompt_override: typing.Optional[str] = None,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
function_name="ExtractContentGraphWithEnvPrompt",
args={
"content": content,
"prompt_override": prompt_override,
},
)
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 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( ctx, result = self.__options.merge_options(baml_options).create_sync_stream(function_name="SummarizeCode", args={
function_name="SummarizeCode", "content": content,
args={ })
"content": content,
},
)
return baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode]( return baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode](
result, result,
lambda x: typing.cast( lambda x: typing.cast(stream_types.SummarizedCode, x.cast_to(types, types, stream_types, True, __runtime__)),
stream_types.SummarizedCode, lambda x: typing.cast(types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)),
x.cast_to(types, types, stream_types, True, __runtime__), ctx,
),
lambda x: typing.cast(
types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
) )
def SummarizeContent(self, content: str,
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( ctx, result = self.__options.merge_options(baml_options).create_sync_stream(function_name="SummarizeContent", args={
function_name="SummarizeContent", "content": content,
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( lambda x: typing.cast(stream_types.SummarizedContent, x.cast_to(types, types, stream_types, True, __runtime__)),
stream_types.SummarizedContent, lambda x: typing.cast(types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)),
x.cast_to(types, types, stream_types, True, __runtime__), ctx,
),
lambda x: typing.cast(
types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
) )
class BamlHttpRequestClient: class BamlHttpRequestClient:
__options: DoNotUseDirectlyCallManager __options: DoNotUseDirectlyCallManager
@ -407,128 +197,42 @@ class BamlHttpRequestClient:
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
def ExtractContentGraph( def ExtractCategories(self, content: str,
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( result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="ExtractCategories", args={
function_name="ExtractContentGraph", "content": content,
args={ }, mode="request")
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_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( result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="ExtractContentGraphGeneric", args={
function_name="ExtractContentGraphGeneric", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ }, mode="request")
"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 ExtractContentGraphWithAnthropic(
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( result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="ExtractDynamicContentGraph", args={
function_name="ExtractContentGraphWithAnthropic", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ }, mode="request")
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="request",
)
return result return result
def SummarizeCode(self, content: str,
def ExtractContentGraphWithEnvPrompt(
self,
content: str,
prompt_override: 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( result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="SummarizeCode", args={
function_name="ExtractContentGraphWithEnvPrompt", "content": content,
args={ }, mode="request")
"content": content,
"prompt_override": prompt_override,
},
mode="request",
)
return result return result
def SummarizeContent(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( result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="SummarizeContent", args={
function_name="SummarizeCode", "content": content,
args={ }, mode="request")
"content": content,
},
mode="request",
)
return result return result
def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(
function_name="SummarizeContent",
args={
"content": content,
},
mode="request",
)
return result
class BamlHttpStreamRequestClient: class BamlHttpStreamRequestClient:
__options: DoNotUseDirectlyCallManager __options: DoNotUseDirectlyCallManager
@ -536,127 +240,41 @@ class BamlHttpStreamRequestClient:
def __init__(self, options: DoNotUseDirectlyCallManager): def __init__(self, options: DoNotUseDirectlyCallManager):
self.__options = options self.__options = options
def ExtractContentGraph( def ExtractCategories(self, content: str,
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( result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="ExtractCategories", args={
function_name="ExtractContentGraph", "content": content,
args={ }, mode="stream")
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_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( result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="ExtractContentGraphGeneric", args={
function_name="ExtractContentGraphGeneric", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ }, mode="stream")
"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 ExtractContentGraphWithAnthropic(
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( result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="ExtractDynamicContentGraph", args={
function_name="ExtractContentGraphWithAnthropic", "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
args={ }, mode="stream")
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="stream",
)
return result return result
def SummarizeCode(self, content: str,
def ExtractContentGraphWithEnvPrompt(
self,
content: str,
prompt_override: 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( result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="SummarizeCode", args={
function_name="ExtractContentGraphWithEnvPrompt", "content": content,
args={ }, mode="stream")
"content": content,
"prompt_override": prompt_override,
},
mode="stream",
)
return result return result
def SummarizeContent(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( result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="SummarizeContent", args={
function_name="SummarizeCode", "content": content,
args={ }, mode="stream")
"content": content,
},
mode="stream",
)
return result return result
def SummarizeContent( b = BamlSyncClient(DoNotUseDirectlyCallManager({}))
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(
function_name="SummarizeContent",
args={
"content": content,
},
mode="stream",
)
return result
b = BamlSyncClient(DoNotUseDirectlyCallManager({}))

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

@ -14,13 +14,9 @@ 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"]

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

@ -15,45 +15,67 @@ 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__( super().__init__(classes=set(
classes=set( ["AudioContent","ContentLabel","DefaultContentPrediction","DynamicKnowledgeGraph","Edge","ImageContent","KnowledgeGraph","Model3DContent","MultimediaContent","Node","ProceduralContent","SummarizedClass","SummarizedCode","SummarizedContent","SummarizedFunction","TextContent","VideoContent",]
[ ), enums=set(
"Edge", []
"KnowledgeGraph", ), runtime=DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME)
"Node",
"SummarizedClass",
"SummarizedCode",
"SummarizedContent",
"SummarizedFunction",
]
),
enums=set([]),
runtime=DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME,
)
# ######################################################################### # #########################################################################
# Generated enums 0 # Generated enums 0
# ######################################################################### # #########################################################################
# ######################################################################### # #########################################################################
# Generated classes 7 # Generated classes 17
# ######################################################################### # #########################################################################
@property
def AudioContent(self) -> "AudioContentViewer":
return AudioContentViewer(self)
@property
def ContentLabel(self) -> "ContentLabelViewer":
return ContentLabelViewer(self)
@property
def DefaultContentPrediction(self) -> "DefaultContentPredictionViewer":
return DefaultContentPredictionViewer(self)
@property
def DynamicKnowledgeGraph(self) -> "DynamicKnowledgeGraphBuilder":
return DynamicKnowledgeGraphBuilder(self)
@property @property
def Edge(self) -> "EdgeViewer": def Edge(self) -> "EdgeViewer":
return EdgeViewer(self) return EdgeViewer(self)
@property
def ImageContent(self) -> "ImageContentViewer":
return ImageContentViewer(self)
@property @property
def KnowledgeGraph(self) -> "KnowledgeGraphViewer": def KnowledgeGraph(self) -> "KnowledgeGraphViewer":
return KnowledgeGraphViewer(self) return KnowledgeGraphViewer(self)
@property
def Model3DContent(self) -> "Model3DContentViewer":
return Model3DContentViewer(self)
@property
def MultimediaContent(self) -> "MultimediaContentViewer":
return MultimediaContentViewer(self)
@property @property
def Node(self) -> "NodeBuilder": def Node(self) -> "NodeBuilder":
return NodeBuilder(self) return NodeBuilder(self)
@property
def ProceduralContent(self) -> "ProceduralContentViewer":
return ProceduralContentViewer(self)
@property @property
def SummarizedClass(self) -> "SummarizedClassViewer": def SummarizedClass(self) -> "SummarizedClassViewer":
return SummarizedClassViewer(self) return SummarizedClassViewer(self)
@ -70,6 +92,15 @@ class TypeBuilder(type_builder.TypeBuilder):
def SummarizedFunction(self) -> "SummarizedFunctionViewer": def SummarizedFunction(self) -> "SummarizedFunctionViewer":
return SummarizedFunctionViewer(self) return SummarizedFunctionViewer(self)
@property
def TextContent(self) -> "TextContentViewer":
return TextContentViewer(self)
@property
def VideoContent(self) -> "VideoContentViewer":
return VideoContentViewer(self)
# ######################################################################### # #########################################################################
# Generated enums 0 # Generated enums 0
@ -77,21 +108,189 @@ class TypeBuilder(type_builder.TypeBuilder):
# ######################################################################### # #########################################################################
# Generated classes 7 # Generated classes 17
# ######################################################################### # #########################################################################
class AudioContentAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("AudioContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ])
self._props = AudioContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "AudioContentProperties":
return self._props
class AudioContentViewer(AudioContentAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
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]
class AudioContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property
def type(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("type"))
@property
def subclass(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
class ContentLabelAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("ContentLabel")
self._properties: typing.Set[str] = set([ "content_type", "type", "subclass", ])
self._props = ContentLabelProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "ContentLabelProperties":
return self._props
class ContentLabelViewer(ContentLabelAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
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]
class ContentLabelProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property
def content_type(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("content_type"))
@property
def type(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("type"))
@property
def subclass(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
class DefaultContentPredictionAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("DefaultContentPrediction")
self._properties: typing.Set[str] = set([ "label", ])
self._props = DefaultContentPredictionProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "DefaultContentPredictionProperties":
return self._props
class DefaultContentPredictionViewer(DefaultContentPredictionAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
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]
class DefaultContentPredictionProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property
def label(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("label"))
class DynamicKnowledgeGraphAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("DynamicKnowledgeGraph")
self._properties: typing.Set[str] = set([ ])
self._props = DynamicKnowledgeGraphProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "DynamicKnowledgeGraphProperties":
return self._props
class DynamicKnowledgeGraphBuilder(DynamicKnowledgeGraphAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
def add_property(self, name: str, type: baml_py.FieldType) -> baml_py.ClassPropertyBuilder:
if name in self._properties:
raise ValueError(f"Property {name} already exists.")
return self._bldr.property(name).type(type)
def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]:
return [(name, self._bldr.property(name)) for name in self._properties]
class DynamicKnowledgeGraphProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
def __getattr__(self, name: str) -> baml_py.ClassPropertyBuilder:
if name not in self.__properties:
raise AttributeError(f"Property {name} not found.")
return self.__bldr.property(name)
class 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( self._properties: typing.Set[str] = set([ "source_node_id", "target_node_id", "relationship_name", ])
[
"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:
@ -106,41 +305,82 @@ 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 [ return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties]
(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:
return type_builder.ClassPropertyViewer(self.__bldr.property("source_node_id")) return type_builder.ClassPropertyViewer(self.__bldr.property("source_node_id"))
@property @property
def target_node_id(self) -> type_builder.ClassPropertyViewer: def target_node_id(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("target_node_id")) return type_builder.ClassPropertyViewer(self.__bldr.property("target_node_id"))
@property @property
def relationship_name(self) -> type_builder.ClassPropertyViewer: def relationship_name(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("relationship_name")) return type_builder.ClassPropertyViewer(self.__bldr.property("relationship_name"))
class ImageContentAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("ImageContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ])
self._props = ImageContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "ImageContentProperties":
return self._props
class ImageContentViewer(ImageContentAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
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]
class ImageContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property
def type(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("type"))
@property
def subclass(self) -> type_builder.ClassPropertyViewer:
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( self._properties: typing.Set[str] = set([ "nodes", "edges", ])
[
"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:
@ -155,39 +395,121 @@ 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 [ return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties]
(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:
return type_builder.ClassPropertyViewer(self.__bldr.property("nodes")) return type_builder.ClassPropertyViewer(self.__bldr.property("nodes"))
@property @property
def edges(self) -> type_builder.ClassPropertyViewer: def edges(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("edges")) return type_builder.ClassPropertyViewer(self.__bldr.property("edges"))
class Model3DContentAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("Model3DContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ])
self._props = Model3DContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "Model3DContentProperties":
return self._props
class Model3DContentViewer(Model3DContentAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
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]
class Model3DContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property
def type(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("type"))
@property
def subclass(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
class MultimediaContentAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("MultimediaContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ])
self._props = MultimediaContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "MultimediaContentProperties":
return self._props
class MultimediaContentViewer(MultimediaContentAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
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]
class MultimediaContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property
def type(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("type"))
@property
def subclass(self) -> type_builder.ClassPropertyViewer:
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( self._properties: typing.Set[str] = set([ "id", "name", "type", "description", ])
[
"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:
@ -202,6 +524,7 @@ 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.")
@ -210,46 +533,88 @@ class NodeBuilder(NodeAst):
def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]: def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]:
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")
@property @property
def name(self) -> baml_py.ClassPropertyBuilder: def name(self) -> baml_py.ClassPropertyBuilder:
return self.__bldr.property("name") return self.__bldr.property("name")
@property @property
def type(self) -> baml_py.ClassPropertyBuilder: def type(self) -> baml_py.ClassPropertyBuilder:
return self.__bldr.property("type") return self.__bldr.property("type")
@property @property
def description(self) -> baml_py.ClassPropertyBuilder: def description(self) -> baml_py.ClassPropertyBuilder:
return self.__bldr.property("description") return self.__bldr.property("description")
class ProceduralContentAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("ProceduralContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ])
self._props = ProceduralContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "ProceduralContentProperties":
return self._props
class ProceduralContentViewer(ProceduralContentAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
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]
class ProceduralContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property
def type(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("type"))
@property
def subclass(self) -> type_builder.ClassPropertyViewer:
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( self._properties: typing.Set[str] = set([ "name", "description", "methods", "decorators", ])
[
"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:
@ -264,50 +629,43 @@ 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 [ return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties]
(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:
return type_builder.ClassPropertyViewer(self.__bldr.property("name")) return type_builder.ClassPropertyViewer(self.__bldr.property("name"))
@property @property
def description(self) -> type_builder.ClassPropertyViewer: def description(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("description")) return type_builder.ClassPropertyViewer(self.__bldr.property("description"))
@property @property
def methods(self) -> type_builder.ClassPropertyViewer: def methods(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("methods")) return type_builder.ClassPropertyViewer(self.__bldr.property("methods"))
@property @property
def decorators(self) -> type_builder.ClassPropertyViewer: def decorators(self) -> type_builder.ClassPropertyViewer:
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( self._properties: typing.Set[str] = set([ "high_level_summary", "key_features", "imports", "constants", "classes", "functions", "workflow_description", ])
[
"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:
@ -322,57 +680,55 @@ 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 [ return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties]
(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:
return type_builder.ClassPropertyViewer(self.__bldr.property("high_level_summary")) return type_builder.ClassPropertyViewer(self.__bldr.property("high_level_summary"))
@property @property
def key_features(self) -> type_builder.ClassPropertyViewer: def key_features(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("key_features")) return type_builder.ClassPropertyViewer(self.__bldr.property("key_features"))
@property @property
def imports(self) -> type_builder.ClassPropertyViewer: def imports(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("imports")) return type_builder.ClassPropertyViewer(self.__bldr.property("imports"))
@property @property
def constants(self) -> type_builder.ClassPropertyViewer: def constants(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("constants")) return type_builder.ClassPropertyViewer(self.__bldr.property("constants"))
@property @property
def classes(self) -> type_builder.ClassPropertyViewer: def classes(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("classes")) return type_builder.ClassPropertyViewer(self.__bldr.property("classes"))
@property @property
def functions(self) -> type_builder.ClassPropertyViewer: def functions(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("functions")) return type_builder.ClassPropertyViewer(self.__bldr.property("functions"))
@property @property
def workflow_description(self) -> type_builder.ClassPropertyViewer: def workflow_description(self) -> type_builder.ClassPropertyViewer:
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( self._properties: typing.Set[str] = set([ "summary", "description", ])
[
"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:
@ -387,40 +743,35 @@ 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 [ return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties]
(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:
return type_builder.ClassPropertyViewer(self.__bldr.property("summary")) return type_builder.ClassPropertyViewer(self.__bldr.property("summary"))
@property @property
def description(self) -> type_builder.ClassPropertyViewer: def description(self) -> type_builder.ClassPropertyViewer:
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( self._properties: typing.Set[str] = set([ "name", "description", "inputs", "outputs", "decorators", ])
[
"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:
@ -435,34 +786,124 @@ 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 [ return [(name, type_builder.ClassPropertyViewer(self._bldr.property(name))) for name in self._properties]
(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:
return type_builder.ClassPropertyViewer(self.__bldr.property("name")) return type_builder.ClassPropertyViewer(self.__bldr.property("name"))
@property @property
def description(self) -> type_builder.ClassPropertyViewer: def description(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("description")) return type_builder.ClassPropertyViewer(self.__bldr.property("description"))
@property @property
def inputs(self) -> type_builder.ClassPropertyViewer: def inputs(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("inputs")) return type_builder.ClassPropertyViewer(self.__bldr.property("inputs"))
@property @property
def outputs(self) -> type_builder.ClassPropertyViewer: def outputs(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("outputs")) return type_builder.ClassPropertyViewer(self.__bldr.property("outputs"))
@property @property
def decorators(self) -> type_builder.ClassPropertyViewer: def decorators(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("decorators")) return type_builder.ClassPropertyViewer(self.__bldr.property("decorators"))
class TextContentAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("TextContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ])
self._props = TextContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "TextContentProperties":
return self._props
class TextContentViewer(TextContentAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
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]
class TextContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property
def type(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("type"))
@property
def subclass(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
class VideoContentAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("VideoContent")
self._properties: typing.Set[str] = set([ "type", "subclass", ])
self._props = VideoContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "VideoContentProperties":
return self._props
class VideoContentViewer(VideoContentAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
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]
class VideoContentProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
@property
def type(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("type"))
@property
def subclass(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))

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

@ -15,18 +15,57 @@ from . import stream_types
type_map = { type_map = {
"types.AudioContent": types.AudioContent,
"stream_types.AudioContent": stream_types.AudioContent,
"types.ContentLabel": types.ContentLabel,
"stream_types.ContentLabel": stream_types.ContentLabel,
"types.DefaultContentPrediction": types.DefaultContentPrediction,
"stream_types.DefaultContentPrediction": stream_types.DefaultContentPrediction,
"types.DynamicKnowledgeGraph": 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,
"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,
"stream_types.Model3DContent": stream_types.Model3DContent,
"types.MultimediaContent": 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,
"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,
"stream_types.TextContent": stream_types.TextContent,
"types.VideoContent": types.VideoContent,
"stream_types.VideoContent": stream_types.VideoContent,
}

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

@ -20,37 +20,44 @@ 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)
# ######################################################################### # #########################################################################
# ######################################################################### # #########################################################################
# Generated classes (7) # Generated classes (17)
# ######################################################################### # #########################################################################
class AudioContent(BaseModel):
type: str
subclass: typing.List[str]
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']]
type: str
subclass: typing.List[str]
class DefaultContentPrediction(BaseModel):
label: "ContentLabel"
class DynamicKnowledgeGraph(BaseModel):
model_config = ConfigDict(extra='allow')
class Edge(BaseModel): class Edge(BaseModel):
# doc string for edge # doc string for edge
@ -60,19 +67,32 @@ class Edge(BaseModel):
target_node_id: str target_node_id: str
relationship_name: str relationship_name: str
class ImageContent(BaseModel):
type: 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):
type: str
subclass: typing.List[str]
class MultimediaContent(BaseModel):
type: 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):
type: str
subclass: typing.List[str]
class SummarizedClass(BaseModel): class SummarizedClass(BaseModel):
name: str name: str
@ -80,7 +100,6 @@ class SummarizedClass(BaseModel):
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]
@ -90,12 +109,10 @@ 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
@ -103,6 +120,13 @@ 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):
type: str
subclass: typing.List[str]
class VideoContent(BaseModel):
type: str
subclass: typing.List[str]
# ######################################################################### # #########################################################################
# Generated type aliases (0) # Generated type aliases (0)

109
cognee/infrastructure/llm/structured_output_framework/baml_src/extract_categories.baml
View file

@ -0,0 +1,109 @@
// Content classification data models - matching shared/data_models.py
class TextContent {
type string
subclass string[]
}
class AudioContent {
type string
subclass string[]
}
class ImageContent {
type string
subclass string[]
}
class VideoContent {
type string
subclass string[]
}
class MultimediaContent {
type string
subclass string[]
}
class Model3DContent {
type string
subclass string[]
}
class ProceduralContent {
type string
subclass string[]
}
class ContentLabel {
content_type "text" | "audio" | "image" | "video" | "multimedia" | "3d_model" | "procedural"
type string
subclass string[]
}
class DefaultContentPrediction {
label ContentLabel
}
// Content classification prompt template
template_string ClassifyContentPrompt() #"
You are a classification engine and should classify content. Make sure to use one of the existing classification options and not invent your own.
Classify the content into one of these main categories and their relevant subclasses:
**TEXT CONTENT** (content_type: "text"):
- type: "TEXTUAL_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
- subclass options: ["Articles, essays, and reports", "Books and manuscripts", "News stories and blog posts", "Research papers and academic publications", "Social media posts and comments", "Website content and product descriptions", "Personal narratives and stories", "Spreadsheets and tables", "Forms and surveys", "Databases and CSV files", "Source code in various programming languages", "Shell commands and scripts", "Markup languages (HTML, XML)", "Stylesheets (CSS) and configuration files (YAML, JSON, INI)", "Chat transcripts and messaging history", "Customer service logs and interactions", "Conversational AI training data", "Textbook content and lecture notes", "Exam questions and academic exercises", "E-learning course materials", "Poetry and prose", "Scripts for plays, movies, and television", "Song lyrics", "Manuals and user guides", "Technical specifications and API documentation", "Helpdesk articles and FAQs", "Contracts and agreements", "Laws, regulations, and legal case documents", "Policy documents and compliance materials", "Clinical trial reports", "Patient records and case notes", "Scientific journal articles", "Financial reports and statements", "Business plans and proposals", "Market research and analysis reports", "Ad copies and marketing slogans", "Product catalogs and brochures", "Press releases and promotional content", "Professional and formal correspondence", "Personal emails and letters", "Image and video captions", "Annotations and metadata for various media", "Vocabulary lists and grammar rules", "Language exercises and quizzes", "Other types of text data"]
**AUDIO CONTENT** (content_type: "audio"):
- type: "AUDIO_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
- subclass options: ["Music tracks and albums", "Podcasts and radio broadcasts", "Audiobooks and audio guides", "Recorded interviews and speeches", "Sound effects and ambient sounds", "Other types of audio recordings"]
**IMAGE CONTENT** (content_type: "image"):
- type: "IMAGE_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
- subclass options: ["Photographs and digital images", "Illustrations, diagrams, and charts", "Infographics and visual data representations", "Artwork and paintings", "Screenshots and graphical user interfaces", "Other types of images"]
**VIDEO CONTENT** (content_type: "video"):
- type: "VIDEO_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
- subclass options: ["Movies and short films", "Documentaries and educational videos", "Video tutorials and how-to guides", "Animated features and cartoons", "Live event recordings and sports broadcasts", "Other types of video content"]
**MULTIMEDIA CONTENT** (content_type: "multimedia"):
- type: "MULTIMEDIA_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
- subclass options: ["Interactive web content and games", "Virtual reality (VR) and augmented reality (AR) experiences", "Mixed media presentations and slide decks", "E-learning modules with integrated multimedia", "Digital exhibitions and virtual tours", "Other types of multimedia content"]
**3D MODEL CONTENT** (content_type: "3d_model"):
- type: "3D_MODEL_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
- subclass options: ["Architectural renderings and building plans", "Product design models and prototypes", "3D animations and character models", "Scientific simulations and visualizations", "Virtual objects for AR/VR applications", "Other types of 3D models"]
**PROCEDURAL CONTENT** (content_type: "procedural"):
- type: "PROCEDURAL_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
- subclass options: ["Tutorials and step-by-step guides", "Workflow and process descriptions", "Simulation and training exercises", "Recipes and crafting instructions", "Other types of procedural content"]
Select the most appropriate content_type, type, and relevant subclasses.
"#
// OpenAI client defined once for all BAML files
// Classification function
function ExtractCategories(content: string) -> DefaultContentPrediction {
client OpenAI
prompt #"
{{ ClassifyContentPrompt() }}
{{ ctx.output_format(prefix="Answer in this schema:\n") }}
{{ _.role('user') }}
{{ content }}
"#
}
// Test case for classification
test ExtractCategoriesExample {
functions [ExtractCategories]
args {
content #"
Natural language processing (NLP) is an interdisciplinary subfield of computer science and information retrieval.
It deals with the interaction between computers and human language, in particular how to program computers to process and analyze large amounts of natural language data.
"#
}
}

164
cognee/infrastructure/llm/structured_output_framework/baml_src/extract_content_graph.baml
View file

@ -50,6 +50,11 @@ class SummarizedCode {
workflow_description string? workflow_description string?
} }
class DynamicKnowledgeGraph {
@@dynamic
}
// Simple template for basic extraction (fast, good quality) // Simple template for basic extraction (fast, good quality)
template_string ExtractContentGraphPrompt() #" template_string ExtractContentGraphPrompt() #"
You are an advanced algorithm that extracts structured data into a knowledge graph. You are an advanced algorithm that extracts structured data into a knowledge graph.
@ -219,31 +224,15 @@ template_string StrictExtractContentGraphPrompt() #"
"# "#
// OpenAI client with environment model selection // OpenAI client with environment model selection
client<llm> OpenAIClientWithEnvModel { client<llm> OpenAI {
provider openai provider openai
options { options {
model env.LLM_MODEL model client_registry.model
api_key env.OPENAI_API_KEY api_key client_registry.api_key
} }
} }
// Anthropic client with environment model selection
client<llm> AnthropicClientWithEnvModel {
provider anthropic
options {
model env.LLM_MODEL
api_key env.ANTHROPIC_API_KEY
}
}
// Default client (maintains backward compatibility)
client<llm> DefaultClient {
provider openai
options {
model "gpt-4o-mini"
api_key env.OPENAI_API_KEY
}
}
// Function that returns raw structured output (for custom objects - to be handled in Python) // Function that returns raw structured output (for custom objects - to be handled in Python)
function ExtractContentGraphGeneric( function ExtractContentGraphGeneric(
@ -251,7 +240,7 @@ function ExtractContentGraphGeneric(
mode: "simple" | "base" | "guided" | "strict" | "custom"?, mode: "simple" | "base" | "guided" | "strict" | "custom"?,
custom_prompt_content: string? custom_prompt_content: string?
) -> KnowledgeGraph { ) -> KnowledgeGraph {
client OpenAIClientWithEnvModel client OpenAI
prompt #" prompt #"
{% if mode == "base" %} {% if mode == "base" %}
@ -281,12 +270,12 @@ function ExtractContentGraphGeneric(
} }
// Backward-compatible function specifically for KnowledgeGraph // Backward-compatible function specifically for KnowledgeGraph
function ExtractContentGraph( function ExtractDynamicContentGraph(
content: string, content: string,
mode: "simple" | "base" | "guided" | "strict" | "custom"?, mode: "simple" | "base" | "guided" | "strict" | "custom"?,
custom_prompt_content: string? custom_prompt_content: string?
) -> KnowledgeGraph { ) -> DynamicKnowledgeGraph {
client OpenAIClientWithEnvModel client OpenAI
prompt #" prompt #"
{% if mode == "base" %} {% if mode == "base" %}
@ -315,72 +304,10 @@ function ExtractContentGraph(
"# "#
} }
// Alternative function that uses environment variable for prompt selection
function ExtractContentGraphWithEnvPrompt(
content: string,
prompt_override: string?
) -> KnowledgeGraph {
client OpenAIClientWithEnvModel
prompt #"
{% if prompt_override %}
{{ prompt_override }}
{% else %}
{{ ExtractContentGraphPrompt() }}
{% endif %}
{{ ctx.output_format(prefix="Answer in this schema:\n") }}
Before answering, briefly describe what you'll extract from the text, then provide the structured output.
Example format:
I'll extract the main entities and their relationships from this text...
{ ... }
{{ _.role('user') }}
{{ content }}
"#
}
// Function that uses Anthropic client
function ExtractContentGraphWithAnthropic(
content: string,
mode: "simple" | "base" | "guided" | "strict" | "custom"?,
custom_prompt_content: string?
) -> KnowledgeGraph {
client AnthropicClientWithEnvModel
prompt #"
{% if mode == "base" %}
{{ DetailedExtractContentGraphPrompt() }}
{% elif mode == "guided" %}
{{ GuidedExtractContentGraphPrompt() }}
{% elif mode == "strict" %}
{{ StrictExtractContentGraphPrompt() }}
{% elif mode == "custom" and custom_prompt_content %}
{{ custom_prompt_content }}
{% else %}
{{ ExtractContentGraphPrompt() }}
{% endif %}
{{ ctx.output_format(prefix="Answer in this schema:\n") }}
Before answering, briefly describe what you'll extract from the text, then provide the structured output.
Example format:
I'll extract the main entities and their relationships from this text...
{ ... }
{{ _.role('user') }}
{{ content }}
"#
}
// Summarization functions // Summarization functions
function SummarizeContent(content: string) -> SummarizedContent { function SummarizeContent(content: string) -> SummarizedContent {
client OpenAIClientWithEnvModel client OpenAI
prompt #" prompt #"
{{ SummarizeContentPrompt() }} {{ SummarizeContentPrompt() }}
@ -393,7 +320,7 @@ function SummarizeContent(content: string) -> SummarizedContent {
} }
function SummarizeCode(content: string) -> SummarizedCode { function SummarizeCode(content: string) -> SummarizedCode {
client OpenAIClientWithEnvModel client OpenAI
prompt #" prompt #"
{{ SummarizeCodePrompt() }} {{ SummarizeCodePrompt() }}
@ -405,71 +332,12 @@ function SummarizeCode(content: string) -> SummarizedCode {
"# "#
} }
test ExtractPersonExample {
functions [ExtractContentGraph]
args {
content #"
My name is Vasiliy. I was born in 1992. I am a software engineer. I work at Google and am based in Berlin.
"#
mode "simple"
}
}
test ExtractGuidedExample {
functions [ExtractContentGraph]
args {
content #"
Apple Inc. was founded by Steve Jobs in 1976. The company is headquartered in Cupertino, California.
Tim Cook is the current CEO of Apple Inc.
"#
mode "guided"
}
}
test ExtractStrictExample { test ExtractStrictExample {
functions [ExtractContentGraph] functions [ExtractContentGraphGeneric]
args { args {
content #" content #"
The Python programming language was created by Guido van Rossum in 1991. The Python programming language was created by Guido van Rossum in 1991.
"# "#
mode "strict" mode "strict"
} }
} }
test ExtractGenericExample {
functions [ExtractContentGraphGeneric]
args {
content #"
React is a JavaScript library for building user interfaces, developed by Facebook.
"#
mode "simple"
}
}
test SummarizeContentExample {
functions [SummarizeContent]
args {
content #"
Natural language processing (NLP) is an interdisciplinary subfield of computer science and information retrieval.
It deals with the interaction between computers and human language, in particular how to program computers to process and analyze large amounts of natural language data.
"#
}
}
test SummarizeCodeExample {
functions [SummarizeCode]
args {
content #"
def fibonacci(n):
if n <= 1:
return n
return fibonacci(n-1) + fibonacci(n-2)
def main():
print(fibonacci(10))
if __name__ == "__main__":
main()
"#
}
}

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

@ -0,0 +1,113 @@
import os
from typing import Type
from pydantic import BaseModel
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.shared.data_models import SummarizedCode
from cognee.shared.logging_utils import get_logger
from baml_py import ClientRegistry
logger = get_logger("extract_summary_baml")
def get_mock_summarized_code():
"""Local mock function to avoid circular imports."""
return SummarizedCode(
high_level_summary="Mock code summary",
key_features=["Mock feature 1", "Mock feature 2"],
imports=["mock_import"],
constants=["MOCK_CONSTANT"],
classes=[],
functions=[],
workflow_description="Mock workflow description",
)
async def extract_summary(content: str, response_model: Type[BaseModel]):
"""
Extract summary using BAML framework.
Args:
content: The content to summarize
response_model: The Pydantic model type for the response
Returns:
BaseModel: The summarized content in the specified format
"""
config = get_llm_config()
baml_registry = ClientRegistry()
baml_registry.add_llm_client(
name="extract_category_client",
provider=config.llm_model,
options={
"model": config.llm_model,
"temperature": config.llm_temperature,
"api_key": config.llm_api_key,
},
)
baml_registry.set_primary("extract_category_client")
# Use BAML's SummarizeContent function
summary_result = await b.SummarizeContent(
content, baml_options={"client_registry": baml_registry}
)
# Convert BAML result to the expected response model
if response_model is SummarizedCode:
# If it's asking for SummarizedCode but we got SummarizedContent,
# we need to use SummarizeCode instead
code_result = await b.SummarizeCode(
content, baml_options={"client_registry": baml_registry}
)
return code_result
else:
# For other models, return the summary result
return summary_result
async def extract_code_summary(content: str):
"""
Extract code summary using BAML framework with mocking support.
Args:
content: The code content to summarize
Returns:
SummarizedCode: The summarized code information
"""
enable_mocking = os.getenv("MOCK_CODE_SUMMARY", "false")
if isinstance(enable_mocking, bool):
enable_mocking = str(enable_mocking).lower()
enable_mocking = enable_mocking in ("true", "1", "yes")
if enable_mocking:
result = get_mock_summarized_code()
return result
else:
try:
config = get_llm_config()
baml_registry = ClientRegistry()
baml_registry.add_llm_client(
name="extract_content_category",
provider=config.llm_provider,
options={
"model": config.llm_model,
"temperature": config.llm_temperature,
"api_key": config.llm_api_key,
},
)
baml_registry.set_primary("extract_content_category")
result = await b.SummarizeCode(content, baml_options={"client_registry": baml_registry})
except Exception as e:
logger.error(
"Failed to extract code summary with BAML, falling back to mock summary", exc_info=e
)
result = get_mock_summarized_code()
return result

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

@ -1,11 +1,13 @@
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
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.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.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
logger = get_logger("extract_summary_baml") logger = get_logger("extract_summary_baml")
@ -35,9 +37,21 @@ async def extract_summary(content: str, response_model: Type[BaseModel]):
""" """
config = get_llm_config() config = get_llm_config()
baml_registry = ClientRegistry()
baml_registry.add_llm_client(
name="def",
provider="openai",
options={
"model": config.llm_model,
"temperature": config.llm_temperature,
"api_key": config.llm_api_key
})
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(
content, baml_options={"client_registry": config.baml_registry} content, baml_options={"client_registry": baml_registry}
) )
# Convert BAML result to the expected response model # Convert BAML result to the expected response model
@ -74,8 +88,20 @@ async def extract_code_summary(content: str):
else: else:
try: try:
config = get_llm_config() config = get_llm_config()
baml_registry = ClientRegistry()
baml_registry.add_llm_client(
name="def",
provider="openai",
options={
"model": config.llm_model,
"temperature": config.llm_temperature,
"api_key": config.llm_api_key
})
baml_registry.set_primary('def')
result = await b.SummarizeCode( result = await b.SummarizeCode(
content, baml_options={"client_registry": config.baml_registry} content, baml_options={"client_registry": baml_registry}
) )
except Exception as e: except Exception as e:
logger.error( logger.error(

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

@ -1,26 +1,50 @@
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.baml_client.async_client import b
from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.type_builder import (
TypeBuilder,
)
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.shared.logging_utils import get_logger, setup_logging from cognee.shared.logging_utils import get_logger, setup_logging
from baml_py import ClientRegistry
async def extract_content_graph(content: str, response_model: Type[BaseModel]): async def extract_content_graph(
# tb = TypeBuilder() content: str, response_model: Type[BaseModel], mode: str = "simple"
):
config = get_llm_config() config = get_llm_config()
setup_logging() setup_logging()
get_logger(level="INFO")
# country = tb.union \
# ([tb.literal_string("USA"), tb.literal_string("UK"), tb.literal_string("Germany"), tb.literal_string("other")])
# tb.Node.add_property("country", country)
get_logger(level="INFO")
baml_registry = ClientRegistry()
baml_registry.add_llm_client(
name="extract_content_client",
provider=config.llm_provider,
options={
"model": config.llm_model,
"temperature": config.llm_temperature,
"api_key": config.llm_api_key,
},
)
baml_registry.set_primary("extract_content_client")
# if response_model:
# # tb = TypeBuilder()
# # country = tb.union \
# # ([tb.literal_string("USA"), tb.literal_string("UK"), tb.literal_string("Germany"), tb.literal_string("other")])
# # tb.Node.add_property("country", country)
#
# graph = await b.ExtractDynamicContentGraph(
# content, mode=mode, baml_options={"client_registry": baml_registry}
# )
#
# return graph
# else:
graph = await b.ExtractContentGraph( graph = await b.ExtractContentGraph(
content, mode="simple", baml_options={"client_registry": config.baml_registry} content, mode=mode, baml_options={"client_registry": baml_registry}
) )
return graph return graph