fix: removing node_type and try except
This commit is contained in:
parent
13396871c9
commit
6f06e4a5eb
3 changed files with 31 additions and 28 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue