9cfdd464a9
* fix(retain): preserve normalized experience fact types and remove deprecated opinion type The ExtractedFactType conversion was re-checking for raw "assistant" fact_type after the parsing layer had already normalized it to "experience". Since fact_from_llm.fact_type was always "experience" (never "assistant"), the ternary always fell through to "world", silently losing experience classification. Also removes the deprecated "opinion" fact type from internal extraction models, database constraints/indexes (via migration), and dead code paths. The public API surface (descriptions, response models, backwards-compat filter) is unchanged. * refactor(retain): drop unused confidence_score column The confidence_score column was only ever non-null for opinion facts (which are now removed). It was always written as NULL and never read back from the database. Remove it from: - DB model and migration (DROP COLUMN) - INSERT queries in fact_storage.py - retain_async/retain_batch_async parameters - RetainContext/RetainResult extension models - RetainBatch dataclass
481 lines
16 KiB
Python
481 lines
16 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 UTC, datetime
|
|
from typing import Any, Literal
|
|
|
|
from .trace import (
|
|
EntryPoint,
|
|
LinkInfo,
|
|
NodeVisit,
|
|
PruningDecision,
|
|
QueryInfo,
|
|
RerankedResult,
|
|
RetrievalMethodResults,
|
|
RetrievalResult,
|
|
RRFMergeResult,
|
|
SearchPhaseMetrics,
|
|
SearchSummary,
|
|
SearchTrace,
|
|
TemporalConstraint,
|
|
WeightComponents,
|
|
)
|
|
|
|
|
|
class SearchTracer:
|
|
"""
|
|
Tracer for collecting detailed search execution information.
|
|
|
|
Usage:
|
|
tracer = SearchTracer(query="Who is Alice?", budget=50, max_tokens=4096)
|
|
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,
|
|
budget: int,
|
|
max_tokens: int,
|
|
tags: list[str] | None = None,
|
|
tags_match: str | None = None,
|
|
):
|
|
"""
|
|
Initialize tracer.
|
|
|
|
Args:
|
|
query: Search query text
|
|
budget: Maximum nodes to explore
|
|
max_tokens: Maximum tokens to return in results
|
|
tags: Tags filter applied to recall
|
|
tags_match: Tags matching mode (any, all, any_strict, all_strict)
|
|
"""
|
|
self.query_text = query
|
|
self.budget = budget
|
|
self.max_tokens = max_tokens
|
|
self.tags = tags
|
|
self.tags_match = tags_match
|
|
|
|
# Trace data
|
|
self.query_embedding: list[float] | None = None
|
|
self.start_time: float | None = None
|
|
self.entry_points: list[EntryPoint] = []
|
|
self.visits: list[NodeVisit] = []
|
|
self.pruned: list[PruningDecision] = []
|
|
self.phase_metrics: list[SearchPhaseMetrics] = []
|
|
|
|
# Temporal constraint detected from query
|
|
self.temporal_constraint: TemporalConstraint | None = None
|
|
|
|
# New 4-way retrieval tracking
|
|
self.retrieval_results: list[RetrievalMethodResults] = []
|
|
self.rrf_merged: list[RRFMergeResult] = []
|
|
self.reranked: list[RerankedResult] = []
|
|
|
|
# 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 record_temporal_constraint(self, start: datetime | None, end: datetime | None):
|
|
"""Record the detected temporal constraint from query analysis."""
|
|
if start is not None or end is not None:
|
|
self.temporal_constraint = TemporalConstraint(start=start, end=end)
|
|
|
|
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)
|
|
"""
|
|
# Clamp similarity to [0.0, 1.0] to handle floating-point precision
|
|
similarity = min(1.0, max(0.0, similarity))
|
|
|
|
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 | None,
|
|
is_entry_point: bool,
|
|
parent_node_id: str | None,
|
|
link_type: Literal["temporal", "semantic", "entity"] | None,
|
|
link_weight: float | None,
|
|
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
|
|
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)
|
|
|
|
# Clamp values to handle floating-point precision issues
|
|
# (sometimes normalization produces values like 1.0000005 instead of 1.0)
|
|
semantic_similarity = min(1.0, max(0.0, semantic_similarity))
|
|
recency = min(1.0, max(0.0, recency))
|
|
frequency = min(1.0, max(0.0, frequency))
|
|
|
|
# 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,
|
|
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: str | None,
|
|
new_activation: float | None,
|
|
followed: bool,
|
|
prune_reason: str | None = None,
|
|
is_supplementary: bool = False,
|
|
):
|
|
"""
|
|
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 (None for supplementary links)
|
|
followed: Whether link was followed
|
|
prune_reason: Why link was not followed (if not followed)
|
|
is_supplementary: Whether this is a supplementary link (multiple connections)
|
|
"""
|
|
# 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,
|
|
is_supplementary=is_supplementary,
|
|
)
|
|
|
|
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: dict[str, Any] | None = 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 add_retrieval_results(
|
|
self,
|
|
method_name: Literal["semantic", "bm25", "graph", "temporal"],
|
|
results: list[tuple], # List of (doc_id, data) tuples
|
|
duration_seconds: float,
|
|
score_field: str, # e.g., "similarity", "bm25_score"
|
|
metadata: dict[str, Any] | None = None,
|
|
fact_type: str | None = None,
|
|
):
|
|
"""
|
|
Record results from a single retrieval method.
|
|
|
|
Args:
|
|
method_name: Name of the retrieval method
|
|
results: List of (doc_id, data) tuples from retrieval
|
|
duration_seconds: Time taken for this retrieval
|
|
score_field: Field name containing the score in data dict
|
|
metadata: Optional metadata about this retrieval method
|
|
fact_type: Fact type this retrieval was for (world, experience)
|
|
"""
|
|
retrieval_results = []
|
|
for rank, (doc_id, data) in enumerate(results, start=1):
|
|
score = data.get(score_field)
|
|
if score is None:
|
|
score = 0.0
|
|
retrieval_results.append(
|
|
RetrievalResult(
|
|
rank=rank,
|
|
node_id=doc_id,
|
|
text=data.get("text") or "",
|
|
context=data.get("context") or "",
|
|
event_date=data.get("event_date"),
|
|
fact_type=data.get("fact_type") or fact_type,
|
|
score=score,
|
|
score_name=score_field,
|
|
)
|
|
)
|
|
|
|
self.retrieval_results.append(
|
|
RetrievalMethodResults(
|
|
method_name=method_name,
|
|
fact_type=fact_type,
|
|
results=retrieval_results,
|
|
duration_seconds=duration_seconds,
|
|
metadata=metadata or {},
|
|
)
|
|
)
|
|
|
|
def add_rrf_merged(self, merged_results: list[tuple]):
|
|
"""
|
|
Record RRF merged results.
|
|
|
|
Args:
|
|
merged_results: List of (doc_id, data, rrf_meta) tuples from RRF merge
|
|
"""
|
|
self.rrf_merged = []
|
|
for rank, (doc_id, data, rrf_meta) in enumerate(merged_results, start=1):
|
|
self.rrf_merged.append(
|
|
RRFMergeResult(
|
|
node_id=doc_id,
|
|
text=data.get("text", ""),
|
|
rrf_score=rrf_meta.get("rrf_score", 0.0),
|
|
source_ranks=rrf_meta.get("source_ranks", {}),
|
|
final_rrf_rank=rank,
|
|
)
|
|
)
|
|
|
|
def add_reranked(self, reranked_results: list[dict[str, Any]], rrf_merged: list):
|
|
"""
|
|
Record reranked results.
|
|
|
|
Args:
|
|
reranked_results: List of result dicts after reranking
|
|
rrf_merged: Original RRF merged results for comparison
|
|
"""
|
|
# Build map of node_id -> rrf_rank
|
|
rrf_rank_map = {}
|
|
for item in self.rrf_merged:
|
|
rrf_rank_map[item.node_id] = item.final_rrf_rank
|
|
|
|
self.reranked = []
|
|
for rank, result in enumerate(reranked_results, start=1):
|
|
node_id = result["id"]
|
|
rrf_rank = rrf_rank_map.get(node_id, len(rrf_merged) + 1)
|
|
rank_change = rrf_rank - rank # Positive = moved up
|
|
|
|
# Extract score components (only include non-None values)
|
|
# Keys from ScoredResult.to_dict(): cross_encoder_score, cross_encoder_score_normalized,
|
|
# rrf_normalized, temporal, recency, combined_score, weight
|
|
score_components = {}
|
|
for key in [
|
|
"cross_encoder_score",
|
|
"cross_encoder_score_normalized",
|
|
"rrf_score",
|
|
"rrf_normalized",
|
|
"temporal",
|
|
"recency",
|
|
"combined_score",
|
|
]:
|
|
if key in result and result[key] is not None:
|
|
score_components[key] = result[key]
|
|
|
|
self.reranked.append(
|
|
RerankedResult(
|
|
node_id=node_id,
|
|
text=result.get("text", ""),
|
|
rerank_score=result.get("weight", 0.0),
|
|
rerank_rank=rank,
|
|
rrf_rank=rrf_rank,
|
|
rank_change=rank_change,
|
|
score_components=score_components,
|
|
)
|
|
)
|
|
|
|
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(UTC),
|
|
budget=self.budget,
|
|
max_tokens=self.max_tokens,
|
|
tags=self.tags,
|
|
tags_match=self.tags_match,
|
|
temporal_constraint=self.temporal_constraint,
|
|
)
|
|
|
|
# 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.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,
|
|
retrieval_results=self.retrieval_results,
|
|
rrf_merged=self.rrf_merged,
|
|
reranked=self.reranked,
|
|
entry_points=self.entry_points,
|
|
visits=self.visits,
|
|
pruned=self.pruned,
|
|
summary=summary,
|
|
final_results=final_results,
|
|
)
|
|
|
|
return trace
|