Merge pull request #385 from topoteretes/feature/cog-920-implement-mock-summaryobject-for-codegraph

Feature/cog 920 implement mock summaryobject for codegraph

Author: Vasilije1990

cognee/modules/data/extraction/extract_summary.py

@@ -1,10 +1,11 @@
 from typing import Type
-
+import os
 from pydantic import BaseModel
 
 from cognee.infrastructure.llm.get_llm_client import get_llm_client
 from cognee.infrastructure.llm.prompts import read_query_prompt
-from cognee.shared.data_models import SummarizedCode
+from cognee.shared.data_models import SummarizedCode, SummarizedClass, SummarizedFunction
+from cognee.tasks.summarization.mock_summary import get_mock_summarized_code
 
 
 async def extract_summary(content: str, response_model: Type[BaseModel]):
@@ -17,5 +18,14 @@ async def extract_summary(content: str, response_model: Type[BaseModel]):
     return llm_output
 
 async def extract_code_summary(content: str):
-    
-    return await extract_summary(content, response_model=SummarizedCode)
+    enable_mocking = os.getenv("MOCK_CODE_SUMMARY", "false")
+    if isinstance(enable_mocking, bool):
+        enable_mocking = str(enable_mocking).lower()
+    enable_mocking = enable_mocking in ("true", "1", "yes")
+
+    if enable_mocking:
+        result = get_mock_summarized_code()
+        return result
+    else:
+        result = await extract_summary(content, response_model=SummarizedCode)
+        return result

cognee/tasks/repo_processor/get_repo_file_dependencies.py

@@ -73,7 +73,7 @@ async def get_repo_file_dependencies(repo_path: str) -> AsyncGenerator[list, Non
 
     yield repo
 
-    with ProcessPoolExecutor() as executor:
+    with ProcessPoolExecutor(max_workers = 12) as executor:
         loop = asyncio.get_event_loop()
 
         tasks = [

cognee/tasks/summarization/mock_summary.py

@@ -0,0 +1,37 @@
+from cognee.shared.data_models import SummarizedCode, SummarizedClass, SummarizedFunction
+
+def get_mock_summarized_code() -> SummarizedCode:
+    return SummarizedCode(
+        file_name="mock_file.py",
+        high_level_summary="This is a mock high-level summary.",
+        key_features=["Mock feature 1", "Mock feature 2"],
+        imports=["mock_import1", "mock_import2"],
+        constants=["MOCK_CONSTANT = 'mock_value'"],
+        classes=[
+            SummarizedClass(
+                name="MockClass",
+                description="This is a mock description of the MockClass.",
+                methods=[
+                    SummarizedFunction(
+                        name="mock_method",
+                        description="This is a description of the mock method.",
+                        docstring="This is a mock method.",
+                        inputs=["mock_input: str"],
+                        outputs=["mock_output: str"],
+                        decorators=None,
+                    )
+                ],
+            )
+        ],
+        functions=[
+            SummarizedFunction(
+                name="mock_function",
+                description="This is a description of the mock function.",
+                docstring="This is a mock function.",
+                inputs=["mock_input: str"],
+                outputs=["mock_output: str"],
+                decorators=None,
+            )
+        ],
+        workflow_description="This is a mock workflow description.",
+    )