added fixes

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.

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

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 (

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],

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],

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