cognee 0.2.1.dev7__py3-none-any.whl → 0.2.2.dev1__py3-none-any.whl

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

@@ -0,0 +1,550 @@
+# ----------------------------------------------------------------------------
+#
+#  Welcome to Baml! To use this generated code, please run the following:
+#
+#  $ pip install baml
+#
+# ----------------------------------------------------------------------------
+
+# This file was generated by BAML: please do not edit it. Instead, edit the
+# BAML files and re-generate this code using: baml-cli generate
+# baml-cli is available with the baml package.
+
+import typing
+import typing_extensions
+import baml_py
+
+from . import stream_types, types, type_builder
+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"
+    __http_request: "BamlHttpRequestClient"
+    __http_stream_request: "BamlHttpStreamRequestClient"
+    __llm_response_parser: LlmResponseParser
+    __llm_stream_parser: LlmStreamParser
+
+    def __init__(self, options: DoNotUseDirectlyCallManager):
+        self.__options = options
+        self.__stream_client = BamlStreamClient(options)
+        self.__http_request = BamlHttpRequestClient(options)
+        self.__http_stream_request = BamlHttpStreamRequestClient(options)
+        self.__llm_response_parser = LlmResponseParser(options)
+        self.__llm_stream_parser = LlmStreamParser(options)
+
+    def __getstate__(self):
+        # Return state needed for pickling
+        return {"options": self.__options}
+
+    def __setstate__(self, state):
+        # Restore state from pickling
+        self.__options = state["options"]
+        self.__stream_client = BamlStreamClient(self.__options)
+        self.__http_request = BamlHttpRequestClient(self.__options)
+        self.__http_stream_request = BamlHttpStreamRequestClient(self.__options)
+        self.__llm_response_parser = LlmResponseParser(self.__options)
+        self.__llm_stream_parser = LlmStreamParser(self.__options)
+
+    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,
+        env: typing.Optional[typing.Dict[str, typing.Optional[str]]] = None,
+    ) -> "BamlSyncClient":
+        options: BamlCallOptions = {}
+        if tb is not None:
+            options["tb"] = tb
+        if client_registry is not None:
+            options["client_registry"] = client_registry
+        if collector is not None:
+            options["collector"] = collector
+        if env is not None:
+            options["env"] = env
+        return BamlSyncClient(self.__options.merge_options(options))
+
+    @property
+    def stream(self):
+        return self.__stream_client
+
+    @property
+    def request(self):
+        return self.__http_request
+
+    @property
+    def stream_request(self):
+        return self.__http_stream_request
+
+    @property
+    def parse(self):
+        return self.__llm_response_parser
+
+    @property
+    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__)
+        )
+
+
+class BamlStreamClient:
+    __options: DoNotUseDirectlyCallManager
+
+    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(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent]:
+        ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
+            function_name="SummarizeContent",
+            args={
+                "content": content,
+            },
+        )
+        return baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent](
+            result,
+            lambda x: typing.cast(
+                stream_types.SummarizedContent,
+                x.cast_to(types, types, stream_types, True, __runtime__),
+            ),
+            lambda x: typing.cast(
+                types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)
+            ),
+            ctx,
+        )
+
+
+class BamlHttpRequestClient:
+    __options: DoNotUseDirectlyCallManager
+
+    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="request",
+        )
+        return result
+
+    def ExtractContentGraphGeneric(
+        self,
+        content: str,
+        mode: typing.Optional[
+            typing.Union[
+                typing_extensions.Literal["simple"],
+                typing_extensions.Literal["base"],
+                typing_extensions.Literal["guided"],
+                typing_extensions.Literal["strict"],
+                typing_extensions.Literal["custom"],
+            ]
+        ] = None,
+        custom_prompt_content: typing.Optional[str] = None,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = self.__options.merge_options(baml_options).create_http_request_sync(
+            function_name="ExtractContentGraphGeneric",
+            args={
+                "content": content,
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+            mode="request",
+        )
+        return result
+
+    def ExtractDynamicContentGraph(
+        self,
+        content: str,
+        mode: typing.Optional[
+            typing.Union[
+                typing_extensions.Literal["simple"],
+                typing_extensions.Literal["base"],
+                typing_extensions.Literal["guided"],
+                typing_extensions.Literal["strict"],
+                typing_extensions.Literal["custom"],
+            ]
+        ] = None,
+        custom_prompt_content: typing.Optional[str] = None,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = self.__options.merge_options(baml_options).create_http_request_sync(
+            function_name="ExtractDynamicContentGraph",
+            args={
+                "content": content,
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+            mode="request",
+        )
+        return result
+
+    def SummarizeCode(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = self.__options.merge_options(baml_options).create_http_request_sync(
+            function_name="SummarizeCode",
+            args={
+                "content": content,
+            },
+            mode="request",
+        )
+        return result
+
+    def SummarizeContent(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = self.__options.merge_options(baml_options).create_http_request_sync(
+            function_name="SummarizeContent",
+            args={
+                "content": content,
+            },
+            mode="request",
+        )
+        return result
+
+
+class BamlHttpStreamRequestClient:
+    __options: DoNotUseDirectlyCallManager
+
+    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(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = self.__options.merge_options(baml_options).create_http_request_sync(
+            function_name="SummarizeContent",
+            args={
+                "content": content,
+            },
+            mode="stream",
+        )
+        return result
+
+
+b = BamlSyncClient(DoNotUseDirectlyCallManager({}))

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

@@ -0,0 +1,26 @@
+# ----------------------------------------------------------------------------
+#
+#  Welcome to Baml! To use this generated code, please run the following:
+#
+#  $ pip install baml
+#
+# ----------------------------------------------------------------------------
+
+# This file was generated by BAML: please do not edit it. Instead, edit the
+# BAML files and re-generate this code using: baml-cli generate
+# baml-cli is available with the baml package.
+
+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"]