376 lines
11 KiB
Python
376 lines
11 KiB
Python
"""
|
|
FastAPI server for memory graph visualization and API.
|
|
|
|
Provides REST API endpoints for memory operations and serves
|
|
the interactive visualization interface.
|
|
"""
|
|
import asyncpg
|
|
import asyncio
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.staticfiles import StaticFiles
|
|
from fastapi.responses import FileResponse
|
|
from pydantic import BaseModel
|
|
from dotenv import load_dotenv
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Optional, List, Dict, Any
|
|
from datetime import datetime
|
|
|
|
# Add parent directory to path for imports
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
from memory import TemporalSemanticMemory
|
|
|
|
import logging
|
|
|
|
load_dotenv()
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
app = FastAPI(title="Memory Graph API", version="1.0.0")
|
|
|
|
# Mount static files
|
|
app.mount("/static", StaticFiles(directory="web/static"), name="static")
|
|
|
|
|
|
class SearchRequest(BaseModel):
|
|
"""Request model for search endpoint."""
|
|
query: str
|
|
agent_id: str = "default"
|
|
thinking_budget: int = 100
|
|
top_k: int = 10
|
|
mmr_lambda: float = 0.5
|
|
trace: bool = False
|
|
|
|
|
|
class MemoryItem(BaseModel):
|
|
"""Single memory item for batch put."""
|
|
content: str
|
|
event_date: Optional[datetime] = None
|
|
context: Optional[str] = None
|
|
|
|
|
|
class BatchPutRequest(BaseModel):
|
|
"""Request model for batch put endpoint."""
|
|
agent_id: str
|
|
items: List[MemoryItem]
|
|
document_id: Optional[str] = None
|
|
document_metadata: Optional[Dict[str, Any]] = None
|
|
upsert: bool = False
|
|
|
|
|
|
async def get_graph_data():
|
|
"""Fetch graph data from database."""
|
|
conn = await asyncpg.connect(
|
|
os.getenv('DATABASE_URL'),
|
|
statement_cache_size=0 # Disable statement caching for pgbouncer compatibility
|
|
)
|
|
|
|
# Get all memory units
|
|
units = await conn.fetch("""
|
|
SELECT id, text, event_date, context
|
|
FROM memory_units
|
|
ORDER BY event_date
|
|
""")
|
|
|
|
# Get all links with weights
|
|
links = await conn.fetch("""
|
|
SELECT
|
|
ml.from_unit_id,
|
|
ml.to_unit_id,
|
|
ml.link_type,
|
|
ml.weight,
|
|
e.canonical_name as entity_name
|
|
FROM memory_links ml
|
|
LEFT JOIN entities e ON ml.entity_id = e.id
|
|
ORDER BY ml.link_type, ml.weight DESC
|
|
""")
|
|
|
|
# Get entity information
|
|
unit_entities = await conn.fetch("""
|
|
SELECT ue.unit_id, e.canonical_name, e.entity_type
|
|
FROM unit_entities ue
|
|
JOIN entities e ON ue.entity_id = e.id
|
|
ORDER BY ue.unit_id
|
|
""")
|
|
|
|
await conn.close()
|
|
|
|
# Build entity mapping
|
|
entity_map = {}
|
|
for row in unit_entities:
|
|
unit_id = row['unit_id']
|
|
entity_name = row['canonical_name']
|
|
entity_type = row['entity_type']
|
|
if unit_id not in entity_map:
|
|
entity_map[unit_id] = []
|
|
entity_map[unit_id].append(f"{entity_name} ({entity_type})")
|
|
|
|
# Build nodes
|
|
nodes = []
|
|
for row in units:
|
|
unit_id = row['id']
|
|
text = row['text']
|
|
event_date = row['event_date']
|
|
context = row['context']
|
|
|
|
entities = entity_map.get(unit_id, [])
|
|
entity_count = len(entities)
|
|
|
|
# Color by entity count
|
|
if entity_count == 0:
|
|
color = "#e0e0e0"
|
|
elif entity_count == 1:
|
|
color = "#90caf9"
|
|
else:
|
|
color = "#42a5f5"
|
|
|
|
nodes.append({
|
|
"data": {
|
|
"id": str(unit_id),
|
|
"label": text[:50] + "..." if len(text) > 50 else text,
|
|
"text": text,
|
|
"context": context,
|
|
"date": str(event_date.date()),
|
|
"entities": ", ".join(entities) if entities else "None",
|
|
"color": color
|
|
}
|
|
})
|
|
|
|
# Build edges
|
|
edges = []
|
|
for row in links:
|
|
from_id = row['from_unit_id']
|
|
to_id = row['to_unit_id']
|
|
link_type = row['link_type']
|
|
weight = row['weight']
|
|
entity_name = row['entity_name']
|
|
|
|
# Set color based on link type
|
|
if link_type == 'temporal':
|
|
color = "#00bcd4"
|
|
line_style = "dashed"
|
|
elif link_type == 'semantic':
|
|
color = "#ff69b4"
|
|
line_style = "solid"
|
|
elif link_type == 'entity':
|
|
color = "#ffd700"
|
|
line_style = "solid"
|
|
else:
|
|
color = "#999999"
|
|
line_style = "solid"
|
|
|
|
edges.append({
|
|
"data": {
|
|
"id": f"{from_id}-{to_id}-{link_type}",
|
|
"source": str(from_id),
|
|
"target": str(to_id),
|
|
"weight": weight,
|
|
"linkType": link_type,
|
|
"entityName": entity_name or "",
|
|
"color": color,
|
|
"lineStyle": line_style
|
|
}
|
|
})
|
|
|
|
# Build table rows
|
|
table_rows = []
|
|
for row in units:
|
|
unit_id = row['id']
|
|
text = row['text']
|
|
event_date = row['event_date']
|
|
context = row['context']
|
|
|
|
entities = entity_map.get(unit_id, [])
|
|
entity_str = ", ".join(entities) if entities else "None"
|
|
table_rows.append({
|
|
"id": str(unit_id)[:8] + "...",
|
|
"text": text,
|
|
"context": context,
|
|
"date": str(event_date.date()),
|
|
"entities": entity_str
|
|
})
|
|
|
|
return {
|
|
"nodes": nodes,
|
|
"edges": edges,
|
|
"table_rows": table_rows,
|
|
"total_units": len(units)
|
|
}
|
|
|
|
memory = TemporalSemanticMemory()
|
|
|
|
@app.get("/")
|
|
async def index():
|
|
"""Serve the visualization page."""
|
|
return FileResponse("web/templates/index.html")
|
|
|
|
|
|
@app.get("/api/graph")
|
|
async def api_graph():
|
|
"""Get graph data from database."""
|
|
try:
|
|
data = await get_graph_data()
|
|
return data
|
|
except Exception as e:
|
|
import traceback
|
|
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
|
print(f"Error in /api/graph: {error_detail}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.post("/api/search")
|
|
async def api_search(request: SearchRequest):
|
|
"""Run a search and return results with trace."""
|
|
try:
|
|
# Initialize memory system
|
|
# Run search with tracing
|
|
results, trace = await memory.search_async(
|
|
agent_id=request.agent_id,
|
|
query=request.query,
|
|
thinking_budget=request.thinking_budget,
|
|
top_k=request.top_k,
|
|
enable_trace=request.trace,
|
|
mmr_lambda=request.mmr_lambda
|
|
)
|
|
|
|
# Convert trace to dict
|
|
trace_dict = trace.to_dict() if trace else None
|
|
|
|
return {
|
|
'results': results,
|
|
'trace': trace_dict
|
|
}
|
|
except Exception as e:
|
|
import traceback
|
|
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
|
print(f"Error in /api/search: {error_detail}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.get("/api/agents")
|
|
async def api_agents():
|
|
"""Get list of available agents from database."""
|
|
try:
|
|
conn = await asyncpg.connect(
|
|
os.getenv('DATABASE_URL'),
|
|
statement_cache_size=0
|
|
)
|
|
|
|
# Get distinct agent IDs from memory_units
|
|
agents = await conn.fetch("""
|
|
SELECT DISTINCT agent_id
|
|
FROM memory_units
|
|
WHERE agent_id IS NOT NULL
|
|
ORDER BY agent_id
|
|
""")
|
|
|
|
await conn.close()
|
|
|
|
agent_list = [row['agent_id'] for row in agents]
|
|
return {"agents": agent_list}
|
|
except Exception as e:
|
|
import traceback
|
|
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
|
print(f"Error in /api/agents: {error_detail}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.post("/api/memories/batch")
|
|
async def api_batch_put(request: BatchPutRequest):
|
|
"""
|
|
Store multiple memories in batch.
|
|
|
|
This endpoint calls put_batch_async to efficiently store multiple memory items.
|
|
Supports document tracking and upsert operations.
|
|
|
|
Example request:
|
|
{
|
|
"agent_id": "user123",
|
|
"items": [
|
|
{"content": "Alice works at Google", "context": "work"},
|
|
{"content": "Bob went hiking yesterday", "event_date": "2024-01-15T10:00:00Z"}
|
|
],
|
|
"document_id": "conversation_123",
|
|
"upsert": false
|
|
}
|
|
"""
|
|
try:
|
|
# Validate agent_id - prevent writing to reserved agents
|
|
RESERVED_AGENT_IDS = {"locomo"}
|
|
if request.agent_id in RESERVED_AGENT_IDS:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail=f"Cannot write to reserved agent_id '{request.agent_id}'. Reserved agents: {', '.join(RESERVED_AGENT_IDS)}"
|
|
)
|
|
|
|
# Initialize memory system
|
|
|
|
|
|
# Prepare contents for put_batch_async
|
|
contents = []
|
|
for item in request.items:
|
|
content_dict = {"content": item.content}
|
|
if item.event_date:
|
|
content_dict["event_date"] = item.event_date
|
|
if item.context:
|
|
content_dict["context"] = item.context
|
|
contents.append(content_dict)
|
|
|
|
# Call put_batch_async
|
|
result = await memory.put_batch_async(
|
|
agent_id=request.agent_id,
|
|
contents=contents,
|
|
document_id=request.document_id,
|
|
document_metadata=request.document_metadata,
|
|
upsert=request.upsert
|
|
)
|
|
|
|
await memory.close()
|
|
|
|
return {
|
|
"success": True,
|
|
"message": f"Successfully stored {len(contents)} memory items",
|
|
"agent_id": request.agent_id,
|
|
"document_id": request.document_id,
|
|
"items_count": len(contents)
|
|
}
|
|
except Exception as e:
|
|
import traceback
|
|
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
|
print(f"Error in /api/memories/batch: {error_detail}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.get("/api/locomo")
|
|
async def api_locomo():
|
|
"""Get Locomo benchmark results."""
|
|
import json
|
|
try:
|
|
results_path = Path(__file__).parent.parent / "benchmarks" / "locomo" / "benchmark_results.json"
|
|
if not results_path.exists():
|
|
raise HTTPException(status_code=404, detail="Benchmark results not found")
|
|
|
|
with open(results_path, 'r') as f:
|
|
data = json.load(f)
|
|
return data
|
|
except FileNotFoundError:
|
|
raise HTTPException(status_code=404, detail="Benchmark results not found")
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
print("\n" + "=" * 80)
|
|
print("Memory Graph API Server")
|
|
print("=" * 80)
|
|
print("\nStarting server at http://localhost:8080")
|
|
print("\nEndpoints:")
|
|
print(" GET / - Visualization UI")
|
|
print(" GET /api/graph - Get graph data")
|
|
print(" POST /api/search - Run search with trace")
|
|
print(" POST /api/memories/batch - Store multiple memories in batch")
|
|
print(" GET /api/agents - List available agents")
|
|
print("\n" + "=" * 80 + "\n")
|
|
|
|
uvicorn.run("server:app", host="0.0.0.0", port=8080, reload=True)
|