From 5e06e06db9df51c3a2e4194c092e1c3455aee0c2 Mon Sep 17 00:00:00 2001 From: vasilije Date: Sun, 6 Jul 2025 18:27:27 +0200 Subject: [PATCH] added fixes --- .../baml_src/extraction/extract_summary.py | 15 +++++++++++++-- .../llitellm_instructor/extraction/__init__.py | 1 + cognee/tasks/graph/extract_graph_from_code.py | 6 ++---- cognee/tasks/graph/extract_graph_from_data.py | 14 +++++++------- cognee/tasks/summarization/summarize_code.py | 4 ++-- cognee/tasks/summarization/summarize_text.py | 11 ++++++----- 6 files changed, 31 insertions(+), 20 deletions(-) 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 1f904a7dd..bc7eac031 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 @@ -4,13 +4,24 @@ 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_src.config import get_llm_config from cognee.shared.data_models import SummarizedCode -from cognee.tasks.summarization.mock_summary import get_mock_summarized_code from cognee.shared.logging_utils import get_logger -from instructor.exceptions import InstructorRetryException logger = get_logger("extract_summary_baml") +def get_mock_summarized_code(): + """Local mock function to avoid circular imports.""" + return SummarizedCode( + high_level_summary="Mock code summary", + key_features=["Mock feature 1", "Mock feature 2"], + imports=["mock_import"], + constants=["MOCK_CONSTANT"], + classes=[], + functions=[], + workflow_description="Mock workflow description", + ) + + async def extract_summary(content: str, response_model: Type[BaseModel]): """ Extract summary using BAML framework. diff --git a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/__init__.py b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/__init__.py index b6419282d..157cbe7e7 100644 --- a/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/__init__.py +++ b/cognee/infrastructure/llm/structured_output_framework/llitellm_instructor/extraction/__init__.py @@ -1 +1,2 @@ from .knowledge_graph.extract_content_graph import extract_content_graph +from .extract_summary import extract_summary, extract_code_summary diff --git a/cognee/tasks/graph/extract_graph_from_code.py b/cognee/tasks/graph/extract_graph_from_code.py index c61c76876..7e1316e02 100644 --- a/cognee/tasks/graph/extract_graph_from_code.py +++ b/cognee/tasks/graph/extract_graph_from_code.py @@ -1,15 +1,13 @@ 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.modules.chunking.models.DocumentChunk import DocumentChunk from cognee.tasks.storage import add_data_points from cognee.base_config import get_base_config +# Framework selection base = get_base_config() - 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 ( diff --git a/cognee/tasks/graph/extract_graph_from_data.py b/cognee/tasks/graph/extract_graph_from_data.py index 9aefa1409..3aeda8b70 100644 --- a/cognee/tasks/graph/extract_graph_from_data.py +++ b/cognee/tasks/graph/extract_graph_from_data.py @@ -7,7 +7,14 @@ 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 +from cognee.modules.graph.utils import ( + expand_with_nodes_and_edges, + retrieve_existing_edges, +) +from cognee.shared.data_models import KnowledgeGraph +from cognee.tasks.storage import add_data_points +# Framework selection base = get_base_config() if base.structured_output_framework == "BAML": print(f"Using BAML framework: {base.structured_output_framework}") @@ -20,13 +27,6 @@ else: extract_content_graph, ) -from cognee.modules.graph.utils import ( - expand_with_nodes_and_edges, - retrieve_existing_edges, -) -from cognee.shared.data_models import KnowledgeGraph -from cognee.tasks.storage import add_data_points - async def integrate_chunk_graphs( data_chunks: list[DocumentChunk], diff --git a/cognee/tasks/summarization/summarize_code.py b/cognee/tasks/summarization/summarize_code.py index 701b4b560..7a5801a83 100644 --- a/cognee/tasks/summarization/summarize_code.py +++ b/cognee/tasks/summarization/summarize_code.py @@ -4,7 +4,9 @@ from uuid import uuid5 from cognee.infrastructure.engine import DataPoint from cognee.base_config import get_base_config +from .models import CodeSummary +# Framework selection base = get_base_config() if base.structured_output_framework == "BAML": print(f"Using BAML framework for code summarization: {base.structured_output_framework}") @@ -19,8 +21,6 @@ else: extract_code_summary, ) -from .models import CodeSummary - async def summarize_code( code_graph_nodes: list[DataPoint], diff --git a/cognee/tasks/summarization/summarize_text.py b/cognee/tasks/summarization/summarize_text.py index b0c7a8ab2..f7f5de939 100644 --- a/cognee/tasks/summarization/summarize_text.py +++ b/cognee/tasks/summarization/summarize_text.py @@ -2,8 +2,13 @@ import asyncio from typing import Type from uuid import uuid5 from pydantic import BaseModel -from cognee.base_config import get_base_config +from cognee.base_config import get_base_config +from cognee.modules.chunking.models.DocumentChunk import DocumentChunk +from cognee.modules.cognify.config import get_cognify_config +from .models import TextSummary + +# Framework selection base = get_base_config() if base.structured_output_framework == "BAML": print(f"Using BAML framework for text summarization: {base.structured_output_framework}") @@ -18,10 +23,6 @@ else: extract_summary, ) -from cognee.modules.chunking.models.DocumentChunk import DocumentChunk -from cognee.modules.cognify.config import get_cognify_config -from .models import TextSummary - async def summarize_text( data_chunks: list[DocumentChunk], summarization_model: Type[BaseModel] = None