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/repo_processor/get_repo_file_dependencies.py b/cognee/tasks/repo_processor/get_repo_file_dependencies.py index b82870796..221af6cf6 100644 --- a/cognee/tasks/repo_processor/get_repo_file_dependencies.py +++ b/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 = [ 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