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

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

@@ -57,6 +57,7 @@ class BamlSyncClient:
             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,
+        on_tick: typing.Optional[typing.Callable[[str, baml_py.baml_py.FunctionLog], None]] = None,
     ) -> "BamlSyncClient":
         options: BamlCallOptions = {}
         if tb is not None:
@@ -67,6 +68,8 @@ class BamlSyncClient:
             options["collector"] = collector
         if env is not None:
             options["env"] = env
+        if on_tick is not None:
+            options["on_tick"] = on_tick
         return BamlSyncClient(self.__options.merge_options(options))
 
     @property
@@ -89,106 +92,30 @@ class BamlSyncClient:
     def parse_stream(self):
         return self.__llm_stream_parser
 
-    def ExtractCategories(
-        self,
-        content: str,
-        baml_options: BamlCallOptions = {},
-    ) -> types.DefaultContentPrediction:
-        result = self.__options.merge_options(baml_options).call_function_sync(
-            function_name="ExtractCategories",
-            args={
-                "content": content,
-            },
-        )
-        return typing.cast(
-            types.DefaultContentPrediction,
-            result.cast_to(types, types, stream_types, False, __runtime__),
-        )
-
-    def ExtractContentGraphGeneric(
-        self,
-        content: str,
-        mode: typing.Optional[
-            typing.Union[
-                typing_extensions.Literal["simple"],
-                typing_extensions.Literal["base"],
-                typing_extensions.Literal["guided"],
-                typing_extensions.Literal["strict"],
-                typing_extensions.Literal["custom"],
-            ]
-        ] = None,
-        custom_prompt_content: typing.Optional[str] = None,
-        baml_options: BamlCallOptions = {},
-    ) -> types.KnowledgeGraph:
-        result = self.__options.merge_options(baml_options).call_function_sync(
-            function_name="ExtractContentGraphGeneric",
-            args={
-                "content": content,
-                "mode": mode,
-                "custom_prompt_content": custom_prompt_content,
-            },
-        )
-        return typing.cast(
-            types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)
-        )
-
-    def ExtractDynamicContentGraph(
-        self,
-        content: str,
-        mode: typing.Optional[
-            typing.Union[
-                typing_extensions.Literal["simple"],
-                typing_extensions.Literal["base"],
-                typing_extensions.Literal["guided"],
-                typing_extensions.Literal["strict"],
-                typing_extensions.Literal["custom"],
-            ]
-        ] = None,
-        custom_prompt_content: typing.Optional[str] = None,
-        baml_options: BamlCallOptions = {},
-    ) -> types.DynamicKnowledgeGraph:
-        result = self.__options.merge_options(baml_options).call_function_sync(
-            function_name="ExtractDynamicContentGraph",
-            args={
-                "content": content,
-                "mode": mode,
-                "custom_prompt_content": custom_prompt_content,
-            },
-        )
-        return typing.cast(
-            types.DynamicKnowledgeGraph,
-            result.cast_to(types, types, stream_types, False, __runtime__),
-        )
-
-    def SummarizeCode(
-        self,
-        content: str,
-        baml_options: BamlCallOptions = {},
-    ) -> types.SummarizedCode:
-        result = self.__options.merge_options(baml_options).call_function_sync(
-            function_name="SummarizeCode",
-            args={
-                "content": content,
-            },
-        )
-        return typing.cast(
-            types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__)
-        )
-
-    def SummarizeContent(
-        self,
-        content: str,
-        baml_options: BamlCallOptions = {},
-    ) -> types.SummarizedContent:
-        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__)
-        )
+    def AcreateStructuredOutput(
+        self,
+        text_input: str,
+        system_prompt: str,
+        baml_options: BamlCallOptions = {},
+    ) -> types.ResponseModel:
+        # Check if on_tick is provided
+        if "on_tick" in baml_options:
+            stream = self.stream.AcreateStructuredOutput(
+                text_input=text_input, system_prompt=system_prompt, baml_options=baml_options
+            )
+            return stream.get_final_response()
+        else:
+            # Original non-streaming code
+            result = self.__options.merge_options(baml_options).call_function_sync(
+                function_name="AcreateStructuredOutput",
+                args={
+                    "text_input": text_input,
+                    "system_prompt": system_prompt,
+                },
+            )
+            return typing.cast(
+                types.ResponseModel, result.cast_to(types, types, stream_types, False, __runtime__)
+            )
 
 
 class BamlStreamClient:
