From 4689e55e68b175df70eaaf9f7ddcb3db82c28a3b Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Wed, 18 Dec 2024 16:42:48 +0100
Subject: [PATCH] feat: Adds mock summary for codegraph pipeline

---
 .../data/extraction/extract_summary.py        | 18 +++++++--
 cognee/tasks/summarization/mock_summary.py    | 37 +++++++++++++++++++
 2 files changed, 51 insertions(+), 4 deletions(-)
 create mode 100644 cognee/tasks/summarization/mock_summary.py

diff --git a/cognee/modules/data/extraction/extract_summary.py b/cognee/modules/data/extraction/extract_summary.py
index 10d429da9..102d5bff4 100644
--- a/cognee/modules/data/extraction/extract_summary.py
+++ b/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
diff --git a/cognee/tasks/summarization/mock_summary.py b/cognee/tasks/summarization/mock_summary.py
new file mode 100644
index 000000000..e61385bb1
--- /dev/null
+++ b/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.",
+    )
\ No newline at end of file