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/async_client.py

@@ -0,0 +1,533 @@
+# ----------------------------------------------------------------------------
+#
+#  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 BamlAsyncClient:
+    __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 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,
+    ) -> "BamlAsyncClient":
+        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 BamlAsyncClient(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
+
+    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={
+                "content": content,
+            },
+        )
+        return typing.cast(
+            types.DefaultContentPrediction,
+            result.cast_to(types, types, stream_types, False, __runtime__),
+        )
+
+    async def ExtractContentGraphGeneric(
+        self,
+        content: str,
+        mode: typing.Optional[
+            typing.Union[
+                typing_extensions.Literal["simple"],
+                typing_extensions.Literal["base"],
+                typing_extensions.Literal["guided"],
+                typing_extensions.Literal["strict"],
+                typing_extensions.Literal["custom"],
+            ]
+        ] = None,
+        custom_prompt_content: typing.Optional[str] = None,
+        baml_options: BamlCallOptions = {},
+    ) -> types.KnowledgeGraph:
+        result = await self.__options.merge_options(baml_options).call_function_async(
+            function_name="ExtractContentGraphGeneric",
+            args={
+                "content": content,
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+        )
+        return typing.cast(
+            types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)
+        )
+
+    async def ExtractDynamicContentGraph(
+        self,
+        content: str,
+        mode: typing.Optional[
+            typing.Union[
+                typing_extensions.Literal["simple"],
+                typing_extensions.Literal["base"],
+                typing_extensions.Literal["guided"],
+                typing_extensions.Literal["strict"],
+                typing_extensions.Literal["custom"],
+            ]
+        ] = None,
+        custom_prompt_content: typing.Optional[str] = None,
+        baml_options: BamlCallOptions = {},
+    ) -> types.DynamicKnowledgeGraph:
+        result = await self.__options.merge_options(baml_options).call_function_async(
+            function_name="ExtractDynamicContentGraph",
+            args={
+                "content": content,
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+        )
+        return typing.cast(
+            types.DynamicKnowledgeGraph,
+            result.cast_to(types, types, stream_types, False, __runtime__),
+        )
+
+    async def SummarizeCode(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> types.SummarizedCode:
+        result = await self.__options.merge_options(baml_options).call_function_async(
+            function_name="SummarizeCode",
+            args={
+                "content": content,
+            },
+        )
+        return typing.cast(
+            types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__)
+        )
+
+    async def SummarizeContent(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> types.SummarizedContent:
+        result = await self.__options.merge_options(baml_options).call_function_async(
+            function_name="SummarizeContent",
+            args={
+                "content": content,
+            },
+        )
+        return typing.cast(
+            types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__)
+        )
+
+
+class BamlStreamClient:
+    __options: DoNotUseDirectlyCallManager
+
+    def __init__(self, options: DoNotUseDirectlyCallManager):
+        self.__options = options
+
+    def ExtractCategories(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.BamlStream[stream_types.DefaultContentPrediction, types.DefaultContentPrediction]:
+        ctx, result = self.__options.merge_options(baml_options).create_async_stream(
+            function_name="ExtractCategories",
+            args={
+                "content": content,
+            },
+        )
+        return baml_py.BamlStream[
+            stream_types.DefaultContentPrediction, types.DefaultContentPrediction
+        ](
+            result,
+            lambda x: typing.cast(
+                stream_types.DefaultContentPrediction,
+                x.cast_to(types, types, stream_types, True, __runtime__),
+            ),
+            lambda x: typing.cast(
+                types.DefaultContentPrediction,
+                x.cast_to(types, types, stream_types, False, __runtime__),
+            ),
+            ctx,
+        )
+
+    def ExtractContentGraphGeneric(
+        self,
+        content: str,
+        mode: typing.Optional[
+            typing.Union[
+                typing_extensions.Literal["simple"],
+                typing_extensions.Literal["base"],
+                typing_extensions.Literal["guided"],
+                typing_extensions.Literal["strict"],
+                typing_extensions.Literal["custom"],
+            ]
+        ] = None,
+        custom_prompt_content: typing.Optional[str] = None,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.BamlStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]:
+        ctx, result = self.__options.merge_options(baml_options).create_async_stream(
+            function_name="ExtractContentGraphGeneric",
+            args={
+                "content": content,
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+        )
+        return baml_py.BamlStream[stream_types.KnowledgeGraph, types.KnowledgeGraph](
+            result,
+            lambda x: typing.cast(
+                stream_types.KnowledgeGraph,
+                x.cast_to(types, types, stream_types, True, __runtime__),
+            ),
+            lambda x: typing.cast(
+                types.KnowledgeGraph, x.cast_to(types, types, stream_types, False, __runtime__)
+            ),
+            ctx,
+        )
+
+    def ExtractDynamicContentGraph(
+        self,
+        content: str,
+        mode: typing.Optional[
+            typing.Union[
+                typing_extensions.Literal["simple"],
+                typing_extensions.Literal["base"],
+                typing_extensions.Literal["guided"],
+                typing_extensions.Literal["strict"],
+                typing_extensions.Literal["custom"],
+            ]
+        ] = None,
+        custom_prompt_content: typing.Optional[str] = None,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.BamlStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph]:
+        ctx, result = self.__options.merge_options(baml_options).create_async_stream(
+            function_name="ExtractDynamicContentGraph",
+            args={
+                "content": content,
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+        )
+        return baml_py.BamlStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph](
+            result,
+            lambda x: typing.cast(
+                stream_types.DynamicKnowledgeGraph,
+                x.cast_to(types, types, stream_types, True, __runtime__),
+            ),
+            lambda x: typing.cast(
+                types.DynamicKnowledgeGraph,
+                x.cast_to(types, types, stream_types, False, __runtime__),
+            ),
+            ctx,
+        )
+
+    def SummarizeCode(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.BamlStream[stream_types.SummarizedCode, types.SummarizedCode]:
+        ctx, result = self.__options.merge_options(baml_options).create_async_stream(
+            function_name="SummarizeCode",
+            args={
+                "content": content,
+            },
+        )
+        return baml_py.BamlStream[stream_types.SummarizedCode, types.SummarizedCode](
+            result,
+            lambda x: typing.cast(
+                stream_types.SummarizedCode,
+                x.cast_to(types, types, stream_types, True, __runtime__),
+            ),
+            lambda x: typing.cast(
+                types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)
+            ),
+            ctx,
+        )
+
+    def SummarizeContent(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.BamlStream[stream_types.SummarizedContent, types.SummarizedContent]:
+        ctx, result = self.__options.merge_options(baml_options).create_async_stream(
+            function_name="SummarizeContent",
+            args={
+                "content": content,
+            },
+        )
+        return baml_py.BamlStream[stream_types.SummarizedContent, types.SummarizedContent](
+            result,
+            lambda x: typing.cast(
+                stream_types.SummarizedContent,
+                x.cast_to(types, types, stream_types, True, __runtime__),
+            ),
+            lambda x: typing.cast(
+                types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)
+            ),
+            ctx,
+        )
+
+
+class BamlHttpRequestClient:
+    __options: DoNotUseDirectlyCallManager
+
+    def __init__(self, options: DoNotUseDirectlyCallManager):
+        self.__options = options
+
+    async def ExtractCategories(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="ExtractCategories",
+            args={
+                "content": content,
+            },
+            mode="request",
+        )
+        return result
+
+    async def ExtractContentGraphGeneric(
+        self,
+        content: str,
+        mode: typing.Optional[
+            typing.Union[
+                typing_extensions.Literal["simple"],
+                typing_extensions.Literal["base"],
+                typing_extensions.Literal["guided"],
+                typing_extensions.Literal["strict"],
+                typing_extensions.Literal["custom"],
+            ]
+        ] = None,
+        custom_prompt_content: typing.Optional[str] = None,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="ExtractContentGraphGeneric",
+            args={
+                "content": content,
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+            mode="request",
+        )
+        return result
+
+    async def ExtractDynamicContentGraph(
+        self,
+        content: str,
+        mode: typing.Optional[
+            typing.Union[
+                typing_extensions.Literal["simple"],
+                typing_extensions.Literal["base"],
+                typing_extensions.Literal["guided"],
+                typing_extensions.Literal["strict"],
+                typing_extensions.Literal["custom"],
+            ]
+        ] = None,
+        custom_prompt_content: typing.Optional[str] = None,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="ExtractDynamicContentGraph",
+            args={
+                "content": content,
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+            mode="request",
+        )
+        return result
+
+    async def SummarizeCode(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="SummarizeCode",
+            args={
+                "content": content,
+            },
+            mode="request",
+        )
+        return result
+
+    async def SummarizeContent(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="SummarizeContent",
+            args={
+                "content": content,
+            },
+            mode="request",
+        )
+        return result
+
+
+class BamlHttpStreamRequestClient:
+    __options: DoNotUseDirectlyCallManager
+
+    def __init__(self, options: DoNotUseDirectlyCallManager):
+        self.__options = options
+
+    async def ExtractCategories(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="ExtractCategories",
+            args={
+                "content": content,
+            },
+            mode="stream",
+        )
+        return result
+
+    async def ExtractContentGraphGeneric(
+        self,
+        content: str,
+        mode: typing.Optional[
+            typing.Union[
+                typing_extensions.Literal["simple"],
+                typing_extensions.Literal["base"],
+                typing_extensions.Literal["guided"],
+                typing_extensions.Literal["strict"],
+                typing_extensions.Literal["custom"],
+            ]
+        ] = None,
+        custom_prompt_content: typing.Optional[str] = None,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="ExtractContentGraphGeneric",
+            args={
+                "content": content,
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+            mode="stream",
+        )
+        return result
+
+    async def ExtractDynamicContentGraph(
+        self,
+        content: str,
+        mode: typing.Optional[
+            typing.Union[
+                typing_extensions.Literal["simple"],
+                typing_extensions.Literal["base"],
+                typing_extensions.Literal["guided"],
+                typing_extensions.Literal["strict"],
+                typing_extensions.Literal["custom"],
+            ]
+        ] = None,
+        custom_prompt_content: typing.Optional[str] = None,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="ExtractDynamicContentGraph",
+            args={
+                "content": content,
+                "mode": mode,
+                "custom_prompt_content": custom_prompt_content,
+            },
+            mode="stream",
+        )
+        return result
+
+    async def SummarizeCode(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="SummarizeCode",
+            args={
+                "content": content,
+            },
+            mode="stream",
+        )
+        return result
+
+    async def SummarizeContent(
+        self,
+        content: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.baml_py.HTTPRequest:
+        result = await self.__options.merge_options(baml_options).create_http_request_async(
+            function_name="SummarizeContent",
+            args={
+                "content": content,
+            },
+            mode="stream",
+        )
+        return result
+
+
+b = BamlAsyncClient(DoNotUseDirectlyCallManager({}))

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

@@ -0,0 +1,94 @@
+# ----------------------------------------------------------------------------
+#
+#  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 __future__ import annotations
+
+import os
+import warnings
+import typing_extensions
+import typing
+import functools
+
+from baml_py.logging import (
+    get_log_level as baml_get_log_level,
+    set_log_level as baml_set_log_level,
+)
+from .globals import reset_baml_env_vars
+
+rT = typing_extensions.TypeVar("rT")  # return type
+pT = typing_extensions.ParamSpec("pT")  # parameters type
+
+
+def _deprecated(message: str):
+    def decorator(func: typing.Callable[pT, rT]) -> typing.Callable[pT, rT]:
+        """Use this decorator to mark functions as deprecated.
+        Every time the decorated function runs, it will emit
+        a "deprecation" warning."""
+
+        @functools.wraps(func)
+        def new_func(*args: pT.args, **kwargs: pT.kwargs):
+            warnings.simplefilter("always", DeprecationWarning)  # turn off filter
+            warnings.warn(
+                "Call to a deprecated function {}.".format(func.__name__) + message,
+                category=DeprecationWarning,
+                stacklevel=2,
+            )
+            warnings.simplefilter("default", DeprecationWarning)  # reset filter
+            return func(*args, **kwargs)
+
+        return new_func
+
+    return decorator
+
+
+@_deprecated("Use os.environ['BAML_LOG'] instead")
+def get_log_level():
+    """
+    Get the log level for the BAML Python client.
+    """
+    return baml_get_log_level()
+
+
+@_deprecated("Use os.environ['BAML_LOG'] instead")
+def set_log_level(
+    level: typing_extensions.Literal["DEBUG", "INFO", "WARN", "ERROR", "OFF"] | str,
+):
+    """
+    Set the log level for the BAML Python client
+    """
+    baml_set_log_level(level)
+    os.environ["BAML_LOG"] = level
+
+
+@_deprecated("Use os.environ['BAML_LOG_JSON_MODE'] instead")
+def set_log_json_mode():
+    """
+    Set the log JSON mode for the BAML Python client.
+    """
+    os.environ["BAML_LOG_JSON_MODE"] = "true"
+
+
+@_deprecated("Use os.environ['BAML_LOG_MAX_CHUNK_LENGTH'] instead")
+def set_log_max_chunk_length():
+    """
+    Set the maximum log chunk length for the BAML Python client.
+    """
+    os.environ["BAML_LOG_MAX_CHUNK_LENGTH"] = "1000"
+
+
+__all__ = [
+    "set_log_level",
+    "get_log_level",
+    "set_log_json_mode",
+    "reset_baml_env_vars",
+    "set_log_max_chunk_length",
+]

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

@@ -0,0 +1,37 @@
+# ----------------------------------------------------------------------------
+#
+#  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 __future__ import annotations
+import os
+import warnings
+
+from baml_py import BamlCtxManager, BamlRuntime
+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", 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
+)
+
+
+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,
+    )
+
+
+__all__ = []