From 6f06e4a5eb1143ddcb2ad08132486630b8a2deae Mon Sep 17 00:00:00 2001
From: chinu0609 <chinmayhbhosale02@gmail.com>
Date: Fri, 31 Oct 2025 00:17:13 +0530
Subject: [PATCH] fix: removing node_type and try except

---
 cognee/modules/retrieval/chunks_retriever.py  |  2 +-
 .../modules/retrieval/summaries_retriever.py  |  2 +-
 .../retrieval/utils/access_tracking.py        | 55 ++++++++++---------
 3 files changed, 31 insertions(+), 28 deletions(-)

diff --git a/cognee/modules/retrieval/chunks_retriever.py b/cognee/modules/retrieval/chunks_retriever.py
index f821fc902..be1f95811 100644
--- a/cognee/modules/retrieval/chunks_retriever.py
+++ b/cognee/modules/retrieval/chunks_retriever.py
@@ -49,7 +49,7 @@ class ChunksRetriever(BaseRetriever):
         try:
             found_chunks = await vector_engine.search("DocumentChunk_text", query, limit=self.top_k)
             logger.info(f"Found {len(found_chunks)} chunks from vector search")
-            await update_node_access_timestamps(found_chunks, "DocumentChunk")
+            await update_node_access_timestamps(found_chunks)
 
         except CollectionNotFoundError as error:
             logger.error("DocumentChunk_text collection not found in vector database")
diff --git a/cognee/modules/retrieval/summaries_retriever.py b/cognee/modules/retrieval/summaries_retriever.py
index 9ac8b096d..0df750d22 100644
--- a/cognee/modules/retrieval/summaries_retriever.py
+++ b/cognee/modules/retrieval/summaries_retriever.py
@@ -56,7 +56,7 @@ class SummariesRetriever(BaseRetriever):
             )
             logger.info(f"Found {len(summaries_results)} summaries from vector search")
             
-            await update_node_access_timestamps(summaries_results, "TextSummary")
+            await update_node_access_timestamps(summaries_results)
         
         except CollectionNotFoundError as error:
             logger.error("TextSummary_text collection not found in vector database")
diff --git a/cognee/modules/retrieval/utils/access_tracking.py b/cognee/modules/retrieval/utils/access_tracking.py
index ca5ed88cd..79afd25db 100644
--- a/cognee/modules/retrieval/utils/access_tracking.py
+++ b/cognee/modules/retrieval/utils/access_tracking.py
@@ -1,4 +1,4 @@
-
+  
 """Utilities for tracking data access in retrievers."""  
   
 import json  
@@ -11,51 +11,54 @@ from cognee.shared.logging_utils import get_logger
 logger = get_logger(__name__)  
   
   
-async def update_node_access_timestamps(items: List[Any], node_type: str):  
+async def update_node_access_timestamps(items: List[Any]):  
     """  
     Update last_accessed_at for nodes in Kuzu graph database.  
+    Automatically determines node type from the graph database.  
       
     Parameters  
     ----------  
     items : List[Any]  
         List of items with payload containing 'id' field (from vector search results)  
-    node_type : str  
-        Type of node to update (e.g., 'DocumentChunk', 'Entity', 'TextSummary')  
     """  
     if not items:  
         return  
       
     graph_engine = await get_graph_engine()  
-    # Convert to milliseconds since epoch (matching the field format)  
     timestamp_ms = int(datetime.now(timezone.utc).timestamp() * 1000)  
       
     for item in items:  
-        # Extract ID from payload (vector search results have this structure)  
+        # Extract ID from payload  
         item_id = item.payload.get("id") if hasattr(item, 'payload') else item.get("id")  
         if not item_id:  
             continue  
               
-        try:  
-            # Get current node properties from Kuzu's Node table  
-            result = await graph_engine.query(  
-                "MATCH (n:Node {id: $id}) WHERE n.type = $node_type RETURN n.properties as props",  
-                {"id": str(item_id), "node_type": node_type}  
+        # try:  
+        # Query to get both node type and properties in one call  
+        result = await graph_engine.query(  
+            "MATCH (n:Node {id: $id}) RETURN n.type as node_type, n.properties as props",  
+            {"id": str(item_id)}  
+        )  
+          
+        if result and len(result) > 0 and result[0]:  
+            node_type = result[0][0]  # First column: node_type  
+            props_json = result[0][1]  # Second column: properties  
+              
+            # Parse existing properties JSON  
+            props = json.loads(props_json) if props_json else {}  
+            # Update last_accessed_at with millisecond timestamp  
+            props["last_accessed_at"] = timestamp_ms  
+              
+            # Write back to graph database  
+            await graph_engine.query(  
+                "MATCH (n:Node {id: $id}) SET n.properties = $props",  
+                {"id": str(item_id), "props": json.dumps(props)}  
             )  
               
-            if result and len(result) > 0 and result[0][0]:  
-                # Parse existing properties JSON  
-                props = json.loads(result[0][0]) if result[0][0] else {}  
-                # Update last_accessed_at with millisecond timestamp  
-                props["last_accessed_at"] = timestamp_ms  
+            logger.debug(f"Updated access timestamp for {node_type} node {item_id}")  
                   
-                # Write back to graph database  
-                await graph_engine.query(  
-                    "MATCH (n:Node {id: $id}) WHERE n.type = $node_type SET n.properties = $props",  
-                    {"id": str(item_id), "node_type": node_type, "props": json.dumps(props)}  
-                )  
-        except Exception as e:  
-            logger.warning(f"Failed to update timestamp for {node_type} {item_id}: {e}")  
-            continue  
+        # except Exception as e:  
+        #     logger.error(f"Failed to update timestamp for node {item_id}: {e}")  
+        #     continue  
       
-    logger.debug(f"Updated access timestamps for {len(items)} {node_type} nodes")
-
+    logger.debug(f"Updated access timestamps for {len(items)} nodes")