cognee 0.3.4.dev4__py3-none-any.whl → 0.3.5__py3-none-any.whl

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

@@ -23,55 +23,15 @@ class LlmResponseParser:
     def __init__(self, options: DoNotUseDirectlyCallManager):
         self.__options = options
 
-    def ExtractCategories(
+    def AcreateStructuredOutput(
         self,
         llm_response: str,
         baml_options: BamlCallOptions = {},
-    ) -> types.DefaultContentPrediction:
+    ) -> types.ResponseModel:
         result = self.__options.merge_options(baml_options).parse_response(
-            function_name="ExtractCategories", llm_response=llm_response, mode="request"
+            function_name="AcreateStructuredOutput", llm_response=llm_response, mode="request"
         )
-        return typing.cast(types.DefaultContentPrediction, result)
-
-    def ExtractContentGraphGeneric(
-        self,
-        llm_response: str,
-        baml_options: BamlCallOptions = {},
-    ) -> types.KnowledgeGraph:
-        result = self.__options.merge_options(baml_options).parse_response(
-            function_name="ExtractContentGraphGeneric", llm_response=llm_response, mode="request"
-        )
-        return typing.cast(types.KnowledgeGraph, result)
-
-    def ExtractDynamicContentGraph(
-        self,
-        llm_response: str,
-        baml_options: BamlCallOptions = {},
-    ) -> types.DynamicKnowledgeGraph:
-        result = self.__options.merge_options(baml_options).parse_response(
-            function_name="ExtractDynamicContentGraph", llm_response=llm_response, mode="request"
-        )
-        return typing.cast(types.DynamicKnowledgeGraph, result)
-
-    def SummarizeCode(
-        self,
-        llm_response: str,
-        baml_options: BamlCallOptions = {},
-    ) -> types.SummarizedCode:
-        result = self.__options.merge_options(baml_options).parse_response(
-            function_name="SummarizeCode", llm_response=llm_response, mode="request"
-        )
-        return typing.cast(types.SummarizedCode, result)
-
-    def SummarizeContent(
-        self,
-        llm_response: str,
-        baml_options: BamlCallOptions = {},
-    ) -> types.SummarizedContent:
-        result = self.__options.merge_options(baml_options).parse_response(
-            function_name="SummarizeContent", llm_response=llm_response, mode="request"
-        )
-        return typing.cast(types.SummarizedContent, result)
+        return typing.cast(types.ResponseModel, result)
 
 
 class LlmStreamParser:
@@ -80,52 +40,12 @@ class LlmStreamParser:
     def __init__(self, options: DoNotUseDirectlyCallManager):
         self.__options = options
 
-    def ExtractCategories(
-        self,
-        llm_response: str,
-        baml_options: BamlCallOptions = {},
-    ) -> stream_types.DefaultContentPrediction:
-        result = self.__options.merge_options(baml_options).parse_response(
-            function_name="ExtractCategories", llm_response=llm_response, mode="stream"
-        )
-        return typing.cast(stream_types.DefaultContentPrediction, result)
-
-    def ExtractContentGraphGeneric(
-        self,
-        llm_response: str,
-        baml_options: BamlCallOptions = {},
-    ) -> stream_types.KnowledgeGraph:
-        result = self.__options.merge_options(baml_options).parse_response(
-            function_name="ExtractContentGraphGeneric", llm_response=llm_response, mode="stream"
-        )
-        return typing.cast(stream_types.KnowledgeGraph, result)
-
-    def ExtractDynamicContentGraph(
-        self,
-        llm_response: str,
-        baml_options: BamlCallOptions = {},
-    ) -> stream_types.DynamicKnowledgeGraph:
-        result = self.__options.merge_options(baml_options).parse_response(
-            function_name="ExtractDynamicContentGraph", llm_response=llm_response, mode="stream"
-        )
-        return typing.cast(stream_types.DynamicKnowledgeGraph, result)
-
-    def SummarizeCode(
-        self,
-        llm_response: str,
-        baml_options: BamlCallOptions = {},
-    ) -> stream_types.SummarizedCode:
-        result = self.__options.merge_options(baml_options).parse_response(
-            function_name="SummarizeCode", llm_response=llm_response, mode="stream"
-        )
-        return typing.cast(stream_types.SummarizedCode, result)
-
-    def SummarizeContent(
+    def AcreateStructuredOutput(
         self,
         llm_response: str,
         baml_options: BamlCallOptions = {},
-    ) -> stream_types.SummarizedContent:
+    ) -> stream_types.ResponseModel:
         result = self.__options.merge_options(baml_options).parse_response(
-            function_name="SummarizeContent", llm_response=llm_response, mode="stream"
+            function_name="AcreateStructuredOutput", llm_response=llm_response, mode="stream"
         )
