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

cognee/modules/retrieval/temporal_retriever.py

@@ -5,6 +5,7 @@ from datetime import datetime
 
 from operator import itemgetter
 from cognee.infrastructure.databases.vector import get_vector_engine
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
 from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
 from cognee.modules.retrieval.utils.session_cache import (
     save_conversation_history,
@@ -49,6 +50,8 @@ class TemporalRetriever(GraphCompletionRetriever):
         node_name: Optional[List[str]] = None,
         wide_search_top_k: Optional[int] = 100,
         triplet_distance_penalty: Optional[float] = 3.5,
+        session_id: Optional[str] = None,
+        response_model: Type = str,
     ):
         super().__init__(
             user_prompt_path=user_prompt_path,
@@ -58,6 +61,8 @@ class TemporalRetriever(GraphCompletionRetriever):
             node_name=node_name,
             wide_search_top_k=wide_search_top_k,
             triplet_distance_penalty=triplet_distance_penalty,
+            session_id=session_id,
+            response_model=response_model,
         )
         self.user_prompt_path = user_prompt_path
         self.system_prompt_path = system_prompt_path
@@ -98,7 +103,7 @@ class TemporalRetriever(GraphCompletionRetriever):
 
     async def filter_top_k_events(self, relevant_events, scored_results):
         # Build a score lookup from vector search results
-        score_lookup = {res.payload["id"]: res.score for res in scored_results}
+        score_lookup = {res.id: res.score for res in scored_results}
 
         events_with_scores = []
         for event in relevant_events[0]["events"]:
@@ -109,9 +114,7 @@ class TemporalRetriever(GraphCompletionRetriever):
 
         return events_with_scores[: self.top_k]
 
-    async def get_context(self, query: str) -> Any:
-        """Retrieves context based on the query."""
-
+    async def get_retrieved_objects(self, query: str) -> dict:
         time_from, time_to = await self.extract_time_from_query(query)
 
         graph_engine = await get_graph_engine()
@@ -127,7 +130,7 @@ class TemporalRetriever(GraphCompletionRetriever):
                 "No timestamps identified based on the query, performing retrieval using triplet search on events and entities."
             )
             triplets = await self.get_triplets(query)
-            return await self.resolve_edges_to_text(triplets)
+            return {"triplets": triplets}
 
         if ids:
             relevant_events = await graph_engine.collect_events(ids=ids)
@@ -136,7 +139,7 @@ class TemporalRetriever(GraphCompletionRetriever):
                 "No events identified based on timestamp filtering, performing retrieval using triplet search on events and entities."
             )
             triplets = await self.get_triplets(query)
-            return await self.resolve_edges_to_text(triplets)
+            return {"triplets": triplets}
 
         vector_engine = get_vector_engine()
         query_vector = (await vector_engine.embedding_engine.embed_text([query]))[0]
@@ -145,16 +148,26 @@ class TemporalRetriever(GraphCompletionRetriever):
             collection_name="Event_name", query_vector=query_vector, limit=None
         )
 
-        top_k_events = await self.filter_top_k_events(relevant_events, vector_search_results)
+        return {"relevant_events": relevant_events, "vector_search_results": vector_search_results}
 
-        return self.descriptions_to_string(top_k_events)
+    async def get_context_from_objects(self, query: str, retrieved_objects: Any) -> Any:
+        """Retrieves context based on the query."""
+        if retrieved_objects.get("relevant_events", None) and retrieved_objects.get(
+            "vector_search_results", None
+        ):
+            top_k_events = await self.filter_top_k_events(
+                retrieved_objects.get("relevant_events"),
+                retrieved_objects.get("vector_search_results", None),
+            )
+            return self.descriptions_to_string(top_k_events)
+        else:
+            # In case no events were found, fall back to triplet context
+            triplets = retrieved_objects.get("triplets", [])
+            context_text = await self.resolve_edges_to_text(triplets)
+            return context_text
 
