cognee 0.5.1.dev0__py3-none-any.whl → 0.5.2.dev0__py3-none-any.whl

cognee/modules/retrieval/graph_completion_retriever.py

@@ -7,7 +7,7 @@ from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
 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.base_retriever import BaseRetriever
 from cognee.modules.retrieval.utils.brute_force_triplet_search import brute_force_triplet_search
 from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
 from cognee.modules.retrieval.utils.session_cache import (
@@ -16,26 +16,24 @@ from cognee.modules.retrieval.utils.session_cache import (
 )
 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.access_tracking import update_node_access_timestamps
 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
+from cognee.modules.graph.utils import get_entity_nodes_from_triplets
 
 logger = get_logger("GraphCompletionRetriever")
 
 
-class GraphCompletionRetriever(BaseGraphRetriever):
+class GraphCompletionRetriever(BaseRetriever):
     """
     Retriever for handling graph-based completion searches.
 
-    This class provides methods to retrieve graph nodes and edges, resolve them into a
-    human-readable format, and generate completions based on graph context. Public methods
-    include:
-    - resolve_edges_to_text
-    - get_triplets
-    - get_context
-    - get_completion
+    This class implements the retrieval pipeline by searching for graph triplets (get_retrieved_objects function),
+    resolving those triplets into human-readable text context (get_context_from_objects function), and generating
+    LLM completions using the retrieved graph data (get_completion_from_context function).
     """
 
     def __init__(
@@ -49,6 +47,8 @@ class GraphCompletionRetriever(BaseGraphRetriever):
         save_interaction: bool = False,
         wide_search_top_k: Optional[int] = 100,
         triplet_distance_penalty: Optional[float] = 3.5,
+        session_id: Optional[str] = None,
+        response_model: Type = str,
     ):
         """Initialize retriever with prompt paths and search parameters."""
         self.save_interaction = save_interaction
@@ -60,6 +60,39 @@ class GraphCompletionRetriever(BaseGraphRetriever):
         self.node_type = node_type
         self.node_name = node_name
         self.triplet_distance_penalty = triplet_distance_penalty
+        # session_id (Optional[str]): Identifier for managing conversation history.
+        self.session_id = session_id
+        # response_model (Type): The Pydantic model or type for the expected response.
+        self.response_model = response_model
+
+    async def get_retrieved_objects(self, query: str) -> List[Edge]:
+        """
+        Performs a brute-force triplet search on the graph and updates access timestamps.
+
+        Args:
+            query (str): The search query to find relevant graph triplets.
+
+        Returns:
+            List[Edge]: A list of retrieved Edge objects (triplets).
+                       Returns an empty list if the graph is empty or no results 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:
+            logger.warning("Empty context was provided to the completion")
+            return []
+        # TODO: Remove when refactor of timestamps tracking is merged
+        entity_nodes = get_entity_nodes_from_triplets(triplets)
+        await update_node_access_timestamps(entity_nodes)
+
+        return triplets
 
     async def resolve_edges_to_text(self, retrieved_edges: list) -> str:
         """
@@ -115,72 +148,54 @@ class GraphCompletionRetriever(BaseGraphRetriever):
 
         return found_triplets
 
-    async def get_context(self, query: str) -> List[Edge]:
+    async def get_context_from_objects(self, query, retrieved_objects) -> str:
         """
-        Retrieves and resolves graph triplets into context based on a query.
+        Transforms raw retrieved graph triplets into a textual context string.
 
-        Parameters:
-        -----------
-
-            - query (str): The query string used to retrieve context from the graph triplets.
+        Args:
+            query (str): The original search query.
+            retrieved_objects (List[Edge]): The raw triplets returned from the search.
+                                            Output of the get_retrieved_objects method.
 
         Returns:
-        --------
+            str: A string representing the resolved graph context.
+                 Returns an empty list (as string) if no triplets are provided.
 
-            - str: A string representing the resolved context from the retrieved triplets, or an
-              empty string if no triplets are found.
+        Note: To avoid duplicate retrievals, ensure that retrieved_objects
+              are provided from get_retrieved_objects method call.
         """
-        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)
+        triplets = retrieved_objects
 
         if len(triplets) == 0:
             logger.warning("Empty context was provided to the completion")
-            return []
+            return ""
 
-        # context = await self.resolve_edges_to_text(triplets)
+        return await self.resolve_edges_to_text(triplets)
 
