Skip to content

Commit

Permalink
Merge branch 'dev' into COG-697-ingest-non-code-files
Browse files Browse the repository at this point in the history
  • Loading branch information
hajdul88 authored Dec 18, 2024
2 parents 3f96416 + 1c8115b commit 8f4d480
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 5 deletions.
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.",
)

0 comments on commit 8f4d480

Please sign in to comment.