322 lines
9.9 KiB
Python
322 lines
9.9 KiB
Python
"""
|
|
Search tracer for collecting detailed search execution traces.
|
|
|
|
The SearchTracer collects comprehensive information about each step
|
|
of the spreading activation search process for debugging and visualization.
|
|
"""
|
|
import time
|
|
from datetime import datetime, timezone
|
|
from typing import List, Optional, Dict, Any, Literal
|
|
|
|
from .search_trace import (
|
|
SearchTrace,
|
|
QueryInfo,
|
|
EntryPoint,
|
|
NodeVisit,
|
|
WeightComponents,
|
|
LinkInfo,
|
|
PruningDecision,
|
|
SearchSummary,
|
|
SearchPhaseMetrics,
|
|
)
|
|
|
|
|
|
class SearchTracer:
|
|
"""
|
|
Tracer for collecting detailed search execution information.
|
|
|
|
Usage:
|
|
tracer = SearchTracer(query="Who is Alice?", thinking_budget=50, top_k=10)
|
|
tracer.start()
|
|
|
|
# During search...
|
|
tracer.record_query_embedding(embedding)
|
|
tracer.add_entry_point(node_id, text, similarity, rank)
|
|
tracer.visit_node(...)
|
|
tracer.prune_node(...)
|
|
|
|
# After search...
|
|
trace = tracer.finalize(final_results)
|
|
json_output = trace.to_json()
|
|
"""
|
|
|
|
def __init__(self, query: str, thinking_budget: int, top_k: int):
|
|
"""
|
|
Initialize tracer.
|
|
|
|
Args:
|
|
query: Search query text
|
|
thinking_budget: Maximum nodes to explore
|
|
top_k: Number of results requested
|
|
"""
|
|
self.query_text = query
|
|
self.thinking_budget = thinking_budget
|
|
self.top_k = top_k
|
|
|
|
# Trace data
|
|
self.query_embedding: Optional[List[float]] = None
|
|
self.start_time: Optional[float] = None
|
|
self.entry_points: List[EntryPoint] = []
|
|
self.visits: List[NodeVisit] = []
|
|
self.pruned: List[PruningDecision] = []
|
|
self.phase_metrics: List[SearchPhaseMetrics] = []
|
|
|
|
# Tracking state
|
|
self.current_step = 0
|
|
self.nodes_visited_set = set() # For quick lookups
|
|
|
|
# Link statistics
|
|
self.temporal_links_followed = 0
|
|
self.semantic_links_followed = 0
|
|
self.entity_links_followed = 0
|
|
|
|
def start(self):
|
|
"""Start timing the search."""
|
|
self.start_time = time.time()
|
|
|
|
def record_query_embedding(self, embedding: List[float]):
|
|
"""Record the query embedding."""
|
|
self.query_embedding = embedding
|
|
|
|
def add_entry_point(self, node_id: str, text: str, similarity: float, rank: int):
|
|
"""
|
|
Record an entry point.
|
|
|
|
Args:
|
|
node_id: Memory unit ID
|
|
text: Memory unit text
|
|
similarity: Cosine similarity to query
|
|
rank: Rank among entry points (1-based)
|
|
"""
|
|
self.entry_points.append(
|
|
EntryPoint(
|
|
node_id=node_id,
|
|
text=text,
|
|
similarity_score=similarity,
|
|
rank=rank,
|
|
)
|
|
)
|
|
|
|
def visit_node(
|
|
self,
|
|
node_id: str,
|
|
text: str,
|
|
context: str,
|
|
event_date: datetime,
|
|
access_count: int,
|
|
is_entry_point: bool,
|
|
parent_node_id: Optional[str],
|
|
link_type: Optional[Literal["temporal", "semantic", "entity"]],
|
|
link_weight: Optional[float],
|
|
activation: float,
|
|
semantic_similarity: float,
|
|
recency: float,
|
|
frequency: float,
|
|
final_weight: float,
|
|
):
|
|
"""
|
|
Record visiting a node.
|
|
|
|
Args:
|
|
node_id: Memory unit ID
|
|
text: Memory unit text
|
|
context: Memory unit context
|
|
event_date: When the memory occurred
|
|
access_count: Access count before this search
|
|
is_entry_point: Whether this is an entry point
|
|
parent_node_id: Node that led here (None for entry points)
|
|
link_type: Type of link from parent
|
|
link_weight: Weight of link from parent
|
|
activation: Activation score
|
|
semantic_similarity: Semantic similarity to query
|
|
recency: Recency weight
|
|
frequency: Frequency weight
|
|
final_weight: Combined final weight
|
|
"""
|
|
self.current_step += 1
|
|
self.nodes_visited_set.add(node_id)
|
|
|
|
# Calculate weight contributions for transparency
|
|
weights = WeightComponents(
|
|
activation=activation,
|
|
semantic_similarity=semantic_similarity,
|
|
recency=recency,
|
|
frequency=frequency,
|
|
final_weight=final_weight,
|
|
activation_contribution=0.3 * activation,
|
|
semantic_contribution=0.3 * semantic_similarity,
|
|
recency_contribution=0.25 * recency,
|
|
frequency_contribution=0.15 * frequency,
|
|
)
|
|
|
|
visit = NodeVisit(
|
|
step=self.current_step,
|
|
node_id=node_id,
|
|
text=text,
|
|
context=context,
|
|
event_date=event_date,
|
|
access_count=access_count,
|
|
is_entry_point=is_entry_point,
|
|
parent_node_id=parent_node_id,
|
|
link_type=link_type,
|
|
link_weight=link_weight,
|
|
weights=weights,
|
|
neighbors_explored=[],
|
|
final_rank=None, # Will be set later
|
|
)
|
|
|
|
self.visits.append(visit)
|
|
|
|
# Track link statistics
|
|
if link_type == "temporal":
|
|
self.temporal_links_followed += 1
|
|
elif link_type == "semantic":
|
|
self.semantic_links_followed += 1
|
|
elif link_type == "entity":
|
|
self.entity_links_followed += 1
|
|
|
|
def add_neighbor_link(
|
|
self,
|
|
from_node_id: str,
|
|
to_node_id: str,
|
|
link_type: Literal["temporal", "semantic", "entity"],
|
|
link_weight: float,
|
|
entity_id: Optional[str],
|
|
new_activation: float,
|
|
followed: bool,
|
|
prune_reason: Optional[str] = None,
|
|
):
|
|
"""
|
|
Record a link to a neighbor (whether followed or not).
|
|
|
|
Args:
|
|
from_node_id: Source node
|
|
to_node_id: Target node
|
|
link_type: Type of link
|
|
link_weight: Weight of link
|
|
entity_id: Entity ID if link is entity-based
|
|
new_activation: Activation passed to neighbor
|
|
followed: Whether link was followed
|
|
prune_reason: Why link was not followed (if not followed)
|
|
"""
|
|
# Find the visit for the source node
|
|
visit = None
|
|
for v in self.visits:
|
|
if v.node_id == from_node_id:
|
|
visit = v
|
|
break
|
|
|
|
if visit is None:
|
|
# Node not found, skip
|
|
return
|
|
|
|
link_info = LinkInfo(
|
|
to_node_id=to_node_id,
|
|
link_type=link_type,
|
|
link_weight=link_weight,
|
|
entity_id=entity_id,
|
|
new_activation=new_activation,
|
|
followed=followed,
|
|
prune_reason=prune_reason,
|
|
)
|
|
|
|
visit.neighbors_explored.append(link_info)
|
|
|
|
def prune_node(
|
|
self,
|
|
node_id: str,
|
|
reason: Literal["already_visited", "activation_too_low", "budget_exhausted"],
|
|
activation: float,
|
|
):
|
|
"""
|
|
Record a node being pruned (not visited).
|
|
|
|
Args:
|
|
node_id: Node that was pruned
|
|
reason: Why it was pruned
|
|
activation: Activation value when pruned
|
|
"""
|
|
self.pruned.append(
|
|
PruningDecision(
|
|
node_id=node_id,
|
|
reason=reason,
|
|
activation=activation,
|
|
would_have_been_step=self.current_step + 1,
|
|
)
|
|
)
|
|
|
|
def add_phase_metric(self, phase_name: str, duration_seconds: float, details: Optional[Dict[str, Any]] = None):
|
|
"""
|
|
Record metrics for a search phase.
|
|
|
|
Args:
|
|
phase_name: Name of the phase
|
|
duration_seconds: Time taken
|
|
details: Additional phase-specific details
|
|
"""
|
|
self.phase_metrics.append(
|
|
SearchPhaseMetrics(
|
|
phase_name=phase_name,
|
|
duration_seconds=duration_seconds,
|
|
details=details or {},
|
|
)
|
|
)
|
|
|
|
def finalize(self, final_results: List[Dict[str, Any]]) -> SearchTrace:
|
|
"""
|
|
Finalize the trace and return the complete SearchTrace object.
|
|
|
|
Args:
|
|
final_results: Final ranked results returned to user
|
|
|
|
Returns:
|
|
Complete SearchTrace object
|
|
"""
|
|
if self.start_time is None:
|
|
raise ValueError("Tracer not started - call start() first")
|
|
|
|
total_duration = time.time() - self.start_time
|
|
|
|
# Set final ranks on visits based on results
|
|
for rank, result in enumerate(final_results, 1):
|
|
result_node_id = result["id"]
|
|
for visit in self.visits:
|
|
if visit.node_id == result_node_id:
|
|
visit.final_rank = rank
|
|
break
|
|
|
|
# Create query info
|
|
query_info = QueryInfo(
|
|
query_text=self.query_text,
|
|
query_embedding=self.query_embedding or [],
|
|
timestamp=datetime.now(timezone.utc),
|
|
thinking_budget=self.thinking_budget,
|
|
top_k=self.top_k,
|
|
)
|
|
|
|
# Create summary
|
|
summary = SearchSummary(
|
|
total_nodes_visited=len(self.visits),
|
|
total_nodes_pruned=len(self.pruned),
|
|
entry_points_found=len(self.entry_points),
|
|
budget_used=len(self.visits),
|
|
budget_remaining=self.thinking_budget - len(self.visits),
|
|
total_duration_seconds=total_duration,
|
|
results_returned=len(final_results),
|
|
temporal_links_followed=self.temporal_links_followed,
|
|
semantic_links_followed=self.semantic_links_followed,
|
|
entity_links_followed=self.entity_links_followed,
|
|
phase_metrics=self.phase_metrics,
|
|
)
|
|
|
|
# Create complete trace
|
|
trace = SearchTrace(
|
|
query=query_info,
|
|
entry_points=self.entry_points,
|
|
visits=self.visits,
|
|
pruned=self.pruned,
|
|
summary=summary,
|
|
final_results=final_results,
|
|
)
|
|
|
|
return trace
|