added files

cognee/infrastructure/llm/structured_output_framework/baml/baml_client/__init__.py

@@ -42,7 +42,6 @@ with EnsureBamlPyImport(__version__) as e:
     from .sync_client import b
 
 
-
 # FOR LEGACY COMPATIBILITY, expose "partial_types" as an alias for "stream_types"
 # WE RECOMMEND USERS TO USE "stream_types" INSTEAD
 partial_types = stream_types

cognee/infrastructure/llm/structured_output_framework/baml/baml_client/async_client.py

@@ -36,10 +36,13 @@ class BamlAsyncClient:
         self.__llm_response_parser = LlmResponseParser(options)
         self.__llm_stream_parser = LlmStreamParser(options)
 
-    def with_options(self,
+    def with_options(
+        self,
         tb: typing.Optional[type_builder.TypeBuilder] = None,
         client_registry: typing.Optional[baml_py.baml_py.ClientRegistry] = None,
-        collector: typing.Optional[typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]] = None,
+        collector: typing.Optional[
+            typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]
+        ] = None,
         env: typing.Optional[typing.Dict[str, typing.Optional[str]]] = None,
     ) -> "BamlAsyncClient":
         options: BamlCallOptions = {}
@@ -73,42 +76,106 @@ class BamlAsyncClient:
     def parse_stream(self):
         return self.__llm_stream_parser
 
