167 lines
8 KiB
Python
167 lines
8 KiB
Python
"""
|
|
Search trace models for debugging and visualization.
|
|
|
|
These Pydantic models define the structure of search traces, capturing
|
|
every step of the spreading activation search process for analysis.
|
|
"""
|
|
from datetime import datetime
|
|
from typing import List, Optional, Dict, Any, Literal
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class QueryInfo(BaseModel):
|
|
"""Information about the search query."""
|
|
query_text: str = Field(description="Original query text")
|
|
query_embedding: List[float] = Field(description="Generated query embedding vector")
|
|
timestamp: datetime = Field(description="When the query was executed")
|
|
thinking_budget: int = Field(description="Maximum nodes to explore")
|
|
top_k: int = Field(description="Number of results requested")
|
|
|
|
|
|
class EntryPoint(BaseModel):
|
|
"""An entry point node selected for search."""
|
|
node_id: str = Field(description="Memory unit ID")
|
|
text: str = Field(description="Memory unit text content")
|
|
similarity_score: float = Field(description="Cosine similarity to query", ge=0.0, le=1.0)
|
|
rank: int = Field(description="Rank among entry points (1-based)")
|
|
|
|
|
|
class WeightComponents(BaseModel):
|
|
"""Breakdown of weight calculation components."""
|
|
activation: float = Field(description="Activation from spreading", ge=0.0, le=1.0)
|
|
semantic_similarity: float = Field(description="Semantic similarity to query", ge=0.0, le=1.0)
|
|
recency: float = Field(description="Recency weight", ge=0.0, le=1.0)
|
|
frequency: float = Field(description="Normalized frequency weight", ge=0.0, le=1.0)
|
|
final_weight: float = Field(description="Combined final weight")
|
|
|
|
# Weight formula components (for transparency)
|
|
activation_contribution: float = Field(description="0.3 * activation")
|
|
semantic_contribution: float = Field(description="0.3 * semantic_similarity")
|
|
recency_contribution: float = Field(description="0.25 * recency")
|
|
frequency_contribution: float = Field(description="0.15 * frequency")
|
|
|
|
|
|
class LinkInfo(BaseModel):
|
|
"""Information about a link to a neighbor."""
|
|
to_node_id: str = Field(description="Target node ID")
|
|
link_type: Literal["temporal", "semantic", "entity"] = Field(description="Type of link")
|
|
link_weight: float = Field(description="Weight of the link (can exceed 1.0 when aggregating multiple connections)", ge=0.0)
|
|
entity_id: Optional[str] = Field(default=None, description="Entity ID if link_type is 'entity'")
|
|
new_activation: Optional[float] = Field(default=None, description="Activation that would be passed to neighbor (None for supplementary links)")
|
|
followed: bool = Field(description="Whether this link was followed (or pruned)")
|
|
prune_reason: Optional[str] = Field(default=None, description="Why link was not followed (if not followed)")
|
|
is_supplementary: bool = Field(default=False, description="Whether this is a supplementary link (multiple connections to same node)")
|
|
|
|
|
|
class NodeVisit(BaseModel):
|
|
"""Information about visiting a node during search."""
|
|
step: int = Field(description="Step number in search (1-based)")
|
|
node_id: str = Field(description="Memory unit ID")
|
|
text: str = Field(description="Memory unit text content")
|
|
context: str = Field(description="Memory unit context")
|
|
event_date: datetime = Field(description="When the memory occurred")
|
|
access_count: int = Field(description="Number of times accessed before this search")
|
|
|
|
# How this node was reached
|
|
is_entry_point: bool = Field(description="Whether this is an entry point")
|
|
parent_node_id: Optional[str] = Field(default=None, description="Node that led to this one")
|
|
link_type: Optional[Literal["temporal", "semantic", "entity"]] = Field(default=None, description="Type of link from parent")
|
|
link_weight: Optional[float] = Field(default=None, description="Weight of link from parent")
|
|
|
|
# Weights
|
|
weights: WeightComponents = Field(description="Weight calculation breakdown")
|
|
|
|
# Neighbors discovered from this node
|
|
neighbors_explored: List[LinkInfo] = Field(default_factory=list, description="Links explored from this node")
|
|
|
|
# Ranking
|
|
final_rank: Optional[int] = Field(default=None, description="Final rank in results (1-based, None if not in top-k)")
|
|
|
|
|
|
class PruningDecision(BaseModel):
|
|
"""Records when a node was considered but not visited."""
|
|
node_id: str = Field(description="Node that was pruned")
|
|
reason: Literal["already_visited", "activation_too_low", "budget_exhausted"] = Field(description="Why it was pruned")
|
|
activation: float = Field(description="Activation value when pruned")
|
|
would_have_been_step: int = Field(description="What step it would have been if visited")
|
|
|
|
|
|
class SearchPhaseMetrics(BaseModel):
|
|
"""Performance metrics for a search phase."""
|
|
phase_name: str = Field(description="Name of the phase")
|
|
duration_seconds: float = Field(description="Time taken in seconds")
|
|
details: Dict[str, Any] = Field(default_factory=dict, description="Additional phase-specific metrics")
|
|
|
|
|
|
class SearchSummary(BaseModel):
|
|
"""Summary statistics about the search."""
|
|
total_nodes_visited: int = Field(description="Total nodes visited")
|
|
total_nodes_pruned: int = Field(description="Total nodes pruned")
|
|
entry_points_found: int = Field(description="Number of entry points")
|
|
budget_used: int = Field(description="How much budget was used")
|
|
budget_remaining: int = Field(description="How much budget remained")
|
|
total_duration_seconds: float = Field(description="Total search duration")
|
|
results_returned: int = Field(description="Number of results returned")
|
|
|
|
# Link statistics
|
|
temporal_links_followed: int = Field(default=0, description="Temporal links followed")
|
|
semantic_links_followed: int = Field(default=0, description="Semantic links followed")
|
|
entity_links_followed: int = Field(default=0, description="Entity links followed")
|
|
|
|
# Phase timings
|
|
phase_metrics: List[SearchPhaseMetrics] = Field(default_factory=list, description="Metrics for each phase")
|
|
|
|
|
|
class SearchTrace(BaseModel):
|
|
"""Complete trace of a search operation."""
|
|
query: QueryInfo = Field(description="Query information")
|
|
entry_points: List[EntryPoint] = Field(description="Entry points selected for search")
|
|
visits: List[NodeVisit] = Field(description="All nodes visited during search (in order)")
|
|
pruned: List[PruningDecision] = Field(default_factory=list, description="Nodes that were pruned")
|
|
summary: SearchSummary = Field(description="Summary statistics")
|
|
|
|
# Final results (for comparison with visits)
|
|
final_results: List[Dict[str, Any]] = Field(description="Final ranked results returned to user")
|
|
|
|
model_config = {
|
|
"json_encoders": {
|
|
datetime: lambda v: v.isoformat()
|
|
}
|
|
}
|
|
|
|
def to_json(self, **kwargs) -> str:
|
|
"""Export trace as JSON string."""
|
|
return self.model_dump_json(indent=2, **kwargs)
|
|
|
|
def to_dict(self) -> dict:
|
|
"""Export trace as dictionary."""
|
|
return self.model_dump()
|
|
|
|
def get_visit_by_node_id(self, node_id: str) -> Optional[NodeVisit]:
|
|
"""Find a visit by node ID."""
|
|
for visit in self.visits:
|
|
if visit.node_id == node_id:
|
|
return visit
|
|
return None
|
|
|
|
def get_search_path_to_node(self, node_id: str) -> List[NodeVisit]:
|
|
"""Get the path from entry point to a specific node."""
|
|
path = []
|
|
current_visit = self.get_visit_by_node_id(node_id)
|
|
|
|
while current_visit:
|
|
path.insert(0, current_visit)
|
|
if current_visit.parent_node_id:
|
|
current_visit = self.get_visit_by_node_id(current_visit.parent_node_id)
|
|
else:
|
|
break
|
|
|
|
return path
|
|
|
|
def get_nodes_by_link_type(self, link_type: Literal["temporal", "semantic", "entity"]) -> List[NodeVisit]:
|
|
"""Get all nodes reached via a specific link type."""
|
|
return [v for v in self.visits if v.link_type == link_type]
|
|
|
|
def get_entry_point_nodes(self) -> List[NodeVisit]:
|
|
"""Get all entry point visits."""
|
|
return [v for v in self.visits if v.is_entry_point]
|