-    async def get_completion(
-        self,
-        query: str,
-        context: Optional[str] = None,
-        session_id: Optional[str] = None,
-        response_model: Type = str,
+    async def get_completion_from_context(
+        self, query: str, retrieved_objects: Any = None, context: Optional[str] = None
     ) -> List[Any]:
         """
         Generates a response using the query and optional context.
@@ -174,45 +187,42 @@ class TemporalRetriever(GraphCompletionRetriever):
 
             - List[str]: A list containing the generated completion.
         """
-        if not context:
-            context = await self.get_context(query=query)
-
-        if context:
-            # 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,
-                        response_model=response_model,
-                    ),
-                )
-            else:
-                completion = await generate_completion(
+
+        # 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=self.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,
-                    response_model=response_model,
-                )
+                    conversation_history=conversation_history,
+                    response_model=self.response_model,
+                ),
+            )
+        else:
+            completion = await generate_completion(
+                query=query,
+                context=context,
+                user_prompt_path=self.user_prompt_path,
+                system_prompt_path=self.system_prompt_path,
+                response_model=self.response_model,
+            )
 
-            if session_save:
-                await save_conversation_history(
-                    query=query,
-                    context_summary=context_summary,
-                    answer=completion,
-                    session_id=session_id,
-                )
+        if session_save:
+            await save_conversation_history(
+                query=query,
+                context_summary=context_summary,
+                answer=completion,
+                session_id=self.session_id,
+            )
 
         return [completion]

cognee/modules/retrieval/triplet_retriever.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Any, Optional, Type, List
+from typing import Any, Optional, Type, List, Union
 
 from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.databases.vector import get_vector_engine
@@ -32,19 +32,23 @@ class TripletRetriever(BaseRetriever):
         system_prompt_path: str = "answer_simple_question.txt",
         system_prompt: Optional[str] = None,
         top_k: Optional[int] = 5,
+        session_id: Optional[str] = None,
+        response_model: Type = str,
     ):
         """Initialize retriever with optional custom prompt paths."""
         self.user_prompt_path = user_prompt_path
         self.system_prompt_path = system_prompt_path
         self.top_k = top_k if top_k is not None else 5
         self.system_prompt = system_prompt
+        self.session_id = session_id
+        self.response_model = response_model
 
-    async def get_context(self, query: str) -> str:
+    async def get_retrieved_objects(self, query: str) -> Any:
         """
-        Retrieves relevant triplets as context.
+        Retrieves relevant triplets.
 
-        Fetches triplets based on a query from a vector engine and combines their text.
-        Returns empty string if no triplets are found. Raises NoDataError if the collection is not
+        Fetches triplets based on a query from a vector engine.
+        Returns empty list if no triplets are found. Raises NoDataError if the collection is not
         found.
 
         Parameters:
@@ -55,8 +59,7 @@ class TripletRetriever(BaseRetriever):
         Returns:
         --------
 
-            - str: A string containing the combined text of the retrieved triplets, or an
-              empty string if none are found.
+            - Any: A list containing the retrieved triplets, or an empty list if none are found.
         """
         vector_engine = get_vector_engine()
 
@@ -67,25 +70,30 @@ class TripletRetriever(BaseRetriever):
                     "In order to use TRIPLET_COMPLETION first use the create_triplet_embeddings memify pipeline. "
                 )
 
-            found_triplets = await vector_engine.search("Triplet_text", query, limit=self.top_k)
+            found_triplets = await vector_engine.search(
+                "Triplet_text", query, limit=self.top_k, include_payload=True
+            )
 
             if len(found_triplets) == 0:
-                return ""
+                return []
 
-            triplets_payload = [found_triplet.payload["text"] for found_triplet in found_triplets]
-            combined_context = "\n".join(triplets_payload)
-            return combined_context
+            return found_triplets
         except CollectionNotFoundError as error:
             logger.error("Triplet_text collection not found")
             raise NoDataError("No data found in the system, please add data first.") from error
 
-    async def get_completion(
-        self,
-        query: str,
-        context: Optional[Any] = None,
-        session_id: Optional[str] = None,
-        response_model: Type = str,
-    ) -> List[Any]:
+    async def get_context_from_objects(self, query: str, retrieved_objects: Any) -> str:
+        if retrieved_objects:
+            triplets_payload = [
+                found_triplet.payload["text"] for found_triplet in retrieved_objects
+            ]
+            combined_context = "\n".join(triplets_payload)
+            return combined_context
+        return ""
+
+    async def get_completion_from_context(
+        self, query: str, retrieved_objects: Any, context: Any
+    ) -> Union[List[str], List[dict]]:
         """
         Generates an LLM completion using the context.
 
@@ -107,9 +115,6 @@ class TripletRetriever(BaseRetriever):
 
             - Any: The generated completion based on the provided query and context.
         """
