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

ruff format

Browse source
This commit is contained in:
hajdul88 2025-10-16 18:25:13 +02:00
parent cdbdbfd755
commit 2595a419b3
5 changed files with 26 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
cognee/modules/retrieval/EntityCompletionRetriever.py
View file

@ -6,7 +6,10 @@ from cognee.infrastructure.entities.BaseEntityExtractor import BaseEntityExtract
from cognee.infrastructure.context.BaseContextProvider import BaseContextProvider
from cognee.modules.retrieval.base_retriever import BaseRetriever
from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
from cognee.modules.retrieval.utils.session_cache import save_to_session_cache, get_conversation_history
from cognee.modules.retrieval.utils.session_cache import (
save_to_session_cache,
get_conversation_history,
)
from cognee.context_global_variables import session_user
from cognee.infrastructure.databases.cache.config import CacheConfig
@ -121,7 +124,7 @@ class EntityCompletionRetriever(BaseRetriever):
if session_save:
conversation_history = await get_conversation_history(session_id=session_id)
context_summary, completion = await asyncio.gather(
summarize_text(str(context)),
generate_completion(

7
cognee/modules/retrieval/completion_retriever.py
View file

@ -4,7 +4,10 @@ from typing import Any, Optional
from cognee.shared.logging_utils import get_logger
from cognee.infrastructure.databases.vector import get_vector_engine
from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
from cognee.modules.retrieval.utils.session_cache import save_to_session_cache, get_conversation_history
from cognee.modules.retrieval.utils.session_cache import (
save_to_session_cache,
get_conversation_history,
)
from cognee.modules.retrieval.base_retriever import BaseRetriever
from cognee.modules.retrieval.exceptions.exceptions import NoDataError
from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
@ -105,7 +108,7 @@ class CompletionRetriever(BaseRetriever):
if session_save:
conversation_history = await get_conversation_history(session_id=session_id)
context_summary, completion = await asyncio.gather(
summarize_text(context),
generate_completion(

7
cognee/modules/retrieval/graph_completion_context_extension_retriever.py
View file

@ -4,7 +4,10 @@ from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
from cognee.shared.logging_utils import get_logger
from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
from cognee.modules.retrieval.utils.session_cache import save_to_session_cache, get_conversation_history
from cognee.modules.retrieval.utils.session_cache import (
save_to_session_cache,
get_conversation_history,
)
from cognee.context_global_variables import session_user
from cognee.infrastructure.databases.cache.config import CacheConfig
@ -130,7 +133,7 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
if session_save:
conversation_history = await get_conversation_history(session_id=session_id)
context_summary, completion = await asyncio.gather(
summarize_text(context_text),
generate_completion(

9
cognee/modules/retrieval/graph_completion_cot_retriever.py
View file

@ -5,7 +5,10 @@ from cognee.shared.logging_utils import get_logger
from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
from cognee.modules.retrieval.utils.session_cache import save_to_session_cache, get_conversation_history
from cognee.modules.retrieval.utils.session_cache import (
save_to_session_cache,
get_conversation_history,
)
from cognee.infrastructure.llm.LLMGateway import LLMGateway
from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt
from cognee.context_global_variables import session_user
@ -92,13 +95,13 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
followup_question = ""
triplets = []
completion = ""
# Retrieve conversation history if session saving is enabled
cache_config = CacheConfig()
user = session_user.get()
user_id = getattr(user, "id", None)
session_save = user_id and cache_config.caching
conversation_history = ""
if session_save:
conversation_history = await get_conversation_history(session_id=session_id)

7
cognee/modules/retrieval/temporal_retriever.py
View file

@ -6,7 +6,10 @@ from typing import Any, Optional, List, Type
from operator import itemgetter
from cognee.infrastructure.databases.vector import get_vector_engine
from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
from cognee.modules.retrieval.utils.session_cache import save_to_session_cache, get_conversation_history
from cognee.modules.retrieval.utils.session_cache import (
save_to_session_cache,
get_conversation_history,
)
from cognee.infrastructure.databases.graph import get_graph_engine
from cognee.infrastructure.llm.prompts import render_prompt
from cognee.infrastructure.llm import LLMGateway
@ -172,7 +175,7 @@ class TemporalRetriever(GraphCompletionRetriever):
if session_save:
conversation_history = await get_conversation_history(session_id=session_id)
context_summary, completion = await asyncio.gather(
summarize_text(context),
generate_completion(