-        return triplets
-
-    async def convert_retrieved_objects_to_context(self, triplets: List[Edge]):
-        context = await self.resolve_edges_to_text(triplets)
-        return context
-
-    async def get_completion(
+    async def get_completion_from_context(
         self,
         query: str,
-        context: Optional[List[Edge]] = None,
-        session_id: Optional[str] = None,
-        response_model: Type = str,
+        retrieved_objects: Optional[List[Edge]],
+        context: str,
     ) -> List[Any]:
         """
-        Generates a completion using graph connections context based on a query.
+        Generates an LLM response based on the query, context, and conversation history.
+        Optionally saves the interaction and updates the session cache.
 
-        Parameters:
-        -----------
-
-            - 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)
+        Args:
+            query (str): The user's question or prompt.
+            retrieved_objects (Optional[List[Edge]]): Raw triplets used for interaction mapping.
+                                                     Output of get_retrieved_objects method.
+            context (str): The text-resolved graph context.
+                           Output of the get_context_from_objects method.
 
         Returns:
-        --------
+            List[Any]: A list containing the generated response (completion).
 
-            - Any: A generated completion based on the query and context provided.
+        Note: To avoid duplicate retrievals, ensure that retrieved_objects and context
+              are provided from previous method calls.
         """
-        triplets = context
-
-        if triplets is None:
-            triplets = await self.get_context(query)
-
-        context_text = await resolve_edges_to_text(triplets)
 
         cache_config = CacheConfig()
         user = session_user.get()
@@ -188,33 +203,33 @@ class GraphCompletionRetriever(BaseGraphRetriever):
         session_save = user_id and cache_config.caching
 
         if session_save:
-            conversation_history = await get_conversation_history(session_id=session_id)
+            conversation_history = await get_conversation_history(session_id=self.session_id)
 
             context_summary, completion = await asyncio.gather(
-                summarize_text(context_text),
+                summarize_text(context),
                 generate_completion(
                     query=query,
-                    context=context_text,
+                    context=context,
                     user_prompt_path=self.user_prompt_path,
                     system_prompt_path=self.system_prompt_path,
                     system_prompt=self.system_prompt,
                     conversation_history=conversation_history,
-                    response_model=response_model,
+                    response_model=self.response_model,
                 ),
             )
         else:
             completion = await generate_completion(
                 query=query,
-                context=context_text,
+                context=context,
                 user_prompt_path=self.user_prompt_path,
                 system_prompt_path=self.system_prompt_path,
                 system_prompt=self.system_prompt,
-                response_model=response_model,
+                response_model=self.response_model,
             )
 
-        if self.save_interaction and context and triplets and completion:
+        if self.save_interaction and retrieved_objects and completion:
             await self.save_qa(
-                question=query, answer=completion, context=context_text, triplets=triplets
+                question=query, answer=completion, context=context, triplets=retrieved_objects
             )
 
         if session_save:
@@ -222,7 +237,7 @@ class GraphCompletionRetriever(BaseGraphRetriever):
                 query=query,
                 context_summary=context_summary,
                 answer=completion,
-                session_id=session_id,
+                session_id=self.session_id,
             )
 
         return [completion]

cognee/modules/retrieval/graph_summary_completion_retriever.py

@@ -28,6 +28,7 @@ class GraphSummaryCompletionRetriever(GraphCompletionRetriever):
         save_interaction: bool = False,
         wide_search_top_k: Optional[int] = 100,
         triplet_distance_penalty: Optional[float] = 3.5,
+        session_id: Optional[str] = None,
     ):
         """Initialize retriever with default prompt paths and search parameters."""
         super().__init__(
@@ -40,6 +41,7 @@ class GraphSummaryCompletionRetriever(GraphCompletionRetriever):
             system_prompt=system_prompt,
             wide_search_top_k=wide_search_top_k,
             triplet_distance_penalty=triplet_distance_penalty,
+            session_id=session_id,
         )
         self.summarize_prompt_path = summarize_prompt_path
 

cognee/modules/retrieval/lexical_retriever.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Optional, List, Union
 from heapq import nlargest
 
 from cognee.infrastructure.databases.graph import get_graph_engine
@@ -72,7 +72,7 @@ class LexicalRetriever(BaseRetriever):
             self._initialized = True
             logger.info("Initialized with %d document chunks", len(self.chunks))
 