@@ -197,149 +124,26 @@ class BamlStreamClient:
     def __init__(self, options: DoNotUseDirectlyCallManager):
         self.__options = options
 
-    def ExtractCategories(
-        self,
-        content: str,
-        baml_options: BamlCallOptions = {},
-    ) -> baml_py.BamlSyncStream[
-        stream_types.DefaultContentPrediction, types.DefaultContentPrediction
-    ]:
-        ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
-            function_name="ExtractCategories",
-            args={
-                "content": content,
-            },
-        )
-        return baml_py.BamlSyncStream[
-            stream_types.DefaultContentPrediction, types.DefaultContentPrediction
-        ](
-            result,
-            lambda x: typing.cast(
-                stream_types.DefaultContentPrediction,
-                x.cast_to(types, types, stream_types, True, __runtime__),
-            ),
-            lambda x: typing.cast(
-                types.DefaultContentPrediction,
-                x.cast_to(types, types, stream_types, False, __runtime__),
-            ),
-            ctx,
-        )
-
-    def ExtractContentGraphGeneric(
-        self,
-        content: str,
-        mode: typing.Optional[
-            typing.Union[
-                typing_extensions.Literal["simple"],
-                typing_extensions.Literal["base"],
-                typing_extensions.Literal["guided"],
-                typing_extensions.Literal["strict"],
-                typing_extensions.Literal["custom"],
-            ]
-        ] = None,
-        custom_prompt_content: typing.Optional[str] = None,
-        baml_options: BamlCallOptions = {},
-    ) -> baml_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]:
-        ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
-            function_name="ExtractContentGraphGeneric",
-            args={
-                "content": content,
-                "mode": mode,
-                "custom_prompt_content": custom_prompt_content,
-            },
-        )
-        return baml_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph](
-            result,
-            lambda x: typing.cast(
-                stream_types.KnowledgeGraph,
-                x.cast_to(types, types, stream_types, True, __runtime__),
-            ),
-            lambda x: typing.cast(
-                types.KnowledgeGraph, x.cast_to(types, types, stream_types, False, __runtime__)
-            ),
-            ctx,
-        )
-
-    def ExtractDynamicContentGraph(
-        self,
-        content: str,
-        mode: typing.Optional[
-            typing.Union[
-                typing_extensions.Literal["simple"],
-                typing_extensions.Literal["base"],
-                typing_extensions.Literal["guided"],
-                typing_extensions.Literal["strict"],
-                typing_extensions.Literal["custom"],
-            ]
-        ] = None,
-        custom_prompt_content: typing.Optional[str] = None,
-        baml_options: BamlCallOptions = {},
-    ) -> baml_py.BamlSyncStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph]:
-        ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
-            function_name="ExtractDynamicContentGraph",
-            args={
-                "content": content,
-                "mode": mode,
-                "custom_prompt_content": custom_prompt_content,
-            },
-        )
-        return baml_py.BamlSyncStream[
-            stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph
-        ](
-            result,
-            lambda x: typing.cast(
-                stream_types.DynamicKnowledgeGraph,
-                x.cast_to(types, types, stream_types, True, __runtime__),
-            ),
-            lambda x: typing.cast(
-                types.DynamicKnowledgeGraph,
-                x.cast_to(types, types, stream_types, False, __runtime__),
-            ),
-            ctx,
-        )
-
-    def SummarizeCode(
-        self,
-        content: str,
-        baml_options: BamlCallOptions = {},
-    ) -> baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode]:
-        ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
-            function_name="SummarizeCode",
-            args={
-                "content": content,
-            },
-        )
-        return baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode](
-            result,
-            lambda x: typing.cast(
-                stream_types.SummarizedCode,
-                x.cast_to(types, types, stream_types, True, __runtime__),
-            ),
-            lambda x: typing.cast(
-                types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)
-            ),
-            ctx,
-        )
-
-    def SummarizeContent(
+    def AcreateStructuredOutput(
         self,
-        content: str,
+        text_input: str,
+        system_prompt: str,
         baml_options: BamlCallOptions = {},
-    ) -> baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent]:
+    ) -> baml_py.BamlSyncStream[stream_types.ResponseModel, types.ResponseModel]:
         ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
