Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 14 additions & 4 deletions cognee/modules/data/extraction/extract_summary.py
Original file line number Diff line number Diff line change
@@ -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]):
Expand All @@ -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
2 changes: 1 addition & 1 deletion cognee/tasks/repo_processor/get_repo_file_dependencies.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down
37 changes: 37 additions & 0 deletions cognee/tasks/summarization/mock_summary.py
Original file line number Diff line number Diff line change
@@ -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.",
)
Loading