-    async def get_context(self, query: str) -> Any:
+    async def get_retrieved_objects(self, query: str) -> Any:
         """Retrieves relevant chunks for the given query."""
         if not self._initialized:
             await self.initialize()
@@ -116,11 +116,36 @@ 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, session_id: Optional[str] = None
-    ) -> Any:
+    async def get_context_from_objects(self, query: str, retrieved_objects: Any) -> str:
         """
-        Returns context for the given query (retrieves if not provided).
+        Retrieves context from retrieved chunks, in text form.
+
+        Parameters:
+        -----------
+
+            - query (str): The query string used to search for relevant document chunk payloads.
+            - retrieved_objects (Any): The retrieved objects to be used for generating textual context.
+
+        Returns:
+        --------
+
+            - str: A string containing the combined text of the retrieved chunk payloads, or an
+              empty string if none are found.
+        """
+        if retrieved_objects:
+            payload_texts = [payload["text"] for payload in retrieved_objects]
+            return "\n".join(payload_texts)
+        else:
+            return ""
+
+    async def get_completion_from_context(
+        self, query: str, retrieved_objects: Any, context: Any
+    ) -> Union[List[str], List[dict]]:
+        """
+        Returns a completion for the given query.
+
+        In case of the Lexical Retriever, we do not generate a completion, we just return
+        the scored chunk payloads, i.e. the retrieved objects.
 
         Parameters:
         -----------
@@ -128,14 +153,11 @@ class LexicalRetriever(BaseRetriever):
             - 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.
+            - List[dict]: The retrieved objects, i.e. the scored payloads.
         """
-        if context is None:
-            context = await self.get_context(query)
-        return context
+        # TODO: Do we want to generate a completion using LLM here?
+        return retrieved_objects

cognee/modules/retrieval/natural_language_retriever.py

@@ -4,7 +4,6 @@ from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.infrastructure.llm.prompts import render_prompt
 from cognee.modules.retrieval.base_retriever import BaseRetriever
-from cognee.modules.retrieval.exceptions import SearchTypeNotSupported
 from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
 
 logger = get_logger("NaturalLanguageRetriever")
@@ -25,10 +24,12 @@ class NaturalLanguageRetriever(BaseRetriever):
         self,
         system_prompt_path: str = "natural_language_retriever_system.txt",
         max_attempts: int = 3,
+        session_id: Optional[str] = None,
     ):
         """Initialize retriever with optional custom prompt paths."""
         self.system_prompt_path = system_prompt_path
         self.max_attempts = max_attempts
+        self.session_id = session_id
 
     async def _get_graph_schema(self, graph_engine) -> tuple:
         """Retrieve the node and edge schemas from the graph database."""
@@ -102,7 +103,17 @@ class NaturalLanguageRetriever(BaseRetriever):
         )
         return []
 
-    async def get_context(self, query: str) -> Optional[Any]:
+    async def get_retrieved_objects(self, query: str) -> Any:
+        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_context_from_objects(self, query: str, retrieved_objects: Any) -> Optional[Any]:
         """
         Retrieves relevant context using a natural language query converted to Cypher.
 
@@ -121,17 +132,11 @@ class NaturalLanguageRetriever(BaseRetriever):
             - Optional[Any]: Returns the context retrieved from the graph database based on the
               query.
         """
-        graph_engine = await get_graph_engine()
-        is_empty = await graph_engine.is_empty()
+        # TODO: Do we want to process retrieved_objects into a context string?
+        return retrieved_objects
 
-        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, session_id: Optional[str] = None
+    async def get_completion_from_context(
+        self, query: str, retrieved_objects: Any, context: Optional[Any] = None
     ) -> Any:
         """
         Returns a completion based on the query and context.
@@ -154,7 +159,5 @@ class NaturalLanguageRetriever(BaseRetriever):
 
             - Any: Returns the completion derived from the given query and context.
         """
-        if context is None:
-            context = await self.get_context(query)
-
+        # TODO: Do we want to generate a completion using LLM here?
         return context

cognee/modules/retrieval/summaries_retriever.py

@@ -1,9 +1,10 @@
-from typing import Any, Optional
+from typing import Any, Optional, List, Union
 
 from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.modules.retrieval.base_retriever import BaseRetriever
 from cognee.modules.retrieval.exceptions.exceptions import NoDataError
