From 53b8df8506f012e7c40c95bb624b4d1d598a9de8 Mon Sep 17 00:00:00 2001 From: vasilije Date: Sun, 6 Jul 2025 18:05:06 +0200 Subject: [PATCH] ruff format --- .../codingagents/coding_rule_associations.py | 8 +- cognee/api/v1/config/config.py | 4 +- .../responses/routers/get_responses_router.py | 4 +- .../evaluation/direct_llm_eval_adapter.py | 9 +- .../embeddings/FastembedEmbeddingEngine.py | 4 +- .../embeddings/LiteLLMEmbeddingEngine.py | 16 +- .../embeddings/OllamaEmbeddingEngine.py | 4 +- .../vector/embeddings/get_embedding_engine.py | 4 +- cognee/infrastructure/llm/__init__.py | 16 +- .../baml/baml_client/__init__.py | 37 +- .../baml/baml_client/async_client.py | 668 ++++++++++++----- .../baml/baml_client/globals.py | 12 +- .../baml/baml_client/inlinedbaml.py | 8 +- .../baml/baml_client/parser.py | 108 ++- .../baml/baml_client/runtime.py | 28 +- .../baml/baml_client/stream_types.py | 16 +- .../baml/baml_client/sync_client.py | 669 +++++++++++++----- .../baml/baml_client/tracing.py | 8 +- .../baml/baml_client/type_builder.py | 228 +++--- .../baml/baml_client/type_map.py | 11 +- .../baml/baml_client/types.py | 21 +- .../baml_src/config.py | 15 +- .../baml_src/extraction/extract_summary.py | 18 +- .../knowledge_graph/extract_content_graph.py | 13 +- .../extraction/extract_categories.py | 8 +- .../extraction/extract_summary.py | 8 +- .../knowledge_graph/extract_content_graph.py | 12 +- .../llm/anthropic/adapter.py | 13 +- .../llm/embedding_rate_limiter.py | 4 +- .../llitellm_instructor/llm/gemini/adapter.py | 8 +- .../llm/generic_llm_api/adapter.py | 9 +- .../llitellm_instructor/llm/get_llm_client.py | 24 +- .../llitellm_instructor/llm/llm_interface.py | 4 +- .../llitellm_instructor/llm/ollama/adapter.py | 4 +- .../llitellm_instructor/llm/openai/adapter.py | 8 +- .../llitellm_instructor/llm/rate_limiter.py | 4 +- .../llm/tokenizer/Gemini/adapter.py | 4 +- .../llitellm_instructor/llm/utils.py | 4 +- .../document_types/AudioDocument.py | 4 +- .../document_types/ImageDocument.py | 4 +- .../modules/pipelines/operations/pipeline.py | 5 +- cognee/modules/retrieval/code_retriever.py | 8 +- .../graph_completion_cot_retriever.py | 9 +- .../retrieval/natural_language_retriever.py | 8 +- cognee/modules/retrieval/utils/completion.py | 9 +- .../utils/description_to_codepart_search.py | 4 +- cognee/shared/data_models.py | 4 +- .../chunk_naive_llm_classifier.py | 4 +- .../entity_extractors/llm_entity_extractor.py | 9 +- ...ct_content_nodes_and_relationship_names.py | 9 +- .../utils/extract_edge_triplets.py | 9 +- .../cascade_extract/utils/extract_nodes.py | 9 +- cognee/tasks/graph/extract_graph_from_code.py | 15 +- cognee/tasks/graph/extract_graph_from_data.py | 11 +- cognee/tasks/graph/infer_data_ontology.py | 8 +- cognee/tasks/summarization/summarize_code.py | 14 +- cognee/tasks/summarization/summarize_text.py | 14 +- .../test_embedding_rate_limiting_realistic.py | 8 +- .../test_rate_limiting_realistic.py | 8 +- .../agentic_reasoning_procurement_example.py | 4 +- examples/python/graphiti_example.py | 9 +- 61 files changed, 1619 insertions(+), 610 deletions(-) diff --git a/cognee-mcp/src/codingagents/coding_rule_associations.py b/cognee-mcp/src/codingagents/coding_rule_associations.py index a7b4390f4..952ddfd66 100644 --- a/cognee-mcp/src/codingagents/coding_rule_associations.py +++ b/cognee-mcp/src/codingagents/coding_rule_associations.py @@ -2,9 +2,13 @@ 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.structured_output_framework.llitellm_instructor.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.structured_output_framework.llitellm_instructor.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 913444aa8..35fcf01ec 100644 --- a/cognee/api/v1/config/config.py +++ b/cognee/api/v1/config/config.py @@ -7,7 +7,9 @@ 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.structured_output_framework.llitellm_instructor.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 6bb293a05..ae0d6caea 100644 --- a/cognee/api/v1/responses/routers/get_responses_router.py +++ b/cognee/api/v1/responses/routers/get_responses_router.py @@ -17,7 +17,9 @@ 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.structured_output_framework.llitellm_instructor.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/eval_framework/evaluation/direct_llm_eval_adapter.py b/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py index 32a5b4926..2359c08ec 100644 --- a/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py +++ b/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py @@ -1,8 +1,13 @@ from typing import Any, Dict, List from pydantic import BaseModel -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.get_llm_client import ( + get_llm_client, +) from cognee.eval_framework.evaluation.base_eval_adapter import BaseEvalAdapter -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.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 bc197fffb..5b45c96b4 100644 --- a/cognee/infrastructure/databases/vector/embeddings/FastembedEmbeddingEngine.py +++ b/cognee/infrastructure/databases/vector/embeddings/FastembedEmbeddingEngine.py @@ -5,7 +5,9 @@ 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.structured_output_framework.llitellm_instructor.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 502ca1fb7..94257699e 100644 --- a/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py +++ b/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py @@ -7,10 +7,18 @@ 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.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.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 4db38648e..24cf58d2b 100644 --- a/cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py +++ b/cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py @@ -7,7 +7,9 @@ import os import aiohttp.http_exceptions from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine -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.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 d97984c96..4fa1473bf 100644 --- a/cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py +++ b/cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py @@ -1,5 +1,7 @@ from cognee.infrastructure.databases.vector.embeddings.config import get_embedding_config -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.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 8a3511003..8e0cce1e6 100644 --- a/cognee/infrastructure/llm/__init__.py +++ b/cognee/infrastructure/llm/__init__.py @@ -1,4 +1,12 @@ -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 +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/structured_output_framework/baml/baml_client/__init__.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/__init__.py index c4fe2bc85..85dfab90c 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/__init__.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/__init__.py @@ -13,9 +13,9 @@ __version__ = "0.201.0" try: - from baml_py.safe_import import EnsureBamlPyImport + from baml_py.safe_import import EnsureBamlPyImport except ImportError: - raise ImportError(f"""Update to baml-py required. + raise ImportError(f"""Update to baml-py required. Version of baml_client generator (see generators.baml): {__version__} Please upgrade baml-py to version "{__version__}". @@ -31,16 +31,15 @@ https://boundaryml.com/discord with EnsureBamlPyImport(__version__) as e: - e.raise_if_incompatible_version(__version__) + 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 - + 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" @@ -48,11 +47,11 @@ with EnsureBamlPyImport(__version__) as e: partial_types = stream_types __all__ = [ - "b", - "stream_types", - "partial_types", - "tracing", - "types", - "reset_baml_env_vars", - "config", -] \ No newline at end of file + "b", + "stream_types", + "partial_types", + "tracing", + "types", + "reset_baml_env_vars", + "config", +] 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 index cb58fb552..17c57f261 100644 --- 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 @@ -36,10 +36,13 @@ class BamlAsyncClient: self.__llm_response_parser = LlmResponseParser(options) self.__llm_stream_parser = LlmStreamParser(options) - def with_options(self, + 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, + 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 = {} @@ -55,67 +58,151 @@ class BamlAsyncClient: @property def stream(self): - return self.__stream_client + return self.__stream_client @property def request(self): - return self.__http_request + return self.__http_request @property def stream_request(self): - return self.__http_stream_request + return self.__http_stream_request @property def parse(self): - return self.__llm_response_parser + 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, + 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, + 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, + 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, + 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__)) - async def SummarizeCode(self, content: str, + 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__) + ) + + async def SummarizeCode( + self, + content: str, baml_options: BamlCallOptions = {}, ) -> types.SummarizedCode: - result = await self.__options.merge_options(baml_options).call_function_async(function_name="SummarizeCode", args={ - "content": content, - }) - return typing.cast(types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__)) - async def SummarizeContent(self, content: str, + result = await self.__options.merge_options(baml_options).call_function_async( + function_name="SummarizeCode", + args={ + "content": content, + }, + ) + return typing.cast( + types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__) + ) + + async def SummarizeContent( + self, + content: str, baml_options: BamlCallOptions = {}, ) -> types.SummarizedContent: - result = await self.__options.merge_options(baml_options).call_function_async(function_name="SummarizeContent", args={ - "content": content, - }) - return typing.cast(types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__)) - + result = await self.__options.merge_options(baml_options).call_function_async( + function_name="SummarizeContent", + args={ + "content": content, + }, + ) + return typing.cast( + types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__) + ) class BamlStreamClient: @@ -124,79 +211,182 @@ class BamlStreamClient: 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, + 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, + 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, + }, ) - 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, + 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, + 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, + }, ) - 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, + 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, + 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, + }, ) - def ExtractContentGraphWithEnvPrompt(self, content: str,prompt_override: typing.Optional[str] = None, + 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, + ctx, result = self.__options.merge_options(baml_options).create_async_stream( + function_name="ExtractContentGraphWithEnvPrompt", + args={ + "content": content, + "prompt_override": prompt_override, + }, ) - def SummarizeCode(self, content: str, + 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 SummarizeCode( + self, + content: str, baml_options: BamlCallOptions = {}, ) -> baml_py.BamlStream[stream_types.SummarizedCode, types.SummarizedCode]: - ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="SummarizeCode", args={ - "content": content, - }) - return baml_py.BamlStream[stream_types.SummarizedCode, types.SummarizedCode]( - result, - lambda x: typing.cast(stream_types.SummarizedCode, x.cast_to(types, types, stream_types, True, __runtime__)), - lambda x: typing.cast(types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)), - ctx, + ctx, result = self.__options.merge_options(baml_options).create_async_stream( + function_name="SummarizeCode", + args={ + "content": content, + }, ) - def SummarizeContent(self, content: str, + return baml_py.BamlStream[stream_types.SummarizedCode, types.SummarizedCode]( + result, + lambda x: typing.cast( + stream_types.SummarizedCode, + x.cast_to(types, types, stream_types, True, __runtime__), + ), + lambda x: typing.cast( + types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__) + ), + ctx, + ) + + def SummarizeContent( + self, + content: str, baml_options: BamlCallOptions = {}, ) -> baml_py.BamlStream[stream_types.SummarizedContent, types.SummarizedContent]: - ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="SummarizeContent", args={ - "content": content, - }) - return baml_py.BamlStream[stream_types.SummarizedContent, types.SummarizedContent]( - result, - lambda x: typing.cast(stream_types.SummarizedContent, x.cast_to(types, types, stream_types, True, __runtime__)), - lambda x: typing.cast(types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)), - ctx, + ctx, result = self.__options.merge_options(baml_options).create_async_stream( + function_name="SummarizeContent", + args={ + "content": content, + }, ) - + return baml_py.BamlStream[stream_types.SummarizedContent, types.SummarizedContent]( + result, + lambda x: typing.cast( + stream_types.SummarizedContent, + x.cast_to(types, types, stream_types, True, __runtime__), + ), + lambda x: typing.cast( + types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__) + ), + ctx, + ) + class BamlHttpRequestClient: __options: DoNotUseDirectlyCallManager @@ -204,49 +394,128 @@ class BamlHttpRequestClient: 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, + 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") + 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, + + 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") + 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, + + 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") + 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, + + 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") + 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 - async def SummarizeCode(self, content: str, + + async def SummarizeCode( + self, + content: str, baml_options: BamlCallOptions = {}, ) -> baml_py.baml_py.HTTPRequest: - result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="SummarizeCode", args={ - "content": content, - }, mode="request") + result = await self.__options.merge_options(baml_options).create_http_request_async( + function_name="SummarizeCode", + args={ + "content": content, + }, + mode="request", + ) return result - async def SummarizeContent(self, content: str, + + async def SummarizeContent( + self, + content: str, baml_options: BamlCallOptions = {}, ) -> baml_py.baml_py.HTTPRequest: - result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="SummarizeContent", args={ - "content": content, - }, mode="request") + result = await self.__options.merge_options(baml_options).create_http_request_async( + function_name="SummarizeContent", + args={ + "content": content, + }, + mode="request", + ) return result - + class BamlHttpStreamRequestClient: __options: DoNotUseDirectlyCallManager @@ -254,48 +523,127 @@ class BamlHttpStreamRequestClient: 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, + 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") + 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 - async def SummarizeCode(self, content: str, - baml_options: BamlCallOptions = {}, - ) -> baml_py.baml_py.HTTPRequest: - result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="SummarizeCode", args={ - "content": content, - }, mode="stream") - return result - async def SummarizeContent(self, content: str, - baml_options: BamlCallOptions = {}, - ) -> baml_py.baml_py.HTTPRequest: - result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="SummarizeContent", args={ - "content": content, - }, mode="stream") - return result - -b = BamlAsyncClient(DoNotUseDirectlyCallManager({})) \ No newline at end of file + 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 + + async def SummarizeCode( + self, + content: str, + baml_options: BamlCallOptions = {}, + ) -> baml_py.baml_py.HTTPRequest: + result = await self.__options.merge_options(baml_options).create_http_request_async( + function_name="SummarizeCode", + args={ + "content": content, + }, + mode="stream", + ) + return result + + async def SummarizeContent( + self, + content: str, + baml_options: BamlCallOptions = {}, + ) -> baml_py.baml_py.HTTPRequest: + result = await self.__options.merge_options(baml_options).create_http_request_async( + function_name="SummarizeContent", + args={ + "content": content, + }, + mode="stream", + ) + return result + + +b = BamlAsyncClient(DoNotUseDirectlyCallManager({})) 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 index 769e055bb..7ef10535c 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/globals.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/globals.py @@ -19,17 +19,19 @@ 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() + "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) +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 + 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 index 09c747670..66ec84fbb 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/inlinedbaml.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/inlinedbaml.py @@ -11,11 +11,11 @@ # 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// Summarization classes\nclass SummarizedContent {\n summary string\n description string\n}\n\nclass SummarizedFunction {\n name string\n description string\n inputs string[]?\n outputs string[]?\n decorators string[]?\n}\n\nclass SummarizedClass {\n name string\n description string\n methods SummarizedFunction[]?\n decorators string[]?\n}\n\nclass SummarizedCode {\n high_level_summary string\n key_features string[]\n imports string[]\n constants string[]\n classes SummarizedClass[]\n functions SummarizedFunction[]\n workflow_description string?\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// Summarization prompt template\ntemplate_string SummarizeContentPrompt() #\"\n You are a top-tier summarization engine. Your task is to summarize text and make it versatile.\n Be brief and concise, but keep the important information and the subject.\n Use synonym words where possible in order to change the wording but keep the meaning.\n\"#\n\n// Code summarization prompt template\ntemplate_string SummarizeCodePrompt() #\"\n You are an expert code analyst. Analyze the provided source code and extract key information:\n\n 1. Provide a high-level summary of what the code does\n 2. List key features and functionality\n 3. Identify imports and dependencies\n 4. List constants and global variables\n 5. Summarize classes with their methods\n 6. Summarize standalone functions\n 7. Describe the overall workflow if applicable\n\n Be precise and technical while remaining clear and concise.\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\n// Summarization functions\nfunction SummarizeContent(content: string) -> SummarizedContent {\n client OpenAIClientWithEnvModel\n\n prompt #\"\n {{ SummarizeContentPrompt() }}\n\n {{ ctx.output_format(prefix=\"Answer in this schema:\\n\") }}\n\n {{ _.role('user') }}\n {{ content }}\n \"#\n}\n\nfunction SummarizeCode(content: string) -> SummarizedCode {\n client OpenAIClientWithEnvModel\n\n prompt #\"\n {{ SummarizeCodePrompt() }}\n\n {{ ctx.output_format(prefix=\"Answer in this schema:\\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\ntest SummarizeContentExample {\n functions [SummarizeContent]\n args {\n content #\"\n Natural language processing (NLP) is an interdisciplinary subfield of computer science and information retrieval.\n It deals with the interaction between computers and human language, in particular how to program computers to process and analyze large amounts of natural language data.\n \"#\n }\n}\n\ntest SummarizeCodeExample {\n functions [SummarizeCode]\n args {\n content #\"\n def fibonacci(n):\n if n <= 1:\n return n\n return fibonacci(n-1) + fibonacci(n-2)\n \n def main():\n print(fibonacci(10))\n \n if __name__ == \"__main__\":\n main()\n \"#\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", + "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// Summarization classes\nclass SummarizedContent {\n summary string\n description string\n}\n\nclass SummarizedFunction {\n name string\n description string\n inputs string[]?\n outputs string[]?\n decorators string[]?\n}\n\nclass SummarizedClass {\n name string\n description string\n methods SummarizedFunction[]?\n decorators string[]?\n}\n\nclass SummarizedCode {\n high_level_summary string\n key_features string[]\n imports string[]\n constants string[]\n classes SummarizedClass[]\n functions SummarizedFunction[]\n workflow_description string?\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// Summarization prompt template\ntemplate_string SummarizeContentPrompt() #"\n You are a top-tier summarization engine. Your task is to summarize text and make it versatile.\n Be brief and concise, but keep the important information and the subject.\n Use synonym words where possible in order to change the wording but keep the meaning.\n"#\n\n// Code summarization prompt template\ntemplate_string SummarizeCodePrompt() #"\n You are an expert code analyst. Analyze the provided source code and extract key information:\n\n 1. Provide a high-level summary of what the code does\n 2. List key features and functionality\n 3. Identify imports and dependencies\n 4. List constants and global variables\n 5. Summarize classes with their methods\n 6. Summarize standalone functions\n 7. Describe the overall workflow if applicable\n\n Be precise and technical while remaining clear and concise.\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\n// Summarization functions\nfunction SummarizeContent(content: string) -> SummarizedContent {\n client OpenAIClientWithEnvModel\n\n prompt #"\n {{ SummarizeContentPrompt() }}\n\n {{ ctx.output_format(prefix="Answer in this schema:\\n") }}\n\n {{ _.role(\'user\') }}\n {{ content }}\n "#\n}\n\nfunction SummarizeCode(content: string) -> SummarizedCode {\n client OpenAIClientWithEnvModel\n\n prompt #"\n {{ SummarizeCodePrompt() }}\n\n {{ ctx.output_format(prefix="Answer in this schema:\\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\ntest SummarizeContentExample {\n functions [SummarizeContent]\n args {\n content #"\n Natural language processing (NLP) is an interdisciplinary subfield of computer science and information retrieval.\n It deals with the interaction between computers and human language, in particular how to program computers to process and analyze large amounts of natural language data.\n "#\n }\n}\n\ntest SummarizeCodeExample {\n functions [SummarizeCode]\n args {\n content #"\n def fibonacci(n):\n if n <= 1:\n return n\n return fibonacci(n-1) + fibonacci(n-2)\n \n def main():\n print(fibonacci(10))\n \n if __name__ == "__main__":\n main()\n "#\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 + return _file_map 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 index 926aca307..f923f8255 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/parser.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/parser.py @@ -16,6 +16,7 @@ import typing_extensions from . import stream_types, types from .runtime import DoNotUseDirectlyCallManager, BamlCallOptions + class LlmResponseParser: __options: DoNotUseDirectlyCallManager @@ -23,42 +24,69 @@ class LlmResponseParser: self.__options = options def ExtractContentGraph( - self, llm_response: str, baml_options: BamlCallOptions = {}, + 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") + 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 = {}, + 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") + 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 = {}, + 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") + 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 = {}, + 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") + result = self.__options.merge_options(baml_options).parse_response( + function_name="ExtractContentGraphWithEnvPrompt", + llm_response=llm_response, + mode="request", + ) return typing.cast(types.KnowledgeGraph, result) def SummarizeCode( - self, llm_response: str, baml_options: BamlCallOptions = {}, + self, + llm_response: str, + baml_options: BamlCallOptions = {}, ) -> types.SummarizedCode: - result = self.__options.merge_options(baml_options).parse_response(function_name="SummarizeCode", llm_response=llm_response, mode="request") + result = self.__options.merge_options(baml_options).parse_response( + function_name="SummarizeCode", llm_response=llm_response, mode="request" + ) return typing.cast(types.SummarizedCode, result) def SummarizeContent( - self, llm_response: str, baml_options: BamlCallOptions = {}, + self, + llm_response: str, + baml_options: BamlCallOptions = {}, ) -> types.SummarizedContent: - result = self.__options.merge_options(baml_options).parse_response(function_name="SummarizeContent", llm_response=llm_response, mode="request") + result = self.__options.merge_options(baml_options).parse_response( + function_name="SummarizeContent", llm_response=llm_response, mode="request" + ) return typing.cast(types.SummarizedContent, result) - class LlmStreamParser: __options: DoNotUseDirectlyCallManager @@ -67,39 +95,65 @@ class LlmStreamParser: self.__options = options def ExtractContentGraph( - self, llm_response: str, baml_options: BamlCallOptions = {}, + 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") + 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 = {}, + 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") + 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 = {}, + 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") + 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 = {}, + 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") + 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) def SummarizeCode( - self, llm_response: str, baml_options: BamlCallOptions = {}, + self, + llm_response: str, + baml_options: BamlCallOptions = {}, ) -> stream_types.SummarizedCode: - result = self.__options.merge_options(baml_options).parse_response(function_name="SummarizeCode", llm_response=llm_response, mode="stream") + result = self.__options.merge_options(baml_options).parse_response( + function_name="SummarizeCode", llm_response=llm_response, mode="stream" + ) return typing.cast(stream_types.SummarizedCode, result) def SummarizeContent( - self, llm_response: str, baml_options: BamlCallOptions = {}, + self, + llm_response: str, + baml_options: BamlCallOptions = {}, ) -> stream_types.SummarizedContent: - result = self.__options.merge_options(baml_options).parse_response(function_name="SummarizeContent", llm_response=llm_response, mode="stream") + result = self.__options.merge_options(baml_options).parse_response( + function_name="SummarizeContent", llm_response=llm_response, mode="stream" + ) return typing.cast(stream_types.SummarizedContent, 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 index c94bb950e..1955e5c14 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/runtime.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/runtime.py @@ -17,7 +17,10 @@ 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__ +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): @@ -48,9 +51,6 @@ class _ResolvedBamlOptions: self.env_vars = env_vars - - - class DoNotUseDirectlyCallManager: def __init__(self, baml_options: BamlCallOptions): self.__baml_options = baml_options @@ -74,7 +74,9 @@ class DoNotUseDirectlyCallManager: collectors_as_list = ( collector if isinstance(collector, list) - else [collector] if collector is not None else [] + else [collector] + if collector is not None + else [] ) env_vars = os.environ.copy() for k, v in self.__baml_options.get("env", {}).items(): @@ -164,7 +166,9 @@ class DoNotUseDirectlyCallManager: *, function_name: str, args: typing.Dict[str, typing.Any], - ) -> typing.Tuple[baml_py.baml_py.RuntimeContextManager, baml_py.baml_py.SyncFunctionResultStream]: + ) -> 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( @@ -172,7 +176,7 @@ class DoNotUseDirectlyCallManager: args, # this is always None, we set this later! # on_event - None, + None, # ctx ctx, # tb @@ -232,7 +236,13 @@ class DoNotUseDirectlyCallManager: mode == "stream", ) - def parse_response(self, *, function_name: str, llm_response: str, mode: typing_extensions.Literal["stream", "request"]) -> typing.Any: + 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, @@ -253,4 +263,4 @@ class DoNotUseDirectlyCallManager: 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 index a3afc0652..0cd28bf45 100644 --- 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 @@ -18,14 +18,19 @@ import baml_py from . import types -StreamStateValueT = typing.TypeVar('StreamStateValueT') +StreamStateValueT = typing.TypeVar("StreamStateValueT") + + class StreamState(BaseModel, typing.Generic[StreamStateValueT]): value: StreamStateValueT state: typing_extensions.Literal["Pending", "Incomplete", "Complete"] + + # ######################################################################### # Generated classes (7) # ######################################################################### + class Edge(BaseModel): # doc string for edge # doc string for source_node_id @@ -34,23 +39,27 @@ class Edge(BaseModel): 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') + model_config = ConfigDict(extra="allow") id: typing.Optional[str] = None name: typing.Optional[str] = None type: typing.Optional[str] = None description: typing.Optional[str] = None + class SummarizedClass(BaseModel): name: typing.Optional[str] = None description: typing.Optional[str] = None methods: typing.Optional[typing.List["SummarizedFunction"]] = None decorators: typing.Optional[typing.List[str]] = None + class SummarizedCode(BaseModel): high_level_summary: typing.Optional[str] = None key_features: typing.List[str] @@ -60,10 +69,12 @@ class SummarizedCode(BaseModel): functions: typing.List["SummarizedFunction"] workflow_description: typing.Optional[str] = None + class SummarizedContent(BaseModel): summary: typing.Optional[str] = None description: typing.Optional[str] = None + class SummarizedFunction(BaseModel): name: typing.Optional[str] = None description: typing.Optional[str] = None @@ -71,6 +82,7 @@ class SummarizedFunction(BaseModel): outputs: typing.Optional[typing.List[str]] = None decorators: typing.Optional[typing.List[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 index 66768467f..f0dbe2144 100644 --- 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 @@ -19,6 +19,7 @@ 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" @@ -48,10 +49,13 @@ class BamlSyncClient: self.__llm_response_parser = LlmResponseParser(self.__options) self.__llm_stream_parser = LlmStreamParser(self.__options) - def with_options(self, + 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, + 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 = {} @@ -67,67 +71,151 @@ class BamlSyncClient: @property def stream(self): - return self.__stream_client + return self.__stream_client @property def request(self): - return self.__http_request + return self.__http_request @property def stream_request(self): - return self.__http_stream_request + return self.__http_stream_request @property def parse(self): - return self.__llm_response_parser + 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, + 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, + 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, + 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, + 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__)) - def SummarizeCode(self, content: str, + 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__) + ) + + def SummarizeCode( + self, + content: str, baml_options: BamlCallOptions = {}, ) -> types.SummarizedCode: - result = self.__options.merge_options(baml_options).call_function_sync(function_name="SummarizeCode", args={ - "content": content, - }) - return typing.cast(types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__)) - def SummarizeContent(self, content: str, + result = self.__options.merge_options(baml_options).call_function_sync( + function_name="SummarizeCode", + args={ + "content": content, + }, + ) + return typing.cast( + types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__) + ) + + def SummarizeContent( + self, + content: str, baml_options: BamlCallOptions = {}, ) -> types.SummarizedContent: - result = self.__options.merge_options(baml_options).call_function_sync(function_name="SummarizeContent", args={ - "content": content, - }) - return typing.cast(types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__)) - + result = self.__options.merge_options(baml_options).call_function_sync( + function_name="SummarizeContent", + args={ + "content": content, + }, + ) + return typing.cast( + types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__) + ) class BamlStreamClient: @@ -136,79 +224,182 @@ class BamlStreamClient: 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, + 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, + 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, + }, ) - 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, + 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, + 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, + }, ) - 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, + 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, + 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, + }, ) - def ExtractContentGraphWithEnvPrompt(self, content: str,prompt_override: typing.Optional[str] = None, + 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, + ctx, result = self.__options.merge_options(baml_options).create_sync_stream( + function_name="ExtractContentGraphWithEnvPrompt", + args={ + "content": content, + "prompt_override": prompt_override, + }, ) - def SummarizeCode(self, content: str, + 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 SummarizeCode( + self, + content: str, baml_options: BamlCallOptions = {}, ) -> baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode]: - ctx, result = self.__options.merge_options(baml_options).create_sync_stream(function_name="SummarizeCode", args={ - "content": content, - }) - return baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode]( - result, - lambda x: typing.cast(stream_types.SummarizedCode, x.cast_to(types, types, stream_types, True, __runtime__)), - lambda x: typing.cast(types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)), - ctx, + ctx, result = self.__options.merge_options(baml_options).create_sync_stream( + function_name="SummarizeCode", + args={ + "content": content, + }, ) - def SummarizeContent(self, content: str, + return baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode]( + result, + lambda x: typing.cast( + stream_types.SummarizedCode, + x.cast_to(types, types, stream_types, True, __runtime__), + ), + lambda x: typing.cast( + types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__) + ), + ctx, + ) + + def SummarizeContent( + self, + content: str, baml_options: BamlCallOptions = {}, ) -> baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent]: - ctx, result = self.__options.merge_options(baml_options).create_sync_stream(function_name="SummarizeContent", args={ - "content": content, - }) - return baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent]( - result, - lambda x: typing.cast(stream_types.SummarizedContent, x.cast_to(types, types, stream_types, True, __runtime__)), - lambda x: typing.cast(types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)), - ctx, + ctx, result = self.__options.merge_options(baml_options).create_sync_stream( + function_name="SummarizeContent", + args={ + "content": content, + }, ) - + return baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent]( + result, + lambda x: typing.cast( + stream_types.SummarizedContent, + x.cast_to(types, types, stream_types, True, __runtime__), + ), + lambda x: typing.cast( + types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__) + ), + ctx, + ) + class BamlHttpRequestClient: __options: DoNotUseDirectlyCallManager @@ -216,49 +407,128 @@ class BamlHttpRequestClient: 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, + 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") + 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, + + 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") + 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, + + 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") + 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, + + 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") + 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 - def SummarizeCode(self, content: str, + + def SummarizeCode( + self, + content: str, baml_options: BamlCallOptions = {}, ) -> baml_py.baml_py.HTTPRequest: - result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="SummarizeCode", args={ - "content": content, - }, mode="request") + result = self.__options.merge_options(baml_options).create_http_request_sync( + function_name="SummarizeCode", + args={ + "content": content, + }, + mode="request", + ) return result - def SummarizeContent(self, content: str, + + def SummarizeContent( + self, + content: str, baml_options: BamlCallOptions = {}, ) -> baml_py.baml_py.HTTPRequest: - result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="SummarizeContent", args={ - "content": content, - }, mode="request") + result = self.__options.merge_options(baml_options).create_http_request_sync( + function_name="SummarizeContent", + args={ + "content": content, + }, + mode="request", + ) return result - + class BamlHttpStreamRequestClient: __options: DoNotUseDirectlyCallManager @@ -266,48 +536,127 @@ class BamlHttpStreamRequestClient: 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, + 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") + 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 - def SummarizeCode(self, content: str, - baml_options: BamlCallOptions = {}, - ) -> baml_py.baml_py.HTTPRequest: - result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="SummarizeCode", args={ - "content": content, - }, mode="stream") - return result - def SummarizeContent(self, content: str, - baml_options: BamlCallOptions = {}, - ) -> baml_py.baml_py.HTTPRequest: - result = self.__options.merge_options(baml_options).create_http_request_sync(function_name="SummarizeContent", args={ - "content": content, - }, mode="stream") - return result - -b = BamlSyncClient(DoNotUseDirectlyCallManager({})) \ No newline at end of file + 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 + + def SummarizeCode( + self, + content: str, + baml_options: BamlCallOptions = {}, + ) -> baml_py.baml_py.HTTPRequest: + result = self.__options.merge_options(baml_options).create_http_request_sync( + function_name="SummarizeCode", + args={ + "content": content, + }, + mode="stream", + ) + return result + + def SummarizeContent( + self, + content: str, + baml_options: BamlCallOptions = {}, + ) -> baml_py.baml_py.HTTPRequest: + result = self.__options.merge_options(baml_options).create_http_request_sync( + function_name="SummarizeContent", + args={ + "content": content, + }, + mode="stream", + ) + return result + + +b = BamlSyncClient(DoNotUseDirectlyCallManager({})) 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 index 06725593c..f925306f9 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/tracing.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/tracing.py @@ -14,9 +14,13 @@ 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() + 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"] +__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 index cbd4e479f..5cdeb0c1b 100644 --- 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 @@ -15,19 +15,29 @@ 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","SummarizedClass","SummarizedCode","SummarizedContent","SummarizedFunction",] - ), enums=set( - [] - ), runtime=DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME) + super().__init__( + classes=set( + [ + "Edge", + "KnowledgeGraph", + "Node", + "SummarizedClass", + "SummarizedCode", + "SummarizedContent", + "SummarizedFunction", + ] + ), + enums=set([]), + runtime=DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME, + ) # ######################################################################### # Generated enums 0 # ######################################################################### - # ######################################################################### # Generated classes 7 # ######################################################################### @@ -61,7 +71,6 @@ class TypeBuilder(type_builder.TypeBuilder): return SummarizedFunctionViewer(self) - # ######################################################################### # Generated enums 0 # ######################################################################### @@ -71,11 +80,18 @@ class TypeBuilder(type_builder.TypeBuilder): # Generated classes 7 # ######################################################################### + class EdgeAst: def __init__(self, tb: type_builder.TypeBuilder): - _tb = tb._tb # type: ignore (we know how to use this private attribute) + _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._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: @@ -90,39 +106,41 @@ 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] - + 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 + 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) + _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._properties: typing.Set[str] = set( + [ + "nodes", + "edges", + ] + ) self._props = KnowledgeGraphProperties(self._bldr, self._properties) def type(self) -> baml_py.FieldType: @@ -137,35 +155,39 @@ 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] - + 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 + 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) + _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._properties: typing.Set[str] = set( + [ + "id", + "name", + "type", + "description", + ] + ) self._props = NodeProperties(self._bldr, self._properties) def type(self) -> baml_py.FieldType: @@ -180,7 +202,6 @@ 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.") @@ -189,45 +210,46 @@ class NodeBuilder(NodeAst): 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 + 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") - - class SummarizedClassAst: def __init__(self, tb: type_builder.TypeBuilder): - _tb = tb._tb # type: ignore (we know how to use this private attribute) + _tb = tb._tb # type: ignore (we know how to use this private attribute) self._bldr = _tb.class_("SummarizedClass") - self._properties: typing.Set[str] = set([ "name", "description", "methods", "decorators", ]) + self._properties: typing.Set[str] = set( + [ + "name", + "description", + "methods", + "decorators", + ] + ) self._props = SummarizedClassProperties(self._bldr, self._properties) def type(self) -> baml_py.FieldType: @@ -242,43 +264,50 @@ class SummarizedClassViewer(SummarizedClassAst): 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] - + return [ + (name, type_builder.ClassPropertyViewer(self._bldr.property(name))) + for name in self._properties + ] class SummarizedClassProperties: 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 + self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 - - @property def name(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("name")) - + @property def description(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("description")) - + @property def methods(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("methods")) - + @property def decorators(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("decorators")) - - class SummarizedCodeAst: def __init__(self, tb: type_builder.TypeBuilder): - _tb = tb._tb # type: ignore (we know how to use this private attribute) + _tb = tb._tb # type: ignore (we know how to use this private attribute) self._bldr = _tb.class_("SummarizedCode") - self._properties: typing.Set[str] = set([ "high_level_summary", "key_features", "imports", "constants", "classes", "functions", "workflow_description", ]) + self._properties: typing.Set[str] = set( + [ + "high_level_summary", + "key_features", + "imports", + "constants", + "classes", + "functions", + "workflow_description", + ] + ) self._props = SummarizedCodeProperties(self._bldr, self._properties) def type(self) -> baml_py.FieldType: @@ -293,55 +322,57 @@ class SummarizedCodeViewer(SummarizedCodeAst): 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] - + return [ + (name, type_builder.ClassPropertyViewer(self._bldr.property(name))) + for name in self._properties + ] class SummarizedCodeProperties: 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 + self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 - - @property def high_level_summary(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("high_level_summary")) - + @property def key_features(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("key_features")) - + @property def imports(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("imports")) - + @property def constants(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("constants")) - + @property def classes(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("classes")) - + @property def functions(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("functions")) - + @property def workflow_description(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("workflow_description")) - - class SummarizedContentAst: def __init__(self, tb: type_builder.TypeBuilder): - _tb = tb._tb # type: ignore (we know how to use this private attribute) + _tb = tb._tb # type: ignore (we know how to use this private attribute) self._bldr = _tb.class_("SummarizedContent") - self._properties: typing.Set[str] = set([ "summary", "description", ]) + self._properties: typing.Set[str] = set( + [ + "summary", + "description", + ] + ) self._props = SummarizedContentProperties(self._bldr, self._properties) def type(self) -> baml_py.FieldType: @@ -356,35 +387,40 @@ class SummarizedContentViewer(SummarizedContentAst): 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] - + return [ + (name, type_builder.ClassPropertyViewer(self._bldr.property(name))) + for name in self._properties + ] class SummarizedContentProperties: 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 + self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 - - @property def summary(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("summary")) - + @property def description(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("description")) - - class SummarizedFunctionAst: def __init__(self, tb: type_builder.TypeBuilder): - _tb = tb._tb # type: ignore (we know how to use this private attribute) + _tb = tb._tb # type: ignore (we know how to use this private attribute) self._bldr = _tb.class_("SummarizedFunction") - self._properties: typing.Set[str] = set([ "name", "description", "inputs", "outputs", "decorators", ]) + self._properties: typing.Set[str] = set( + [ + "name", + "description", + "inputs", + "outputs", + "decorators", + ] + ) self._props = SummarizedFunctionProperties(self._bldr, self._properties) def type(self) -> baml_py.FieldType: @@ -399,38 +435,34 @@ class SummarizedFunctionViewer(SummarizedFunctionAst): 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] - + return [ + (name, type_builder.ClassPropertyViewer(self._bldr.property(name))) + for name in self._properties + ] class SummarizedFunctionProperties: 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 + self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 - - @property def name(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("name")) - + @property def description(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("description")) - + @property def inputs(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("inputs")) - + @property def outputs(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("outputs")) - + @property def decorators(self) -> type_builder.ClassPropertyViewer: return type_builder.ClassPropertyViewer(self.__bldr.property("decorators")) - - - 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 index 787a62d94..09dfb30cb 100644 --- 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 @@ -15,27 +15,18 @@ 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, - "types.SummarizedClass": types.SummarizedClass, "stream_types.SummarizedClass": stream_types.SummarizedClass, - "types.SummarizedCode": types.SummarizedCode, "stream_types.SummarizedCode": stream_types.SummarizedCode, - "types.SummarizedContent": types.SummarizedContent, "stream_types.SummarizedContent": stream_types.SummarizedContent, - "types.SummarizedFunction": types.SummarizedFunction, "stream_types.SummarizedFunction": stream_types.SummarizedFunction, - - -} \ 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 index 72c1e59c9..f7b8b8f1b 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/types.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/types.py @@ -20,22 +20,29 @@ from pydantic import BaseModel, ConfigDict import baml_py -CheckT = typing_extensions.TypeVar('CheckT') -CheckName = typing_extensions.TypeVar('CheckName', bound=str) +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) # ######################################################################### @@ -44,6 +51,7 @@ def all_succeeded(checks: typing.Dict[CheckName, Check]) -> bool: # Generated classes (7) # ######################################################################### + class Edge(BaseModel): # doc string for edge # doc string for source_node_id @@ -52,23 +60,27 @@ class Edge(BaseModel): 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') + model_config = ConfigDict(extra="allow") id: str name: str type: str description: str + class SummarizedClass(BaseModel): name: str description: str methods: typing.Optional[typing.List["SummarizedFunction"]] = None decorators: typing.Optional[typing.List[str]] = None + class SummarizedCode(BaseModel): high_level_summary: str key_features: typing.List[str] @@ -78,10 +90,12 @@ class SummarizedCode(BaseModel): functions: typing.List["SummarizedFunction"] workflow_description: typing.Optional[str] = None + class SummarizedContent(BaseModel): summary: str description: str + class SummarizedFunction(BaseModel): name: str description: str @@ -89,6 +103,7 @@ class SummarizedFunction(BaseModel): outputs: typing.Optional[typing.List[str]] = None decorators: typing.Optional[typing.List[str]] = None + # ######################################################################### # 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 index 153507e8a..2a305cf9e 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml_src/config.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml_src/config.py @@ -5,6 +5,7 @@ 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. @@ -54,11 +55,15 @@ class LLMConfig(BaseSettings): def model_post_init(self, __context) -> None: """Initialize the BAML registry after the model is created.""" - self.baml_registry.add_llm_client(name=self.llm_provider, provider=self.llm_provider, options={ - "model": self.llm_model, - "temperature": self.llm_temperature, - "api_key": self.llm_api_key - }) + self.baml_registry.add_llm_client( + name=self.llm_provider, + provider=self.llm_provider, + options={ + "model": self.llm_model, + "temperature": self.llm_temperature, + "api_key": self.llm_api_key, + }, + ) # Sets the primary client self.baml_registry.set_primary(self.llm_provider) diff --git a/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_summary.py b/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_summary.py index 2417eaa1e..1f904a7dd 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_summary.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_summary.py @@ -14,22 +14,22 @@ logger = get_logger("extract_summary_baml") async def extract_summary(content: str, response_model: Type[BaseModel]): """ Extract summary using BAML framework. - + Args: content: The content to summarize response_model: The Pydantic model type for the response - + Returns: BaseModel: The summarized content in the specified format """ config = get_llm_config() - + # Use BAML's SummarizeContent function summary_result = await b.SummarizeContent(content, baml_options={"tb": config.baml_registry}) - + # Convert BAML result to the expected response model if response_model is SummarizedCode: - # If it's asking for SummarizedCode but we got SummarizedContent, + # If it's asking for SummarizedCode but we got SummarizedContent, # we need to use SummarizeCode instead code_result = await b.SummarizeCode(content, baml_options={"tb": config.baml_registry}) return code_result @@ -41,10 +41,10 @@ async def extract_summary(content: str, response_model: Type[BaseModel]): async def extract_code_summary(content: str): """ Extract code summary using BAML framework with mocking support. - + Args: content: The code content to summarize - + Returns: SummarizedCode: The summarized code information """ @@ -61,7 +61,9 @@ async def extract_code_summary(content: str): config = get_llm_config() result = await b.SummarizeCode(content, baml_options={"tb": config.baml_registry}) except Exception as e: - logger.error("Failed to extract code summary with BAML, falling back to mock summary", exc_info=e) + logger.error( + "Failed to extract code summary with BAML, falling back to mock summary", exc_info=e + ) result = get_mock_summarized_code() return result 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 index 63cc01705..894eb387f 100644 --- 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 @@ -2,14 +2,14 @@ import os from typing import Type from pydantic import BaseModel from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.async_client import b -from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.type_builder import TypeBuilder +from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.type_builder import ( + TypeBuilder, +) from cognee.infrastructure.llm.structured_output_framework.baml_src.config import get_llm_config from cognee.shared.logging_utils import get_logger, setup_logging - - async def extract_content_graph(content: str, response_model: Type[BaseModel]): # tb = TypeBuilder() config = get_llm_config() @@ -19,9 +19,8 @@ async def extract_content_graph(content: str, response_model: Type[BaseModel]): # ([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}) + graph = await b.ExtractContentGraph( + content, mode="simple", baml_options={"tb": config.baml_registry} + ) return graph - diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_categories.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_categories.py index 8b7e3db6d..cea066db1 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_categories.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_categories.py @@ -1,7 +1,11 @@ from typing import Type from pydantic import BaseModel -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.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/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_summary.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_summary.py index 6935e9c44..a18ae8b29 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_summary.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/extract_summary.py @@ -5,8 +5,12 @@ from typing import Type from instructor.exceptions import InstructorRetryException from pydantic import BaseModel -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.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/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/knowledge_graph/extract_content_graph.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/knowledge_graph/extract_content_graph.py index d25afef76..aa243246d 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/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,15 @@ import os from typing import Type from pydantic import BaseModel -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 +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/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/anthropic/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/anthropic/adapter.py index 7b9710675..5e8a41894 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/anthropic/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/anthropic/adapter.py @@ -3,9 +3,16 @@ from pydantic import BaseModel import instructor from cognee.exceptions import InvalidValueError -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 +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/structured_output_framework/llitellm_instructor/llm/embedding_rate_limiter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/embedding_rate_limiter.py index 0aed181e4..145bc3457 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/embedding_rate_limiter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/embedding_rate_limiter.py @@ -6,7 +6,9 @@ import time import asyncio import random from cognee.shared.logging_utils import get_logger -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.config import ( + get_llm_config, +) logger = get_logger() diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/gemini/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/gemini/adapter.py index e2f3b16e4..8adf3cde7 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/gemini/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/gemini/adapter.py @@ -6,8 +6,12 @@ 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.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.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/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/generic_llm_api/adapter.py index 3a1a7b2c3..421334cd0 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/generic_llm_api/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/generic_llm_api/adapter.py @@ -4,8 +4,13 @@ from typing import Type from pydantic import BaseModel import instructor -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 +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/structured_output_framework/llitellm_instructor/llm/get_llm_client.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/get_llm_client.py index 4738aa71d..3f9bc3994 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/get_llm_client.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/get_llm_client.py @@ -4,7 +4,9 @@ from enum import Enum from cognee.exceptions import InvalidValueError from cognee.infrastructure.llm import get_llm_config -from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.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 @@ -59,7 +61,9 @@ def get_llm_client(): if llm_config.llm_api_key is None: raise InvalidValueError(message="LLM API key is not set.") - from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.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 +79,9 @@ def get_llm_client(): if llm_config.llm_api_key is None: raise InvalidValueError(message="LLM API key is not set.") - from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.generic_llm_api 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 +92,9 @@ def get_llm_client(): ) elif provider == LLMProvider.ANTHROPIC: - from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.anthropic 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 +102,9 @@ def get_llm_client(): if llm_config.llm_api_key is None: raise InvalidValueError(message="LLM API key is not set.") - from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.generic_llm_api 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 +118,9 @@ def get_llm_client(): if llm_config.llm_api_key is None: raise InvalidValueError(message="LLM API key is not set.") - from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.gemini 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/structured_output_framework/llitellm_instructor/llm/llm_interface.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/llm_interface.py index e5b4b5f0d..f0882c1b9 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/llm_interface.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/llm_interface.py @@ -3,7 +3,9 @@ from typing import Type, Protocol from abc import abstractmethod from pydantic import BaseModel -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.prompts import ( + read_query_prompt, +) class LLMInterface(Protocol): diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/ollama/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/ollama/adapter.py index fc2d35431..bcabb6adf 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/ollama/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/ollama/adapter.py @@ -1,7 +1,9 @@ from typing import Type from pydantic import BaseModel import instructor -from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.llm.llm_interface import LLMInterface +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, diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/openai/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/openai/adapter.py index 35c0a5101..c1406d571 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/openai/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/openai/adapter.py @@ -7,8 +7,12 @@ 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.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.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, diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/rate_limiter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/rate_limiter.py index 9e8e6012a..bb39713f0 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/rate_limiter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/rate_limiter.py @@ -49,7 +49,9 @@ 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.structured_output_framework.llitellm_instructor.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/tokenizer/Gemini/adapter.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/Gemini/adapter.py index 02419f3ab..7db010b1d 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/Gemini/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/tokenizer/Gemini/adapter.py @@ -24,7 +24,9 @@ class GeminiTokenizer(TokenizerInterface): # Get LLM API key from config from cognee.infrastructure.databases.vector.embeddings.config import get_embedding_config - 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.config import ( + get_llm_config, + ) config = get_embedding_config() llm_config = get_llm_config() diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/utils.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/utils.py index 9f6e92d5c..e66743762 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/utils.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/llm/utils.py @@ -1,6 +1,8 @@ import litellm -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.get_llm_client import ( + get_llm_client, +) from cognee.shared.logging_utils import get_logger logger = get_logger() diff --git a/cognee/modules/data/processing/document_types/AudioDocument.py b/cognee/modules/data/processing/document_types/AudioDocument.py index 763061c8c..9af822c5b 100644 --- a/cognee/modules/data/processing/document_types/AudioDocument.py +++ b/cognee/modules/data/processing/document_types/AudioDocument.py @@ -1,4 +1,6 @@ -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.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 c8400a7fc..27cf9bb63 100644 --- a/cognee/modules/data/processing/document_types/ImageDocument.py +++ b/cognee/modules/data/processing/document_types/ImageDocument.py @@ -1,4 +1,6 @@ -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.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 03945f37a..1aad038b4 100644 --- a/cognee/modules/pipelines/operations/pipeline.py +++ b/cognee/modules/pipelines/operations/pipeline.py @@ -69,7 +69,10 @@ async def cognee_pipeline( cognee_pipeline.first_run = True if cognee_pipeline.first_run: - from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.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 33c8254a0..6986862a1 100644 --- a/cognee/modules/retrieval/code_retriever.py +++ b/cognee/modules/retrieval/code_retriever.py @@ -6,8 +6,12 @@ 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.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.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/graph_completion_cot_retriever.py b/cognee/modules/retrieval/graph_completion_cot_retriever.py index c61f72d2c..722f0e04c 100644 --- a/cognee/modules/retrieval/graph_completion_cot_retriever.py +++ b/cognee/modules/retrieval/graph_completion_cot_retriever.py @@ -1,9 +1,14 @@ from typing import Any, Optional, List, Type from cognee.shared.logging_utils import get_logger -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.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.structured_output_framework.llitellm_instructor.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 f0afddc19..0a4990201 100644 --- a/cognee/modules/retrieval/natural_language_retriever.py +++ b/cognee/modules/retrieval/natural_language_retriever.py @@ -2,8 +2,12 @@ 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.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.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 7613eccc8..22a7fce49 100644 --- a/cognee/modules/retrieval/utils/completion.py +++ b/cognee/modules/retrieval/utils/completion.py @@ -1,5 +1,10 @@ -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 +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 57804dbd4..3470968f8 100644 --- a/cognee/modules/retrieval/utils/description_to_codepart_search.py +++ b/cognee/modules/retrieval/utils/description_to_codepart_search.py @@ -9,7 +9,9 @@ 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.structured_output_framework.llitellm_instructor.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 8adced97e..d5093d0cb 100644 --- a/cognee/shared/data_models.py +++ b/cognee/shared/data_models.py @@ -4,7 +4,9 @@ from enum import Enum, auto from typing import Any, Dict, List, Optional, Union from pydantic import BaseModel, Field -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.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 bba413d69..43cd38a5c 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,9 @@ 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.infrastructure.llm.structured_output_framework.llitellm_instructor.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 1ec22a02d..c3d89905a 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,13 @@ 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.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.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 ca9f40e81..105062cb5 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,13 @@ from typing import List, Tuple from pydantic import BaseModel -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.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 7c7d0819e..119eb6661 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,11 @@ 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.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 e5749ba50..647be9de7 100644 --- a/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py +++ b/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py @@ -1,8 +1,13 @@ from typing import List from pydantic import BaseModel -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.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 d74099ed6..c61c76876 100644 --- a/cognee/tasks/graph/extract_graph_from_code.py +++ b/cognee/tasks/graph/extract_graph_from_code.py @@ -1,19 +1,26 @@ import asyncio from typing import Type, List from pydantic import BaseModel -from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.extraction 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': +if base.structured_output_framework == "BAML": print(f"Using BAML framework: {base.structured_output_framework}") - from cognee.infrastructure.llm.structured_output_framework.baml_src.extraction import extract_content_graph + from cognee.infrastructure.llm.structured_output_framework.baml_src.extraction import ( + extract_content_graph, + ) else: print(f"Using llitellm_instructor framework: {base.structured_output_framework}") - from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.extraction import extract_content_graph + 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 a13b7967d..9aefa1409 100644 --- a/cognee/tasks/graph/extract_graph_from_data.py +++ b/cognee/tasks/graph/extract_graph_from_data.py @@ -7,13 +7,18 @@ 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.base_config import get_base_config + base = get_base_config() -if base.structured_output_framework == 'BAML': +if base.structured_output_framework == "BAML": print(f"Using BAML framework: {base.structured_output_framework}") - from cognee.infrastructure.llm.structured_output_framework.baml_src.extraction import extract_content_graph + from cognee.infrastructure.llm.structured_output_framework.baml_src.extraction import ( + extract_content_graph, + ) else: print(f"Using llitellm_instructor framework: {base.structured_output_framework}") - from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.extraction import extract_content_graph + 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, diff --git a/cognee/tasks/graph/infer_data_ontology.py b/cognee/tasks/graph/infer_data_ontology.py index 6c4e6a8b1..45ad74c04 100644 --- a/cognee/tasks/graph/infer_data_ontology.py +++ b/cognee/tasks/graph/infer_data_ontology.py @@ -15,8 +15,12 @@ from pydantic import BaseModel from cognee.modules.graph.exceptions import EntityNotFoundError from cognee.modules.ingestion.exceptions import IngestionError -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.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 diff --git a/cognee/tasks/summarization/summarize_code.py b/cognee/tasks/summarization/summarize_code.py index be2eee570..701b4b560 100644 --- a/cognee/tasks/summarization/summarize_code.py +++ b/cognee/tasks/summarization/summarize_code.py @@ -6,12 +6,18 @@ from cognee.infrastructure.engine import DataPoint from cognee.base_config import get_base_config base = get_base_config() -if base.structured_output_framework == 'BAML': +if base.structured_output_framework == "BAML": print(f"Using BAML framework for code summarization: {base.structured_output_framework}") - from cognee.infrastructure.llm.structured_output_framework.baml_src.extraction import extract_code_summary + from cognee.infrastructure.llm.structured_output_framework.baml_src.extraction import ( + extract_code_summary, + ) else: - print(f"Using llitellm_instructor framework for code summarization: {base.structured_output_framework}") - from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.extraction import extract_code_summary + print( + f"Using llitellm_instructor framework for code summarization: {base.structured_output_framework}" + ) + 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 2785f4a68..b0c7a8ab2 100644 --- a/cognee/tasks/summarization/summarize_text.py +++ b/cognee/tasks/summarization/summarize_text.py @@ -5,12 +5,18 @@ from pydantic import BaseModel from cognee.base_config import get_base_config base = get_base_config() -if base.structured_output_framework == 'BAML': +if base.structured_output_framework == "BAML": print(f"Using BAML framework for text summarization: {base.structured_output_framework}") - from cognee.infrastructure.llm.structured_output_framework.baml_src.extraction import extract_summary + from cognee.infrastructure.llm.structured_output_framework.baml_src.extraction import ( + extract_summary, + ) else: - print(f"Using llitellm_instructor framework for text summarization: {base.structured_output_framework}") - from cognee.infrastructure.llm.structured_output_framework.llitellm_instructor.extraction import extract_summary + print( + f"Using llitellm_instructor framework for text summarization: {base.structured_output_framework}" + ) + 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 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 9426b6d69..7d18bb6b9 100644 --- a/cognee/tests/unit/infrastructure/test_embedding_rate_limiting_realistic.py +++ b/cognee/tests/unit/infrastructure/test_embedding_rate_limiting_realistic.py @@ -3,8 +3,12 @@ import time import asyncio import logging -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.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 0405704b3..ec64bc1a4 100644 --- a/cognee/tests/unit/infrastructure/test_rate_limiting_realistic.py +++ b/cognee/tests/unit/infrastructure/test_rate_limiting_realistic.py @@ -2,8 +2,12 @@ import asyncio import os from unittest.mock import patch from cognee.shared.logging_utils import get_logger -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 +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/examples/python/agentic_reasoning_procurement_example.py b/examples/python/agentic_reasoning_procurement_example.py index 5fd8a1e12..4c99ab8bb 100644 --- a/examples/python/agentic_reasoning_procurement_example.py +++ b/examples/python/agentic_reasoning_procurement_example.py @@ -3,7 +3,9 @@ import logging import cognee import asyncio -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.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 28d73c23f..c372830b3 100644 --- a/examples/python/graphiti_example.py +++ b/examples/python/graphiti_example.py @@ -12,8 +12,13 @@ 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.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.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 = [