-        return typing.cast(stream_types.SummarizedContent, result)
+        return typing.cast(stream_types.ResponseModel, result)

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

@@ -30,6 +30,10 @@ class BamlCallOptions(typing.TypedDict, total=False):
     collector: typing_extensions.NotRequired[
         typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]
     ]
+    abort_controller: typing_extensions.NotRequired[baml_py.baml_py.AbortController]
+    on_tick: typing_extensions.NotRequired[
+        typing.Callable[[str, baml_py.baml_py.FunctionLog], None]
+    ]
 
 
 class _ResolvedBamlOptions:
@@ -37,6 +41,8 @@ class _ResolvedBamlOptions:
     client_registry: typing.Optional[baml_py.baml_py.ClientRegistry]
     collectors: typing.List[baml_py.baml_py.Collector]
     env_vars: typing.Dict[str, str]
+    abort_controller: typing.Optional[baml_py.baml_py.AbortController]
+    on_tick: typing.Optional[typing.Callable[[], None]]
 
     def __init__(
         self,
@@ -44,11 +50,15 @@ class _ResolvedBamlOptions:
         client_registry: typing.Optional[baml_py.baml_py.ClientRegistry],
         collectors: typing.List[baml_py.baml_py.Collector],
         env_vars: typing.Dict[str, str],
+        abort_controller: typing.Optional[baml_py.baml_py.AbortController],
+        on_tick: typing.Optional[typing.Callable[[], None]],
     ):
         self.tb = tb
         self.client_registry = client_registry
         self.collectors = collectors
         self.env_vars = env_vars
+        self.abort_controller = abort_controller
+        self.on_tick = on_tick
 
 
 class DoNotUseDirectlyCallManager:
@@ -85,11 +95,27 @@ class DoNotUseDirectlyCallManager:
             else:
                 env_vars.pop(k, None)
 
+        abort_controller = self.__baml_options.get("abort_controller")
+
+        on_tick = self.__baml_options.get("on_tick")
+        if on_tick is not None:
+            collector = baml_py.baml_py.Collector("on-tick-collector")
+            collectors_as_list.append(collector)
+
+            def on_tick_wrapper():
+                log = collector.last
+                if log is not None:
+                    on_tick("Unknown", log)
+        else:
+            on_tick_wrapper = None
+
         return _ResolvedBamlOptions(
             baml_tb,
             client_registry,
             collectors_as_list,
             env_vars,
+            abort_controller,
+            on_tick_wrapper,
         )
 
     def merge_options(self, options: BamlCallOptions) -> "DoNotUseDirectlyCallManager":
@@ -99,6 +125,14 @@ class DoNotUseDirectlyCallManager:
         self, *, function_name: str, args: typing.Dict[str, typing.Any]
     ) -> baml_py.baml_py.FunctionResult:
         resolved_options = self.__resolve()
+
+        # Check if already aborted
+        if (
+            resolved_options.abort_controller is not None
+            and resolved_options.abort_controller.aborted
+        ):
+            raise Exception("BamlAbortError: Operation was aborted")
+
         return await __runtime__.call_function(
             function_name,
             args,
@@ -112,12 +146,22 @@ class DoNotUseDirectlyCallManager:
             resolved_options.collectors,
             # env_vars
             resolved_options.env_vars,
+            # abort_controller
+            resolved_options.abort_controller,
         )
 
     def call_function_sync(
         self, *, function_name: str, args: typing.Dict[str, typing.Any]
     ) -> baml_py.baml_py.FunctionResult:
         resolved_options = self.__resolve()
+
+        # Check if already aborted
+        if (
+            resolved_options.abort_controller is not None
+            and resolved_options.abort_controller.aborted
+        ):
+            raise Exception("BamlAbortError: Operation was aborted")
+
         ctx = __ctx__manager__.get()
         return __runtime__.call_function_sync(
             function_name,
@@ -132,6 +176,8 @@ class DoNotUseDirectlyCallManager:
             resolved_options.collectors,
             # env_vars
             resolved_options.env_vars,
+            # abort_controller
+            resolved_options.abort_controller,
         )
 
     def create_async_stream(
@@ -158,6 +204,8 @@ class DoNotUseDirectlyCallManager:
             resolved_options.collectors,
             # env_vars
             resolved_options.env_vars,
+            # on_tick
+            resolved_options.on_tick,
         )
         return ctx, result
 
@@ -170,6 +218,10 @@ class DoNotUseDirectlyCallManager:
         baml_py.baml_py.RuntimeContextManager, baml_py.baml_py.SyncFunctionResultStream
     ]:
         resolved_options = self.__resolve()