+from cognee.modules.retrieval.utils.access_tracking import update_node_access_timestamps
 from cognee.infrastructure.databases.vector.exceptions.exceptions import CollectionNotFoundError
 
 logger = get_logger("SummariesRetriever")
@@ -22,13 +23,14 @@ class SummariesRetriever(BaseRetriever):
     - top_k: int - Number of top summaries to retrieve.
     """
 
-    def __init__(self, top_k: int = 5):
+    def __init__(self, top_k: int = 5, session_id: Optional[str] = None):
         """Initialize retriever with search parameters."""
         self.top_k = top_k
+        self.session_id = session_id
 
-    async def get_context(self, query: str) -> Any:
+    async def get_retrieved_objects(self, query: str) -> Any:
         """
-        Retrieves summary context based on the query.
+        Retrieves text summary objects based on the query.
 
         On encountering a missing collection, raises NoDataError with a message to add data
         first.
@@ -41,7 +43,7 @@ class SummariesRetriever(BaseRetriever):
         Returns:
         --------
 
-            - Any: A list of payloads from the retrieved summaries.
+            - Any: A list of text summaries retrieved from the search.
         """
         logger.info(
             f"Starting summary retrieval for query: '{query[:100]}{'...' if len(query) > 100 else ''}'"
@@ -51,51 +53,66 @@ class SummariesRetriever(BaseRetriever):
 
         try:
             summaries_results = await vector_engine.search(
-                "TextSummary_text", query, limit=self.top_k
+                "TextSummary_text", query, limit=self.top_k, include_payload=True
             )
             logger.info(f"Found {len(summaries_results)} summaries from vector search")
+
+            await update_node_access_timestamps(summaries_results)
+
+            return summaries_results
         except CollectionNotFoundError as error:
             logger.error("TextSummary_text collection not found in vector database")
             raise NoDataError("No data found in the system, please add data first.") from error
 
-        summary_payloads = [summary.payload for summary in summaries_results]
-        logger.info(f"Returning {len(summary_payloads)} summary payloads")
-        return summary_payloads
-
-    async def get_completion(
-        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None, **kwargs
-    ) -> Any:
+    async def get_context_from_objects(self, query: str, retrieved_objects: Any) -> str:
         """
-        Generates a completion using summaries context.
+        Retrieves relevant summaries as context.
 
-        If no context is provided, retrieves context using the query. Returns the provided
-        context or the retrieved context if none was given.
+        Fetches text summaries based on a query from a vector engine and combines their text.
+        Returns empty string if no summaries are found. Raises NoDataError if the collection is not
+        found.
 
         Parameters:
         -----------
 
-            - 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)
+            - query (str): The query string used to search for relevant text summaries.
 
         Returns:
         --------
 
-            - Any: The generated completion context, which is either provided or retrieved.
+            - str: A string containing the combined text of the retrieved summaries, or an
+              empty string if none are found.
         """
-        logger.info(
-            f"Starting completion generation for query: '{query[:100]}{'...' if len(query) > 100 else ''}'"
-        )
-
-        if context is None:
-            logger.debug("No context provided, retrieving context from vector database")
-            context = await self.get_context(query)
+        if retrieved_objects:
+            summary_payload_texts = [summary.payload["text"] for summary in retrieved_objects]
+            return "\n".join(summary_payload_texts)
         else:
-            logger.debug("Using provided context")
+            return ""
 
-        logger.info(
-            f"Returning context with {len(context) if isinstance(context, list) else 1} item(s)"
-        )
-        return context
+    async def get_completion_from_context(
+        self, query: str, retrieved_objects: Any, context: Any
+    ) -> Union[List[str], List[dict]]:
+        """
+        Generates a completion using text summaries.
+        In case of the Summaries Retriever, we do not generate a completion, we just return
+        the payloads of found summaries.
+
+        Parameters:
+        -----------
+
+            - query (str): The query string to be used for generating a completion.
+            - retrieved_objects (Any): The retrieved objects to be used for generating a completion.
+            - context (Any): The context to be used for generating a completion.
+
+        Returns:
+        --------
+
+            - List[dict]: A list of payloads of found summaries.
+        """
+        # TODO: Do we want to generate a completion using LLM here?
+        if retrieved_objects:
+            summary_payloads = [summary.payload for summary in retrieved_objects]
+            logger.info(f"Returning {len(summary_payloads)} summary payloads")
+            return summary_payloads
+        else:
+            return []