-            function_name="SummarizeContent",
+            function_name="AcreateStructuredOutput",
             args={
-                "content": content,
+                "text_input": text_input,
+                "system_prompt": system_prompt,
             },
         )
-        return baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent](
+        return baml_py.BamlSyncStream[stream_types.ResponseModel, types.ResponseModel](
             result,
             lambda x: typing.cast(
-                stream_types.SummarizedContent,
-                x.cast_to(types, types, stream_types, True, __runtime__),
+                stream_types.ResponseModel, x.cast_to(types, types, stream_types, True, __runtime__)
             ),
             lambda x: typing.cast(
-                types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)
+                types.ResponseModel, x.cast_to(types, types, stream_types, False, __runtime__)
             ),
             ctx,
         )
@@ -351,95 +155,17 @@ class BamlHttpRequestClient:
     def __init__(self, options: DoNotUseDirectlyCallManager):
         self.__options = options
 
-    def ExtractCategories(
+    def AcreateStructuredOutput(
         self,
-        content: str,
+        text_input: str,
+        system_prompt: str,
         baml_options: BamlCallOptions = {},
     ) -> baml_py.baml_py.HTTPRequest:
         result = self.__options.merge_options(baml_options).create_http_request_sync(
-            function_name="ExtractCategories",
+            function_name="AcreateStructuredOutput",
             args={
-                "content": content,
-            },
-            mode="request",
-        )
-        return result
-
-    def ExtractContentGraphGeneric(
-        self,
-        content: str,
-        mode: typing.Optional[
-            typing.Union[
-                typing_extensions.Literal["simple"],
-                typing_extensions.Literal["base"],
-                typing_extensions.Literal["guided"],
-                typing_extensions.Literal["strict"],
-                typing_extensions.Literal["custom"],
-            ]
-        ] = None,
-        custom_prompt_content: typing.Optional[str] = None,
-        baml_options: BamlCallOptions = {},
-    ) -> baml_py.baml_py.HTTPRequest:
-        result = self.__options.merge_options(baml_options).create_http_request_sync(
-            function_name="ExtractContentGraphGeneric",
-            args={
-                "content": content,
-                "mode": mode,
-                "custom_prompt_content": custom_prompt_content,
-            },
-            mode="request",
-        )
-        return result
-
-    def ExtractDynamicContentGraph(
-        self,
-        content: str,
-        mode: typing.Optional[
-            typing.Union[
-                typing_extensions.Literal["simple"],
-                typing_extensions.Literal["base"],
-                typing_extensions.Literal["guided"],
-                typing_extensions.Literal["strict"],
-                typing_extensions.Literal["custom"],
-            ]
-        ] = None,
-        custom_prompt_content: typing.Optional[str] = None,
-        baml_options: BamlCallOptions = {},
-    ) -> baml_py.baml_py.HTTPRequest:
-        result = self.__options.merge_options(baml_options).create_http_request_sync(
-            function_name="ExtractDynamicContentGraph",
-            args={
-                "content": content,
-                "mode": mode,
-                "custom_prompt_content": custom_prompt_content,
-            },
-            mode="request",
-        )
-        return result
-
-    def SummarizeCode(
-        self,
-        content: str,
-        baml_options: BamlCallOptions = {},
-    ) -> baml_py.baml_py.HTTPRequest:
-        result = self.__options.merge_options(baml_options).create_http_request_sync(
-            function_name="SummarizeCode",
-            args={
-                "content": content,
-            },
-            mode="request",
-        )
-        return result
-
-    def SummarizeContent(
-        self,
-        content: str,
-        baml_options: BamlCallOptions = {},
-    ) -> baml_py.baml_py.HTTPRequest:
-        result = self.__options.merge_options(baml_options).create_http_request_sync(
-            function_name="SummarizeContent",
-            args={
-                "content": content,
+                "text_input": text_input,
+                "system_prompt": system_prompt,
             },
             mode="request",
         )
@@ -452,95 +178,17 @@ class BamlHttpStreamRequestClient:
     def __init__(self, options: DoNotUseDirectlyCallManager):
         self.__options = options
 
-    def ExtractCategories(
-        self,
-        content: str,
-        baml_options: BamlCallOptions = {},
-    ) -> baml_py.baml_py.HTTPRequest:
-        result = self.__options.merge_options(baml_options).create_http_request_sync(
-            function_name="ExtractCategories",
-            args={
-                "content": content,
-            },
-            mode="stream",
-        )
-        return result
-
-    def ExtractContentGraphGeneric(
-        self,
-        content: str,
-        mode: typing.Optional[
-            typing.Union[
-                typing_extensions.Literal["simple"],
-                typing_extensions.Literal["base"],
-                typing_extensions.Literal["guided"],
-                typing_extensions.Literal["strict"],
-                typing_extensions.Literal["custom"],
-            ]
-        ] = None,
-        custom_prompt_content: typing.Optional[str] = None,
-        baml_options: BamlCallOptions = {},
-    ) -> baml_py.baml_py.HTTPRequest:
-        result = self.__options.merge_options(baml_options).create_http_request_sync(
-            function_name="ExtractContentGraphGeneric",
-            args={
-                "content": content,
-                "mode": mode,
-                "custom_prompt_content": custom_prompt_content,
-            },
-            mode="stream",
-        )
-        return result
-
-    def ExtractDynamicContentGraph(
-        self,
-        content: str,
-        mode: typing.Optional[
-            typing.Union[
-                typing_extensions.Literal["simple"],
-                typing_extensions.Literal["base"],
-                typing_extensions.Literal["guided"],
-                typing_extensions.Literal["strict"],
-                typing_extensions.Literal["custom"],
-            ]
-        ] = None,
-        custom_prompt_content: typing.Optional[str] = None,
-        baml_options: BamlCallOptions = {},
-    ) -> baml_py.baml_py.HTTPRequest:
-        result = self.__options.merge_options(baml_options).create_http_request_sync(
-            function_name="ExtractDynamicContentGraph",
-            args={
-                "content": content,
-                "mode": mode,
-                "custom_prompt_content": custom_prompt_content,
-            },
-            mode="stream",
-        )
-        return result
-
-    def SummarizeCode(
-        self,
-        content: str,
-        baml_options: BamlCallOptions = {},
-    ) -> baml_py.baml_py.HTTPRequest:
-        result = self.__options.merge_options(baml_options).create_http_request_sync(
-            function_name="SummarizeCode",
-            args={
-                "content": content,
-            },
-            mode="stream",
-        )
-        return result
-
-    def SummarizeContent(
+    def AcreateStructuredOutput(
         self,
-        content: str,
+        text_input: str,
+        system_prompt: str,
         baml_options: BamlCallOptions = {},
     ) -> baml_py.baml_py.HTTPRequest:
         result = self.__options.merge_options(baml_options).create_http_request_sync(
-            function_name="SummarizeContent",
+            function_name="AcreateStructuredOutput",
             args={
-                "content": content,
+                "text_input": text_input,
+                "system_prompt": system_prompt,
             },
             mode="stream",
         )