cognee 0.3.5__py3-none-any.whl → 0.3.7__py3-none-any.whl

cognee/modules/retrieval/graph_completion_retriever.py

@@ -1,20 +1,26 @@
+import asyncio
 from typing import Any, Optional, Type, List
 from uuid import NAMESPACE_OID, uuid5
 
 from cognee.infrastructure.engine import DataPoint
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
-from cognee.modules.users.methods import get_default_user
 from cognee.tasks.storage import add_data_points
 from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.modules.graph.utils.convert_node_to_data_point import get_all_subclasses
 from cognee.modules.retrieval.base_graph_retriever import BaseGraphRetriever
 from cognee.modules.retrieval.utils.brute_force_triplet_search import brute_force_triplet_search
-from cognee.modules.retrieval.utils.completion import generate_completion
+from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
+from cognee.modules.retrieval.utils.session_cache import (
+    save_conversation_history,
+    get_conversation_history,
+)
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.retrieval.utils.extract_uuid_from_node import extract_uuid_from_node
 from cognee.modules.retrieval.utils.models import CogneeUserInteraction
 from cognee.modules.engine.models.node_set import NodeSet
 from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.context_global_variables import session_user
+from cognee.infrastructure.databases.cache.config import CacheConfig
 
 logger = get_logger("GraphCompletionRetriever")
 
@@ -118,6 +124,13 @@ class GraphCompletionRetriever(BaseGraphRetriever):
             - str: A string representing the resolved context from the retrieved triplets, or an
               empty string if no triplets are found.
         """
+        graph_engine = await get_graph_engine()
+        is_empty = await graph_engine.is_empty()
+
+        if is_empty:
+            logger.warning("Search attempt on an empty knowledge graph")
+            return []
+
         triplets = await self.get_triplets(query)
 
         if len(triplets) == 0:
@@ -132,6 +145,7 @@ class GraphCompletionRetriever(BaseGraphRetriever):
         self,
         query: str,
         context: Optional[List[Edge]] = None,
+        session_id: Optional[str] = None,
     ) -> List[str]:
         """
         Generates a completion using graph connections context based on a query.
@@ -142,6 +156,8 @@ class GraphCompletionRetriever(BaseGraphRetriever):
             - query (str): The query string for which a completion is generated.
             - context (Optional[Any]): Optional context to use for generating the completion; if
               not provided, context is retrieved based on the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
 
         Returns:
         --------
@@ -155,19 +171,47 @@ class GraphCompletionRetriever(BaseGraphRetriever):
 
         context_text = await resolve_edges_to_text(triplets)
 
-        completion = await generate_completion(
-            query=query,
-            context=context_text,
-            user_prompt_path=self.user_prompt_path,
-            system_prompt_path=self.system_prompt_path,
-            system_prompt=self.system_prompt,
-        )
+        cache_config = CacheConfig()
+        user = session_user.get()
+        user_id = getattr(user, "id", None)
+        session_save = user_id and cache_config.caching
+
+        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(
+                    query=query,
+                    context=context_text,
+                    user_prompt_path=self.user_prompt_path,
+                    system_prompt_path=self.system_prompt_path,
+                    system_prompt=self.system_prompt,
+                    conversation_history=conversation_history,
+                ),
+            )
+        else:
+            completion = await generate_completion(
+                query=query,
+                context=context_text,
+                user_prompt_path=self.user_prompt_path,
+                system_prompt_path=self.system_prompt_path,
+                system_prompt=self.system_prompt,
+            )
 
         if self.save_interaction and context and triplets and completion:
             await self.save_qa(
                 question=query, answer=completion, context=context_text, triplets=triplets
             )
 
+        if session_save:
+            await save_conversation_history(
+                query=query,
+                context_summary=context_summary,
+                answer=completion,
+                session_id=session_id,
+            )
+
         return [completion]
 
     async def save_qa(self, question: str, answer: str, context: str, triplets: List) -> None:
@@ -194,7 +238,7 @@ class GraphCompletionRetriever(BaseGraphRetriever):
             belongs_to_set=interactions_node_set,
         )
 
-        await add_data_points(data_points=[cognee_user_interaction], update_edge_collection=False)
+        await add_data_points(data_points=[cognee_user_interaction])
 
         relationships = []
         relationship_name = "used_graph_element_to_answer"

cognee/modules/retrieval/lexical_retriever.py

@@ -116,8 +116,26 @@ class LexicalRetriever(BaseRetriever):
         else:
             return [self.payloads[chunk_id] for chunk_id, _ in top_results]
 
-    async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
-        """Returns context for the given query (retrieves if not provided)."""
+    async def get_completion(
+        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None
+    ) -> Any:
+        """
+        Returns context for the given query (retrieves if not provided).
+
+        Parameters:
+        -----------
+
+            - query (str): The query string to retrieve context for.
+            - context (Optional[Any]): Optional pre-fetched context; if None, it retrieves
+              the context for the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
+
+        Returns:
+        --------
+
+            - Any: The context, either provided or retrieved.
+        """
         if context is None:
             context = await self.get_context(query)
         return context

cognee/modules/retrieval/natural_language_retriever.py

@@ -122,10 +122,17 @@ class NaturalLanguageRetriever(BaseRetriever):
               query.
         """
         graph_engine = await get_graph_engine()
+        is_empty = await graph_engine.is_empty()
+
+        if is_empty:
+            logger.warning("Search attempt on an empty knowledge graph")
+            return []
 
         return await self._execute_cypher_query(query, graph_engine)
 
-    async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
+    async def get_completion(
+        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None
+    ) -> Any:
         """
         Returns a completion based on the query and context.
 
@@ -139,6 +146,8 @@ class NaturalLanguageRetriever(BaseRetriever):
             - query (str): The natural language query to get a completion from.
             - context (Optional[Any]): The context in which to base the completion; if not
               provided, it will be retrieved using the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
 
         Returns:
         --------

cognee/modules/retrieval/summaries_retriever.py

@@ -62,7 +62,9 @@ class SummariesRetriever(BaseRetriever):
         logger.info(f"Returning {len(summary_payloads)} summary payloads")
         return summary_payloads
 
