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

refactor: Have a dynamic model generator for python for baml responses

Browse source
This commit is contained in:
Igor Ilic 2025-09-08 22:58:52 +02:00
parent befb8ac237
commit 2f59e6ee08
12 changed files with 151 additions and 2771 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -105,153 +105,6 @@ class BamlAsyncClient:
types.ResponseModel, result.cast_to(types, types, stream_types, False, __runtime__)
)
async def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> types.DefaultContentPrediction:
# Check if on_tick is provided
if "on_tick" in baml_options:
# Use streaming internally when on_tick is provided
stream = self.stream.ExtractCategories(content=content, baml_options=baml_options)
return await stream.get_final_response()
else:
# Original non-streaming code
result = await self.__options.merge_options(baml_options).call_function_async(
function_name="ExtractCategories",
args={
"content": content,
},
)
return typing.cast(
types.DefaultContentPrediction,
result.cast_to(types, types, stream_types, False, __runtime__),
)
async def ExtractContentGraphGeneric(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {},
) -> types.KnowledgeGraph:
# Check if on_tick is provided
if "on_tick" in baml_options:
# Use streaming internally when on_tick is provided
stream = self.stream.ExtractContentGraphGeneric(
content=content,
mode=mode,
custom_prompt_content=custom_prompt_content,
baml_options=baml_options,
)
return await stream.get_final_response()
else:
# Original non-streaming code
result = await self.__options.merge_options(baml_options).call_function_async(
function_name="ExtractContentGraphGeneric",
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 ExtractDynamicContentGraph(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {},
) -> types.DynamicKnowledgeGraph:
# Check if on_tick is provided
if "on_tick" in baml_options:
# Use streaming internally when on_tick is provided
stream = self.stream.ExtractDynamicContentGraph(
content=content,
mode=mode,
custom_prompt_content=custom_prompt_content,
baml_options=baml_options,
)
return await stream.get_final_response()
else:
# Original non-streaming code
result = await self.__options.merge_options(baml_options).call_function_async(
function_name="ExtractDynamicContentGraph",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.DynamicKnowledgeGraph,
result.cast_to(types, types, stream_types, False, __runtime__),
)
async def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> types.SummarizedCode:
# Check if on_tick is provided
if "on_tick" in baml_options:
# Use streaming internally when on_tick is provided
stream = self.stream.SummarizeCode(content=content, baml_options=baml_options)
return await stream.get_final_response()
else:
# Original non-streaming code
result = await self.__options.merge_options(baml_options).call_function_async(
function_name="SummarizeCode",
args={
"content": content,
},
)
return typing.cast(
types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__)
)
async def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> types.SummarizedContent:
# Check if on_tick is provided
if "on_tick" in baml_options:
# Use streaming internally when on_tick is provided
stream = self.stream.SummarizeContent(content=content, baml_options=baml_options)
return await stream.get_final_response()
else:
# Original non-streaming code
result = await self.__options.merge_options(baml_options).call_function_async(
function_name="SummarizeContent",
args={
"content": content,
},
)
return typing.cast(
types.SummarizedContent,
result.cast_to(types, types, stream_types, False, __runtime__),
)
class BamlStreamClient:
__options: DoNotUseDirectlyCallManager
@ -283,149 +136,6 @@ class BamlStreamClient:
ctx,
)
def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlStream[stream_types.DefaultContentPrediction, types.DefaultContentPrediction]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream(
function_name="ExtractCategories",
args={
"content": content,
},
)
return baml_py.BamlStream[
stream_types.DefaultContentPrediction, types.DefaultContentPrediction
](
result,
lambda x: typing.cast(
stream_types.DefaultContentPrediction,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.DefaultContentPrediction,
x.cast_to(types, types, stream_types, False, __runtime__),
),
ctx,
)
def ExtractContentGraphGeneric(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream(
function_name="ExtractContentGraphGeneric",
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 ExtractDynamicContentGraph(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream(
function_name="ExtractDynamicContentGraph",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return baml_py.BamlStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph](
result,
lambda x: typing.cast(
stream_types.DynamicKnowledgeGraph,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.DynamicKnowledgeGraph,
x.cast_to(types, types, stream_types, False, __runtime__),
),
ctx,
)
def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlStream[stream_types.SummarizedCode, types.SummarizedCode]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream(
function_name="SummarizeCode",
args={
"content": content,
},
)
return baml_py.BamlStream[stream_types.SummarizedCode, types.SummarizedCode](
result,
lambda x: typing.cast(
stream_types.SummarizedCode,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
)
def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlStream[stream_types.SummarizedContent, types.SummarizedContent]:
ctx, result = self.__options.merge_options(baml_options).create_async_stream(
function_name="SummarizeContent",
args={
"content": content,
},
)
return baml_py.BamlStream[stream_types.SummarizedContent, types.SummarizedContent](
result,
lambda x: typing.cast(
stream_types.SummarizedContent,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
)
class BamlHttpRequestClient:
__options: DoNotUseDirectlyCallManager
@ -449,100 +159,6 @@ class BamlHttpRequestClient:
)
return result
async def ExtractCategories(
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="ExtractCategories",
args={
"content": content,
},
mode="request",
)
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,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(
function_name="ExtractContentGraphGeneric",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="request",
)
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,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(
function_name="ExtractDynamicContentGraph",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="request",
)
return result
async def SummarizeCode(
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="SummarizeCode",
args={
"content": content,
},
mode="request",
)
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:
__options: DoNotUseDirectlyCallManager
@ -566,99 +182,5 @@ class BamlHttpStreamRequestClient:
)
return result
async def ExtractCategories(
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="ExtractCategories",
args={
"content": content,
},
mode="stream",
)
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,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(
function_name="ExtractContentGraphGeneric",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="stream",
)
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,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = await self.__options.merge_options(baml_options).create_http_request_async(
function_name="ExtractDynamicContentGraph",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="stream",
)
return result
async def SummarizeCode(
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="SummarizeCode",
args={
"content": content,
},
mode="stream",
)
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="stream",
)
return result
b = BamlAsyncClient(DoNotUseDirectlyCallManager({}))

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

File diff suppressed because one or more lines are too long

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

@ -33,56 +33,6 @@ class LlmResponseParser:
)
return typing.cast(types.ResponseModel, result)
def ExtractCategories(
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.DefaultContentPrediction:
result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractCategories", llm_response=llm_response, mode="request"
)
return typing.cast(types.DefaultContentPrediction, result)
def ExtractContentGraphGeneric(
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.KnowledgeGraph:
result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractContentGraphGeneric", llm_response=llm_response, mode="request"
)
return typing.cast(types.KnowledgeGraph, result)
def ExtractDynamicContentGraph(
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.DynamicKnowledgeGraph:
result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractDynamicContentGraph", llm_response=llm_response, mode="request"
)
return typing.cast(types.DynamicKnowledgeGraph, result)
def SummarizeCode(
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.SummarizedCode:
result = self.__options.merge_options(baml_options).parse_response(
function_name="SummarizeCode", llm_response=llm_response, mode="request"
)
return typing.cast(types.SummarizedCode, result)
def SummarizeContent(
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> types.SummarizedContent:
result = self.__options.merge_options(baml_options).parse_response(
function_name="SummarizeContent", llm_response=llm_response, mode="request"
)
return typing.cast(types.SummarizedContent, result)
class LlmStreamParser:
__options: DoNotUseDirectlyCallManager
@ -99,53 +49,3 @@ class LlmStreamParser:
function_name="AcreateStructuredOutput", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.ResponseModel, result)
def ExtractCategories(
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.DefaultContentPrediction:
result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractCategories", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.DefaultContentPrediction, result)
def ExtractContentGraphGeneric(
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.KnowledgeGraph:
result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractContentGraphGeneric", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.KnowledgeGraph, result)
def ExtractDynamicContentGraph(
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.DynamicKnowledgeGraph:
result = self.__options.merge_options(baml_options).parse_response(
function_name="ExtractDynamicContentGraph", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.DynamicKnowledgeGraph, result)
def SummarizeCode(
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.SummarizedCode:
result = self.__options.merge_options(baml_options).parse_response(
function_name="SummarizeCode", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.SummarizedCode, result)
def SummarizeContent(
self,
llm_response: str,
baml_options: BamlCallOptions = {},
) -> stream_types.SummarizedContent:
result = self.__options.merge_options(baml_options).parse_response(
function_name="SummarizeContent", llm_response=llm_response, mode="stream"
)
return typing.cast(stream_types.SummarizedContent, result)

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

@ -27,115 +27,14 @@ class StreamState(BaseModel, typing.Generic[StreamStateValueT]):
# #########################################################################
# Generated classes (18)
# Generated classes (1)
# #########################################################################
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):
# doc string for edge
# doc string for source_node_id
source_node_id: typing.Optional[str] = None
target_node_id: 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):
nodes: typing.List["types.Node"]
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):
model_config = ConfigDict(extra="allow")
id: typing.Optional[str] = None
name: typing.Optional[str] = None
type: typing.Optional[str] = None
description: typing.Optional[str] = None
class ProceduralContent(BaseModel):
type: typing.Optional[str] = None
subclass: typing.List[str]
class ResponseModel(BaseModel):
model_config = ConfigDict(extra="allow")
class SummarizedClass(BaseModel):
name: typing.Optional[str] = None
description: typing.Optional[str] = None
methods: typing.Optional[typing.List["SummarizedFunction"]] = None
decorators: typing.Optional[typing.List[str]] = None
class SummarizedCode(BaseModel):
high_level_summary: typing.Optional[str] = None
key_features: typing.List[str]
imports: typing.List[str]
constants: typing.List[str]
classes: typing.List["SummarizedClass"]
functions: typing.List["SummarizedFunction"]
workflow_description: typing.Optional[str] = None
class SummarizedContent(BaseModel):
summary: typing.Optional[str] = None
description: typing.Optional[str] = None
class SummarizedFunction(BaseModel):
name: typing.Optional[str] = None
description: typing.Optional[str] = None
inputs: typing.Optional[typing.List[str]] = None
outputs: 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)
# #########################################################################

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

@ -117,148 +117,6 @@ class BamlSyncClient:
types.ResponseModel, result.cast_to(types, types, stream_types, False, __runtime__)
)
def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> types.DefaultContentPrediction:
# Check if on_tick is provided
if "on_tick" in baml_options:
stream = self.stream.ExtractCategories(content=content, baml_options=baml_options)
return stream.get_final_response()
else:
# Original non-streaming code
result = self.__options.merge_options(baml_options).call_function_sync(
function_name="ExtractCategories",
args={
"content": content,
},
)
return typing.cast(
types.DefaultContentPrediction,
result.cast_to(types, types, stream_types, False, __runtime__),
)
def ExtractContentGraphGeneric(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {},
) -> types.KnowledgeGraph:
# Check if on_tick is provided
if "on_tick" in baml_options:
stream = self.stream.ExtractContentGraphGeneric(
content=content,
mode=mode,
custom_prompt_content=custom_prompt_content,
baml_options=baml_options,
)
return stream.get_final_response()
else:
# Original non-streaming code
result = self.__options.merge_options(baml_options).call_function_sync(
function_name="ExtractContentGraphGeneric",
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 ExtractDynamicContentGraph(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {},
) -> types.DynamicKnowledgeGraph:
# Check if on_tick is provided
if "on_tick" in baml_options:
stream = self.stream.ExtractDynamicContentGraph(
content=content,
mode=mode,
custom_prompt_content=custom_prompt_content,
baml_options=baml_options,
)
return stream.get_final_response()
else:
# Original non-streaming code
result = self.__options.merge_options(baml_options).call_function_sync(
function_name="ExtractDynamicContentGraph",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.DynamicKnowledgeGraph,
result.cast_to(types, types, stream_types, False, __runtime__),
)
def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> types.SummarizedCode:
# Check if on_tick is provided
if "on_tick" in baml_options:
stream = self.stream.SummarizeCode(content=content, baml_options=baml_options)
return stream.get_final_response()
else:
# Original non-streaming code
result = self.__options.merge_options(baml_options).call_function_sync(
function_name="SummarizeCode",
args={
"content": content,
},
)
return typing.cast(
types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__)
)
def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> types.SummarizedContent:
# Check if on_tick is provided
if "on_tick" in baml_options:
stream = self.stream.SummarizeContent(content=content, baml_options=baml_options)
return stream.get_final_response()
else:
# Original non-streaming code
result = self.__options.merge_options(baml_options).call_function_sync(
function_name="SummarizeContent",
args={
"content": content,
},
)
return typing.cast(
types.SummarizedContent,
result.cast_to(types, types, stream_types, False, __runtime__),
)
class BamlStreamClient:
__options: DoNotUseDirectlyCallManager
@ -290,153 +148,6 @@ class BamlStreamClient:
ctx,
)
def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlSyncStream[
stream_types.DefaultContentPrediction, types.DefaultContentPrediction
]:
ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
function_name="ExtractCategories",
args={
"content": content,
},
)
return baml_py.BamlSyncStream[
stream_types.DefaultContentPrediction, types.DefaultContentPrediction
](
result,
lambda x: typing.cast(
stream_types.DefaultContentPrediction,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.DefaultContentPrediction,
x.cast_to(types, types, stream_types, False, __runtime__),
),
ctx,
)
def ExtractContentGraphGeneric(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
function_name="ExtractContentGraphGeneric",
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 ExtractDynamicContentGraph(
self,
content: str,
mode: typing.Optional[
typing.Union[
typing_extensions.Literal["simple"],
typing_extensions.Literal["base"],
typing_extensions.Literal["guided"],
typing_extensions.Literal["strict"],
typing_extensions.Literal["custom"],
]
] = None,
custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlSyncStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph]:
ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
function_name="ExtractDynamicContentGraph",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return baml_py.BamlSyncStream[
stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph
](
result,
lambda x: typing.cast(
stream_types.DynamicKnowledgeGraph,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.DynamicKnowledgeGraph,
x.cast_to(types, types, stream_types, False, __runtime__),
),
ctx,
)
def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode]:
ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
function_name="SummarizeCode",
args={
"content": content,
},
)
return baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode](
result,
lambda x: typing.cast(
stream_types.SummarizedCode,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
)
def SummarizeContent(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent]:
ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
function_name="SummarizeContent",
args={
"content": content,
},
)
return baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent](
result,
lambda x: typing.cast(
stream_types.SummarizedContent,
x.cast_to(types, types, stream_types, True, __runtime__),
),
lambda x: typing.cast(
types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)
),
ctx,
)
class BamlHttpRequestClient:
__options: DoNotUseDirectlyCallManager
@ -460,100 +171,6 @@ class BamlHttpRequestClient:
)
return result
def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(
function_name="ExtractCategories",
args={
"content": content,
},
mode="request",
)
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,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(
function_name="ExtractContentGraphGeneric",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="request",
)
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,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(
function_name="ExtractDynamicContentGraph",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="request",
)
return result
def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(
function_name="SummarizeCode",
args={
"content": content,
},
mode="request",
)
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:
__options: DoNotUseDirectlyCallManager
@ -577,99 +194,5 @@ class BamlHttpStreamRequestClient:
)
return result
def ExtractCategories(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(
function_name="ExtractCategories",
args={
"content": content,
},
mode="stream",
)
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,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(
function_name="ExtractContentGraphGeneric",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="stream",
)
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,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(
function_name="ExtractDynamicContentGraph",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
mode="stream",
)
return result
def SummarizeCode(
self,
content: str,
baml_options: BamlCallOptions = {},
) -> baml_py.baml_py.HTTPRequest:
result = self.__options.merge_options(baml_options).create_http_request_sync(
function_name="SummarizeCode",
args={
"content": content,
},
mode="stream",
)
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="stream",
)
return result
b = BamlSyncClient(DoNotUseDirectlyCallManager({}))

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