+        if resolved_options.on_tick is not None:
+            raise ValueError(
+                "on_tick is not supported for sync streams. Please use async streams instead."
+            )
         ctx = __ctx__manager__.get()
         result = __runtime__.stream_function_sync(
             function_name,
@@ -187,6 +239,9 @@ class DoNotUseDirectlyCallManager:
             resolved_options.collectors,
             # env_vars
             resolved_options.env_vars,
+            # on_tick
+            # always None! sync streams don't support on_tick
+            None,
         )
         return ctx, result
 
@@ -264,3 +319,26 @@ class DoNotUseDirectlyCallManager:
             # env_vars
             resolved_options.env_vars,
         )
+
+
+def disassemble(function: typing.Callable) -> None:
+    import inspect
+    from . import b
+
+    if not callable(function):
+        print(f"disassemble: object {function} is not a Baml function")
+        return
+
+    is_client_method = False
+
+    for method_name, _ in inspect.getmembers(b, predicate=inspect.ismethod):
+        if method_name == function.__name__:
+            is_client_method = True
+            break
+
+    if not is_client_method:
+        print(f"disassemble: function {function.__name__} is not a Baml function")
+        return
+
+    print(f"----- function {function.__name__} -----")
+    __runtime__.disassemble(function.__name__)

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

@@ -27,109 +27,12 @@ class StreamState(BaseModel, typing.Generic[StreamStateValueT]):
 
 
 # #########################################################################
-# Generated classes (17)
+# Generated classes (1)
 # #########################################################################
 
 
-class AudioContent(BaseModel):
-    type: typing.Optional[str] = None
-    subclass: typing.List[str]
-
-
-class ContentLabel(BaseModel):
-    content_type: typing.Optional[typing.Union[str, str, str, str, str, str, str]] = None
-    type: typing.Optional[str] = None
-    subclass: typing.List[str]
-
-
-class DefaultContentPrediction(BaseModel):
-    label: typing.Optional["ContentLabel"] = None
-
-
-class DynamicKnowledgeGraph(BaseModel):
-    model_config = ConfigDict(extra="allow")
-
-
-class Edge(BaseModel):
-    # doc string for edge
-    # doc string for source_node_id
-
-    source_node_id: typing.Optional[str] = None
-    target_node_id: typing.Optional[str] = None
-    relationship_name: typing.Optional[str] = None
-
-
-class ImageContent(BaseModel):
-    type: typing.Optional[str] = None
-    subclass: typing.List[str]
-
-
-class KnowledgeGraph(BaseModel):
-    nodes: typing.List["types.Node"]
-    edges: typing.List["Edge"]
-
-
-class Model3DContent(BaseModel):
-    type: typing.Optional[str] = None
-    subclass: typing.List[str]
-
-
-class MultimediaContent(BaseModel):
-    type: typing.Optional[str] = None
-    subclass: typing.List[str]
-
-
-class Node(BaseModel):
+class ResponseModel(BaseModel):
     model_config = ConfigDict(extra="allow")
-    id: typing.Optional[str] = None
-    name: typing.Optional[str] = None
-    type: typing.Optional[str] = None
-    description: typing.Optional[str] = None
-
-
-class ProceduralContent(BaseModel):
-    type: typing.Optional[str] = None
-    subclass: typing.List[str]
-
-
-class SummarizedClass(BaseModel):
-    name: typing.Optional[str] = None
-    description: typing.Optional[str] = None
-    methods: typing.Optional[typing.List["SummarizedFunction"]] = None
-    decorators: typing.Optional[typing.List[str]] = None
-
-
-class SummarizedCode(BaseModel):
-    high_level_summary: typing.Optional[str] = None
-    key_features: typing.List[str]
-    imports: typing.List[str]
-    constants: typing.List[str]
-    classes: typing.List["SummarizedClass"]
-    functions: typing.List["SummarizedFunction"]
-    workflow_description: typing.Optional[str] = None
-
-
-class SummarizedContent(BaseModel):
-    summary: typing.Optional[str] = None
-    description: typing.Optional[str] = None
-
-
-class SummarizedFunction(BaseModel):
-    name: typing.Optional[str] = None
-    description: typing.Optional[str] = None
-    inputs: typing.Optional[typing.List[str]] = None
-    outputs: typing.Optional[typing.List[str]] = None
-    decorators: typing.Optional[typing.List[str]] = None
-
-
-class TextContent(BaseModel):
-    type: typing.Optional[str] = None
-    subclass: typing.List[str]
-
-
-class VideoContent(BaseModel):
-    type: typing.Optional[str] = None
-    subclass: typing.List[str]
 
 
 # #########################################################################