feat: Adds mock summary for codegraph pipeline
This commit is contained in:
parent
852532fcad
commit
4689e55e68
2 changed files with 51 additions and 4 deletions
|
|
@ -1,10 +1,11 @@
|
||||||
from typing import Type
|
from typing import Type
|
||||||
|
import os
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
from cognee.infrastructure.llm.get_llm_client import get_llm_client
|
from cognee.infrastructure.llm.get_llm_client import get_llm_client
|
||||||
from cognee.infrastructure.llm.prompts import read_query_prompt
|
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]):
|
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
|
return llm_output
|
||||||
|
|
||||||
async def extract_code_summary(content: str):
|
async def extract_code_summary(content: str):
|
||||||
|
enable_mocking = os.getenv("MOCK_CODE_SUMMARY", "false")
|
||||||
return await extract_summary(content, response_model=SummarizedCode)
|
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
|
||||||
|
|
|
||||||
37
cognee/tasks/summarization/mock_summary.py
Normal file
37
cognee/tasks/summarization/mock_summary.py
Normal file
|
|
@ -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…
Add table
Reference in a new issue