@ -24,24 +24,7 @@ class TypeBuilder(type_builder.TypeBuilder):
super().__init__(
classes=set(
[
"AudioContent",
"ContentLabel",
"DefaultContentPrediction",
"DynamicKnowledgeGraph",
"Edge",
"ImageContent",
"KnowledgeGraph",
"Model3DContent",
"MultimediaContent",
"Node",
"ProceduralContent",
"ResponseModel",
"SummarizedClass",
"SummarizedCode",
"SummarizedContent",
"SummarizedFunction",
"TextContent",
"VideoContent",
]
),
enums=set([]),
@ -53,81 +36,13 @@ class TypeBuilder(type_builder.TypeBuilder):
# #########################################################################
# #########################################################################
# Generated classes 18
# Generated classes 1
# #########################################################################
@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
def Edge(self) -> "EdgeViewer":
return EdgeViewer(self)
@property
def ImageContent(self) -> "ImageContentViewer":
return ImageContentViewer(self)
@property
def KnowledgeGraph(self) -> "KnowledgeGraphViewer":
return KnowledgeGraphViewer(self)
@property
def Model3DContent(self) -> "Model3DContentViewer":
return Model3DContentViewer(self)
@property
def MultimediaContent(self) -> "MultimediaContentViewer":
return MultimediaContentViewer(self)
@property
def Node(self) -> "NodeBuilder":
return NodeBuilder(self)
@property
def ProceduralContent(self) -> "ProceduralContentViewer":
return ProceduralContentViewer(self)
@property
def ResponseModel(self) -> "ResponseModelBuilder":
return ResponseModelBuilder(self)
@property
def SummarizedClass(self) -> "SummarizedClassViewer":
return SummarizedClassViewer(self)
@property
def SummarizedCode(self) -> "SummarizedCodeViewer":
return SummarizedCodeViewer(self)
@property
def SummarizedContent(self) -> "SummarizedContentViewer":
return SummarizedContentViewer(self)
@property
def SummarizedFunction(self) -> "SummarizedFunctionViewer":
return SummarizedFunctionViewer(self)
@property
def TextContent(self) -> "TextContentViewer":
return TextContentViewer(self)
@property
def VideoContent(self) -> "VideoContentViewer":
return VideoContentViewer(self)
# #########################################################################
# Generated enums 0
@ -135,533 +50,10 @@ class TypeBuilder(type_builder.TypeBuilder):
# #########################################################################
# Generated classes 18
# Generated classes 1
# #########################################################################
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 self._bldr.list_properties()
def remove_property(self, name: str) -> None:
self._bldr.remove_property(name)
def reset(self) -> None:
self._bldr.reset()
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:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("Edge")
self._properties: typing.Set[str] = set(
[
"source_node_id",
"target_node_id",
"relationship_name",
]
)
self._props = EdgeProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "EdgeProperties":
return self._props
class EdgeViewer(EdgeAst):
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 EdgeProperties:
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 source_node_id(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("source_node_id"))
@property
def target_node_id(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("target_node_id"))
@property
def relationship_name(self) -> type_builder.ClassPropertyViewer:
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:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("KnowledgeGraph")
self._properties: typing.Set[str] = set(
[
"nodes",
"edges",
]
)
self._props = KnowledgeGraphProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "KnowledgeGraphProperties":
return self._props
class KnowledgeGraphViewer(KnowledgeGraphAst):
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 KnowledgeGraphProperties:
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 nodes(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("nodes"))
@property
def edges(self) -> type_builder.ClassPropertyViewer:
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:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("Node")
self._properties: typing.Set[str] = set(
[
"id",
"name",
"type",
"description",
]
)
self._props = NodeProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "NodeProperties":
return self._props
class NodeBuilder(NodeAst):
def __init__(self, tb: type_builder.TypeBuilder):
super().__init__(tb)
def add_property(self, name: str, type: baml_py.FieldType) -> baml_py.ClassPropertyBuilder:
if name in self._properties:
raise ValueError(f"Property {name} already exists.")
return self._bldr.property(name).type(type)
def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]:
return self._bldr.list_properties()
def remove_property(self, name: str) -> None:
self._bldr.remove_property(name)
def reset(self) -> None:
self._bldr.reset()
class NodeProperties:
def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]):
self.__bldr = bldr
self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821
def __getattr__(self, name: str) -> baml_py.ClassPropertyBuilder:
if name not in self.__properties:
raise AttributeError(f"Property {name} not found.")
return self.__bldr.property(name)
@property
def id(self) -> baml_py.ClassPropertyBuilder:
return self.__bldr.property("id")
@property
def name(self) -> baml_py.ClassPropertyBuilder:
return self.__bldr.property("name")
@property
def type(self) -> baml_py.ClassPropertyBuilder:
return self.__bldr.property("type")
@property
def description(self) -> baml_py.ClassPropertyBuilder:
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 ResponseModelAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
@ -705,323 +97,3 @@ class ResponseModelProperties:
if name not in self.__properties:
raise AttributeError(f"Property {name} not found.")
return self.__bldr.property(name)
class SummarizedClassAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("SummarizedClass")
self._properties: typing.Set[str] = set(
[
"name",
"description",
"methods",
"decorators",
]
)
self._props = SummarizedClassProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "SummarizedClassProperties":
return self._props
class SummarizedClassViewer(SummarizedClassAst):
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 SummarizedClassProperties:
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 name(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("name"))
@property
def description(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("description"))
@property
def methods(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("methods"))
@property
def decorators(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("decorators"))
class SummarizedCodeAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("SummarizedCode")
self._properties: typing.Set[str] = set(
[
"high_level_summary",
"key_features",
"imports",
"constants",
"classes",
"functions",
"workflow_description",
]
)
self._props = SummarizedCodeProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "SummarizedCodeProperties":
return self._props
class SummarizedCodeViewer(SummarizedCodeAst):
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 SummarizedCodeProperties:
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 high_level_summary(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("high_level_summary"))
@property
def key_features(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("key_features"))
@property
def imports(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("imports"))
@property
def constants(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("constants"))
@property
def classes(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("classes"))
@property
def functions(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("functions"))
@property
def workflow_description(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("workflow_description"))
class SummarizedContentAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("SummarizedContent")
self._properties: typing.Set[str] = set(
[
"summary",
"description",
]
)
self._props = SummarizedContentProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "SummarizedContentProperties":
return self._props
class SummarizedContentViewer(SummarizedContentAst):
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 SummarizedContentProperties:
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 summary(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("summary"))
@property
def description(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("description"))
class SummarizedFunctionAst:
def __init__(self, tb: type_builder.TypeBuilder):
_tb = tb._tb # type: ignore (we know how to use this private attribute)
self._bldr = _tb.class_("SummarizedFunction")
self._properties: typing.Set[str] = set(
[
"name",
"description",
"inputs",
"outputs",
"decorators",
]
)
self._props = SummarizedFunctionProperties(self._bldr, self._properties)
def type(self) -> baml_py.FieldType:
return self._bldr.field()
@property
def props(self) -> "SummarizedFunctionProperties":
return self._props
class SummarizedFunctionViewer(SummarizedFunctionAst):
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 SummarizedFunctionProperties:
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 name(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("name"))
@property
def description(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("description"))
@property
def inputs(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("inputs"))
@property
def outputs(self) -> type_builder.ClassPropertyViewer:
return type_builder.ClassPropertyViewer(self.__bldr.property("outputs"))
@property
def decorators(self) -> type_builder.ClassPropertyViewer:
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"))

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

