gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'dev' into langfuse_generation

Browse source
This commit is contained in:
hajdul88 2024-12-20 16:21:12 +01:00 • committed by GitHub
commit 0813f9e7ac
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 14 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
cognee/modules/data/extraction/extract_summary.py
View file

@ -1,12 +1,17 @@
from typing import Type import logging
import os import os
from typing import Type
from instructor.exceptions import InstructorRetryException
from pydantic import BaseModel from pydantic import BaseModel
from tenacity import RetryError
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, SummarizedClass, SummarizedFunction from cognee.shared.data_models import SummarizedCode
from cognee.tasks.summarization.mock_summary import get_mock_summarized_code from cognee.tasks.summarization.mock_summary import get_mock_summarized_code
logger = logging.getLogger("extract_summary")
async def extract_summary(content: str, response_model: Type[BaseModel]): async def extract_summary(content: str, response_model: Type[BaseModel]):
llm_client = get_llm_client() llm_client = get_llm_client()
@ -14,7 +19,7 @@ async def extract_summary(content: str, response_model: Type[BaseModel]):
system_prompt = read_query_prompt("summarize_content.txt") system_prompt = read_query_prompt("summarize_content.txt")
llm_output = await llm_client.acreate_structured_output(content, system_prompt, response_model) llm_output = await llm_client.acreate_structured_output(content, system_prompt, response_model)
return llm_output return llm_output
async def extract_code_summary(content: str): async def extract_code_summary(content: str):
@ -27,5 +32,10 @@ async def extract_code_summary(content: str):
result = get_mock_summarized_code() result = get_mock_summarized_code()
return result return result
else: else:
result = await extract_summary(content, response_model=SummarizedCode) try:
result = await extract_summary(content, response_model=SummarizedCode)
except (RetryError, InstructorRetryException) as e:
logger.error("Failed to extract code summary, falling back to mock summary", exc_info=e)
result = get_mock_summarized_code()
return result return result