From 119a6532b08d5e9823a51db3377d917ecdaaa1e3 Mon Sep 17 00:00:00 2001 From: vasilije Date: Tue, 29 Jul 2025 09:59:34 +0200 Subject: [PATCH] made sure base tests pass --- cognee/base_config.py | 6 ++++++ .../extraction/knowledge_graph/extract_content_graph.py | 2 +- .../llitellm_instructor/llm/openai/adapter.py | 4 +++- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/cognee/base_config.py b/cognee/base_config.py index f4a5c00d2..0fe132843 100644 --- a/cognee/base_config.py +++ b/cognee/base_config.py @@ -8,18 +8,24 @@ from pydantic_settings import BaseSettings, SettingsConfigDict class BaseConfig(BaseSettings): data_root_directory: str = get_absolute_path(".data_storage") + system_root_directory: str = get_absolute_path(".data_storage") monitoring_tool: object = Observer.LANGFUSE langfuse_public_key: Optional[str] = os.getenv("LANGFUSE_PUBLIC_KEY") langfuse_secret_key: Optional[str] = os.getenv("LANGFUSE_SECRET_KEY") langfuse_host: Optional[str] = os.getenv("LANGFUSE_HOST") default_user_email: Optional[str] = os.getenv("DEFAULT_USER_EMAIL") default_user_password: Optional[str] = os.getenv("DEFAULT_USER_PASSWORD") + structured_output_framework: str = os.getenv( + "STRUCTURED_OUTPUT_FRAMEWORK", "llitellm_instructor" + ) model_config = SettingsConfigDict(env_file=".env", extra="allow") def to_dict(self) -> dict: return { "data_root_directory": self.data_root_directory, + "system_root_directory": self.system_root_directory, "monitoring_tool": self.monitoring_tool, + "structured_output_framework": self.structured_output_framework, } 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 6ddfdbde2..465b4843c 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 @@ -42,7 +42,7 @@ async def extract_content_graph( # return graph # else: - graph = await b.ExtractContentGraph( + graph = await b.ExtractContentGraphGeneric( content, mode=mode, baml_options={"client_registry": baml_registry} ) 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 0bf527a62..55a919795 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 @@ -8,7 +8,9 @@ from litellm.exceptions import ContentPolicyViolationError from instructor.exceptions import InstructorRetryException from cognee.exceptions import InvalidValueError -from cognee.infrastructure.llm.prompts import read_query_prompt +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, )