diff --git a/cognee-mcp/src/codingagents/coding_rule_associations.py b/cognee-mcp/src/codingagents/coding_rule_associations.py index 15ed8d43a..a7b4390f4 100644 --- a/cognee-mcp/src/codingagents/coding_rule_associations.py +++ b/cognee-mcp/src/codingagents/coding_rule_associations.py @@ -2,9 +2,9 @@ from uuid import NAMESPACE_OID, uuid5 from cognee.infrastructure.databases.graph import get_graph_engine from cognee.infrastructure.databases.vector import get_vector_engine -from cognee.infrastructure.llm.prompts import render_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import render_prompt from cognee.low_level import DataPoint -from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client from cognee.shared.logging_utils import get_logger from cognee.modules.engine.models import NodeSet from cognee.tasks.storage import add_data_points, index_graph_edges diff --git a/cognee/api/v1/config/config.py b/cognee/api/v1/config/config.py index 792d5e200..913444aa8 100644 --- a/cognee/api/v1/config/config.py +++ b/cognee/api/v1/config/config.py @@ -7,7 +7,7 @@ from cognee.modules.cognify.config import get_cognify_config from cognee.infrastructure.data.chunking.config import get_chunk_config from cognee.infrastructure.databases.vector import get_vectordb_config from cognee.infrastructure.databases.graph.config import get_graph_config -from cognee.infrastructure.llm.config import get_llm_config +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.config import get_llm_config from cognee.infrastructure.databases.relational import get_relational_config, get_migration_config from cognee.infrastructure.files.storage import LocalStorage diff --git a/cognee/api/v1/responses/routers/get_responses_router.py b/cognee/api/v1/responses/routers/get_responses_router.py index 16e025403..6bb293a05 100644 --- a/cognee/api/v1/responses/routers/get_responses_router.py +++ b/cognee/api/v1/responses/routers/get_responses_router.py @@ -17,7 +17,7 @@ from cognee.api.v1.responses.models import ( ) from cognee.api.v1.responses.dispatch_function import dispatch_function from cognee.api.v1.responses.default_tools import DEFAULT_TOOLS -from cognee.infrastructure.llm.config import get_llm_config +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.config import get_llm_config from cognee.modules.users.models import User from cognee.modules.users.methods import get_authenticated_user diff --git a/cognee/base_config.py b/cognee/base_config.py index 837834273..2f9ca8cdb 100644 --- a/cognee/base_config.py +++ b/cognee/base_config.py @@ -9,6 +9,7 @@ from pydantic_settings import BaseSettings, SettingsConfigDict class BaseConfig(BaseSettings): data_root_directory: str = get_absolute_path(".data_storage") monitoring_tool: object = Observer.LANGFUSE + structured_output_framework: str = os.getenv("STRUCTURED_OUTPUT_FRAMEWORK") graphistry_username: Optional[str] = os.getenv("GRAPHISTRY_USERNAME") graphistry_password: Optional[str] = os.getenv("GRAPHISTRY_PASSWORD") langfuse_public_key: Optional[str] = os.getenv("LANGFUSE_PUBLIC_KEY") diff --git a/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py b/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py index 00ee7e101..32a5b4926 100644 --- a/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py +++ b/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py @@ -1,8 +1,8 @@ from typing import Any, Dict, List from pydantic import BaseModel -from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client from cognee.eval_framework.evaluation.base_eval_adapter import BaseEvalAdapter -from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import read_query_prompt, render_prompt from cognee.eval_framework.eval_config import EvalConfig diff --git a/cognee/infrastructure/databases/vector/embeddings/FastembedEmbeddingEngine.py b/cognee/infrastructure/databases/vector/embeddings/FastembedEmbeddingEngine.py index 99bb4ec8a..bc197fffb 100644 --- a/cognee/infrastructure/databases/vector/embeddings/FastembedEmbeddingEngine.py +++ b/cognee/infrastructure/databases/vector/embeddings/FastembedEmbeddingEngine.py @@ -5,7 +5,7 @@ import litellm import os from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine from cognee.infrastructure.databases.exceptions.EmbeddingException import EmbeddingException -from cognee.infrastructure.llm.tokenizer.TikToken import TikTokenTokenizer +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.tokenizer.TikToken import TikTokenTokenizer litellm.set_verbose = False logger = get_logger("FastembedEmbeddingEngine") diff --git a/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py b/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py index b51d397ed..502ca1fb7 100644 --- a/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py +++ b/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py @@ -7,11 +7,11 @@ import litellm import os from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine from cognee.infrastructure.databases.exceptions.EmbeddingException import EmbeddingException -from cognee.infrastructure.llm.tokenizer.Gemini import GeminiTokenizer -from cognee.infrastructure.llm.tokenizer.HuggingFace import HuggingFaceTokenizer -from cognee.infrastructure.llm.tokenizer.Mistral import MistralTokenizer -from cognee.infrastructure.llm.tokenizer.TikToken import TikTokenTokenizer -from cognee.infrastructure.llm.embedding_rate_limiter import ( +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.tokenizer.Gemini import GeminiTokenizer +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.tokenizer.HuggingFace import HuggingFaceTokenizer +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.tokenizer import MistralTokenizer +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.tokenizer.TikToken import TikTokenTokenizer +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.embedding_rate_limiter import ( embedding_rate_limit_async, embedding_sleep_and_retry_async, ) diff --git a/cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py b/cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py index 31f615477..4db38648e 100644 --- a/cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py +++ b/cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py @@ -7,9 +7,8 @@ import os import aiohttp.http_exceptions from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine -from cognee.infrastructure.databases.exceptions.EmbeddingException import EmbeddingException -from cognee.infrastructure.llm.tokenizer.HuggingFace import HuggingFaceTokenizer -from cognee.infrastructure.llm.embedding_rate_limiter import ( +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.tokenizer.HuggingFace import HuggingFaceTokenizer +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.embedding_rate_limiter import ( embedding_rate_limit_async, embedding_sleep_and_retry_async, ) diff --git a/cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py b/cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py index ef06eb7b2..d97984c96 100644 --- a/cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py +++ b/cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py @@ -1,5 +1,5 @@ from cognee.infrastructure.databases.vector.embeddings.config import get_embedding_config -from cognee.infrastructure.llm.config import get_llm_config +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.config import get_llm_config from .EmbeddingEngine import EmbeddingEngine from functools import lru_cache diff --git a/cognee/infrastructure/llm/__init__.py b/cognee/infrastructure/llm/__init__.py index b1609b524..8a3511003 100644 --- a/cognee/infrastructure/llm/__init__.py +++ b/cognee/infrastructure/llm/__init__.py @@ -1,4 +1,4 @@ -from .config import get_llm_config -from .utils import get_max_chunk_tokens -from .utils import test_llm_connection -from .utils import test_embedding_connection +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.config import get_llm_config +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.utils import get_max_chunk_tokens +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.utils import test_llm_connection +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.utils import test_embedding_connection diff --git a/cognee/infrastructure/llm/anthropic/__init__.py b/cognee/infrastructure/llm/structured_output_framework/__init__.py similarity index 100% rename from cognee/infrastructure/llm/anthropic/__init__.py rename to cognee/infrastructure/llm/structured_output_framework/__init__.py diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/__init__.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/__init__.py new file mode 100644 index 000000000..c4fe2bc85 --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/__init__.py @@ -0,0 +1,58 @@ +# ---------------------------------------------------------------------------- +# +# 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. + +__version__ = "0.201.0" + +try: + from baml_py.safe_import import EnsureBamlPyImport +except ImportError: + raise ImportError(f"""Update to baml-py required. +Version of baml_client generator (see generators.baml): {__version__} + +Please upgrade baml-py to version "{__version__}". + +$ pip install baml-py=={__version__} +$ uv add baml-py=={__version__} + +If nothing else works, please ask for help: + +https://github.com/boundaryml/baml/issues +https://boundaryml.com/discord +""") from None + + +with EnsureBamlPyImport(__version__) as e: + e.raise_if_incompatible_version(__version__) + + from . import types + from . import tracing + from . import stream_types + from . import config + from .config import reset_baml_env_vars + + 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 + +__all__ = [ + "b", + "stream_types", + "partial_types", + "tracing", + "types", + "reset_baml_env_vars", + "config", +] \ No newline at end of file diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/async_client.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/async_client.py new file mode 100644 index 000000000..dd837cf5a --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/async_client.py @@ -0,0 +1,235 @@ +# ---------------------------------------------------------------------------- +# +# 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 ExtractContentGraph(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="ExtractContentGraph", 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 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 ExtractContentGraphWithAnthropic(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="ExtractContentGraphWithAnthropic", 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 ExtractContentGraphWithEnvPrompt(self, content: str,prompt_override: typing.Optional[str] = None, + baml_options: BamlCallOptions = {}, + ) -> types.KnowledgeGraph: + result = await self.__options.merge_options(baml_options).call_function_async(function_name="ExtractContentGraphWithEnvPrompt", args={ + "content": content,"prompt_override": prompt_override, + }) + return typing.cast(types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)) + + + +class BamlStreamClient: + __options: DoNotUseDirectlyCallManager + + def __init__(self, options: DoNotUseDirectlyCallManager): + self.__options = options + + def ExtractContentGraph(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="ExtractContentGraph", 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 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 ExtractContentGraphWithAnthropic(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="ExtractContentGraphWithAnthropic", 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 ExtractContentGraphWithEnvPrompt(self, content: str,prompt_override: 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="ExtractContentGraphWithEnvPrompt", args={ + "content": content,"prompt_override": prompt_override, + }) + 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, + ) + + +class BamlHttpRequestClient: + __options: DoNotUseDirectlyCallManager + + def __init__(self, options: DoNotUseDirectlyCallManager): + self.__options = options + + async def ExtractContentGraph(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="ExtractContentGraph", args={ + "content": content,"mode": mode,"custom_prompt_content": custom_prompt_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 ExtractContentGraphWithAnthropic(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="ExtractContentGraphWithAnthropic", args={ + "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, + }, mode="request") + return result + async def ExtractContentGraphWithEnvPrompt(self, content: str,prompt_override: 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="ExtractContentGraphWithEnvPrompt", args={ + "content": content,"prompt_override": prompt_override, + }, mode="request") + return result + + +class BamlHttpStreamRequestClient: + __options: DoNotUseDirectlyCallManager + + def __init__(self, options: DoNotUseDirectlyCallManager): + self.__options = options + + async def ExtractContentGraph(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="ExtractContentGraph", args={ + "content": content,"mode": mode,"custom_prompt_content": custom_prompt_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 ExtractContentGraphWithAnthropic(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="ExtractContentGraphWithAnthropic", args={ + "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, + }, mode="stream") + return result + async def ExtractContentGraphWithEnvPrompt(self, content: str,prompt_override: 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="ExtractContentGraphWithEnvPrompt", args={ + "content": content,"prompt_override": prompt_override, + }, mode="stream") + return result + + +b = BamlAsyncClient(DoNotUseDirectlyCallManager({})) \ No newline at end of file diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/config.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/config.py new file mode 100644 index 000000000..998bd61bd --- /dev/null +++ b/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", +] diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/globals.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/globals.py new file mode 100644 index 000000000..769e055bb --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/globals.py @@ -0,0 +1,35 @@ +# ---------------------------------------------------------------------------- +# +# 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__ = [] diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/inlinedbaml.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/inlinedbaml.py new file mode 100644 index 000000000..4e5b22a42 --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/inlinedbaml.py @@ -0,0 +1,21 @@ +# ---------------------------------------------------------------------------- +# +# 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. + +_file_map = { + + "extract_categories.baml": "", + "extract_content_graph.baml": "class Node {\n id string\n name string\n type string\n description string\n @@dynamic\n}\n\n/// doc string for edge\nclass Edge {\n /// doc string for source_node_id\n source_node_id string\n target_node_id string\n relationship_name string\n}\n\nclass KnowledgeGraph {\n nodes (Node @stream.done)[]\n edges Edge[]\n}\n\n// Simple template for basic extraction (fast, good quality)\ntemplate_string ExtractContentGraphPrompt() #\"\n You are an advanced algorithm that extracts structured data into a knowledge graph.\n\n - **Nodes**: Entities/concepts (like Wikipedia articles).\n - **Edges**: Relationships (like Wikipedia links). Use snake_case (e.g., `acted_in`).\n\n **Rules:**\n\n 1. **Node Labeling & IDs**\n - Use basic types only (e.g., \"Person\", \"Date\", \"Organization\").\n - Avoid overly specific or generic terms (e.g., no \"Mathematician\" or \"Entity\").\n - Node IDs must be human-readable names from the text (no numbers).\n\n 2. **Dates & Numbers**\n - Label dates as **\"Date\"** in \"YYYY-MM-DD\" format (use available parts if incomplete).\n - Properties are key-value pairs; do not use escaped quotes.\n\n 3. **Coreference Resolution**\n - Use a single, complete identifier for each entity (e.g., always \"John Doe\" not \"Joe\" or \"he\").\n\n 4. **Relationship Labels**:\n - Use descriptive, lowercase, snake_case names for edges.\n - *Example*: born_in, married_to, invented_by.\n - Avoid vague or generic labels like isA, relatesTo, has.\n - Avoid duplicated relationships like produces, produced by.\n\n 5. **Strict Compliance**\n - Follow these rules exactly. Non-compliance results in termination.\n\"#\n\n// Detailed template for complex extraction (slower, higher quality)\ntemplate_string DetailedExtractContentGraphPrompt() #\"\n You are a top-tier algorithm designed for extracting information in structured formats to build a knowledge graph.\n **Nodes** represent entities and concepts. They're akin to Wikipedia nodes.\n **Edges** represent relationships between concepts. They're akin to Wikipedia links.\n\n The aim is to achieve simplicity and clarity in the knowledge graph.\n\n # 1. Labeling Nodes\n **Consistency**: Ensure you use basic or elementary types for node labels.\n - For example, when you identify an entity representing a person, always label it as **\"Person\"**.\n - Avoid using more specific terms like \"Mathematician\" or \"Scientist\", keep those as \"profession\" property.\n - Don't use too generic terms like \"Entity\".\n **Node IDs**: Never utilize integers as node IDs.\n - Node IDs should be names or human-readable identifiers found in the text.\n\n # 2. Handling Numerical Data and Dates\n - For example, when you identify an entity representing a date, make sure it has type **\"Date\"**.\n - Extract the date in the format \"YYYY-MM-DD\"\n - If not possible to extract the whole date, extract month or year, or both if available.\n - **Property Format**: Properties must be in a key-value format.\n - **Quotation Marks**: Never use escaped single or double quotes within property values.\n - **Naming Convention**: Use snake_case for relationship names, e.g., `acted_in`.\n\n # 3. Coreference Resolution\n - **Maintain Entity Consistency**: When extracting entities, it's vital to ensure consistency.\n If an entity, such as \"John Doe\", is mentioned multiple times in the text but is referred to by different names or pronouns (e.g., \"Joe\", \"he\"),\n always use the most complete identifier for that entity throughout the knowledge graph. In this example, use \"John Doe\" as the Person's ID.\n Remember, the knowledge graph should be coherent and easily understandable, so maintaining consistency in entity references is crucial.\n\n # 4. Strict Compliance\n Adhere to the rules strictly. Non-compliance will result in termination.\n\"#\n\n// Guided template with step-by-step instructions\ntemplate_string GuidedExtractContentGraphPrompt() #\"\n You are an advanced algorithm designed to extract structured information to build a clean, consistent, and human-readable knowledge graph.\n\n **Objective**:\n - Nodes represent entities and concepts, similar to Wikipedia articles.\n - Edges represent typed relationships between nodes, similar to Wikipedia hyperlinks.\n - The graph must be clear, minimal, consistent, and semantically precise.\n\n **Node Guidelines**:\n\n 1. **Label Consistency**:\n - Use consistent, basic types for all node labels.\n - Do not switch between granular or vague labels for the same kind of entity.\n - Pick one label for each category and apply it uniformly.\n - Each entity type should be in a singular form and in a case of multiple words separated by whitespaces\n\n 2. **Node Identifiers**:\n - Node IDs must be human-readable and derived directly from the text.\n - Prefer full names and canonical terms.\n - Never use integers or autogenerated IDs.\n - *Example*: Use \"Marie Curie\", \"Theory of Evolution\", \"Google\".\n\n 3. **Coreference Resolution**:\n - Maintain one consistent node ID for each real-world entity.\n - Resolve aliases, acronyms, and pronouns to the most complete form.\n - *Example*: Always use \"John Doe\" even if later referred to as \"Doe\" or \"he\".\n\n **Edge Guidelines**:\n\n 4. **Relationship Labels**:\n - Use descriptive, lowercase, snake_case names for edges.\n - *Example*: born_in, married_to, invented_by.\n - Avoid vague or generic labels like isA, relatesTo, has.\n\n 5. **Relationship Direction**:\n - Edges must be directional and logically consistent.\n - *Example*:\n - \"Marie Curie\" —[born_in]→ \"Warsaw\"\n - \"Radioactivity\" —[discovered_by]→ \"Marie Curie\"\n\n **Compliance**:\n Strict adherence to these guidelines is required. Any deviation will result in immediate termination of the task.\n\"#\n\n// Strict template with zero-tolerance rules\ntemplate_string StrictExtractContentGraphPrompt() #\"\n You are a top-tier algorithm for **extracting structured information** from unstructured text to build a **knowledge graph**.\n\n Your primary goal is to extract:\n - **Nodes**: Representing **entities** and **concepts** (like Wikipedia nodes).\n - **Edges**: Representing **relationships** between those concepts (like Wikipedia links).\n\n The resulting knowledge graph must be **simple, consistent, and human-readable**.\n\n ## 1. Node Labeling and Identification\n\n ### Node Types\n Use **basic atomic types** for node labels. Always prefer general types over specific roles or professions:\n - \"Person\" for any human.\n - \"Organization\" for companies, institutions, etc.\n - \"Location\" for geographic or place entities.\n - \"Date\" for any temporal expression.\n - \"Event\" for historical or scheduled occurrences.\n - \"Work\" for books, films, artworks, or research papers.\n - \"Concept\" for abstract notions or ideas.\n\n ### Node IDs\n - Always assign **human-readable and unambiguous identifiers**.\n - Never use numeric or autogenerated IDs.\n - Prioritize **most complete form** of entity names for consistency.\n\n ## 2. Relationship Handling\n - Use **snake_case** for all relationship (edge) types.\n - Keep relationship types semantically clear and consistent.\n - Avoid vague relation names like \"related_to\" unless no better alternative exists.\n\n ## 3. Strict Compliance\n Follow all rules exactly. Any deviation may lead to rejection or incorrect graph construction.\n\"#\n\n// OpenAI client with environment model selection\nclient OpenAIClientWithEnvModel {\n provider openai\n options {\n model env.LLM_MODEL\n api_key env.OPENAI_API_KEY\n }\n}\n\n// Anthropic client with environment model selection\nclient AnthropicClientWithEnvModel {\n provider anthropic\n options {\n model env.LLM_MODEL\n api_key env.ANTHROPIC_API_KEY\n }\n}\n\n// Default client (maintains backward compatibility)\nclient DefaultClient {\n provider openai\n options {\n model \"gpt-4o-mini\"\n api_key env.OPENAI_API_KEY\n }\n}\n\n// Function that returns raw structured output (for custom objects - to be handled in Python)\nfunction ExtractContentGraphGeneric(\n content: string,\n mode: \"simple\" | \"base\" | \"guided\" | \"strict\" | \"custom\"?,\n custom_prompt_content: string?\n) -> KnowledgeGraph {\n client OpenAIClientWithEnvModel\n\n prompt #\"\n {% if mode == \"base\" %}\n {{ DetailedExtractContentGraphPrompt() }}\n {% elif mode == \"guided\" %}\n {{ GuidedExtractContentGraphPrompt() }}\n {% elif mode == \"strict\" %}\n {{ StrictExtractContentGraphPrompt() }}\n {% elif mode == \"custom\" and custom_prompt_content %}\n {{ custom_prompt_content }}\n {% else %}\n {{ ExtractContentGraphPrompt() }}\n {% endif %}\n\n {{ ctx.output_format(prefix=\"Answer in this schema:\\n\") }}\n\n Before answering, briefly describe what you'll extract from the text, then provide the structured output.\n\n Example format:\n I'll extract the main entities and their relationships from this text...\n\n { ... }\n\n {{ _.role('user') }}\n {{ content }}\n \"#\n}\n\n// Backward-compatible function specifically for KnowledgeGraph\nfunction ExtractContentGraph(\n content: string,\n mode: \"simple\" | \"base\" | \"guided\" | \"strict\" | \"custom\"?,\n custom_prompt_content: string?\n) -> KnowledgeGraph {\n client OpenAIClientWithEnvModel\n\n prompt #\"\n {% if mode == \"base\" %}\n {{ DetailedExtractContentGraphPrompt() }}\n {% elif mode == \"guided\" %}\n {{ GuidedExtractContentGraphPrompt() }}\n {% elif mode == \"strict\" %}\n {{ StrictExtractContentGraphPrompt() }}\n {% elif mode == \"custom\" and custom_prompt_content %}\n {{ custom_prompt_content }}\n {% else %}\n {{ ExtractContentGraphPrompt() }}\n {% endif %}\n\n {{ ctx.output_format(prefix=\"Answer in this schema:\\n\") }}\n\n Before answering, briefly describe what you'll extract from the text, then provide the structured output.\n\n Example format:\n I'll extract the main entities and their relationships from this text...\n\n { ... }\n\n {{ _.role('user') }}\n {{ content }}\n \"#\n}\n\n// Alternative function that uses environment variable for prompt selection\nfunction ExtractContentGraphWithEnvPrompt(\n content: string,\n prompt_override: string?\n) -> KnowledgeGraph {\n client OpenAIClientWithEnvModel\n\n prompt #\"\n {% if prompt_override %}\n {{ prompt_override }}\n {% else %}\n {{ ExtractContentGraphPrompt() }}\n {% endif %}\n\n {{ ctx.output_format(prefix=\"Answer in this schema:\\n\") }}\n\n Before answering, briefly describe what you'll extract from the text, then provide the structured output.\n\n Example format:\n I'll extract the main entities and their relationships from this text...\n\n { ... }\n\n {{ _.role('user') }}\n {{ content }}\n \"#\n}\n\n// Function that uses Anthropic client\nfunction ExtractContentGraphWithAnthropic(\n content: string,\n mode: \"simple\" | \"base\" | \"guided\" | \"strict\" | \"custom\"?,\n custom_prompt_content: string?\n) -> KnowledgeGraph {\n client AnthropicClientWithEnvModel\n\n prompt #\"\n {% if mode == \"base\" %}\n {{ DetailedExtractContentGraphPrompt() }}\n {% elif mode == \"guided\" %}\n {{ GuidedExtractContentGraphPrompt() }}\n {% elif mode == \"strict\" %}\n {{ StrictExtractContentGraphPrompt() }}\n {% elif mode == \"custom\" and custom_prompt_content %}\n {{ custom_prompt_content }}\n {% else %}\n {{ ExtractContentGraphPrompt() }}\n {% endif %}\n\n {{ ctx.output_format(prefix=\"Answer in this schema:\\n\") }}\n\n Before answering, briefly describe what you'll extract from the text, then provide the structured output.\n\n Example format:\n I'll extract the main entities and their relationships from this text...\n\n { ... }\n\n {{ _.role('user') }}\n {{ content }}\n \"#\n}\n\ntest ExtractPersonExample {\n functions [ExtractContentGraph]\n args {\n content #\"\n My name is Vasiliy. I was born in 1992. I am a software engineer. I work at Google and am based in Berlin.\n \"#\n mode \"simple\"\n }\n}\n\ntest ExtractGuidedExample {\n functions [ExtractContentGraph]\n args {\n content #\"\n Apple Inc. was founded by Steve Jobs in 1976. The company is headquartered in Cupertino, California.\n Tim Cook is the current CEO of Apple Inc.\n \"#\n mode \"guided\"\n }\n}\n\ntest ExtractStrictExample {\n functions [ExtractContentGraph]\n args {\n content #\"\n The Python programming language was created by Guido van Rossum in 1991.\n \"#\n mode \"strict\"\n }\n}\n\ntest ExtractGenericExample {\n functions [ExtractContentGraphGeneric]\n args {\n content #\"\n React is a JavaScript library for building user interfaces, developed by Facebook.\n \"#\n mode \"simple\"\n }\n}\n", + "generators.baml": "// This helps use auto generate libraries you can use in the language of\n// your choice. You can have multiple generators if you use multiple languages.\n// Just ensure that the output_dir is different for each generator.\ngenerator target {\n // Valid values: \"python/pydantic\", \"typescript\", \"ruby/sorbet\", \"rest/openapi\"\n output_type \"python/pydantic\"\n\n // Where the generated code will be saved (relative to baml_src/)\n output_dir \"../baml/\"\n\n // The version of the BAML package you have installed (e.g. same version as your baml-py or @boundaryml/baml).\n // The BAML VSCode extension version should also match this version.\n version \"0.201.0\"\n\n // Valid values: \"sync\", \"async\"\n // This controls what `b.FunctionName()` will be (sync or async).\n default_client_mode sync\n}\n", +} + +def get_baml_files(): + return _file_map \ No newline at end of file diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/parser.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/parser.py new file mode 100644 index 000000000..c389b8c0a --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/parser.py @@ -0,0 +1,81 @@ +# ---------------------------------------------------------------------------- +# +# 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 + +from . import stream_types, types +from .runtime import DoNotUseDirectlyCallManager, BamlCallOptions + +class LlmResponseParser: + __options: DoNotUseDirectlyCallManager + + def __init__(self, options: DoNotUseDirectlyCallManager): + self.__options = options + + def ExtractContentGraph( + self, llm_response: str, baml_options: BamlCallOptions = {}, + ) -> types.KnowledgeGraph: + result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractContentGraph", llm_response=llm_response, mode="request") + return typing.cast(types.KnowledgeGraph, 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 ExtractContentGraphWithAnthropic( + self, llm_response: str, baml_options: BamlCallOptions = {}, + ) -> types.KnowledgeGraph: + result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractContentGraphWithAnthropic", llm_response=llm_response, mode="request") + return typing.cast(types.KnowledgeGraph, result) + + def ExtractContentGraphWithEnvPrompt( + self, llm_response: str, baml_options: BamlCallOptions = {}, + ) -> types.KnowledgeGraph: + result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractContentGraphWithEnvPrompt", llm_response=llm_response, mode="request") + return typing.cast(types.KnowledgeGraph, result) + + + +class LlmStreamParser: + __options: DoNotUseDirectlyCallManager + + def __init__(self, options: DoNotUseDirectlyCallManager): + self.__options = options + + def ExtractContentGraph( + self, llm_response: str, baml_options: BamlCallOptions = {}, + ) -> stream_types.KnowledgeGraph: + result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractContentGraph", llm_response=llm_response, mode="stream") + return typing.cast(stream_types.KnowledgeGraph, 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 ExtractContentGraphWithAnthropic( + self, llm_response: str, baml_options: BamlCallOptions = {}, + ) -> stream_types.KnowledgeGraph: + result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractContentGraphWithAnthropic", llm_response=llm_response, mode="stream") + return typing.cast(stream_types.KnowledgeGraph, result) + + def ExtractContentGraphWithEnvPrompt( + self, llm_response: str, baml_options: BamlCallOptions = {}, + ) -> stream_types.KnowledgeGraph: + result = self.__options.merge_options(baml_options).parse_response(function_name="ExtractContentGraphWithEnvPrompt", llm_response=llm_response, mode="stream") + return typing.cast(stream_types.KnowledgeGraph, result) + + \ No newline at end of file diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/runtime.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/runtime.py new file mode 100644 index 000000000..c94bb950e --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/runtime.py @@ -0,0 +1,256 @@ +# ---------------------------------------------------------------------------- +# +# 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 os +import typing +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__ + + +class BamlCallOptions(typing.TypedDict, total=False): + tb: typing_extensions.NotRequired[type_builder.TypeBuilder] + client_registry: typing_extensions.NotRequired[baml_py.baml_py.ClientRegistry] + env: typing_extensions.NotRequired[typing.Dict[str, typing.Optional[str]]] + collector: typing_extensions.NotRequired[ + typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]] + ] + + +class _ResolvedBamlOptions: + tb: typing.Optional[baml_py.baml_py.TypeBuilder] + client_registry: typing.Optional[baml_py.baml_py.ClientRegistry] + collectors: typing.List[baml_py.baml_py.Collector] + env_vars: typing.Dict[str, str] + + def __init__( + self, + tb: typing.Optional[baml_py.baml_py.TypeBuilder], + client_registry: typing.Optional[baml_py.baml_py.ClientRegistry], + collectors: typing.List[baml_py.baml_py.Collector], + env_vars: typing.Dict[str, str], + ): + self.tb = tb + self.client_registry = client_registry + self.collectors = collectors + self.env_vars = env_vars + + + + + +class DoNotUseDirectlyCallManager: + def __init__(self, baml_options: BamlCallOptions): + self.__baml_options = baml_options + + def __getstate__(self): + # Return state needed for pickling + return {"baml_options": self.__baml_options} + + def __setstate__(self, state): + # Restore state from pickling + self.__baml_options = state["baml_options"] + + def __resolve(self) -> _ResolvedBamlOptions: + tb = self.__baml_options.get("tb") + if tb is not None: + baml_tb = tb._tb # type: ignore (we know how to use this private attribute) + else: + baml_tb = None + client_registry = self.__baml_options.get("client_registry") + collector = self.__baml_options.get("collector") + collectors_as_list = ( + collector + if isinstance(collector, list) + else [collector] if collector is not None else [] + ) + env_vars = os.environ.copy() + for k, v in self.__baml_options.get("env", {}).items(): + if v is not None: + env_vars[k] = v + else: + env_vars.pop(k, None) + + return _ResolvedBamlOptions( + baml_tb, + client_registry, + collectors_as_list, + env_vars, + ) + + def merge_options(self, options: BamlCallOptions) -> "DoNotUseDirectlyCallManager": + return DoNotUseDirectlyCallManager({**self.__baml_options, **options}) + + async def call_function_async( + self, *, function_name: str, args: typing.Dict[str, typing.Any] + ) -> baml_py.baml_py.FunctionResult: + resolved_options = self.__resolve() + return await __runtime__.call_function( + function_name, + args, + # ctx + __ctx__manager__.clone_context(), + # tb + resolved_options.tb, + # cr + resolved_options.client_registry, + # collectors + resolved_options.collectors, + # env_vars + resolved_options.env_vars, + ) + + def call_function_sync( + self, *, function_name: str, args: typing.Dict[str, typing.Any] + ) -> baml_py.baml_py.FunctionResult: + resolved_options = self.__resolve() + ctx = __ctx__manager__.get() + return __runtime__.call_function_sync( + function_name, + args, + # ctx + ctx, + # tb + resolved_options.tb, + # cr + resolved_options.client_registry, + # collectors + resolved_options.collectors, + # env_vars + resolved_options.env_vars, + ) + + def create_async_stream( + self, + *, + function_name: str, + args: typing.Dict[str, typing.Any], + ) -> typing.Tuple[baml_py.baml_py.RuntimeContextManager, baml_py.baml_py.FunctionResultStream]: + resolved_options = self.__resolve() + ctx = __ctx__manager__.clone_context() + result = __runtime__.stream_function( + function_name, + args, + # this is always None, we set this later! + # on_event + None, + # ctx + ctx, + # tb + resolved_options.tb, + # cr + resolved_options.client_registry, + # collectors + resolved_options.collectors, + # env_vars + resolved_options.env_vars, + ) + return ctx, result + + def create_sync_stream( + self, + *, + function_name: str, + args: typing.Dict[str, typing.Any], + ) -> 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( + function_name, + args, + # this is always None, we set this later! + # on_event + None, + # ctx + ctx, + # tb + resolved_options.tb, + # cr + resolved_options.client_registry, + # collectors + resolved_options.collectors, + # env_vars + resolved_options.env_vars, + ) + return ctx, result + + async def create_http_request_async( + self, + *, + function_name: str, + args: typing.Dict[str, typing.Any], + mode: typing_extensions.Literal["stream", "request"], + ) -> baml_py.baml_py.HTTPRequest: + resolved_options = self.__resolve() + return await __runtime__.build_request( + function_name, + args, + # ctx + __ctx__manager__.clone_context(), + # tb + resolved_options.tb, + # cr + resolved_options.client_registry, + # env_vars + resolved_options.env_vars, + # is_stream + mode == "stream", + ) + + def create_http_request_sync( + self, + *, + function_name: str, + args: typing.Dict[str, typing.Any], + mode: typing_extensions.Literal["stream", "request"], + ) -> baml_py.baml_py.HTTPRequest: + resolved_options = self.__resolve() + return __runtime__.build_request_sync( + function_name, + args, + # ctx + __ctx__manager__.get(), + # tb + resolved_options.tb, + # cr + resolved_options.client_registry, + # env_vars + resolved_options.env_vars, + # is_stream + mode == "stream", + ) + + 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, + llm_response, + # enum_module + types, + # cls_module + types, + # partial_cls_module + stream_types, + # allow_partials + mode == "stream", + # ctx + __ctx__manager__.get(), + # tb + resolved_options.tb, + # cr + resolved_options.client_registry, + # env_vars + resolved_options.env_vars, + ) \ No newline at end of file diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/stream_types.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/stream_types.py new file mode 100644 index 000000000..feefe5a89 --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/stream_types.py @@ -0,0 +1,50 @@ +# ---------------------------------------------------------------------------- +# +# 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 +from pydantic import BaseModel, ConfigDict + +import baml_py + +from . import types + +StreamStateValueT = typing.TypeVar('StreamStateValueT') +class StreamState(BaseModel, typing.Generic[StreamStateValueT]): + value: StreamStateValueT + state: typing_extensions.Literal["Pending", "Incomplete", "Complete"] +# ######################################################################### +# Generated classes (3) +# ######################################################################### + +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 KnowledgeGraph(BaseModel): + nodes: typing.List["types.Node"] + edges: typing.List["Edge"] + +class Node(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 + +# ######################################################################### +# Generated type aliases (0) +# ######################################################################### diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/sync_client.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/sync_client.py new file mode 100644 index 000000000..b4fa1c558 --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/sync_client.py @@ -0,0 +1,247 @@ +# ---------------------------------------------------------------------------- +# +# 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 ExtractContentGraph(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="ExtractContentGraph", 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 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 ExtractContentGraphWithAnthropic(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="ExtractContentGraphWithAnthropic", 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 ExtractContentGraphWithEnvPrompt(self, content: str,prompt_override: typing.Optional[str] = None, + baml_options: BamlCallOptions = {}, + ) -> types.KnowledgeGraph: + result = self.__options.merge_options(baml_options).call_function_sync(function_name="ExtractContentGraphWithEnvPrompt", args={ + "content": content,"prompt_override": prompt_override, + }) + return typing.cast(types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)) + + + +class BamlStreamClient: + __options: DoNotUseDirectlyCallManager + + def __init__(self, options: DoNotUseDirectlyCallManager): + self.__options = options + + def ExtractContentGraph(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="ExtractContentGraph", 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 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 ExtractContentGraphWithAnthropic(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="ExtractContentGraphWithAnthropic", 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 ExtractContentGraphWithEnvPrompt(self, content: str,prompt_override: 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="ExtractContentGraphWithEnvPrompt", args={ + "content": content,"prompt_override": prompt_override, + }) + 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, + ) + + +class BamlHttpRequestClient: + __options: DoNotUseDirectlyCallManager + + def __init__(self, options: DoNotUseDirectlyCallManager): + self.__options = options + + def ExtractContentGraph(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="ExtractContentGraph", args={ + "content": content,"mode": mode,"custom_prompt_content": custom_prompt_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 ExtractContentGraphWithAnthropic(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="ExtractContentGraphWithAnthropic", args={ + "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, + }, mode="request") + return result + def ExtractContentGraphWithEnvPrompt(self, content: str,prompt_override: 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="ExtractContentGraphWithEnvPrompt", args={ + "content": content,"prompt_override": prompt_override, + }, mode="request") + return result + + +class BamlHttpStreamRequestClient: + __options: DoNotUseDirectlyCallManager + + def __init__(self, options: DoNotUseDirectlyCallManager): + self.__options = options + + def ExtractContentGraph(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="ExtractContentGraph", args={ + "content": content,"mode": mode,"custom_prompt_content": custom_prompt_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 ExtractContentGraphWithAnthropic(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="ExtractContentGraphWithAnthropic", args={ + "content": content,"mode": mode,"custom_prompt_content": custom_prompt_content, + }, mode="stream") + return result + def ExtractContentGraphWithEnvPrompt(self, content: str,prompt_override: 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="ExtractContentGraphWithEnvPrompt", args={ + "content": content,"prompt_override": prompt_override, + }, mode="stream") + return result + + +b = BamlSyncClient(DoNotUseDirectlyCallManager({})) \ No newline at end of file diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/tracing.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/tracing.py new file mode 100644 index 000000000..06725593c --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/tracing.py @@ -0,0 +1,22 @@ +# ---------------------------------------------------------------------------- +# +# 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"] diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_builder.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_builder.py new file mode 100644 index 000000000..3964e4721 --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_builder.py @@ -0,0 +1,208 @@ +# ---------------------------------------------------------------------------- +# +# 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 +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( + ["Edge","KnowledgeGraph","Node",] + ), enums=set( + [] + ), runtime=DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME) + + # ######################################################################### + # Generated enums 0 + # ######################################################################### + + + # ######################################################################### + # Generated classes 3 + # ######################################################################### + + @property + def Edge(self) -> "EdgeViewer": + return EdgeViewer(self) + + @property + def KnowledgeGraph(self) -> "KnowledgeGraphViewer": + return KnowledgeGraphViewer(self) + + @property + def Node(self) -> "NodeBuilder": + return NodeBuilder(self) + + + +# ######################################################################### +# Generated enums 0 +# ######################################################################### + + +# ######################################################################### +# Generated classes 3 +# ######################################################################### + +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._props = EdgeProperties(self._bldr, self._properties) + + def type(self) -> baml_py.FieldType: + return self._bldr.field() + + @property + def props(self) -> "EdgeProperties": + return self._props + + +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] + + + +class EdgeProperties: + 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 + + + + @property + def source_node_id(self) -> type_builder.ClassPropertyViewer: + return type_builder.ClassPropertyViewer(self.__bldr.property("source_node_id")) + + @property + def target_node_id(self) -> type_builder.ClassPropertyViewer: + return type_builder.ClassPropertyViewer(self.__bldr.property("target_node_id")) + + @property + def relationship_name(self) -> type_builder.ClassPropertyViewer: + return type_builder.ClassPropertyViewer(self.__bldr.property("relationship_name")) + + + + +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._props = KnowledgeGraphProperties(self._bldr, self._properties) + + def type(self) -> baml_py.FieldType: + return self._bldr.field() + + @property + def props(self) -> "KnowledgeGraphProperties": + return self._props + + +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] + + + +class KnowledgeGraphProperties: + 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 + + + + @property + def nodes(self) -> type_builder.ClassPropertyViewer: + return type_builder.ClassPropertyViewer(self.__bldr.property("nodes")) + + @property + def edges(self) -> type_builder.ClassPropertyViewer: + return type_builder.ClassPropertyViewer(self.__bldr.property("edges")) + + + + +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._props = NodeProperties(self._bldr, self._properties) + + def type(self) -> baml_py.FieldType: + return self._bldr.field() + + @property + def props(self) -> "NodeProperties": + return self._props + + +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.") + return self._bldr.property(name).type(type) + + def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]: + 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") + + @property + def name(self) -> baml_py.ClassPropertyBuilder: + return self.__bldr.property("name") + + @property + def type(self) -> baml_py.ClassPropertyBuilder: + return self.__bldr.property("type") + + @property + def description(self) -> baml_py.ClassPropertyBuilder: + return self.__bldr.property("description") + + + diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_map.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_map.py new file mode 100644 index 000000000..61f0c3bab --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_map.py @@ -0,0 +1,29 @@ +# ---------------------------------------------------------------------------- +# +# 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 . import types +from . import stream_types + + +type_map = { + + "types.Edge": types.Edge, + "stream_types.Edge": stream_types.Edge, + + "types.KnowledgeGraph": types.KnowledgeGraph, + "stream_types.KnowledgeGraph": stream_types.KnowledgeGraph, + + "types.Node": types.Node, + "stream_types.Node": stream_types.Node, + + +} \ No newline at end of file diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/types.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/types.py new file mode 100644 index 000000000..09b090f4e --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/types.py @@ -0,0 +1,68 @@ +# ---------------------------------------------------------------------------- +# +# 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 +from enum import Enum + + +from pydantic import BaseModel, ConfigDict + + +import baml_py + +CheckT = typing_extensions.TypeVar('CheckT') +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) +# ######################################################################### + +# ######################################################################### +# Generated classes (3) +# ######################################################################### + +class Edge(BaseModel): + # doc string for edge + # doc string for source_node_id + + source_node_id: str + target_node_id: str + relationship_name: str + +class KnowledgeGraph(BaseModel): + nodes: typing.List["Node"] + edges: typing.List["Edge"] + +class Node(BaseModel): + model_config = ConfigDict(extra='allow') + id: str + name: str + type: str + description: str + +# ######################################################################### +# Generated type aliases (0) +# ######################################################################### diff --git a/cognee/infrastructure/llm/structured_output_framework/baml_src/config.py b/cognee/infrastructure/llm/structured_output_framework/baml_src/config.py new file mode 100644 index 000000000..91f26cb11 --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml_src/config.py @@ -0,0 +1,179 @@ +import os +from typing import Optional +from functools import lru_cache +from pydantic_settings import BaseSettings, SettingsConfigDict +from pydantic import model_validator +from baml_py import ClientRegistry + +class LLMConfig(BaseSettings): + """ + Configuration settings for the LLM (Large Language Model) provider and related options. + + Public instance variables include: + - llm_provider + - llm_model + - llm_endpoint + - llm_api_key + - llm_api_version + - llm_temperature + - llm_streaming + - llm_max_tokens + - transcription_model + - graph_prompt_path + - llm_rate_limit_enabled + - llm_rate_limit_requests + - llm_rate_limit_interval + - embedding_rate_limit_enabled + - embedding_rate_limit_requests + - embedding_rate_limit_interval + + Public methods include: + - ensure_env_vars_for_ollama + - to_dict + """ + + llm_provider: str = "openai" + llm_model: str = "gpt-4o-mini" + llm_endpoint: str = "" + llm_api_key: Optional[str] = None + llm_api_version: Optional[str] = None + llm_temperature: float = 0.0 + llm_streaming: bool = False + llm_max_tokens: int = 16384 + transcription_model: str = "whisper-1" + graph_prompt_path: str = "generate_graph_prompt.txt" + llm_rate_limit_enabled: bool = False + llm_rate_limit_requests: int = 60 + llm_rate_limit_interval: int = 60 # in seconds (default is 60 requests per minute) + embedding_rate_limit_enabled: bool = False + embedding_rate_limit_requests: int = 60 + embedding_rate_limit_interval: int = 60 # in seconds (default is 60 requests per minute) + baml_registry = ClientRegistry() + + model_config = SettingsConfigDict(env_file=".env", extra="allow") + + + baml_registry.add_llm_client(name=llm_provider, provider=llm_provider, options={ + "model": llm_model, + "temperature": llm_temperature, + "api_key": llm_api_key + }) + # Sets MyAmazingClient as the primary client + baml_registry.set_primary('openai') + + @model_validator(mode="after") + def ensure_env_vars_for_ollama(self) -> "LLMConfig": + """ + Validate required environment variables for the 'ollama' LLM provider. + + Raises ValueError if some required environment variables are set without the others. + Only checks are performed when 'llm_provider' is set to 'ollama'. + + Returns: + -------- + + - 'LLMConfig': The instance of LLMConfig after validation. + """ + + if self.llm_provider != "ollama": + # Skip checks unless provider is "ollama" + return self + + def is_env_set(var_name: str) -> bool: + """ + Check if a given environment variable is set and non-empty. + + Parameters: + ----------- + + - var_name (str): The name of the environment variable to check. + + Returns: + -------- + + - bool: True if the environment variable exists and is not empty, otherwise False. + """ + val = os.environ.get(var_name) + return val is not None and val.strip() != "" + + # + # 1. Check LLM environment variables + # + llm_env_vars = { + "LLM_MODEL": is_env_set("LLM_MODEL"), + "LLM_ENDPOINT": is_env_set("LLM_ENDPOINT"), + "LLM_API_KEY": is_env_set("LLM_API_KEY"), + } + if any(llm_env_vars.values()) and not all(llm_env_vars.values()): + missing_llm = [key for key, is_set in llm_env_vars.items() if not is_set] + raise ValueError( + "You have set some but not all of the required environment variables " + f"for LLM usage (LLM_MODEL, LLM_ENDPOINT, LLM_API_KEY). Missing: {missing_llm}" + ) + + # + # 2. Check embedding environment variables + # + embedding_env_vars = { + "EMBEDDING_PROVIDER": is_env_set("EMBEDDING_PROVIDER"), + "EMBEDDING_MODEL": is_env_set("EMBEDDING_MODEL"), + "EMBEDDING_DIMENSIONS": is_env_set("EMBEDDING_DIMENSIONS"), + "HUGGINGFACE_TOKENIZER": is_env_set("HUGGINGFACE_TOKENIZER"), + } + if any(embedding_env_vars.values()) and not all(embedding_env_vars.values()): + missing_embed = [key for key, is_set in embedding_env_vars.items() if not is_set] + raise ValueError( + "You have set some but not all of the required environment variables " + "for embeddings (EMBEDDING_PROVIDER, EMBEDDING_MODEL, " + "EMBEDDING_DIMENSIONS, HUGGINGFACE_TOKENIZER). Missing: " + f"{missing_embed}" + ) + + return self + + def to_dict(self) -> dict: + """ + Convert the LLMConfig instance into a dictionary representation. + + Returns: + -------- + + - dict: A dictionary containing the configuration settings of the LLMConfig + instance. + """ + return { + "provider": self.llm_provider, + "model": self.llm_model, + "endpoint": self.llm_endpoint, + "api_key": self.llm_api_key, + "api_version": self.llm_api_version, + "temperature": self.llm_temperature, + "streaming": self.llm_streaming, + "max_tokens": self.llm_max_tokens, + "transcription_model": self.transcription_model, + "graph_prompt_path": self.graph_prompt_path, + "rate_limit_enabled": self.llm_rate_limit_enabled, + "rate_limit_requests": self.llm_rate_limit_requests, + "rate_limit_interval": self.llm_rate_limit_interval, + "embedding_rate_limit_enabled": self.embedding_rate_limit_enabled, + "embedding_rate_limit_requests": self.embedding_rate_limit_requests, + "embedding_rate_limit_interval": self.embedding_rate_limit_interval, + } + + +@lru_cache +def get_llm_config(): + """ + Retrieve and cache the LLM configuration. + + This function returns an instance of the LLMConfig class. It leverages + caching to ensure that repeated calls do not create new instances, + but instead return the already created configuration object. + + Returns: + -------- + + - LLMConfig: An instance of the LLMConfig class containing the configuration for the + LLM. + """ + return LLMConfig() diff --git a/cognee/infrastructure/llm/gemini/__init__.py b/cognee/infrastructure/llm/structured_output_framework/baml_src/extract_categories.baml similarity index 100% rename from cognee/infrastructure/llm/gemini/__init__.py rename to cognee/infrastructure/llm/structured_output_framework/baml_src/extract_categories.baml diff --git a/cognee/infrastructure/llm/structured_output_framework/baml_src/extract_content_graph.baml b/cognee/infrastructure/llm/structured_output_framework/baml_src/extract_content_graph.baml new file mode 100644 index 000000000..6b4c957fe --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml_src/extract_content_graph.baml @@ -0,0 +1,367 @@ +class Node { + id string + name string + type string + description string + @@dynamic +} + +/// doc string for edge +class Edge { + /// doc string for source_node_id + source_node_id string + target_node_id string + relationship_name string +} + +class KnowledgeGraph { + nodes (Node @stream.done)[] + edges Edge[] +} + +// Simple template for basic extraction (fast, good quality) +template_string ExtractContentGraphPrompt() #" + You are an advanced algorithm that extracts structured data into a knowledge graph. + + - **Nodes**: Entities/concepts (like Wikipedia articles). + - **Edges**: Relationships (like Wikipedia links). Use snake_case (e.g., `acted_in`). + + **Rules:** + + 1. **Node Labeling & IDs** + - Use basic types only (e.g., "Person", "Date", "Organization"). + - Avoid overly specific or generic terms (e.g., no "Mathematician" or "Entity"). + - Node IDs must be human-readable names from the text (no numbers). + + 2. **Dates & Numbers** + - Label dates as **"Date"** in "YYYY-MM-DD" format (use available parts if incomplete). + - Properties are key-value pairs; do not use escaped quotes. + + 3. **Coreference Resolution** + - Use a single, complete identifier for each entity (e.g., always "John Doe" not "Joe" or "he"). + + 4. **Relationship Labels**: + - Use descriptive, lowercase, snake_case names for edges. + - *Example*: born_in, married_to, invented_by. + - Avoid vague or generic labels like isA, relatesTo, has. + - Avoid duplicated relationships like produces, produced by. + + 5. **Strict Compliance** + - Follow these rules exactly. Non-compliance results in termination. +"# + +// Detailed template for complex extraction (slower, higher quality) +template_string DetailedExtractContentGraphPrompt() #" + You are a top-tier algorithm designed for extracting information in structured formats to build a knowledge graph. + **Nodes** represent entities and concepts. They're akin to Wikipedia nodes. + **Edges** represent relationships between concepts. They're akin to Wikipedia links. + + The aim is to achieve simplicity and clarity in the knowledge graph. + + # 1. Labeling Nodes + **Consistency**: Ensure you use basic or elementary types for node labels. + - For example, when you identify an entity representing a person, always label it as **"Person"**. + - Avoid using more specific terms like "Mathematician" or "Scientist", keep those as "profession" property. + - Don't use too generic terms like "Entity". + **Node IDs**: Never utilize integers as node IDs. + - Node IDs should be names or human-readable identifiers found in the text. + + # 2. Handling Numerical Data and Dates + - For example, when you identify an entity representing a date, make sure it has type **"Date"**. + - Extract the date in the format "YYYY-MM-DD" + - If not possible to extract the whole date, extract month or year, or both if available. + - **Property Format**: Properties must be in a key-value format. + - **Quotation Marks**: Never use escaped single or double quotes within property values. + - **Naming Convention**: Use snake_case for relationship names, e.g., `acted_in`. + + # 3. Coreference Resolution + - **Maintain Entity Consistency**: When extracting entities, it's vital to ensure consistency. + If an entity, such as "John Doe", is mentioned multiple times in the text but is referred to by different names or pronouns (e.g., "Joe", "he"), + always use the most complete identifier for that entity throughout the knowledge graph. In this example, use "John Doe" as the Person's ID. + Remember, the knowledge graph should be coherent and easily understandable, so maintaining consistency in entity references is crucial. + + # 4. Strict Compliance + Adhere to the rules strictly. Non-compliance will result in termination. +"# + +// Guided template with step-by-step instructions +template_string GuidedExtractContentGraphPrompt() #" + You are an advanced algorithm designed to extract structured information to build a clean, consistent, and human-readable knowledge graph. + + **Objective**: + - Nodes represent entities and concepts, similar to Wikipedia articles. + - Edges represent typed relationships between nodes, similar to Wikipedia hyperlinks. + - The graph must be clear, minimal, consistent, and semantically precise. + + **Node Guidelines**: + + 1. **Label Consistency**: + - Use consistent, basic types for all node labels. + - Do not switch between granular or vague labels for the same kind of entity. + - Pick one label for each category and apply it uniformly. + - Each entity type should be in a singular form and in a case of multiple words separated by whitespaces + + 2. **Node Identifiers**: + - Node IDs must be human-readable and derived directly from the text. + - Prefer full names and canonical terms. + - Never use integers or autogenerated IDs. + - *Example*: Use "Marie Curie", "Theory of Evolution", "Google". + + 3. **Coreference Resolution**: + - Maintain one consistent node ID for each real-world entity. + - Resolve aliases, acronyms, and pronouns to the most complete form. + - *Example*: Always use "John Doe" even if later referred to as "Doe" or "he". + + **Edge Guidelines**: + + 4. **Relationship Labels**: + - Use descriptive, lowercase, snake_case names for edges. + - *Example*: born_in, married_to, invented_by. + - Avoid vague or generic labels like isA, relatesTo, has. + + 5. **Relationship Direction**: + - Edges must be directional and logically consistent. + - *Example*: + - "Marie Curie" —[born_in]→ "Warsaw" + - "Radioactivity" —[discovered_by]→ "Marie Curie" + + **Compliance**: + Strict adherence to these guidelines is required. Any deviation will result in immediate termination of the task. +"# + +// Strict template with zero-tolerance rules +template_string StrictExtractContentGraphPrompt() #" + You are a top-tier algorithm for **extracting structured information** from unstructured text to build a **knowledge graph**. + + Your primary goal is to extract: + - **Nodes**: Representing **entities** and **concepts** (like Wikipedia nodes). + - **Edges**: Representing **relationships** between those concepts (like Wikipedia links). + + The resulting knowledge graph must be **simple, consistent, and human-readable**. + + ## 1. Node Labeling and Identification + + ### Node Types + Use **basic atomic types** for node labels. Always prefer general types over specific roles or professions: + - "Person" for any human. + - "Organization" for companies, institutions, etc. + - "Location" for geographic or place entities. + - "Date" for any temporal expression. + - "Event" for historical or scheduled occurrences. + - "Work" for books, films, artworks, or research papers. + - "Concept" for abstract notions or ideas. + + ### Node IDs + - Always assign **human-readable and unambiguous identifiers**. + - Never use numeric or autogenerated IDs. + - Prioritize **most complete form** of entity names for consistency. + + ## 2. Relationship Handling + - Use **snake_case** for all relationship (edge) types. + - Keep relationship types semantically clear and consistent. + - Avoid vague relation names like "related_to" unless no better alternative exists. + + ## 3. Strict Compliance + Follow all rules exactly. Any deviation may lead to rejection or incorrect graph construction. +"# + +// OpenAI client with environment model selection +client OpenAIClientWithEnvModel { + provider openai + options { + model env.LLM_MODEL + api_key env.OPENAI_API_KEY + } +} + +// Anthropic client with environment model selection +client AnthropicClientWithEnvModel { + provider anthropic + options { + model env.LLM_MODEL + api_key env.ANTHROPIC_API_KEY + } +} + +// Default client (maintains backward compatibility) +client DefaultClient { + provider openai + options { + model "gpt-4o-mini" + api_key env.OPENAI_API_KEY + } +} + +// Function that returns raw structured output (for custom objects - to be handled in Python) +function ExtractContentGraphGeneric( + content: string, + mode: "simple" | "base" | "guided" | "strict" | "custom"?, + custom_prompt_content: string? +) -> KnowledgeGraph { + client OpenAIClientWithEnvModel + + prompt #" + {% if mode == "base" %} + {{ DetailedExtractContentGraphPrompt() }} + {% elif mode == "guided" %} + {{ GuidedExtractContentGraphPrompt() }} + {% elif mode == "strict" %} + {{ StrictExtractContentGraphPrompt() }} + {% elif mode == "custom" and custom_prompt_content %} + {{ custom_prompt_content }} + {% else %} + {{ ExtractContentGraphPrompt() }} + {% endif %} + + {{ ctx.output_format(prefix="Answer in this schema:\n") }} + + Before answering, briefly describe what you'll extract from the text, then provide the structured output. + + Example format: + I'll extract the main entities and their relationships from this text... + + { ... } + + {{ _.role('user') }} + {{ content }} + "# +} + +// Backward-compatible function specifically for KnowledgeGraph +function ExtractContentGraph( + content: string, + mode: "simple" | "base" | "guided" | "strict" | "custom"?, + custom_prompt_content: string? +) -> KnowledgeGraph { + client OpenAIClientWithEnvModel + + prompt #" + {% if mode == "base" %} + {{ DetailedExtractContentGraphPrompt() }} + {% elif mode == "guided" %} + {{ GuidedExtractContentGraphPrompt() }} + {% elif mode == "strict" %} + {{ StrictExtractContentGraphPrompt() }} + {% elif mode == "custom" and custom_prompt_content %} + {{ custom_prompt_content }} + {% else %} + {{ ExtractContentGraphPrompt() }} + {% endif %} + + {{ ctx.output_format(prefix="Answer in this schema:\n") }} + + Before answering, briefly describe what you'll extract from the text, then provide the structured output. + + Example format: + I'll extract the main entities and their relationships from this text... + + { ... } + + {{ _.role('user') }} + {{ content }} + "# +} + +// Alternative function that uses environment variable for prompt selection +function ExtractContentGraphWithEnvPrompt( + content: string, + prompt_override: string? +) -> KnowledgeGraph { + client OpenAIClientWithEnvModel + + prompt #" + {% if prompt_override %} + {{ prompt_override }} + {% else %} + {{ ExtractContentGraphPrompt() }} + {% endif %} + + {{ ctx.output_format(prefix="Answer in this schema:\n") }} + + Before answering, briefly describe what you'll extract from the text, then provide the structured output. + + Example format: + I'll extract the main entities and their relationships from this text... + + { ... } + + {{ _.role('user') }} + {{ content }} + "# +} + +// Function that uses Anthropic client +function ExtractContentGraphWithAnthropic( + content: string, + mode: "simple" | "base" | "guided" | "strict" | "custom"?, + custom_prompt_content: string? +) -> KnowledgeGraph { + client AnthropicClientWithEnvModel + + prompt #" + {% if mode == "base" %} + {{ DetailedExtractContentGraphPrompt() }} + {% elif mode == "guided" %} + {{ GuidedExtractContentGraphPrompt() }} + {% elif mode == "strict" %} + {{ StrictExtractContentGraphPrompt() }} + {% elif mode == "custom" and custom_prompt_content %} + {{ custom_prompt_content }} + {% else %} + {{ ExtractContentGraphPrompt() }} + {% endif %} + + {{ ctx.output_format(prefix="Answer in this schema:\n") }} + + Before answering, briefly describe what you'll extract from the text, then provide the structured output. + + Example format: + I'll extract the main entities and their relationships from this text... + + { ... } + + {{ _.role('user') }} + {{ content }} + "# +} + +test ExtractPersonExample { + functions [ExtractContentGraph] + args { + content #" + My name is Vasiliy. I was born in 1992. I am a software engineer. I work at Google and am based in Berlin. + "# + mode "simple" + } +} + +test ExtractGuidedExample { + functions [ExtractContentGraph] + args { + content #" + Apple Inc. was founded by Steve Jobs in 1976. The company is headquartered in Cupertino, California. + Tim Cook is the current CEO of Apple Inc. + "# + mode "guided" + } +} + +test ExtractStrictExample { + functions [ExtractContentGraph] + args { + content #" + The Python programming language was created by Guido van Rossum in 1991. + "# + mode "strict" + } +} + +test ExtractGenericExample { + functions [ExtractContentGraphGeneric] + args { + content #" + React is a JavaScript library for building user interfaces, developed by Facebook. + "# + mode "simple" + } +} diff --git a/cognee/modules/data/extraction/__init__.py b/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/__init__.py similarity index 100% rename from cognee/modules/data/extraction/__init__.py rename to cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/__init__.py diff --git a/cognee/infrastructure/llm/generic_llm_api/__init__.py b/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_categories.py similarity index 100% rename from cognee/infrastructure/llm/generic_llm_api/__init__.py rename to cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_categories.py diff --git a/cognee/infrastructure/llm/ollama/__init__.py b/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_summary.py similarity index 100% rename from cognee/infrastructure/llm/ollama/__init__.py rename to cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_summary.py diff --git a/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/knowledge_graph/extract_content_graph.py b/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/knowledge_graph/extract_content_graph.py new file mode 100644 index 000000000..d81254351 --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/knowledge_graph/extract_content_graph.py @@ -0,0 +1,18 @@ +import os +from typing import Type +from pydantic import BaseModel +from cognee.infrastructure.llm.structured_output_framework.baml.async_client import b +from cognee.infrastructure.llm.structured_output_framework.baml.type_builder import TypeBuilder +from cognee.infrastructure.llm.structured_output_framework.baml_src.config import get_llm_config + + +async def extract_content_graph(content: str, response_model: Type[BaseModel]): + # tb = TypeBuilder() + config = get_llm_config() + # country = tb.union \ + # ([tb.literal_string("USA"), tb.literal_string("UK"), tb.literal_string("Germany"), tb.literal_string("other")]) + # tb.Node.add_property("country", country) + + graph = await b.ExtractContentGraph(content, mode="simple", baml_options={ "tb": config.baml_registry}) + return graph + diff --git a/cognee/infrastructure/llm/structured_output_framework/baml_src/generators.baml b/cognee/infrastructure/llm/structured_output_framework/baml_src/generators.baml new file mode 100644 index 000000000..ac4f08efb --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/baml_src/generators.baml @@ -0,0 +1,18 @@ +// This helps use auto generate libraries you can use in the language of +// your choice. You can have multiple generators if you use multiple languages. +// Just ensure that the output_dir is different for each generator. +generator target { + // Valid values: "python/pydantic", "typescript", "ruby/sorbet", "rest/openapi" + output_type "python/pydantic" + + // Where the generated code will be saved (relative to baml_src/) + output_dir "../baml/" + + // The version of the BAML package you have installed (e.g. same version as your baml-py or @boundaryml/baml). + // The BAML VSCode extension version should also match this version. + version "0.201.0" + + // Valid values: "sync", "async" + // This controls what `b.FunctionName()` will be (sync or async). + default_client_mode sync +} diff --git a/cognee/infrastructure/llm/openai/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/__init__.py similarity index 100% rename from cognee/infrastructure/llm/openai/__init__.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/__init__.py diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/__init__.py new file mode 100644 index 000000000..b6419282d --- /dev/null +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/__init__.py @@ -0,0 +1 @@ +from .knowledge_graph.extract_content_graph import extract_content_graph diff --git a/cognee/modules/data/extraction/extract_categories.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_categories.py similarity index 56% rename from cognee/modules/data/extraction/extract_categories.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_categories.py index 3de26e59d..8b7e3db6d 100644 --- a/cognee/modules/data/extraction/extract_categories.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_categories.py @@ -1,7 +1,7 @@ -from typing import Type, List +from typing import Type from pydantic import BaseModel -from cognee.infrastructure.llm.prompts import read_query_prompt -from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import read_query_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client async def extract_categories(content: str, response_model: Type[BaseModel]): diff --git a/cognee/modules/data/extraction/extract_summary.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_summary.py similarity index 84% rename from cognee/modules/data/extraction/extract_summary.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_summary.py index 3e1195434..6935e9c44 100644 --- a/cognee/modules/data/extraction/extract_summary.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_summary.py @@ -5,8 +5,8 @@ from typing import Type from instructor.exceptions import InstructorRetryException from pydantic import BaseModel -from cognee.infrastructure.llm.get_llm_client import get_llm_client -from cognee.infrastructure.llm.prompts import read_query_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import read_query_prompt from cognee.shared.data_models import SummarizedCode from cognee.tasks.summarization.mock_summary import get_mock_summarized_code diff --git a/cognee/modules/data/extraction/knowledge_graph/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/knowledge_graph/__init__.py similarity index 100% rename from cognee/modules/data/extraction/knowledge_graph/__init__.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/knowledge_graph/__init__.py diff --git a/cognee/modules/data/extraction/knowledge_graph/extract_content_graph.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/knowledge_graph/extract_content_graph.py similarity index 70% rename from cognee/modules/data/extraction/knowledge_graph/extract_content_graph.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/knowledge_graph/extract_content_graph.py index f477a4b72..d25afef76 100644 --- a/cognee/modules/data/extraction/knowledge_graph/extract_content_graph.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/knowledge_graph/extract_content_graph.py @@ -1,9 +1,9 @@ import os from typing import Type from pydantic import BaseModel -from cognee.infrastructure.llm.get_llm_client import get_llm_client -from cognee.infrastructure.llm.prompts import render_prompt -from cognee.infrastructure.llm.config import get_llm_config +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import render_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.config import get_llm_config async def extract_content_graph(content: str, response_model: Type[BaseModel]): diff --git a/cognee/modules/data/extraction/texts.json b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/texts.json similarity index 100% rename from cognee/modules/data/extraction/texts.json rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/texts.json diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/anthropic/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/anthropic/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cognee/infrastructure/llm/anthropic/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/anthropic/adapter.py similarity index 87% rename from cognee/infrastructure/llm/anthropic/adapter.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/anthropic/adapter.py index ec2f35f7b..7b9710675 100644 --- a/cognee/infrastructure/llm/anthropic/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/anthropic/adapter.py @@ -3,9 +3,9 @@ from pydantic import BaseModel import instructor from cognee.exceptions import InvalidValueError -from cognee.infrastructure.llm.llm_interface import LLMInterface -from cognee.infrastructure.llm.prompts import read_query_prompt -from cognee.infrastructure.llm.rate_limiter import rate_limit_async, sleep_and_retry_async +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.llm_interface import LLMInterface +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import read_query_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.rate_limiter import rate_limit_async, sleep_and_retry_async class AnthropicAdapter(LLMInterface): diff --git a/cognee/infrastructure/llm/config.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/config.py similarity index 100% rename from cognee/infrastructure/llm/config.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/config.py diff --git a/cognee/infrastructure/llm/embedding_rate_limiter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/embedding_rate_limiter.py similarity index 99% rename from cognee/infrastructure/llm/embedding_rate_limiter.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/embedding_rate_limiter.py index 24312dab1..0aed181e4 100644 --- a/cognee/infrastructure/llm/embedding_rate_limiter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/embedding_rate_limiter.py @@ -6,7 +6,7 @@ import time import asyncio import random from cognee.shared.logging_utils import get_logger -from cognee.infrastructure.llm.config import get_llm_config +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.config import get_llm_config logger = get_logger() diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/gemini/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/gemini/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cognee/infrastructure/llm/gemini/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/gemini/adapter.py similarity index 93% rename from cognee/infrastructure/llm/gemini/adapter.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/gemini/adapter.py index db11a5ab4..e2f3b16e4 100644 --- a/cognee/infrastructure/llm/gemini/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/gemini/adapter.py @@ -1,5 +1,4 @@ import litellm -import logging from pydantic import BaseModel from typing import Type, Optional from litellm import acompletion, JSONSchemaValidationError @@ -7,9 +6,9 @@ from litellm import acompletion, JSONSchemaValidationError from cognee.shared.logging_utils import get_logger from cognee.modules.observability.get_observe import get_observe from cognee.exceptions import InvalidValueError -from cognee.infrastructure.llm.llm_interface import LLMInterface -from cognee.infrastructure.llm.prompts import read_query_prompt -from cognee.infrastructure.llm.rate_limiter import ( +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.llm_interface import LLMInterface +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import read_query_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.rate_limiter import ( rate_limit_async, sleep_and_retry_async, ) diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/generic_llm_api/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/generic_llm_api/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cognee/infrastructure/llm/generic_llm_api/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/generic_llm_api/adapter.py similarity index 90% rename from cognee/infrastructure/llm/generic_llm_api/adapter.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/generic_llm_api/adapter.py index e74a4eb03..3a1a7b2c3 100644 --- a/cognee/infrastructure/llm/generic_llm_api/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/generic_llm_api/adapter.py @@ -1,14 +1,11 @@ """Adapter for Generic API LLM provider API""" -import logging from typing import Type from pydantic import BaseModel import instructor -from cognee.infrastructure.llm.llm_interface import LLMInterface -from cognee.infrastructure.llm.config import get_llm_config -from cognee.infrastructure.llm.rate_limiter import rate_limit_async, sleep_and_retry_async -from cognee.shared.logging_utils import get_logger +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.llm_interface import LLMInterface +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.rate_limiter import rate_limit_async, sleep_and_retry_async import litellm diff --git a/cognee/infrastructure/llm/get_llm_client.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/get_llm_client.py similarity index 80% rename from cognee/infrastructure/llm/get_llm_client.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/get_llm_client.py index a5a0b1f1e..4738aa71d 100644 --- a/cognee/infrastructure/llm/get_llm_client.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/get_llm_client.py @@ -4,7 +4,7 @@ from enum import Enum from cognee.exceptions import InvalidValueError from cognee.infrastructure.llm import get_llm_config -from cognee.infrastructure.llm.ollama.adapter import OllamaAPIAdapter +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.ollama.adapter import OllamaAPIAdapter # Define an Enum for LLM Providers @@ -48,7 +48,7 @@ def get_llm_client(): # Check if max_token value is defined in liteLLM for given model # if not use value from cognee configuration - from cognee.infrastructure.llm.utils import ( + from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.utils import ( get_model_max_tokens, ) # imported here to avoid circular imports @@ -59,7 +59,7 @@ def get_llm_client(): if llm_config.llm_api_key is None: raise InvalidValueError(message="LLM API key is not set.") - from .openai.adapter import OpenAIAdapter + from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.openai.adapter import OpenAIAdapter return OpenAIAdapter( api_key=llm_config.llm_api_key, @@ -75,7 +75,7 @@ def get_llm_client(): if llm_config.llm_api_key is None: raise InvalidValueError(message="LLM API key is not set.") - from .generic_llm_api.adapter import GenericAPIAdapter + from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.generic_llm_api import GenericAPIAdapter return OllamaAPIAdapter( llm_config.llm_endpoint, @@ -86,7 +86,7 @@ def get_llm_client(): ) elif provider == LLMProvider.ANTHROPIC: - from .anthropic.adapter import AnthropicAdapter + from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.anthropic import AnthropicAdapter return AnthropicAdapter(max_tokens=max_tokens, model=llm_config.llm_model) @@ -94,7 +94,7 @@ def get_llm_client(): if llm_config.llm_api_key is None: raise InvalidValueError(message="LLM API key is not set.") - from .generic_llm_api.adapter import GenericAPIAdapter + from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.generic_llm_api import GenericAPIAdapter return GenericAPIAdapter( llm_config.llm_endpoint, @@ -108,7 +108,7 @@ def get_llm_client(): if llm_config.llm_api_key is None: raise InvalidValueError(message="LLM API key is not set.") - from .gemini.adapter import GeminiAdapter + from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.gemini import GeminiAdapter return GeminiAdapter( api_key=llm_config.llm_api_key, diff --git a/cognee/infrastructure/llm/llm_interface.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/llm_interface.py similarity index 94% rename from cognee/infrastructure/llm/llm_interface.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/llm_interface.py index 609bb0a9f..e5b4b5f0d 100644 --- a/cognee/infrastructure/llm/llm_interface.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/llm_interface.py @@ -3,7 +3,7 @@ from typing import Type, Protocol from abc import abstractmethod from pydantic import BaseModel -from cognee.infrastructure.llm.prompts import read_query_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import read_query_prompt class LLMInterface(Protocol): diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/ollama/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/ollama/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cognee/infrastructure/llm/ollama/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/ollama/adapter.py similarity index 96% rename from cognee/infrastructure/llm/ollama/adapter.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/ollama/adapter.py index 82f50ede5..fc2d35431 100644 --- a/cognee/infrastructure/llm/ollama/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/ollama/adapter.py @@ -1,9 +1,8 @@ from typing import Type from pydantic import BaseModel import instructor -from cognee.infrastructure.llm.llm_interface import LLMInterface -from cognee.infrastructure.llm.config import get_llm_config -from cognee.infrastructure.llm.rate_limiter import ( +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.llm_interface import LLMInterface +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.rate_limiter import ( rate_limit_async, rate_limit_sync, sleep_and_retry_async, diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/openai/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/openai/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cognee/infrastructure/llm/openai/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/openai/adapter.py similarity index 96% rename from cognee/infrastructure/llm/openai/adapter.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/openai/adapter.py index a816e3e3f..35c0a5101 100644 --- a/cognee/infrastructure/llm/openai/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/openai/adapter.py @@ -7,9 +7,9 @@ from pydantic import BaseModel from cognee.modules.data.processing.document_types.open_data_file import open_data_file from cognee.exceptions import InvalidValueError -from cognee.infrastructure.llm.llm_interface import LLMInterface -from cognee.infrastructure.llm.prompts import read_query_prompt -from cognee.infrastructure.llm.rate_limiter import ( +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.llm_interface import LLMInterface +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import read_query_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.rate_limiter import ( rate_limit_async, rate_limit_sync, sleep_and_retry_async, diff --git a/cognee/infrastructure/llm/prompts/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/__init__.py similarity index 100% rename from cognee/infrastructure/llm/prompts/__init__.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/__init__.py diff --git a/cognee/infrastructure/llm/prompts/answer_hotpot_question.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_hotpot_question.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/answer_hotpot_question.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_hotpot_question.txt diff --git a/cognee/infrastructure/llm/prompts/answer_hotpot_using_cognee_search.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_hotpot_using_cognee_search.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/answer_hotpot_using_cognee_search.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_hotpot_using_cognee_search.txt diff --git a/cognee/infrastructure/llm/prompts/answer_simple_question.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_simple_question.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/answer_simple_question.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_simple_question.txt diff --git a/cognee/infrastructure/llm/prompts/answer_simple_question_benchmark.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_simple_question_benchmark.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/answer_simple_question_benchmark.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_simple_question_benchmark.txt diff --git a/cognee/infrastructure/llm/prompts/answer_simple_question_benchmark2.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_simple_question_benchmark2.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/answer_simple_question_benchmark2.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_simple_question_benchmark2.txt diff --git a/cognee/infrastructure/llm/prompts/answer_simple_question_benchmark3.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_simple_question_benchmark3.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/answer_simple_question_benchmark3.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_simple_question_benchmark3.txt diff --git a/cognee/infrastructure/llm/prompts/answer_simple_question_benchmark4.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_simple_question_benchmark4.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/answer_simple_question_benchmark4.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_simple_question_benchmark4.txt diff --git a/cognee/infrastructure/llm/prompts/answer_simple_question_restricted.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_simple_question_restricted.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/answer_simple_question_restricted.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/answer_simple_question_restricted.txt diff --git a/cognee/infrastructure/llm/prompts/categorize_categories.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/categorize_categories.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/categorize_categories.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/categorize_categories.txt diff --git a/cognee/infrastructure/llm/prompts/categorize_summary.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/categorize_summary.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/categorize_summary.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/categorize_summary.txt diff --git a/cognee/infrastructure/llm/prompts/classify_content.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/classify_content.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/classify_content.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/classify_content.txt diff --git a/cognee/infrastructure/llm/prompts/codegraph_retriever_system.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/codegraph_retriever_system.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/codegraph_retriever_system.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/codegraph_retriever_system.txt diff --git a/cognee/infrastructure/llm/prompts/coding_rule_association_agent_system.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/coding_rule_association_agent_system.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/coding_rule_association_agent_system.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/coding_rule_association_agent_system.txt diff --git a/cognee/infrastructure/llm/prompts/coding_rule_association_agent_user.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/coding_rule_association_agent_user.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/coding_rule_association_agent_user.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/coding_rule_association_agent_user.txt diff --git a/cognee/infrastructure/llm/prompts/context_for_question.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/context_for_question.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/context_for_question.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/context_for_question.txt diff --git a/cognee/infrastructure/llm/prompts/cot_followup_system_prompt.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/cot_followup_system_prompt.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/cot_followup_system_prompt.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/cot_followup_system_prompt.txt diff --git a/cognee/infrastructure/llm/prompts/cot_followup_user_prompt.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/cot_followup_user_prompt.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/cot_followup_user_prompt.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/cot_followup_user_prompt.txt diff --git a/cognee/infrastructure/llm/prompts/cot_validation_system_prompt.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/cot_validation_system_prompt.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/cot_validation_system_prompt.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/cot_validation_system_prompt.txt diff --git a/cognee/infrastructure/llm/prompts/cot_validation_user_prompt.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/cot_validation_user_prompt.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/cot_validation_user_prompt.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/cot_validation_user_prompt.txt diff --git a/cognee/infrastructure/llm/prompts/direct_llm_eval_prompt.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/direct_llm_eval_prompt.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/direct_llm_eval_prompt.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/direct_llm_eval_prompt.txt diff --git a/cognee/infrastructure/llm/prompts/direct_llm_eval_system.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/direct_llm_eval_system.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/direct_llm_eval_system.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/direct_llm_eval_system.txt diff --git a/cognee/infrastructure/llm/prompts/extract_entities_system.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/extract_entities_system.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/extract_entities_system.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/extract_entities_system.txt diff --git a/cognee/infrastructure/llm/prompts/extract_entities_user.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/extract_entities_user.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/extract_entities_user.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/extract_entities_user.txt diff --git a/cognee/infrastructure/llm/prompts/extract_ontology.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/extract_ontology.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/extract_ontology.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/extract_ontology.txt diff --git a/cognee/infrastructure/llm/prompts/generate_graph_prompt.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/generate_graph_prompt.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/generate_graph_prompt.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/generate_graph_prompt.txt diff --git a/cognee/infrastructure/llm/prompts/generate_graph_prompt_guided.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/generate_graph_prompt_guided.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/generate_graph_prompt_guided.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/generate_graph_prompt_guided.txt diff --git a/cognee/infrastructure/llm/prompts/generate_graph_prompt_oneshot.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/generate_graph_prompt_oneshot.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/generate_graph_prompt_oneshot.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/generate_graph_prompt_oneshot.txt diff --git a/cognee/infrastructure/llm/prompts/generate_graph_prompt_simple.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/generate_graph_prompt_simple.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/generate_graph_prompt_simple.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/generate_graph_prompt_simple.txt diff --git a/cognee/infrastructure/llm/prompts/generate_graph_prompt_strict.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/generate_graph_prompt_strict.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/generate_graph_prompt_strict.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/generate_graph_prompt_strict.txt diff --git a/cognee/infrastructure/llm/prompts/graph_context_for_question.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/graph_context_for_question.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/graph_context_for_question.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/graph_context_for_question.txt diff --git a/cognee/infrastructure/llm/prompts/llm_judge_prompts.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/llm_judge_prompts.py similarity index 100% rename from cognee/infrastructure/llm/prompts/llm_judge_prompts.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/llm_judge_prompts.py diff --git a/cognee/infrastructure/llm/prompts/natural_language_retriever_system.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/natural_language_retriever_system.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/natural_language_retriever_system.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/natural_language_retriever_system.txt diff --git a/cognee/infrastructure/llm/prompts/patch_gen_instructions.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/patch_gen_instructions.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/patch_gen_instructions.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/patch_gen_instructions.txt diff --git a/cognee/infrastructure/llm/prompts/patch_gen_kg_instructions.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/patch_gen_kg_instructions.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/patch_gen_kg_instructions.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/patch_gen_kg_instructions.txt diff --git a/cognee/infrastructure/llm/prompts/read_query_prompt.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/read_query_prompt.py similarity index 100% rename from cognee/infrastructure/llm/prompts/read_query_prompt.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/read_query_prompt.py diff --git a/cognee/infrastructure/llm/prompts/render_prompt.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/render_prompt.py similarity index 100% rename from cognee/infrastructure/llm/prompts/render_prompt.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/render_prompt.py diff --git a/cognee/infrastructure/llm/prompts/summarize_code.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/summarize_code.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/summarize_code.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/summarize_code.txt diff --git a/cognee/infrastructure/llm/prompts/summarize_content.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/summarize_content.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/summarize_content.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/summarize_content.txt diff --git a/cognee/infrastructure/llm/prompts/summarize_search_results.txt b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/summarize_search_results.txt similarity index 100% rename from cognee/infrastructure/llm/prompts/summarize_search_results.txt rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/prompts/summarize_search_results.txt diff --git a/cognee/infrastructure/llm/rate_limiter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/rate_limiter.py similarity index 99% rename from cognee/infrastructure/llm/rate_limiter.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/rate_limiter.py index bfc9bcc09..9e8e6012a 100644 --- a/cognee/infrastructure/llm/rate_limiter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/rate_limiter.py @@ -49,12 +49,7 @@ from functools import wraps from limits import RateLimitItemPerMinute, storage from limits.strategies import MovingWindowRateLimiter from cognee.shared.logging_utils import get_logger -from cognee.infrastructure.llm.config import get_llm_config -import threading -import logging -import functools -import openai -import os +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.config import get_llm_config logger = get_logger() diff --git a/cognee/infrastructure/llm/tokenizer/Gemini/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/Gemini/__init__.py similarity index 100% rename from cognee/infrastructure/llm/tokenizer/Gemini/__init__.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/Gemini/__init__.py diff --git a/cognee/infrastructure/llm/tokenizer/Gemini/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/Gemini/adapter.py similarity index 93% rename from cognee/infrastructure/llm/tokenizer/Gemini/adapter.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/Gemini/adapter.py index 1b7e5a6e6..02419f3ab 100644 --- a/cognee/infrastructure/llm/tokenizer/Gemini/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/Gemini/adapter.py @@ -1,4 +1,4 @@ -from typing import List, Any, Union +from typing import List, Any from ..tokenizer_interface import TokenizerInterface @@ -24,7 +24,7 @@ class GeminiTokenizer(TokenizerInterface): # Get LLM API key from config from cognee.infrastructure.databases.vector.embeddings.config import get_embedding_config - from cognee.infrastructure.llm.config import get_llm_config + from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.config import get_llm_config config = get_embedding_config() llm_config = get_llm_config() diff --git a/cognee/infrastructure/llm/tokenizer/HuggingFace/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/HuggingFace/__init__.py similarity index 100% rename from cognee/infrastructure/llm/tokenizer/HuggingFace/__init__.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/HuggingFace/__init__.py diff --git a/cognee/infrastructure/llm/tokenizer/HuggingFace/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/HuggingFace/adapter.py similarity index 100% rename from cognee/infrastructure/llm/tokenizer/HuggingFace/adapter.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/HuggingFace/adapter.py diff --git a/cognee/infrastructure/llm/tokenizer/Mistral/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/Mistral/__init__.py similarity index 100% rename from cognee/infrastructure/llm/tokenizer/Mistral/__init__.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/Mistral/__init__.py diff --git a/cognee/infrastructure/llm/tokenizer/Mistral/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/Mistral/adapter.py similarity index 100% rename from cognee/infrastructure/llm/tokenizer/Mistral/adapter.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/Mistral/adapter.py diff --git a/cognee/infrastructure/llm/tokenizer/TikToken/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/TikToken/__init__.py similarity index 100% rename from cognee/infrastructure/llm/tokenizer/TikToken/__init__.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/TikToken/__init__.py diff --git a/cognee/infrastructure/llm/tokenizer/TikToken/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/TikToken/adapter.py similarity index 100% rename from cognee/infrastructure/llm/tokenizer/TikToken/adapter.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/TikToken/adapter.py diff --git a/cognee/infrastructure/llm/tokenizer/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/__init__.py similarity index 100% rename from cognee/infrastructure/llm/tokenizer/__init__.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/__init__.py diff --git a/cognee/infrastructure/llm/tokenizer/tokenizer_interface.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/tokenizer_interface.py similarity index 100% rename from cognee/infrastructure/llm/tokenizer/tokenizer_interface.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/tokenizer_interface.py diff --git a/cognee/infrastructure/llm/utils.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/utils.py similarity index 96% rename from cognee/infrastructure/llm/utils.py rename to cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/utils.py index fdc8c521c..9f6e92d5c 100644 --- a/cognee/infrastructure/llm/utils.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/utils.py @@ -1,6 +1,6 @@ import litellm -from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client from cognee.shared.logging_utils import get_logger logger = get_logger() diff --git a/cognee/modules/data/extraction/knowledge_graph/add_model_class_to_graph.py b/cognee/modules/data/methods/add_model_class_to_graph.py similarity index 100% rename from cognee/modules/data/extraction/knowledge_graph/add_model_class_to_graph.py rename to cognee/modules/data/methods/add_model_class_to_graph.py diff --git a/cognee/modules/data/processing/document_types/AudioDocument.py b/cognee/modules/data/processing/document_types/AudioDocument.py index a254d85e8..763061c8c 100644 --- a/cognee/modules/data/processing/document_types/AudioDocument.py +++ b/cognee/modules/data/processing/document_types/AudioDocument.py @@ -1,4 +1,4 @@ -from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client from cognee.modules.chunking.Chunker import Chunker from .Document import Document diff --git a/cognee/modules/data/processing/document_types/ImageDocument.py b/cognee/modules/data/processing/document_types/ImageDocument.py index ae288c34c..c8400a7fc 100644 --- a/cognee/modules/data/processing/document_types/ImageDocument.py +++ b/cognee/modules/data/processing/document_types/ImageDocument.py @@ -1,4 +1,4 @@ -from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client from cognee.modules.chunking.Chunker import Chunker from .Document import Document diff --git a/cognee/modules/pipelines/operations/pipeline.py b/cognee/modules/pipelines/operations/pipeline.py index ab2ccef67..03945f37a 100644 --- a/cognee/modules/pipelines/operations/pipeline.py +++ b/cognee/modules/pipelines/operations/pipeline.py @@ -1,6 +1,6 @@ import asyncio from typing import Union -from uuid import NAMESPACE_OID, uuid5, UUID +from uuid import UUID from cognee.shared.logging_utils import get_logger from cognee.modules.data.methods.get_dataset_data import get_dataset_data @@ -69,7 +69,7 @@ async def cognee_pipeline( cognee_pipeline.first_run = True if cognee_pipeline.first_run: - from cognee.infrastructure.llm.utils import test_llm_connection, test_embedding_connection + from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.utils import test_llm_connection, test_embedding_connection # Test LLM and Embedding configuration once before running Cognee await test_llm_connection() diff --git a/cognee/modules/retrieval/code_retriever.py b/cognee/modules/retrieval/code_retriever.py index 185ba3e8b..33c8254a0 100644 --- a/cognee/modules/retrieval/code_retriever.py +++ b/cognee/modules/retrieval/code_retriever.py @@ -6,8 +6,8 @@ from pydantic import BaseModel from cognee.modules.retrieval.base_retriever import BaseRetriever from cognee.infrastructure.databases.graph import get_graph_engine from cognee.infrastructure.databases.vector import get_vector_engine -from cognee.infrastructure.llm.get_llm_client import get_llm_client -from cognee.infrastructure.llm.prompts import read_query_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import read_query_prompt class CodeRetriever(BaseRetriever): diff --git a/cognee/modules/retrieval/context_providers/TripletSearchContextProvider.py b/cognee/modules/retrieval/context_providers/TripletSearchContextProvider.py index e8b08c9bd..ac29231ee 100644 --- a/cognee/modules/retrieval/context_providers/TripletSearchContextProvider.py +++ b/cognee/modules/retrieval/context_providers/TripletSearchContextProvider.py @@ -4,8 +4,6 @@ import asyncio from cognee.infrastructure.context.BaseContextProvider import BaseContextProvider from cognee.infrastructure.engine import DataPoint from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph -from cognee.infrastructure.llm.get_llm_client import get_llm_client -from cognee.infrastructure.llm.prompts import read_query_prompt from cognee.modules.retrieval.utils.brute_force_triplet_search import ( brute_force_triplet_search, format_triplets, diff --git a/cognee/modules/retrieval/graph_completion_context_extension_retriever.py b/cognee/modules/retrieval/graph_completion_context_extension_retriever.py index 2479a454d..4027646c1 100644 --- a/cognee/modules/retrieval/graph_completion_context_extension_retriever.py +++ b/cognee/modules/retrieval/graph_completion_context_extension_retriever.py @@ -1,9 +1,7 @@ from typing import Any, Optional, List, Type from cognee.shared.logging_utils import get_logger -from cognee.infrastructure.llm.get_llm_client import get_llm_client from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever from cognee.modules.retrieval.utils.completion import generate_completion -from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt logger = get_logger() diff --git a/cognee/modules/retrieval/graph_completion_cot_retriever.py b/cognee/modules/retrieval/graph_completion_cot_retriever.py index 95eb1a9b6..c61f72d2c 100644 --- a/cognee/modules/retrieval/graph_completion_cot_retriever.py +++ b/cognee/modules/retrieval/graph_completion_cot_retriever.py @@ -1,9 +1,9 @@ from typing import Any, Optional, List, Type from cognee.shared.logging_utils import get_logger -from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever from cognee.modules.retrieval.utils.completion import generate_completion -from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import read_query_prompt, render_prompt logger = get_logger() diff --git a/cognee/modules/retrieval/natural_language_retriever.py b/cognee/modules/retrieval/natural_language_retriever.py index 9211c286b..f0afddc19 100644 --- a/cognee/modules/retrieval/natural_language_retriever.py +++ b/cognee/modules/retrieval/natural_language_retriever.py @@ -2,8 +2,8 @@ from typing import Any, Optional import logging from cognee.infrastructure.databases.graph import get_graph_engine from cognee.infrastructure.databases.graph.networkx.adapter import NetworkXAdapter -from cognee.infrastructure.llm.get_llm_client import get_llm_client -from cognee.infrastructure.llm.prompts import render_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import render_prompt from cognee.modules.retrieval.base_retriever import BaseRetriever from cognee.modules.retrieval.exceptions import SearchTypeNotSupported from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface diff --git a/cognee/modules/retrieval/utils/completion.py b/cognee/modules/retrieval/utils/completion.py index ed4bdcefc..7613eccc8 100644 --- a/cognee/modules/retrieval/utils/completion.py +++ b/cognee/modules/retrieval/utils/completion.py @@ -1,5 +1,5 @@ -from cognee.infrastructure.llm.get_llm_client import get_llm_client -from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import read_query_prompt, render_prompt async def generate_completion( diff --git a/cognee/modules/retrieval/utils/description_to_codepart_search.py b/cognee/modules/retrieval/utils/description_to_codepart_search.py index 63256891d..57804dbd4 100644 --- a/cognee/modules/retrieval/utils/description_to_codepart_search.py +++ b/cognee/modules/retrieval/utils/description_to_codepart_search.py @@ -9,7 +9,7 @@ from cognee.modules.users.methods import get_default_user from cognee.modules.users.models import User from cognee.shared.utils import send_telemetry from cognee.modules.search.methods import search -from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client logger = get_logger(level=ERROR) diff --git a/cognee/shared/data_models.py b/cognee/shared/data_models.py index fb9386dcf..8adced97e 100644 --- a/cognee/shared/data_models.py +++ b/cognee/shared/data_models.py @@ -4,7 +4,7 @@ from enum import Enum, auto from typing import Any, Dict, List, Optional, Union from pydantic import BaseModel, Field -from cognee.infrastructure.llm.config import get_llm_config +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.config import get_llm_config if get_llm_config().llm_provider.lower() == "gemini": """ diff --git a/cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py b/cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py index c9d7bf1a1..bba413d69 100644 --- a/cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py +++ b/cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py @@ -7,7 +7,7 @@ from pydantic import BaseModel from cognee.infrastructure.databases.graph import get_graph_engine from cognee.infrastructure.databases.vector import get_vector_engine from cognee.infrastructure.engine.models import DataPoint -from cognee.modules.data.extraction.extract_categories import extract_categories +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.extraction.extract_categories import extract_categories from cognee.modules.chunking.models.DocumentChunk import DocumentChunk diff --git a/cognee/tasks/entity_completion/entity_extractors/llm_entity_extractor.py b/cognee/tasks/entity_completion/entity_extractors/llm_entity_extractor.py index 1512d01d5..1ec22a02d 100644 --- a/cognee/tasks/entity_completion/entity_extractors/llm_entity_extractor.py +++ b/cognee/tasks/entity_completion/entity_extractors/llm_entity_extractor.py @@ -6,8 +6,8 @@ from pydantic import BaseModel from cognee.infrastructure.entities.BaseEntityExtractor import BaseEntityExtractor from cognee.modules.engine.models import Entity from cognee.modules.engine.models.EntityType import EntityType -from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt -from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import read_query_prompt, render_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client logger = get_logger("llm_entity_extractor") diff --git a/cognee/tasks/graph/cascade_extract/utils/extract_content_nodes_and_relationship_names.py b/cognee/tasks/graph/cascade_extract/utils/extract_content_nodes_and_relationship_names.py index 731c35b37..ca9f40e81 100644 --- a/cognee/tasks/graph/cascade_extract/utils/extract_content_nodes_and_relationship_names.py +++ b/cognee/tasks/graph/cascade_extract/utils/extract_content_nodes_and_relationship_names.py @@ -1,8 +1,8 @@ from typing import List, Tuple from pydantic import BaseModel -from cognee.infrastructure.llm.get_llm_client import get_llm_client -from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import render_prompt, read_query_prompt from cognee.root_dir import get_absolute_path diff --git a/cognee/tasks/graph/cascade_extract/utils/extract_edge_triplets.py b/cognee/tasks/graph/cascade_extract/utils/extract_edge_triplets.py index 3a382e5df..7c7d0819e 100644 --- a/cognee/tasks/graph/cascade_extract/utils/extract_edge_triplets.py +++ b/cognee/tasks/graph/cascade_extract/utils/extract_edge_triplets.py @@ -1,6 +1,6 @@ -from typing import List, Tuple -from cognee.infrastructure.llm.get_llm_client import get_llm_client -from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt +from typing import List +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import render_prompt, read_query_prompt from cognee.shared.data_models import KnowledgeGraph from cognee.root_dir import get_absolute_path diff --git a/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py b/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py index 6946d5ae8..e5749ba50 100644 --- a/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py +++ b/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py @@ -1,9 +1,8 @@ from typing import List from pydantic import BaseModel -from cognee.modules.chunking.models.DocumentChunk import DocumentChunk -from cognee.infrastructure.llm.get_llm_client import get_llm_client -from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import render_prompt, read_query_prompt from cognee.root_dir import get_absolute_path diff --git a/cognee/tasks/graph/extract_graph_from_code.py b/cognee/tasks/graph/extract_graph_from_code.py index 42375a683..d38572b30 100644 --- a/cognee/tasks/graph/extract_graph_from_code.py +++ b/cognee/tasks/graph/extract_graph_from_code.py @@ -1,10 +1,17 @@ import asyncio from typing import Type, List from pydantic import BaseModel -from cognee.modules.data.extraction.knowledge_graph import extract_content_graph +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.extraction import extract_content_graph from cognee.modules.chunking.models.DocumentChunk import DocumentChunk from cognee.tasks.storage import add_data_points +from cognee.base_config import get_base_config +base = get_base_config() + +if base.structured_output_framework == 'BAML': + from cognee.infrastructure.llm.structured_output_framework.baml_src.extraction import extract_content_graph +else: + from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.extraction import extract_content_graph async def extract_graph_from_code( data_chunks: list[DocumentChunk], graph_model: Type[BaseModel] diff --git a/cognee/tasks/graph/extract_graph_from_data.py b/cognee/tasks/graph/extract_graph_from_data.py index 01d8bb618..707601a99 100644 --- a/cognee/tasks/graph/extract_graph_from_data.py +++ b/cognee/tasks/graph/extract_graph_from_data.py @@ -1,12 +1,18 @@ import asyncio -from typing import Type, List, Optional +from typing import Type, List from pydantic import BaseModel from cognee.infrastructure.databases.graph import get_graph_engine from cognee.modules.ontology.rdf_xml.OntologyResolver import OntologyResolver from cognee.modules.chunking.models.DocumentChunk import DocumentChunk -from cognee.modules.data.extraction.knowledge_graph import extract_content_graph +from cognee.base_config import get_base_config +base = get_base_config() +if base.structured_output_framework == 'BAML': + from cognee.infrastructure.llm.structured_output_framework.baml_src.extraction import extract_content_graph +else: + from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.extraction import extract_content_graph + from cognee.modules.graph.utils import ( expand_with_nodes_and_edges, retrieve_existing_edges, diff --git a/cognee/tasks/graph/infer_data_ontology.py b/cognee/tasks/graph/infer_data_ontology.py index 23ede6656..6c4e6a8b1 100644 --- a/cognee/tasks/graph/infer_data_ontology.py +++ b/cognee/tasks/graph/infer_data_ontology.py @@ -13,16 +13,16 @@ import aiofiles import pandas as pd from pydantic import BaseModel -from cognee.modules.graph.exceptions import EntityNotFoundError, EntityAlreadyExistsError +from cognee.modules.graph.exceptions import EntityNotFoundError from cognee.modules.ingestion.exceptions import IngestionError -from cognee.infrastructure.llm.prompts import read_query_prompt -from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import read_query_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client from cognee.infrastructure.data.chunking.config import get_chunk_config from cognee.infrastructure.data.chunking.get_chunking_engine import get_chunk_engine from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine from cognee.infrastructure.files.utils.extract_text_from_file import extract_text_from_file from cognee.infrastructure.files.utils.guess_file_type import guess_file_type, FileTypeException -from cognee.modules.data.extraction.knowledge_graph.add_model_class_to_graph import ( +from cognee.modules.data.methods.add_model_class_to_graph import ( add_model_class_to_graph, ) from cognee.tasks.graph.models import NodeModel, GraphOntology diff --git a/cognee/tasks/summarization/summarize_code.py b/cognee/tasks/summarization/summarize_code.py index 9efc5b6ca..a3807152f 100644 --- a/cognee/tasks/summarization/summarize_code.py +++ b/cognee/tasks/summarization/summarize_code.py @@ -3,7 +3,7 @@ from typing import AsyncGenerator, Union from uuid import uuid5 from cognee.infrastructure.engine import DataPoint -from cognee.modules.data.extraction.extract_summary import extract_code_summary +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.extraction import extract_code_summary from .models import CodeSummary diff --git a/cognee/tasks/summarization/summarize_text.py b/cognee/tasks/summarization/summarize_text.py index 9a8b7cbd7..8aff298b4 100644 --- a/cognee/tasks/summarization/summarize_text.py +++ b/cognee/tasks/summarization/summarize_text.py @@ -2,7 +2,7 @@ import asyncio from typing import Type from uuid import uuid5 from pydantic import BaseModel -from cognee.modules.data.extraction.extract_summary import extract_summary +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.extraction import extract_summary from cognee.modules.chunking.models.DocumentChunk import DocumentChunk from cognee.modules.cognify.config import get_cognify_config from .models import TextSummary diff --git a/cognee/tests/unit/infrastructure/databases/test_rate_limiter.py b/cognee/tests/unit/infrastructure/databases/test_rate_limiter.py index 1e9dbe055..2ab8e4dac 100644 --- a/cognee/tests/unit/infrastructure/databases/test_rate_limiter.py +++ b/cognee/tests/unit/infrastructure/databases/test_rate_limiter.py @@ -1,10 +1,8 @@ """Tests for the LLM rate limiter.""" import pytest -import asyncio -import time from unittest.mock import patch -from cognee.infrastructure.llm.rate_limiter import ( +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.rate_limiter import ( llm_rate_limiter, rate_limit_async, rate_limit_sync, diff --git a/cognee/tests/unit/infrastructure/mock_embedding_engine.py b/cognee/tests/unit/infrastructure/mock_embedding_engine.py index a94fafab3..e600430bc 100644 --- a/cognee/tests/unit/infrastructure/mock_embedding_engine.py +++ b/cognee/tests/unit/infrastructure/mock_embedding_engine.py @@ -4,7 +4,7 @@ from typing import List from cognee.infrastructure.databases.vector.embeddings.LiteLLMEmbeddingEngine import ( LiteLLMEmbeddingEngine, ) -from cognee.infrastructure.llm.embedding_rate_limiter import ( +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.embedding_rate_limiter import ( embedding_rate_limit_async, embedding_sleep_and_retry_async, ) diff --git a/cognee/tests/unit/infrastructure/test_embedding_rate_limiting_realistic.py b/cognee/tests/unit/infrastructure/test_embedding_rate_limiting_realistic.py index a2d61308a..9426b6d69 100644 --- a/cognee/tests/unit/infrastructure/test_embedding_rate_limiting_realistic.py +++ b/cognee/tests/unit/infrastructure/test_embedding_rate_limiting_realistic.py @@ -1,14 +1,10 @@ import os import time import asyncio -from functools import lru_cache import logging -from cognee.infrastructure.llm.config import LLMConfig, get_llm_config -from cognee.infrastructure.llm.embedding_rate_limiter import EmbeddingRateLimiter -from cognee.infrastructure.databases.vector.embeddings.LiteLLMEmbeddingEngine import ( - LiteLLMEmbeddingEngine, -) +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.config import get_llm_config +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.embedding_rate_limiter import EmbeddingRateLimiter from cognee.tests.unit.infrastructure.mock_embedding_engine import MockEmbeddingEngine # Configure logging diff --git a/cognee/tests/unit/infrastructure/test_rate_limiting_realistic.py b/cognee/tests/unit/infrastructure/test_rate_limiting_realistic.py index 25eda7f15..0405704b3 100644 --- a/cognee/tests/unit/infrastructure/test_rate_limiting_realistic.py +++ b/cognee/tests/unit/infrastructure/test_rate_limiting_realistic.py @@ -1,11 +1,9 @@ import asyncio -import time import os -from functools import lru_cache from unittest.mock import patch from cognee.shared.logging_utils import get_logger -from cognee.infrastructure.llm.rate_limiter import llm_rate_limiter -from cognee.infrastructure.llm.config import get_llm_config +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.rate_limiter import llm_rate_limiter +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.config import get_llm_config async def test_rate_limiting_realistic(): diff --git a/cognee/tests/unit/infrastructure/test_rate_limiting_retry.py b/cognee/tests/unit/infrastructure/test_rate_limiting_retry.py index 4338ab2ea..eef14405e 100644 --- a/cognee/tests/unit/infrastructure/test_rate_limiting_retry.py +++ b/cognee/tests/unit/infrastructure/test_rate_limiting_retry.py @@ -1,7 +1,7 @@ import time import asyncio from cognee.shared.logging_utils import get_logger -from cognee.infrastructure.llm.rate_limiter import ( +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.rate_limiter import ( sleep_and_retry_sync, sleep_and_retry_async, is_rate_limit_error, diff --git a/cognee_demo/pyproject.toml b/cognee_demo/pyproject.toml new file mode 100644 index 000000000..b47d5dd1d --- /dev/null +++ b/cognee_demo/pyproject.toml @@ -0,0 +1,16 @@ +[project] +name = "cognee-demo" +version = "0.1.0" +description = "" +authors = [ + {name = "vasilije",email = "vas.markovic@gmail.com"} +] +readme = "README.md" +requires-python = ">=3.11" +dependencies = [ +] + + +[build-system] +requires = ["poetry-core>=2.0.0,<3.0.0"] +build-backend = "poetry.core.masonry.api" diff --git a/examples/python/agentic_reasoning_procurement_example.py b/examples/python/agentic_reasoning_procurement_example.py index 4aa99cc88..5fd8a1e12 100644 --- a/examples/python/agentic_reasoning_procurement_example.py +++ b/examples/python/agentic_reasoning_procurement_example.py @@ -2,9 +2,8 @@ import os import logging import cognee import asyncio -import json -from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client from dotenv import load_dotenv from cognee.api.v1.search import SearchType from cognee.modules.engine.models import NodeSet diff --git a/examples/python/graphiti_example.py b/examples/python/graphiti_example.py index 4167e8391..28d73c23f 100644 --- a/examples/python/graphiti_example.py +++ b/examples/python/graphiti_example.py @@ -12,8 +12,8 @@ from cognee.tasks.temporal_awareness.index_graphiti_objects import ( ) from cognee.modules.retrieval.utils.brute_force_triplet_search import brute_force_triplet_search from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever -from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt -from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.prompts import read_query_prompt, render_prompt +from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.get_llm_client import get_llm_client from cognee.modules.users.methods import get_default_user text_list = [ diff --git a/poetry.lock b/poetry.lock index 728f05c4d..29e8ea50f 100644 --- a/poetry.lock +++ b/poetry.lock @@ -629,6 +629,24 @@ files = [ [package.extras] extras = ["regex"] +[[package]] +name = "baml-py" +version = "0.201.0" +description = "BAML python bindings (pyproject.toml)" +optional = false +python-versions = "*" +groups = ["main"] +files = [ + {file = "baml_py-0.201.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:83228d2af2b0e845bbbb4e14f7cbd3376cec385aee01210ac522ab6076e07bec"}, + {file = "baml_py-0.201.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:2a9d016139e3ae5b5ce98c7b05b5fbd53d5d38f04dc810ec4d70fb17dd6c10e4"}, + {file = "baml_py-0.201.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b5058505b1a3c5f04fc1679aec4d730fa9bef2cbd96209b3ed50152f60b96baf"}, + {file = "baml_py-0.201.0-cp38-abi3-manylinux_2_24_aarch64.whl", hash = "sha256:36289d548581ba4accd5eaaab3246872542dd32dc6717e537654fa0cad884071"}, + {file = "baml_py-0.201.0-cp38-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:5ab70e7bd6481d71edca8a33313347b29faccec78b9960138aa437522813ac9a"}, + {file = "baml_py-0.201.0-cp38-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:7efc5c693a7142c230a4f3d6700415127fee0b9f5fdbb36db63e04e27ac4c0f1"}, + {file = "baml_py-0.201.0-cp38-abi3-win_amd64.whl", hash = "sha256:56499857b7a27ae61a661c8ce0dddd0fb567a45c0b826157e44048a14cf586f9"}, + {file = "baml_py-0.201.0-cp38-abi3-win_arm64.whl", hash = "sha256:1e52dc1151db84a302b746590fe2bc484bdd794f83fa5da7216d9394c559f33a"}, +] + [[package]] name = "banks" version = "2.1.3" @@ -12100,4 +12118,4 @@ weaviate = ["weaviate-client"] [metadata] lock-version = "2.1" python-versions = ">=3.10,<=3.13" -content-hash = "c33c4bc9654072b74c2f768c6d4bf9da7b1a814b50046539fceb78a5541d30ee" +content-hash = "87a986dbcdedbbe7e00ccd052af0f74faa195b466247c0c003c49efdc942ec7b" diff --git a/pyproject.toml b/pyproject.toml index bb0459bb0..f3e8af267 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,8 @@ dependencies = [ "structlog>=25.2.0,<26", "onnxruntime<=1.21.1", "pylance==0.22.0", - "kuzu==0.9.0" + "kuzu==0.9.0", + "baml-py (>=0.201.0,<0.202.0)" ] [project.optional-dependencies]