-    async def get_completion(self, query: str, context: Optional[Any] = None, **kwargs) -> Any:
+    async def get_completion(
+        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None, **kwargs
+    ) -> Any:
         """
         Generates a completion using summaries context.
 
@@ -75,6 +77,8 @@ class SummariesRetriever(BaseRetriever):
             - query (str): The search query for generating the completion.
             - context (Optional[Any]): Optional context for the completion; if not provided,
               will be retrieved based on the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
 
         Returns:
         --------

cognee/modules/retrieval/temporal_retriever.py

@@ -1,16 +1,22 @@
 import os
+import asyncio
 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
+from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
+from cognee.modules.retrieval.utils.session_cache import (
+    save_conversation_history,
+    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
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
 from cognee.shared.logging_utils import get_logger
-
+from cognee.context_global_variables import session_user
+from cognee.infrastructure.databases.cache.config import CacheConfig
 
 from cognee.tasks.temporal_graph.models import QueryInterval
 
@@ -137,17 +143,63 @@ class TemporalRetriever(GraphCompletionRetriever):
 
         return self.descriptions_to_string(top_k_events)
 
-    async def get_completion(self, query: str, context: Optional[str] = None) -> List[str]:
-        """Generates a response using the query and optional context."""
+    async def get_completion(
+        self, query: str, context: Optional[str] = None, session_id: Optional[str] = None
+    ) -> List[str]:
+        """
+        Generates a response using the query and optional context.
+
+        Parameters:
+        -----------
+
+            - query (str): The query string for which a completion is generated.
+            - context (Optional[str]): Optional context to use; if None, it will be
+              retrieved based on the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
+
+        Returns:
+        --------
+
+            - List[str]: A list containing the generated completion.
+        """
         if not context:
             context = await self.get_context(query=query)
 
         if context:
-            completion = await generate_completion(
-                query=query,
-                context=context,
-                user_prompt_path=self.user_prompt_path,
-                system_prompt_path=self.system_prompt_path,
-            )
+            # Check if we need to generate context summary for caching
+            cache_config = CacheConfig()
+            user = session_user.get()
+            user_id = getattr(user, "id", None)
+            session_save = user_id and cache_config.caching
+
+            if session_save:
+                conversation_history = await get_conversation_history(session_id=session_id)
+
+                context_summary, completion = await asyncio.gather(
+                    summarize_text(context),
+                    generate_completion(
+                        query=query,
+                        context=context,
+                        user_prompt_path=self.user_prompt_path,
+                        system_prompt_path=self.system_prompt_path,
+                        conversation_history=conversation_history,
+                    ),
+                )
+            else:
+                completion = await generate_completion(
+                    query=query,
+                    context=context,
+                    user_prompt_path=self.user_prompt_path,
+                    system_prompt_path=self.system_prompt_path,
+                )
+
+            if session_save:
+                await save_conversation_history(
+                    query=query,
+                    context_summary=context_summary,
+                    answer=completion,
+                    session_id=session_id,
+                )
 
         return [completion]

cognee/modules/retrieval/user_qa_feedback.py

@@ -8,7 +8,7 @@ from cognee.shared.logging_utils import get_logger
 from cognee.modules.retrieval.base_feedback import BaseFeedback
 from cognee.modules.retrieval.utils.models import CogneeUserFeedback
 from cognee.modules.retrieval.utils.models import UserFeedbackEvaluation
-from cognee.tasks.storage import add_data_points
+from cognee.tasks.storage import add_data_points, index_graph_edges
 
 logger = get_logger("CompletionRetriever")
 
@@ -47,7 +47,7 @@ class UserQAFeedback(BaseFeedback):
             belongs_to_set=feedbacks_node_set,
         )
 
-        await add_data_points(data_points=[cognee_user_feedback], update_edge_collection=False)
+        await add_data_points(data_points=[cognee_user_feedback])
 
         relationships = []
         relationship_name = "gives_feedback_to"
@@ -76,6 +76,7 @@ class UserQAFeedback(BaseFeedback):
         if len(relationships) > 0:
             graph_engine = await get_graph_engine()
             await graph_engine.add_edges(relationships)
+            await index_graph_edges(relationships)
             await graph_engine.apply_feedback_weight(
                 node_ids=to_node_ids, weight=feedback_sentiment.score
             )

cognee/modules/retrieval/utils/completion.py

@@ -9,12 +9,17 @@ async def generate_completion(
     user_prompt_path: str,
     system_prompt_path: str,
     system_prompt: Optional[str] = None,
+    conversation_history: Optional[str] = None,
 ) -> str:
     """Generates a completion using LLM with given context and prompts."""
     args = {"question": query, "context": context}
     user_prompt = render_prompt(user_prompt_path, args)
     system_prompt = system_prompt if system_prompt else read_query_prompt(system_prompt_path)
 
+    if conversation_history:
+        #:TODO: I would separate the history and put it into the system prompt but we have to test what works best with longer convos
+        system_prompt = conversation_history + "\nTASK:" + system_prompt
+
     return await LLMGateway.acreate_structured_output(
         text_input=user_prompt,
         system_prompt=system_prompt,

cognee/modules/retrieval/utils/description_to_codepart_search.py

@@ -62,7 +62,7 @@ async def code_description_to_code_part(
 
     try:
         if include_docs:
-            search_results = await search(query_text=query, query_type="INSIGHTS")
+            search_results = await search(query_text=query, query_type="GRAPH_COMPLETION")
 
             concatenated_descriptions = " ".join(
                 obj["description"]

cognee/modules/retrieval/utils/session_cache.py

@@ -0,0 +1,156 @@
+from typing import Optional, List, Dict, Any
+from cognee.context_global_variables import session_user
+from cognee.infrastructure.databases.cache.config import CacheConfig
+from cognee.infrastructure.databases.exceptions import CacheConnectionError
+from cognee.shared.logging_utils import get_logger
+
+logger = get_logger("session_cache")
+
+
+async def save_conversation_history(
+    query: str,
+    context_summary: str,
+    answer: str,
+    session_id: Optional[str] = None,
+) -> bool:
+    """
+    Saves Q&A interaction to the session cache if user is authenticated and caching is enabled.
+
+    Handles cache unavailability gracefully by logging warnings instead of failing.
+
+    Parameters:
+    -----------
+
+        - query (str): The user's query/question.
+        - context_summary (str): Summarized context used for generating the answer.
+        - answer (str): The generated answer/completion.
+        - session_id (Optional[str]): Session identifier. Defaults to 'default_session' if None.
+
+    Returns:
+    --------
+
+        - bool: True if successfully saved to cache, False otherwise.
+    """
+    try:
+        cache_config = CacheConfig()
+        user = session_user.get()
+        user_id = getattr(user, "id", None)
+
+        if not (user_id and cache_config.caching):
+            logger.debug("Session caching disabled or user not authenticated")
+            return False
+
+        if session_id is None:
+            session_id = "default_session"
+
+        from cognee.infrastructure.databases.cache.get_cache_engine import get_cache_engine
+
+        cache_engine = get_cache_engine()
+
+        if cache_engine is None:
+            logger.warning("Cache engine not available, skipping session save")
+            return False
+
+        await cache_engine.add_qa(
+            str(user_id),
+            session_id=session_id,
+            question=query,
+            context=context_summary,
+            answer=answer,
+        )
+
+        logger.info(
+            f"Successfully saved Q&A to session cache: user_id={user_id}, session_id={session_id}"
+        )
+        return True
+
+    except CacheConnectionError as e:
+        logger.warning(f"Cache unavailable, continuing without session save: {e.message}")
+        return False
+
+    except Exception as e:
+        logger.error(
+            f"Unexpected error saving to session cache: {type(e).__name__}: {str(e)}. Continuing without caching."
+        )
+        return False
+
+
+async def get_conversation_history(
+    session_id: Optional[str] = None,
+) -> str:
+    """
+    Retrieves conversation history from cache and formats it as text.
+
+    Returns formatted conversation history with time, question, context, and answer
+    for the last N Q&A pairs (N is determined by cache engine default).
+
+    Parameters:
+    -----------
+
+        - session_id (Optional[str]): Session identifier. Defaults to 'default_session' if None.
+
+    Returns:
+    --------
+
+        - str: Formatted conversation history string, or empty string if no history or error.
+
+    Format:
+    -------
+
+        Previous conversation:
+
+        [2024-01-15 10:30:45]
+        QUESTION: What is X?
+        CONTEXT: X is a concept...
+        ANSWER: X is...
+
+        [2024-01-15 10:31:20]
+        QUESTION: How does Y work?
+        CONTEXT: Y is related to...
+        ANSWER: Y works by...
+    """
+    try:
+        cache_config = CacheConfig()
+        user = session_user.get()
+        user_id = getattr(user, "id", None)
+
+        if not (user_id and cache_config.caching):
+            logger.debug("Session caching disabled or user not authenticated")
+            return ""
+
+        if session_id is None:
+            session_id = "default_session"
+
+        from cognee.infrastructure.databases.cache.get_cache_engine import get_cache_engine
+
+        cache_engine = get_cache_engine()
+
+        if cache_engine is None:
+            logger.warning("Cache engine not available, skipping conversation history retrieval")
+            return ""
+
+        history_entries = await cache_engine.get_latest_qa(str(user_id), session_id)
+
+        if not history_entries:
+            logger.debug("No conversation history found")
+            return ""
+
+        history_text = "Previous conversation:\n\n"
+        for entry in history_entries:
+            history_text += f"[{entry.get('time', 'Unknown time')}]\n"
+            history_text += f"QUESTION: {entry.get('question', '')}\n"
+            history_text += f"CONTEXT: {entry.get('context', '')}\n"
+            history_text += f"ANSWER: {entry.get('answer', '')}\n\n"
+
+        logger.debug(f"Retrieved {len(history_entries)} conversation history entries")
+        return history_text
+
+    except CacheConnectionError as e:
+        logger.warning(f"Cache unavailable, continuing without conversation history: {e.message}")
+        return ""
+
+    except Exception as e:
+        logger.warning(
+            f"Unexpected error retrieving conversation history: {type(e).__name__}: {str(e)}"
+        )
+        return ""

cognee/modules/search/methods/get_search_type_tools.py

@@ -9,7 +9,6 @@ from cognee.modules.search.exceptions import UnsupportedSearchTypeError
 # Retrievers
 from cognee.modules.retrieval.user_qa_feedback import UserQAFeedback
 from cognee.modules.retrieval.chunks_retriever import ChunksRetriever
-from cognee.modules.retrieval.insights_retriever import InsightsRetriever
 from cognee.modules.retrieval.summaries_retriever import SummariesRetriever
 from cognee.modules.retrieval.completion_retriever import CompletionRetriever
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
@@ -44,10 +43,6 @@ async def get_search_type_tools(
             SummariesRetriever(top_k=top_k).get_completion,
             SummariesRetriever(top_k=top_k).get_context,
         ],
-        SearchType.INSIGHTS: [
-            InsightsRetriever(top_k=top_k).get_completion,
-            InsightsRetriever(top_k=top_k).get_context,
-        ],
         SearchType.CHUNKS: [
             ChunksRetriever(top_k=top_k).get_completion,
             ChunksRetriever(top_k=top_k).get_context,

cognee/modules/search/methods/no_access_control_search.py

@@ -1,12 +1,16 @@
 from typing import Any, List, Optional, Tuple, Type, Union
 
+from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.modules.data.models.Dataset import Dataset
 from cognee.modules.engine.models.node_set import NodeSet
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
 from cognee.modules.search.types import SearchType
+from cognee.shared.logging_utils import get_logger
 
 from .get_search_type_tools import get_search_type_tools
 
+logger = get_logger()
+
 
 async def no_access_control_search(
     query_type: SearchType,
@@ -19,6 +23,7 @@ async def no_access_control_search(
     save_interaction: bool = False,
     last_k: Optional[int] = None,
     only_context: bool = False,
+    session_id: Optional[str] = None,
 ) -> Tuple[Any, Union[str, List[Edge]], List[Dataset]]:
     search_tools = await get_search_type_tools(
         query_type=query_type,
@@ -31,6 +36,12 @@ async def no_access_control_search(
         save_interaction=save_interaction,
         last_k=last_k,
     )
+    graph_engine = await get_graph_engine()
+    is_empty = await graph_engine.is_empty()
+
+    if is_empty:
+        # TODO: we can log here, but not all search types use graph. Still keeping this here for reviewer input
+        logger.warning("Search attempt on an empty knowledge graph")
     if len(search_tools) == 2:
         [get_completion, get_context] = search_tools
 
@@ -38,7 +49,7 @@ async def no_access_control_search(
             return None, await get_context(query_text), []
 
         context = await get_context(query_text)
-        result = await get_completion(query_text, context)
+        result = await get_completion(query_text, context, session_id=session_id)
     else:
         unknown_tool = search_tools[0]
         result = await unknown_tool(query_text)