-        if context is None:
-            context = await self.get_context(query)
-
         cache_config = CacheConfig()
         user = session_user.get()
         user_id = getattr(user, "id", None)
@@ -119,14 +124,11 @@ class TripletRetriever(BaseRetriever):
             completion = await self._get_completion_with_session(
                 query=query,
                 context=context,
-                session_id=session_id,
-                response_model=response_model,
             )
         else:
             completion = await self._get_completion_without_session(
                 query=query,
                 context=context,
-                response_model=response_model,
             )
 
         return [completion]
@@ -135,11 +137,9 @@ class TripletRetriever(BaseRetriever):
         self,
         query: str,
         context: str,
-        session_id: Optional[str],
-        response_model: Type,
     ) -> Any:
         """Generate completion with session history and caching."""
-        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),
@@ -150,7 +150,7 @@ class TripletRetriever(BaseRetriever):
                 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,
             ),
         )
 
@@ -158,7 +158,7 @@ class TripletRetriever(BaseRetriever):
             query=query,
             context_summary=context_summary,
             answer=completion,
-            session_id=session_id,
+            session_id=self.session_id,
         )
 
         return completion
@@ -167,7 +167,6 @@ class TripletRetriever(BaseRetriever):
         self,
         query: str,
         context: str,
-        response_model: Type,
     ) -> Any:
         """Generate completion without session history."""
         completion = await generate_completion(
@@ -176,7 +175,7 @@ class TripletRetriever(BaseRetriever):
             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,
         )
 
         return completion

cognee/modules/retrieval/utils/access_tracking.py

@@ -0,0 +1,88 @@
+"""Utilities for tracking data access in retrievers."""
+
+import json
+from datetime import datetime, timezone
+from typing import List, Any
+from uuid import UUID
+import os
+from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.data.models import Data
+from cognee.shared.logging_utils import get_logger
+from sqlalchemy import update
+from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph
+
+logger = get_logger(__name__)
+
+
+async def update_node_access_timestamps(items: List[Any]):
+    if os.getenv("ENABLE_LAST_ACCESSED", "false").lower() != "true":
+        return
+
+    if not items:
+        return
+
+    graph_engine = await get_graph_engine()
+    timestamp_dt = datetime.now(timezone.utc)
+
+    # Extract node IDs
+    node_ids = []
+    for item in items:
+        item_id = item.payload.get("id") if hasattr(item, "payload") else item.get("id")
+        if item_id:
+            node_ids.append(str(item_id))
+
+    if not node_ids:
+        return
+
+    # Focus on document-level tracking via projection
+    try:
+        doc_ids = await _find_origin_documents_via_projection(graph_engine, node_ids)
+        if doc_ids:
+            await _update_sql_records(doc_ids, timestamp_dt)
+    except Exception as e:
+        logger.error(f"Failed to update SQL timestamps: {e}")
+        raise
+
+
+async def _find_origin_documents_via_projection(graph_engine, node_ids):
+    """Find origin documents using graph projection instead of DB queries"""
+    # Project the entire graph with necessary properties
+    memory_fragment = CogneeGraph()
+    await memory_fragment.project_graph_from_db(
+        graph_engine,
+        node_properties_to_project=["id", "type"],
+        edge_properties_to_project=["relationship_name"],
+    )
+
+    # Find origin documents by traversing the in-memory graph
+    doc_ids = set()
+    for node_id in node_ids:
+        node = memory_fragment.get_node(node_id)
+        if node and node.get_attribute("type") == "DocumentChunk":
+            # Traverse edges to find connected documents
+            for edge in node.get_skeleton_edges():
+                # Get the neighbor node
+                neighbor = (
+                    edge.get_destination_node()
+                    if edge.get_source_node().id == node_id
+                    else edge.get_source_node()
+                )
+                if neighbor and neighbor.get_attribute("type") in ["TextDocument", "Document"]:
+                    doc_ids.add(neighbor.id)
+
+    return list(doc_ids)
+
+
+async def _update_sql_records(doc_ids, timestamp_dt):
+    """Update SQL Data table (same for all providers)"""
+    db_engine = get_relational_engine()
+    async with db_engine.get_async_session() as session:
+        stmt = (
+            update(Data)
+            .where(Data.id.in_([UUID(doc_id) for doc_id in doc_ids]))
+            .values(last_accessed=timestamp_dt)
+        )
+
+        await session.execute(stmt)
+        await session.commit()