180 lines
7.1 KiB
Python
180 lines
7.1 KiB
Python
"""
|
|
Example demonstrating search tracing functionality.
|
|
|
|
This script shows how to:
|
|
1. Enable search tracing
|
|
2. Retrieve the trace object
|
|
3. Export trace to JSON for visualization
|
|
"""
|
|
import asyncio
|
|
import json
|
|
from datetime import datetime, timezone
|
|
from memory import TemporalSemanticMemory
|
|
|
|
|
|
async def main():
|
|
"""Run the trace example."""
|
|
# Initialize memory system
|
|
memory = TemporalSemanticMemory()
|
|
|
|
try:
|
|
# Create a test agent
|
|
agent_id = f"trace_demo_{datetime.now(timezone.utc).timestamp()}"
|
|
|
|
print("=" * 70)
|
|
print("SEARCH TRACE EXAMPLE")
|
|
print("=" * 70)
|
|
|
|
# Store some test memories
|
|
print("\n1. Storing test memories...")
|
|
await memory.put_async(
|
|
agent_id=agent_id,
|
|
content="Alice works at Google as a software engineer in Mountain View",
|
|
context="conversation",
|
|
)
|
|
await memory.put_async(
|
|
agent_id=agent_id,
|
|
content="Bob also works at Google but in the New York office",
|
|
context="conversation",
|
|
)
|
|
await memory.put_async(
|
|
agent_id=agent_id,
|
|
content="Charlie founded TechCorp, a startup in San Francisco",
|
|
context="conversation",
|
|
)
|
|
await memory.put_async(
|
|
agent_id=agent_id,
|
|
content="Alice and Bob met at a Google conference last year",
|
|
context="conversation",
|
|
)
|
|
print(" ✓ 4 memories stored")
|
|
|
|
# Perform search with tracing enabled
|
|
print("\n2. Searching with trace enabled...")
|
|
query = "Who works at Google?"
|
|
|
|
results, trace = await memory.search_async(
|
|
agent_id=agent_id,
|
|
query=query,
|
|
thinking_budget=30,
|
|
top_k=5,
|
|
enable_trace=True,
|
|
)
|
|
|
|
print(f" ✓ Search completed")
|
|
|
|
# Display trace summary
|
|
print("\n3. Trace Summary:")
|
|
print(f" - Query: {trace.query.query_text}")
|
|
print(f" - Thinking budget: {trace.query.thinking_budget}")
|
|
print(f" - Entry points found: {len(trace.entry_points)}")
|
|
print(f" - Total nodes visited: {trace.summary.total_nodes_visited}")
|
|
print(f" - Total nodes pruned: {trace.summary.total_nodes_pruned}")
|
|
print(f" - Budget used: {trace.summary.budget_used}")
|
|
print(f" - Budget remaining: {trace.summary.budget_remaining}")
|
|
print(f" - Results returned: {trace.summary.results_returned}")
|
|
print(f" - Total duration: {trace.summary.total_duration_seconds:.3f}s")
|
|
print(f" - Temporal links followed: {trace.summary.temporal_links_followed}")
|
|
print(f" - Semantic links followed: {trace.summary.semantic_links_followed}")
|
|
print(f" - Entity links followed: {trace.summary.entity_links_followed}")
|
|
|
|
# Show entry points
|
|
print("\n4. Entry Points:")
|
|
for ep in trace.entry_points:
|
|
print(f" [{ep.rank}] {ep.text[:60]}... (similarity: {ep.similarity_score:.3f})")
|
|
|
|
# Show visited nodes with their paths
|
|
print("\n5. Search Path (First 5 visits):")
|
|
for i, visit in enumerate(trace.visits[:5], 1):
|
|
indent = " "
|
|
if visit.is_entry_point:
|
|
print(f"{indent}[{i}] ENTRY POINT: {visit.text[:60]}...")
|
|
else:
|
|
parent = f"from {visit.parent_node_id[:8]}" if visit.parent_node_id else "?"
|
|
link_info = f"via {visit.link_type}" if visit.link_type else ""
|
|
print(f"{indent}[{i}] {parent} {link_info}: {visit.text[:60]}...")
|
|
|
|
print(f"{indent} - Activation: {visit.weights.activation:.3f}")
|
|
print(f"{indent} - Semantic sim: {visit.weights.semantic_similarity:.3f}")
|
|
print(f"{indent} - Recency: {visit.weights.recency:.3f}")
|
|
print(f"{indent} - Final weight: {visit.weights.final_weight:.3f}")
|
|
|
|
if visit.neighbors_explored:
|
|
followed = sum(1 for n in visit.neighbors_explored if n.followed)
|
|
pruned = len(visit.neighbors_explored) - followed
|
|
print(f"{indent} - Neighbors: {followed} followed, {pruned} pruned")
|
|
|
|
# Show pruning decisions
|
|
if trace.pruned:
|
|
print(f"\n6. Pruning Decisions (showing first 5 of {len(trace.pruned)}):")
|
|
for prune in trace.pruned[:5]:
|
|
print(f" - Node {prune.node_id[:8]}: {prune.reason} (activation: {prune.activation:.3f})")
|
|
|
|
# Show phase metrics
|
|
print("\n7. Phase Metrics:")
|
|
for pm in trace.summary.phase_metrics:
|
|
print(f" - {pm.phase_name}: {pm.duration_seconds:.3f}s")
|
|
if pm.details:
|
|
for key, value in pm.details.items():
|
|
if isinstance(value, float):
|
|
print(f" • {key}: {value:.3f}")
|
|
else:
|
|
print(f" • {key}: {value}")
|
|
|
|
# Export to JSON
|
|
print("\n8. Exporting trace to JSON...")
|
|
trace_json = trace.to_json()
|
|
output_file = f"trace_{agent_id}.json"
|
|
with open(output_file, "w") as f:
|
|
f.write(trace_json)
|
|
print(f" ✓ Trace saved to: {output_file}")
|
|
print(f" ✓ JSON size: {len(trace_json):,} bytes")
|
|
|
|
# Show search results
|
|
print("\n9. Search Results:")
|
|
for i, result in enumerate(results, 1):
|
|
print(f" [{i}] {result['text'][:70]}...")
|
|
print(f" Weight: {result['weight']:.3f} "
|
|
f"(act: {result['activation']:.2f}, "
|
|
f"sem: {result['semantic_similarity']:.2f}, "
|
|
f"rec: {result['recency']:.2f})")
|
|
|
|
# Test helper methods
|
|
print("\n10. Testing Helper Methods:")
|
|
|
|
# Get path to first result
|
|
if results:
|
|
first_result_id = results[0]['id']
|
|
path = trace.get_search_path_to_node(first_result_id)
|
|
print(f" - Path to top result has {len(path)} steps")
|
|
|
|
# Count nodes by link type
|
|
temporal_nodes = trace.get_nodes_by_link_type("temporal")
|
|
semantic_nodes = trace.get_nodes_by_link_type("semantic")
|
|
entity_nodes = trace.get_nodes_by_link_type("entity")
|
|
print(f" - Nodes reached via temporal links: {len(temporal_nodes)}")
|
|
print(f" - Nodes reached via semantic links: {len(semantic_nodes)}")
|
|
print(f" - Nodes reached via entity links: {len(entity_nodes)}")
|
|
|
|
print("\n" + "=" * 70)
|
|
print("TRACE EXAMPLE COMPLETE!")
|
|
print("=" * 70)
|
|
print(f"\nYou can now build a visualization using the trace data in:")
|
|
print(f" {output_file}")
|
|
print("\nThe trace contains:")
|
|
print(f" - Complete search path with all nodes visited")
|
|
print(f" - Weight calculations for each node")
|
|
print(f" - Link information (type, weight, whether followed)")
|
|
print(f" - Pruning decisions with reasons")
|
|
print(f" - Performance metrics for each phase")
|
|
|
|
# Cleanup
|
|
print("\nCleaning up test agent...")
|
|
await memory.delete_agent(agent_id)
|
|
|
|
finally:
|
|
await memory.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|