-    async def ExtractCategories(self, content: str,
+    async def ExtractCategories(
+        self,
+        content: str,
         baml_options: BamlCallOptions = {},
     ) -> types.DefaultContentPrediction:
-        result = await self.__options.merge_options(baml_options).call_function_async(function_name="ExtractCategories", args={
+        result = await self.__options.merge_options(baml_options).call_function_async(
+            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,
+        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:
-        result = await self.__options.merge_options(baml_options).call_function_async(function_name="ExtractContentGraphGeneric", args={
+        result = await self.__options.merge_options(baml_options).call_function_async(
-            "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
+            function_name="ExtractContentGraphGeneric",
-        })
+            args={
-        return typing.cast(types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__))
+                "content": content,
-    async def ExtractDynamicContentGraph(self, content: str,mode: typing.Optional[typing.Union[typing_extensions.Literal['simple'], typing_extensions.Literal['base'], typing_extensions.Literal['guided'], typing_extensions.Literal['strict'], typing_extensions.Literal['custom']]] = None,custom_prompt_content: typing.Optional[str] = None,
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+        )
+        return typing.cast(
+            types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)
+        )
+
+    async def ExtractDynamicContentGraph(
+        self,
+        content: str,
+        mode: typing.Optional[
+            typing.Union[
+                typing_extensions.Literal["simple"],
+                typing_extensions.Literal["base"],
+                typing_extensions.Literal["guided"],
+                typing_extensions.Literal["strict"],
+                typing_extensions.Literal["custom"],
+            ]
+        ] = None,
+        custom_prompt_content: typing.Optional[str] = None,
         baml_options: BamlCallOptions = {},
     ) -> types.DynamicKnowledgeGraph:
-        result = await self.__options.merge_options(baml_options).call_function_async(function_name="ExtractDynamicContentGraph", args={
+        result = await self.__options.merge_options(baml_options).call_function_async(
-            "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
+            function_name="ExtractDynamicContentGraph",
-        })
+            args={
-        return typing.cast(types.DynamicKnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__))
+                "content": content,
-    async def SummarizeCode(self, content: str,
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+        )
+        return typing.cast(
+            types.DynamicKnowledgeGraph,
+            result.cast_to(types, types, stream_types, False, __runtime__),
+        )
+
+    async def SummarizeCode(
+        self,
+        content: str,
         baml_options: BamlCallOptions = {},
     ) -> types.SummarizedCode:
-        result = await self.__options.merge_options(baml_options).call_function_async(function_name="SummarizeCode", args={
+        result = await self.__options.merge_options(baml_options).call_function_async(
+            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,
+        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:
-        result = await self.__options.merge_options(baml_options).call_function_async(function_name="SummarizeContent", args={
+        result = await self.__options.merge_options(baml_options).call_function_async(
+            function_name="SummarizeContent",
+            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:
@@ -117,64 +184,146 @@ class BamlStreamClient:
     def __init__(self, options: DoNotUseDirectlyCallManager):
         self.__options = options
 
-    def ExtractCategories(self, content: str,
+    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={
+        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](
+        )
+        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(
-          lambda x: typing.cast(types.DefaultContentPrediction, x.cast_to(types, types, stream_types, False, __runtime__)),
+                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,
+
+    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={
+        ctx, result = self.__options.merge_options(baml_options).create_async_stream(
-            "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
+            function_name="ExtractContentGraphGeneric",
-        })
+            args={
+                "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(
-          lambda x: typing.cast(types.KnowledgeGraph, x.cast_to(types, types, stream_types, False, __runtime__)),
+                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,
+
+    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={
+        ctx, result = self.__options.merge_options(baml_options).create_async_stream(
-            "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
+            function_name="ExtractDynamicContentGraph",
-        })
+            args={
+                "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(
-          lambda x: typing.cast(types.DynamicKnowledgeGraph, x.cast_to(types, types, stream_types, False, __runtime__)),
+                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,
+
+    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={
+        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(
-          lambda x: typing.cast(types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)),
+                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,
+
+    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={
+        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(
-          lambda x: typing.cast(types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)),
+                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,
         )
 
@@ -185,40 +334,98 @@ class BamlHttpRequestClient:
     def __init__(self, options: DoNotUseDirectlyCallManager):
         self.__options = options
 
-    async def ExtractCategories(self, content: str,
+    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={
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="ExtractCategories",
+            args={
                 "content": content,
-        }, mode="request")
+            },
+            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,
+
+    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={
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
-            "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
+            function_name="ExtractContentGraphGeneric",
-        }, mode="request")
+            args={
-        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")
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+            mode="request",
+        )
         return result
-    async def SummarizeContent(self, content: str,
+
+    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="SummarizeContent", args={
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="ExtractDynamicContentGraph",
+            args={
                 "content": content,
-        }, mode="request")
+                "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
 
 
@@ -228,40 +435,98 @@ class BamlHttpStreamRequestClient:
     def __init__(self, options: DoNotUseDirectlyCallManager):
         self.__options = options
 
-    async def ExtractCategories(self, content: str,
+    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={
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="ExtractCategories",
+            args={
                 "content": content,
-        }, mode="stream")
+            },
+            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,
+
+    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={
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
-            "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content,
+            function_name="ExtractContentGraphGeneric",
-        }, mode="stream")
+            args={
-        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")
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+            mode="stream",
+        )
         return result
-    async def SummarizeContent(self, content: str,
+
+    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="SummarizeContent", args={
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="ExtractDynamicContentGraph",
+            args={
                 "content": content,
-        }, mode="stream")
+                "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
 
 

cognee/infrastructure/llm/structured_output_framework/baml/baml_client/globals.py

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

cognee/infrastructure/llm/structured_output_framework/baml/baml_client/parser.py

@@ -16,6 +16,7 @@ import typing_extensions
 from . import stream_types, types
 from .runtime import DoNotUseDirectlyCallManager, BamlCallOptions
 
+
 class LlmResponseParser:
     __options: DoNotUseDirectlyCallManager
 
@@ -23,37 +24,56 @@ class LlmResponseParser:
         self.__options = options
 
     def ExtractCategories(
-        self, llm_response: str, baml_options: BamlCallOptions = {},
+        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")
+        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 = {},
+        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")
+        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 = {},
+        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")
+        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 = {},
+        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")
+        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 = {},
+        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")
+        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
 
@@ -61,33 +81,51 @@ class LlmStreamParser:
         self.__options = options
 
     def ExtractCategories(
-        self, llm_response: str, baml_options: BamlCallOptions = {},
+        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")
+        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 = {},
+        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")
+        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 = {},
+        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")
+        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 = {},
+        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")
+        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 = {},
+        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")
+        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)
-
-    

cognee/infrastructure/llm/structured_output_framework/baml/baml_client/runtime.py

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

cognee/infrastructure/llm/structured_output_framework/baml/baml_client/stream_types.py

@@ -18,28 +18,37 @@ import baml_py
 
 from . import types
 
-StreamStateValueT = typing.TypeVar('StreamStateValueT')
+StreamStateValueT = typing.TypeVar("StreamStateValueT")
+
+
 class StreamState(BaseModel, typing.Generic[StreamStateValueT]):
     value: StreamStateValueT
     state: typing_extensions.Literal["Pending", "Incomplete", "Complete"]
+
+
 # #########################################################################
 # 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')
+    model_config = ConfigDict(extra="allow")
+
 
 class Edge(BaseModel):
     # doc string for edge
@@ -49,39 +58,47 @@ class Edge(BaseModel):
     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')
+    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 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]
@@ -91,10 +108,12 @@ class SummarizedCode(BaseModel):
     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
@@ -102,14 +121,17 @@ class SummarizedFunction(BaseModel):
     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)
 # #########################################################################

cognee/infrastructure/llm/structured_output_framework/baml/baml_client/sync_client.py

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

cognee/infrastructure/llm/structured_output_framework/baml/baml_client/tracing.py

@@ -14,9 +14,13 @@ from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX
 
 trace = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.trace_fn
 set_tags = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.upsert_tags
+
+
 def 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
 
 
-__all__ = ['trace', 'set_tags', "flush", "on_log_event"]
+__all__ = ["trace", "set_tags", "flush", "on_log_event"]

cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_builder.py

@@ -15,19 +15,39 @@ from baml_py import type_builder
 from baml_py import baml_py
 from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME
 
+
 class TypeBuilder(type_builder.TypeBuilder):
     def __init__(self):
-        super().__init__(classes=set(
+        super().__init__(
-          ["AudioContent","ContentLabel","DefaultContentPrediction","DynamicKnowledgeGraph","Edge","ImageContent","KnowledgeGraph","Model3DContent","MultimediaContent","Node","ProceduralContent","SummarizedClass","SummarizedCode","SummarizedContent","SummarizedFunction","TextContent","VideoContent",]
+            classes=set(
-        ), enums=set(
+                [
-          []
+                    "AudioContent",
-        ), runtime=DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME)
+                    "ContentLabel",
+                    "DefaultContentPrediction",
+                    "DynamicKnowledgeGraph",
+                    "Edge",
+                    "ImageContent",
+                    "KnowledgeGraph",
+                    "Model3DContent",
+                    "MultimediaContent",
+                    "Node",
+                    "ProceduralContent",
+                    "SummarizedClass",
+                    "SummarizedCode",
+                    "SummarizedContent",
+                    "SummarizedFunction",
+                    "TextContent",
+                    "VideoContent",
+                ]
+            ),
+            enums=set([]),
+            runtime=DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME,
+        )
 
     # #########################################################################
     # Generated enums 0
     # #########################################################################
 
-
     # #########################################################################
     # Generated classes 17
     # #########################################################################
@@ -101,7 +121,6 @@ class TypeBuilder(type_builder.TypeBuilder):
         return VideoContentViewer(self)
 
 
-
 # #########################################################################
 # Generated enums 0
 # #########################################################################
@@ -111,11 +130,17 @@ class TypeBuilder(type_builder.TypeBuilder):
 # 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._properties: typing.Set[str] = set(
+            [
+                "type",
+                "subclass",
+            ]
+        )
         self._props = AudioContentProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -130,10 +155,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class AudioContentProperties:
@@ -141,8 +167,6 @@ class AudioContentProperties:
         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"))
@@ -152,13 +176,17 @@ class AudioContentProperties:
         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._properties: typing.Set[str] = set(
+            [
+                "content_type",
+                "type",
+                "subclass",
+            ]
+        )
         self._props = ContentLabelProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -173,10 +201,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class ContentLabelProperties:
@@ -184,8 +213,6 @@ class ContentLabelProperties:
         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"))
@@ -199,13 +226,15 @@ class ContentLabelProperties:
         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._properties: typing.Set[str] = set(
+            [
+                "label",
+            ]
+        )
         self._props = DefaultContentPredictionProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -220,10 +249,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class DefaultContentPredictionProperties:
@@ -231,20 +261,16 @@ class DefaultContentPredictionProperties:
         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._properties: typing.Set[str] = set([])
         self._props = DynamicKnowledgeGraphProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -259,7 +285,6 @@ 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.")
@@ -269,28 +294,28 @@ class DynamicKnowledgeGraphBuilder(DynamicKnowledgeGraphAst):
         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:
     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._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:
@@ -305,10 +330,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class EdgeProperties:
@@ -316,8 +342,6 @@ class EdgeProperties:
         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"))
@@ -331,13 +355,16 @@ class EdgeProperties:
         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._properties: typing.Set[str] = set(
+            [
+                "type",
+                "subclass",
+            ]
+        )
         self._props = ImageContentProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -352,10 +379,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class ImageContentProperties:
@@ -363,8 +391,6 @@ class ImageContentProperties:
         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"))
@@ -374,13 +400,16 @@ class ImageContentProperties:
         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._properties: typing.Set[str] = set(
+            [
+                "nodes",
+                "edges",
+            ]
+        )
         self._props = KnowledgeGraphProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -395,10 +424,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class KnowledgeGraphProperties:
@@ -406,8 +436,6 @@ class KnowledgeGraphProperties:
         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"))
@@ -417,13 +445,16 @@ class KnowledgeGraphProperties:
         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._properties: typing.Set[str] = set(
+            [
+                "type",
+                "subclass",
+            ]
+        )
         self._props = Model3DContentProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -438,10 +469,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class Model3DContentProperties:
@@ -449,8 +481,6 @@ class Model3DContentProperties:
         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"))
@@ -460,13 +490,16 @@ class Model3DContentProperties:
         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._properties: typing.Set[str] = set(
+            [
+                "type",
+                "subclass",
+            ]
+        )
         self._props = MultimediaContentProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -481,10 +514,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class MultimediaContentProperties:
@@ -492,8 +526,6 @@ class MultimediaContentProperties:
         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"))
@@ -503,13 +535,18 @@ class MultimediaContentProperties:
         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._properties: typing.Set[str] = set(
+            [
+                "id",
+                "name",
+                "type",
+                "description",
+            ]
+        )
         self._props = NodeProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -524,7 +561,6 @@ 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.")
@@ -534,20 +570,16 @@ class NodeBuilder(NodeAst):
         return [(name, self._bldr.property(name)) for name in self._properties]
 
 
-
-
 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")
@@ -565,13 +597,16 @@ class NodeProperties:
         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._properties: typing.Set[str] = set(
+            [
+                "type",
+                "subclass",
+            ]
+        )
         self._props = ProceduralContentProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -586,10 +621,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class ProceduralContentProperties:
@@ -597,8 +633,6 @@ class ProceduralContentProperties:
         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"))
@@ -608,13 +642,18 @@ class ProceduralContentProperties:
         return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
 
 
-
-
 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._properties: typing.Set[str] = set(
+            [
+                "name",
+                "description",
+                "methods",
+                "decorators",
+            ]
+        )
         self._props = SummarizedClassProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -629,10 +668,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class SummarizedClassProperties:
@@ -640,8 +680,6 @@ class SummarizedClassProperties:
         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"))
@@ -659,13 +697,21 @@ class SummarizedClassProperties:
         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._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:
@@ -680,10 +726,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class SummarizedCodeProperties:
@@ -691,8 +738,6 @@ class SummarizedCodeProperties:
         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"))
@@ -722,13 +767,16 @@ class SummarizedCodeProperties:
         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._properties: typing.Set[str] = set(
+            [
+                "summary",
+                "description",
+            ]
+        )
         self._props = SummarizedContentProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -743,10 +791,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class SummarizedContentProperties:
@@ -754,8 +803,6 @@ class SummarizedContentProperties:
         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"))
@@ -765,13 +812,19 @@ class SummarizedContentProperties:
         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._properties: typing.Set[str] = set(
+            [
+                "name",
+                "description",
+                "inputs",
+                "outputs",
+                "decorators",
+            ]
+        )
         self._props = SummarizedFunctionProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -786,10 +839,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class SummarizedFunctionProperties:
@@ -797,8 +851,6 @@ class SummarizedFunctionProperties:
         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"))
@@ -820,13 +872,16 @@ class SummarizedFunctionProperties:
         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._properties: typing.Set[str] = set(
+            [
+                "type",
+                "subclass",
+            ]
+        )
         self._props = TextContentProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -841,10 +896,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class TextContentProperties:
@@ -852,8 +908,6 @@ class TextContentProperties:
         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"))
@@ -863,13 +917,16 @@ class TextContentProperties:
         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._properties: typing.Set[str] = set(
+            [
+                "type",
+                "subclass",
+            ]
+        )
         self._props = VideoContentProperties(self._bldr, self._properties)
 
     def type(self) -> baml_py.FieldType:
@@ -884,10 +941,11 @@ 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]
+        return [
-    
+            (name, type_builder.ClassPropertyViewer(self._bldr.property(name)))
+            for name in self._properties
+        ]
 
 
 class VideoContentProperties:
@@ -895,8 +953,6 @@ class VideoContentProperties:
         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"))
@@ -904,6 +960,3 @@ class VideoContentProperties:
     @property
     def subclass(self) -> type_builder.ClassPropertyViewer:
         return type_builder.ClassPropertyViewer(self.__bldr.property("subclass"))
-    
-    
-

cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_map.py

@@ -15,57 +15,38 @@ 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.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,
-
-
 }

cognee/infrastructure/llm/structured_output_framework/baml/baml_client/types.py

@@ -20,22 +20,29 @@ from pydantic import BaseModel, ConfigDict
 
 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):
     name: str
     expression: str
     status: str
+
+
 class Checked(BaseModel, typing.Generic[CheckT, CheckName]):
     value: CheckT
     checks: typing.Dict[CheckName, Check]
 
+
 def get_checks(checks: typing.Dict[CheckName, Check]) -> typing.List[Check]:
     return list(checks.values())
 
+
 def all_succeeded(checks: typing.Dict[CheckName, Check]) -> bool:
     return all(check.status == "succeeded" for check in get_checks(checks))
+
+
 # #########################################################################
 # Generated enums (0)
 # #########################################################################
@@ -44,20 +51,33 @@ def all_succeeded(checks: typing.Dict[CheckName, Check]) -> bool:
 # 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']]
+    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')
+    model_config = ConfigDict(extra="allow")
+
 
 class Edge(BaseModel):
     # doc string for edge
@@ -67,39 +87,47 @@ class Edge(BaseModel):
     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')
+    model_config = ConfigDict(extra="allow")
     id: str
     name: str
     type: str
     description: str
 
+
 class ProceduralContent(BaseModel):
     type: str
     subclass: typing.List[str]
 
+
 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]
@@ -109,10 +137,12 @@ class SummarizedCode(BaseModel):
     functions: typing.List["SummarizedFunction"]
     workflow_description: typing.Optional[str] = None
 
+
 class SummarizedContent(BaseModel):
     summary: str
     description: str
 
+
 class SummarizedFunction(BaseModel):
     name: str
     description: str
@@ -120,14 +150,17 @@ class SummarizedFunction(BaseModel):
     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)
 # #########################################################################

cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_categories.py

@@ -2,13 +2,14 @@ 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
 
+config = get_llm_config()
+
+
 logger = get_logger("extract_summary_baml")
 
 

cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_summary.py

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

cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/knowledge_graph/extract_content_graph.py

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

cognee/tasks/graph/extract_graph_from_data.py

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

pyproject.toml

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