@ -15,40 +15,6 @@ from . import stream_types
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,
"stream_types.Edge": stream_types.Edge,
"types.ImageContent": types.ImageContent,
"stream_types.ImageContent": stream_types.ImageContent,
"types.KnowledgeGraph": 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,
"stream_types.Node": stream_types.Node,
"types.ProceduralContent": types.ProceduralContent,
"stream_types.ProceduralContent": stream_types.ProceduralContent,
"types.ResponseModel": types.ResponseModel,
"stream_types.ResponseModel": stream_types.ResponseModel,
"types.SummarizedClass": types.SummarizedClass,
"stream_types.SummarizedClass": stream_types.SummarizedClass,
"types.SummarizedCode": types.SummarizedCode,
"stream_types.SummarizedCode": stream_types.SummarizedCode,
"types.SummarizedContent": types.SummarizedContent,
"stream_types.SummarizedContent": stream_types.SummarizedContent,
"types.SummarizedFunction": 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,
}

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

@ -48,123 +48,14 @@ def all_succeeded(checks: typing.Dict[CheckName, Check]) -> bool:
# #########################################################################
# #########################################################################
# Generated classes (18)
# Generated classes (1)
# #########################################################################
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):
# doc string for edge
# doc string for source_node_id
source_node_id: str
target_node_id: str
relationship_name: str
class ImageContent(BaseModel):
type: str
subclass: typing.List[str]
class KnowledgeGraph(BaseModel):
nodes: typing.List["Node"]
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):
model_config = ConfigDict(extra="allow")
id: str
name: str
type: str
description: str
class ProceduralContent(BaseModel):
type: str
subclass: typing.List[str]
class ResponseModel(BaseModel):
model_config = ConfigDict(extra="allow")
class SummarizedClass(BaseModel):
name: str
description: str
methods: typing.Optional[typing.List["SummarizedFunction"]] = None
decorators: typing.Optional[typing.List[str]] = None
class SummarizedCode(BaseModel):
high_level_summary: str
key_features: typing.List[str]
imports: typing.List[str]
constants: typing.List[str]
classes: typing.List["SummarizedClass"]
functions: typing.List["SummarizedFunction"]
workflow_description: typing.Optional[str] = None
class SummarizedContent(BaseModel):
summary: str
description: str
class SummarizedFunction(BaseModel):
name: str
description: str
inputs: typing.Optional[typing.List[str]] = None
outputs: 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)
# #########################################################################

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

@ -2,6 +2,15 @@ class ResponseModel {
@@dynamic
}
// OpenAI client with environment model selection
client<llm> OpenAI {
provider openai
options {
model client_registry.model
api_key client_registry.api_key
}
}
function AcreateStructuredOutput(
text_input: string,
system_prompt: string,

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

@ -1,109 +0,0 @@
// 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.
"#
}
}

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

@ -1,343 +0,0 @@
class Node {
id string
name string
type string
description string
@@dynamic
}
/// doc string for edge
class Edge {
/// doc string for source_node_id
source_node_id string
target_node_id string
relationship_name string
}
class KnowledgeGraph {
nodes (Node @stream.done)[]
edges Edge[]
}
// Summarization classes
class SummarizedContent {
summary string
description string
}
class SummarizedFunction {
name string
description string
inputs string[]?
outputs string[]?
decorators string[]?
}
class SummarizedClass {
name string
description string
methods SummarizedFunction[]?
decorators string[]?
}
class SummarizedCode {
high_level_summary string
key_features string[]
imports string[]
constants string[]
classes SummarizedClass[]
functions SummarizedFunction[]
workflow_description string?
}
class DynamicKnowledgeGraph {
@@dynamic
}
// Simple template for basic extraction (fast, good quality)
template_string ExtractContentGraphPrompt() #"
You are an advanced algorithm that extracts structured data into a knowledge graph.
- **Nodes**: Entities/concepts (like Wikipedia articles).
- **Edges**: Relationships (like Wikipedia links). Use snake_case (e.g., `acted_in`).
**Rules:**
1. **Node Labeling & IDs**
- Use basic types only (e.g., "Person", "Date", "Organization").
- Avoid overly specific or generic terms (e.g., no "Mathematician" or "Entity").
- Node IDs must be human-readable names from the text (no numbers).
2. **Dates & Numbers**
- Label dates as **"Date"** in "YYYY-MM-DD" format (use available parts if incomplete).
- Properties are key-value pairs; do not use escaped quotes.
3. **Coreference Resolution**
- Use a single, complete identifier for each entity (e.g., always "John Doe" not "Joe" or "he").
4. **Relationship Labels**:
- Use descriptive, lowercase, snake_case names for edges.
- *Example*: born_in, married_to, invented_by.
- Avoid vague or generic labels like isA, relatesTo, has.
- Avoid duplicated relationships like produces, produced by.
5. **Strict Compliance**
- Follow these rules exactly. Non-compliance results in termination.
"#
// Summarization prompt template
template_string SummarizeContentPrompt() #"
You are a top-tier summarization engine. Your task is to summarize text and make it versatile.
Be brief and concise, but keep the important information and the subject.
Use synonym words where possible in order to change the wording but keep the meaning.
"#
// Code summarization prompt template
template_string SummarizeCodePrompt() #"
You are an expert code analyst. Analyze the provided source code and extract key information:
1. Provide a high-level summary of what the code does
2. List key features and functionality
3. Identify imports and dependencies
4. List constants and global variables
5. Summarize classes with their methods
6. Summarize standalone functions
7. Describe the overall workflow if applicable
Be precise and technical while remaining clear and concise.
"#
// Detailed template for complex extraction (slower, higher quality)
template_string DetailedExtractContentGraphPrompt() #"
You are a top-tier algorithm designed for extracting information in structured formats to build a knowledge graph.
**Nodes** represent entities and concepts. They're akin to Wikipedia nodes.
**Edges** represent relationships between concepts. They're akin to Wikipedia links.
The aim is to achieve simplicity and clarity in the knowledge graph.
# 1. Labeling Nodes
**Consistency**: Ensure you use basic or elementary types for node labels.
- For example, when you identify an entity representing a person, always label it as **"Person"**.
- Avoid using more specific terms like "Mathematician" or "Scientist", keep those as "profession" property.
- Don't use too generic terms like "Entity".
**Node IDs**: Never utilize integers as node IDs.
- Node IDs should be names or human-readable identifiers found in the text.
# 2. Handling Numerical Data and Dates
- For example, when you identify an entity representing a date, make sure it has type **"Date"**.
- Extract the date in the format "YYYY-MM-DD"
- If not possible to extract the whole date, extract month or year, or both if available.
- **Property Format**: Properties must be in a key-value format.
- **Quotation Marks**: Never use escaped single or double quotes within property values.
- **Naming Convention**: Use snake_case for relationship names, e.g., `acted_in`.
# 3. Coreference Resolution
- **Maintain Entity Consistency**: When extracting entities, it's vital to ensure consistency.
If an entity, such as "John Doe", is mentioned multiple times in the text but is referred to by different names or pronouns (e.g., "Joe", "he"),
always use the most complete identifier for that entity throughout the knowledge graph. In this example, use "John Doe" as the Person's ID.
Remember, the knowledge graph should be coherent and easily understandable, so maintaining consistency in entity references is crucial.
# 4. Strict Compliance
Adhere to the rules strictly. Non-compliance will result in termination.
"#
// Guided template with step-by-step instructions
template_string GuidedExtractContentGraphPrompt() #"
You are an advanced algorithm designed to extract structured information to build a clean, consistent, and human-readable knowledge graph.
**Objective**:
- Nodes represent entities and concepts, similar to Wikipedia articles.
- Edges represent typed relationships between nodes, similar to Wikipedia hyperlinks.
- The graph must be clear, minimal, consistent, and semantically precise.
**Node Guidelines**:
1. **Label Consistency**:
- Use consistent, basic types for all node labels.
- Do not switch between granular or vague labels for the same kind of entity.
- Pick one label for each category and apply it uniformly.
- Each entity type should be in a singular form and in a case of multiple words separated by whitespaces
2. **Node Identifiers**:
- Node IDs must be human-readable and derived directly from the text.
- Prefer full names and canonical terms.
- Never use integers or autogenerated IDs.
- *Example*: Use "Marie Curie", "Theory of Evolution", "Google".
3. **Coreference Resolution**:
- Maintain one consistent node ID for each real-world entity.
- Resolve aliases, acronyms, and pronouns to the most complete form.
- *Example*: Always use "John Doe" even if later referred to as "Doe" or "he".
**Edge Guidelines**:
4. **Relationship Labels**:
- Use descriptive, lowercase, snake_case names for edges.
- *Example*: born_in, married_to, invented_by.
- Avoid vague or generic labels like isA, relatesTo, has.
5. **Relationship Direction**:
- Edges must be directional and logically consistent.
- *Example*:
- "Marie Curie" —[born_in]→ "Warsaw"
- "Radioactivity" —[discovered_by]→ "Marie Curie"
**Compliance**:
Strict adherence to these guidelines is required. Any deviation will result in immediate termination of the task.
"#
// Strict template with zero-tolerance rules
template_string StrictExtractContentGraphPrompt() #"
You are a top-tier algorithm for **extracting structured information** from unstructured text to build a **knowledge graph**.
Your primary goal is to extract:
- **Nodes**: Representing **entities** and **concepts** (like Wikipedia nodes).
- **Edges**: Representing **relationships** between those concepts (like Wikipedia links).
The resulting knowledge graph must be **simple, consistent, and human-readable**.
## 1. Node Labeling and Identification
### Node Types
Use **basic atomic types** for node labels. Always prefer general types over specific roles or professions:
- "Person" for any human.
- "Organization" for companies, institutions, etc.
- "Location" for geographic or place entities.
- "Date" for any temporal expression.
- "Event" for historical or scheduled occurrences.
- "Work" for books, films, artworks, or research papers.
- "Concept" for abstract notions or ideas.
### Node IDs
- Always assign **human-readable and unambiguous identifiers**.
- Never use numeric or autogenerated IDs.
- Prioritize **most complete form** of entity names for consistency.
## 2. Relationship Handling
- Use **snake_case** for all relationship (edge) types.
- Keep relationship types semantically clear and consistent.
- Avoid vague relation names like "related_to" unless no better alternative exists.
## 3. Strict Compliance
Follow all rules exactly. Any deviation may lead to rejection or incorrect graph construction.
"#
// OpenAI client with environment model selection
client<llm> OpenAI {
provider openai
options {
model client_registry.model
api_key client_registry.api_key
}
}
// Function that returns raw structured output (for custom objects - to be handled in Python)
function ExtractContentGraphGeneric(
content: string,
mode: "simple" | "base" | "guided" | "strict" | "custom"?,
custom_prompt_content: string?
) -> KnowledgeGraph {
client OpenAI
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 }}
"#
}
// Backward-compatible function specifically for KnowledgeGraph
function ExtractDynamicContentGraph(
content: string,
mode: "simple" | "base" | "guided" | "strict" | "custom"?,
custom_prompt_content: string?
) -> DynamicKnowledgeGraph {
client OpenAI
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
function SummarizeContent(content: string) -> SummarizedContent {
client OpenAI
prompt #"
{{ SummarizeContentPrompt() }}
{{ ctx.output_format(prefix="Answer in this schema:\n") }}
{{ _.role('user') }}
{{ content }}
"#
}
function SummarizeCode(content: string) -> SummarizedCode {
client OpenAI
prompt #"
{{ SummarizeCodePrompt() }}
{{ ctx.output_format(prefix="Answer in this schema:\n") }}
{{ _.role('user') }}
{{ content }}
"#
}
test ExtractStrictExample {
functions [ExtractContentGraphGeneric]
args {
content #"
The Python programming language was created by Guido van Rossum in 1991.
"#
mode "strict"
}
}

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

@ -1,7 +1,7 @@
import asyncio
from typing import Type
from cognee.shared.logging_utils import get_logger
from cognee.shared.data_models import SummarizedCode
from cognee.infrastructure.llm.config import get_llm_config
from typing import List, Dict, Union, Optional, Literal
@ -14,94 +14,145 @@ from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.type
)
from cognee.infrastructure.llm.structured_output_framework.baml.baml_client import b
from pydantic import BaseModel
from typing import get_origin, get_args
logger = get_logger("extract_summary_baml")
def create_dynamic_baml_type(pydantic_model):
class SummarizedFunction(BaseModel):
name: str
description: str
inputs: Optional[List[str]] = None
outputs: Optional[List[str]] = None
decorators: Optional[List[str]] = None
class SummarizedClass(BaseModel):
name: str
description: str
methods: Optional[List[SummarizedFunction]] = None
decorators: Optional[List[str]] = None
class SummarizedCode(BaseModel):
high_level_summary: str
key_features: List[str]
imports: List[str] = []
constants: List[str] = []
classes: List[SummarizedClass] = []
functions: List[SummarizedFunction] = []
workflow_description: Optional[str] = None
def create_dynamic_baml_type(baml_model, pydantic_model):
tb = TypeBuilder()
# if pydantic_model == str:
# b.ResponseModel.add_property("text", tb.string())
# return tb
#
# def map_type(field_type, field_info):
# # Handle Optional/Union types
# if getattr(field_type, "__origin__", None) == Union:
# # Extract types from Union
# types = field_type.__args__
# # Handle Optional (Union with NoneType)
# if type(None) in types:
# inner_type = next(t for t in types if t != type(None))
# return map_type(inner_type, field_info).optional()
# # Handle regular Union
# mapped_types = [map_type(t, field_info) for t in types]
# return tb.union(*mapped_types)
#
# # Handle Lists
# if getattr(field_type, "__origin__", None) == list:
# inner_type = field_type.__args__[0]
# return map_type(inner_type, field_info).list()
#
# # Handle Maps/Dictionaries
# if getattr(field_type, "__origin__", None) == dict:
# key_type, value_type = field_type.__args__
# # BAML only supports string or enum keys in maps
# if key_type not in [str, Enum]:
# raise ValueError("Map keys must be strings or enums in BAML")
# return tb.map(map_type(key_type, field_info), map_type(value_type, field_info))
#
# # Handle Literal types
# if getattr(field_type, "__origin__", None) == Literal:
# literal_values = field_type.__args__
# return tb.union(*[tb.literal(val) for val in literal_values])
#
# # Handle Enums
# if isinstance(field_type, type) and issubclass(field_type, Enum):
# enum_type = tb.add_enum(field_type.__name__)
# for member in field_type:
# enum_type.add_value(member.name)
# return enum_type.type()
#
# # Handle primitive and special types
# type_mapping = {
# str: tb.string(),
# int: tb.int(),
# float: tb.float(),
# bool: tb.bool(),
# Image: tb.image(),
# Audio: tb.audio(),
# Video: tb.video(),
# Pdf: tb.pdf(),
# # datetime is not natively supported in BAML, map to string
# datetime: tb.string(),
# }
#
# # Handle nested BaseModel classes
# if isinstance(field_type, type) and issubclass(field_type, BaseModel):
# nested_tb = create_dynamic_baml_type(field_type)
# # Get the last created class from the nested TypeBuilder
# return nested_tb.get_last_class().type()
#
# if field_type in type_mapping:
# return type_mapping[field_type]
#
# raise ValueError(f"Unsupported type: {field_type}")
#
# fields = pydantic_model.model_fields
#
# # Add fields
# for field_name, field_info in fields.items():
# field_type = field_info.annotation
# baml_type = map_type(field_type, field_info)
#
# # Add property with type
# prop = b.ResponseModel.add_property(field_name, baml_type)
#
# # Add description if available
# if field_info.description:
# prop.description(field_info.description)
if pydantic_model is str:
baml_model.add_property("text", tb.string())
return tb
def map_type(field_type, field_info):
"""
Convert a Python / Pydantic type -> BAML TypeBuilder representation.
"""
origin = get_origin(field_type) # e.g. list[…] -> list
args = get_args(field_type) # e.g. list[int] -> (int,)
# ------------------------------------------------------------------
# 1. Optional / Union ------------------------------------------------
# ------------------------------------------------------------------
if origin is Union:
non_none_args = [t for t in args if t is not type(None)]
# Optional[T] ⇢ exactly (T, NoneType)
if len(args) == 2 and len(non_none_args) == 1:
return map_type(non_none_args[0], field_info).optional()
# Plain Union[A, B, …]
return tb.union(*(map_type(t, field_info) for t in args))
# ------------------------------------------------------------------
# 2. List / Sequence -------------------------------------------------
# ------------------------------------------------------------------
if origin in (list,):
(inner_type,) = args # list has exactly one parameter
return map_type(inner_type, field_info).list()
# ------------------------------------------------------------------
# 3. Dict / Map -------------------------------------------------------
# ------------------------------------------------------------------
def _is_enum_subclass(tp) -> bool:
"""Guarded issubclass – returns False when tp is not a class."""
return isinstance(tp, type) and issubclass(tp, Enum)
if origin in (dict,):
key_type, value_type = args or (str, object)
if key_type is not str and not _is_enum_subclass(key_type):
raise ValueError("BAML maps only allow 'str' or Enum subclasses as keys")
return tb.map(
map_type(key_type, field_info), # mostly tb.string() or enum
map_type(value_type, field_info),
)
# ------------------------------------------------------------------
# 4. Literal ----------------------------------------------------------
# ------------------------------------------------------------------
if origin is Literal:
return tb.union(*(tb.literal(v) for v in args))
# ------------------------------------------------------------------
# 5. Enum -------------------------------------------------------------
# ------------------------------------------------------------------
if _is_enum_subclass(field_type):
enum_builder = tb.add_enum(field_type.__name__)
for member in field_type:
enum_builder.add_value(member.name)
return enum_builder.type()
# ------------------------------------------------------------------
# 6. Nested Pydantic model -------------------------------------------
# ------------------------------------------------------------------
from pydantic import BaseModel # local import
if isinstance(field_type, type) and issubclass(field_type, BaseModel):
# Create nested class if it doesn't exist
try:
nested_class = tb.add_class(field_type.__name__)
except ValueError:
pass
# Find dynamic types of nested class
create_dynamic_baml_type(nested_class, field_type)
# Return nested class model
return nested_class.type()
primitive_map = {
str: tb.string(),
int: tb.int(),
float: tb.float(),
bool: tb.bool(),
datetime: tb.string(), # BAML has no native datetime
}
if field_type in primitive_map:
return primitive_map[field_type]
raise ValueError(f"Unsupported type for BAML mapping: {field_type}")
fields = pydantic_model.model_fields
# Add fields
for field_name, field_info in fields.items():
field_type = field_info.annotation
baml_type = map_type(field_type, field_info)
# Add property with type
prop = baml_model.add_property(field_name, baml_type)
# Add description if available
if field_info.description:
prop.description(field_info.description)
return tb
@ -120,8 +171,9 @@ async def acreate_structured_output(
BaseModel: The summarized content in the specified format
"""
config = get_llm_config()
tb = TypeBuilder()
type_builder = create_dynamic_baml_type(response_model)
type_builder = create_dynamic_baml_type(tb.ResponseModel, SummarizedCode)
result = await b.AcreateStructuredOutput(
text_input=text_input,
@ -136,6 +188,6 @@ if __name__ == "__main__":
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
loop.run_until_complete(acreate_structured_output("TEST", SummarizedCode))
loop.run_until_complete(acreate_structured_output("TEST", "THIS IS A TEST", SummarizedCode))
finally:
loop.run_until_complete(loop.shutdown_asyncgens())