add pending ops

This commit is contained in:
Nicolò Boschi 2025-11-07 15:12:37 +01:00
parent 0f33c809e7
commit bff72ca6ae
13 changed files with 2414 additions and 15179 deletions

View file

@ -0,0 +1,48 @@
"""add_async_operations_table
Revision ID: 0e96398aae9e
Revises: 1a35a4fa1950
Create Date: 2025-11-07 14:54:21.224968
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '0e96398aae9e'
down_revision: Union[str, Sequence[str], None] = '1a35a4fa1950'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# Create async_operations table
op.execute("""
CREATE TABLE async_operations (
id UUID PRIMARY KEY,
agent_id TEXT NOT NULL,
task_type TEXT NOT NULL,
items_count INTEGER NOT NULL,
document_id TEXT,
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
)
""")
# Create index on agent_id for fast lookups by agent
op.execute("""
CREATE INDEX idx_async_operations_agent_id
ON async_operations(agent_id)
""")
def downgrade() -> None:
"""Downgrade schema."""
# Drop index
op.execute("DROP INDEX IF EXISTS idx_async_operations_agent_id")
# Drop table
op.execute("DROP TABLE IF EXISTS async_operations")

View file

@ -281,6 +281,11 @@ class BenchmarkRunner:
contents=batch_contents
)
# If using remote API, wait for this batch to complete before continuing
from memora.remote_client import RemoteMemoryClient
if isinstance(self.memory, RemoteMemoryClient):
await self.memory.wait_for_backlog_completion(agent_id, verbose=False)
return len(batch_contents)
async def answer_question(
@ -702,6 +707,10 @@ class BenchmarkRunner:
More realistic scenario where agent accumulates memories over time.
"""
# Check if using remote API client
from memora.remote_client import RemoteMemoryClient
is_remote = isinstance(self.memory, RemoteMemoryClient)
# Phase 1: Ingestion
if not skip_ingestion:
# Calculate and display data statistics
@ -718,23 +727,48 @@ class BenchmarkRunner:
await self.memory.delete_agent(agent_id)
console.print(f" [green]✓[/green] Cleared agent data")
# Collect all sessions from all items into one large batch
console.print(f" [yellow]Collecting sessions from all items...[/yellow]")
all_sessions = []
for item in items:
item_sessions = self.dataset.prepare_sessions_for_ingestion(item)
all_sessions.extend(item_sessions)
if is_remote:
# For remote API: send one request per instance, then poll
console.print(f" [yellow]Sending {len(items)} instances (one request per instance)...[/yellow]")
total_sessions = 0
console.print(f" [cyan]Collected {len(all_sessions)} sessions from {len(items)} items[/cyan]")
console.print(f" [yellow]Ingesting in one batch (auto-chunks if needed)...[/yellow]")
for i, item in enumerate(items, 1):
item_sessions = self.dataset.prepare_sessions_for_ingestion(item)
total_sessions += len(item_sessions)
# Ingest all sessions in one batch call (will auto-chunk if too large)
await self.memory.put_batch_async(
agent_id=agent_id,
contents=all_sessions
)
if item_sessions:
await self.memory.put_batch_async(
agent_id=agent_id,
contents=item_sessions
)
console.print(f" [green]✓[/green] Ingested {len(all_sessions)} sessions from {len(items)} items")
if i % 10 == 0 or i == len(items):
console.print(f" Sent {i}/{len(items)} instances ({total_sessions} sessions so far)")
console.print(f" [green]✓[/green] Sent all {len(items)} instances ({total_sessions} sessions total)")
# Wait for all background processing to complete
console.print(f" [yellow]Waiting for background processing to complete...[/yellow]")
await self.memory.wait_for_backlog_completion(agent_id, verbose=False)
console.print(f" [green]✓[/green] Background processing complete")
else:
# For local memory: collect all and send in one batch (faster with auto-chunking)
console.print(f" [yellow]Collecting sessions from all items...[/yellow]")
all_sessions = []
for item in items:
item_sessions = self.dataset.prepare_sessions_for_ingestion(item)
all_sessions.extend(item_sessions)
console.print(f" [cyan]Collected {len(all_sessions)} sessions from {len(items)} items[/cyan]")
console.print(f" [yellow]Ingesting in one batch (auto-chunks if needed)...[/yellow]")
# Ingest all sessions in one batch call (will auto-chunk if too large)
await self.memory.put_batch_async(
agent_id=agent_id,
contents=all_sessions
)
console.print(f" [green]✓[/green] Ingested {len(all_sessions)} sessions from {len(items)} items")
else:
console.print(f"\n[3] Skipping ingestion (using existing data)")

View file

@ -314,7 +314,8 @@ async def run_benchmark(
max_questions_per_conv: int = None,
skip_ingestion: bool = False,
use_think: bool = False,
conversation: str = None
conversation: str = None,
api_url: str = None
):
"""
Run the LoComo benchmark.
@ -325,16 +326,23 @@ async def run_benchmark(
skip_ingestion: Whether to skip ingestion and use existing data
use_think: Whether to use the think API instead of search + LLM
conversation: Specific conversation ID to run (e.g., "conv-26")
api_url: Optional API URL to connect to (default: use local memory)
"""
# Initialize components
dataset = LoComoDataset()
memory = TemporalSemanticMemory(
db_url=os.getenv("DATABASE_URL"),
memory_llm_provider=os.getenv("MEMORY_LLM_PROVIDER", "groq"),
memory_llm_api_key=os.getenv("MEMORY_LLM_API_KEY"),
memory_llm_model=os.getenv("MEMORY_LLM_MODEL", "openai/gpt-oss-120b"),
memory_llm_base_url=os.getenv("MEMORY_LLM_BASE_URL") or None, # Use None to get provider defaults
)
# Use remote API client if api_url is provided, otherwise use local memory
if api_url:
from memora.remote_client import RemoteMemoryClient
memory = RemoteMemoryClient(base_url=api_url)
else:
memory = TemporalSemanticMemory(
db_url=os.getenv("DATABASE_URL"),
memory_llm_provider=os.getenv("MEMORY_LLM_PROVIDER", "groq"),
memory_llm_api_key=os.getenv("MEMORY_LLM_API_KEY"),
memory_llm_model=os.getenv("MEMORY_LLM_MODEL", "openai/gpt-oss-120b"),
memory_llm_base_url=os.getenv("MEMORY_LLM_BASE_URL") or None, # Use None to get provider defaults
)
await memory.initialize()
if use_think:
@ -466,6 +474,7 @@ if __name__ == "__main__":
parser.add_argument('--skip-ingestion', action='store_true', help='Skip ingestion and use existing data')
parser.add_argument('--use-think', action='store_true', help='Use think API instead of search + LLM')
parser.add_argument('--conversation', type=str, default=None, help='Run only specific conversation (e.g., "conv-26")')
parser.add_argument('--api-url', type=str, default=None, help='Memora API URL (default: use local memory, example: http://localhost:8000)')
args = parser.parse_args()
@ -474,5 +483,6 @@ if __name__ == "__main__":
max_questions_per_conv=args.max_questions,
skip_ingestion=args.skip_ingestion,
use_think=args.use_think,
conversation=args.conversation
conversation=args.conversation,
api_url=args.api_url
))

View file

@ -186,7 +186,8 @@ async def run_benchmark(
max_questions_per_instance: int = None,
thinking_budget: int = 100,
max_tokens: int = 4096,
skip_ingestion: bool = False
skip_ingestion: bool = False,
api_url: str = None
):
"""
Run the LongMemEval benchmark.
@ -197,6 +198,7 @@ async def run_benchmark(
thinking_budget: Thinking budget for spreading activation search
max_tokens: Maximum tokens to retrieve from memories
skip_ingestion: Whether to skip ingestion and use existing data
api_url: Optional API URL to connect to (default: use local memory)
"""
from rich.console import Console
console = Console()
@ -213,13 +215,19 @@ async def run_benchmark(
dataset = LongMemEvalDataset()
answer_generator = LongMemEvalAnswerGenerator()
answer_evaluator = LLMAnswerEvaluator()
memory = TemporalSemanticMemory(
db_url=os.getenv("DATABASE_URL"),
memory_llm_provider=os.getenv("MEMORY_LLM_PROVIDER", "groq"),
memory_llm_api_key=os.getenv("MEMORY_LLM_API_KEY"),
memory_llm_model=os.getenv("MEMORY_LLM_MODEL", "openai/gpt-oss-120b"),
memory_llm_base_url=os.getenv("MEMORY_LLM_BASE_URL") or None, # Use None to get provider defaults
)
# Use remote API client if api_url is provided, otherwise use local memory
if api_url:
from memora.remote_client import RemoteMemoryClient
memory = RemoteMemoryClient(base_url=api_url)
else:
memory = TemporalSemanticMemory(
db_url=os.getenv("DATABASE_URL"),
memory_llm_provider=os.getenv("MEMORY_LLM_PROVIDER", "groq"),
memory_llm_api_key=os.getenv("MEMORY_LLM_API_KEY"),
memory_llm_model=os.getenv("MEMORY_LLM_MODEL", "openai/gpt-oss-120b"),
memory_llm_base_url=os.getenv("MEMORY_LLM_BASE_URL") or None, # Use None to get provider defaults
)
# Create benchmark runner
runner = BenchmarkRunner(
@ -374,6 +382,12 @@ if __name__ == "__main__":
action="store_true",
help="Skip ingestion and use existing data"
)
parser.add_argument(
"--api-url",
type=str,
default=None,
help="Memora API URL (default: use local memory, example: http://localhost:8000)"
)
args = parser.parse_args()
@ -382,5 +396,6 @@ if __name__ == "__main__":
max_questions_per_instance=args.max_questions,
thinking_budget=args.thinking_budget,
max_tokens=args.max_tokens,
skip_ingestion=args.skip_ingestion
skip_ingestion=args.skip_ingestion,
api_url=args.api_url
))

File diff suppressed because it is too large Load diff

View file

@ -12,18 +12,17 @@ from pathlib import Path
# Add parent directory to path to import memory module
sys.path.insert(0, str(Path(__file__).parent))
from memora.web.server import create_app
from memora.api import create_app
from memora import TemporalSemanticMemory
def generate_openapi_spec(output_path: str = "openapi.json"):
"""Generate OpenAPI spec and save to file."""
# Create a temporary memory instance for OpenAPI generation
_memory = TemporalSemanticMemory(
db_url=os.getenv("DATABASE_URL", "postgresql://user:pass@localhost:5432/memora"),
memory_llm_provider=os.getenv("MEMORY_LLM_PROVIDER", "groq"),
memory_llm_api_key=os.getenv("MEMORY_LLM_API_KEY", "dummy"),
memory_llm_model=os.getenv("MEMORY_LLM_MODEL", "openai/gpt-oss-120b"),
memory_llm_base_url=os.getenv("MEMORY_LLM_BASE_URL") or None,
db_url="mock",
memory_llm_provider="ollama",
memory_llm_api_key="mock",
memory_llm_model="mock",
)
app = create_app(_memory)

View file

@ -5,6 +5,7 @@ This module provides the create_app function to create and configure
the FastAPI application with all API endpoints.
"""
import logging
import uuid
from pathlib import Path
from typing import Optional, List, Dict, Any
from datetime import datetime
@ -20,23 +21,23 @@ from memora import TemporalSemanticMemory
class SearchRequest(BaseModel):
"""Request model for search endpoint."""
query: str
fact_type: str
agent_id: str = "default"
thinking_budget: int = 100
max_tokens: int = 4096
reranker: str = "heuristic"
trace: bool = False
fact_type: Optional[str] = None
class Config:
json_schema_extra = {
"example": {
"query": "What did Alice say about machine learning?",
"fact_type": "world",
"agent_id": "user123",
"thinking_budget": 100,
"max_tokens": 4096,
"reranker": "heuristic",
"trace": True,
"fact_type": "world"
"trace": True
}
}
@ -336,11 +337,26 @@ def _register_routes(app: FastAPI):
response_model=SearchResponse,
tags=["Search"],
summary="Search memory",
description="Search memory using semantic similarity and spreading activation. Optionally filter by fact_type (world, agent, opinion)"
description="""
Search memory using semantic similarity and spreading activation.
The fact_type parameter is required and must be one of:
- 'world': General knowledge about people, places, events, and things that happen
- 'agent': Memories about what the AI agent did, actions taken, and tasks performed
- 'opinion': The agent's formed beliefs, perspectives, and viewpoints
"""
)
async def api_search(request: SearchRequest):
"""Run a search and return results with trace."""
try:
# Validate fact_type
valid_fact_types = ["world", "agent", "opinion"]
if request.fact_type not in valid_fact_types:
raise HTTPException(
status_code=400,
detail=f"Invalid fact_type '{request.fact_type}'. Must be one of: {', '.join(valid_fact_types)}"
)
# Run search with tracing
results, trace = await app.state.memory.search_async(
agent_id=request.agent_id,
@ -359,6 +375,8 @@ def _register_routes(app: FastAPI):
results=results,
trace=trace_dict
)
except HTTPException:
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
@ -366,111 +384,6 @@ def _register_routes(app: FastAPI):
raise HTTPException(status_code=500, detail=str(e))
@app.post(
"/api/world_search",
response_model=SearchResponse,
tags=["Search"],
summary="Search world facts",
description="Search only world facts - general knowledge about people, places, events, and things that happen"
)
async def api_world_search(request: SearchRequest):
"""Search only world facts (general knowledge about the world)."""
try:
# Run search with fact_type filter for 'world'
results, trace = await app.state.memory.search_async(
agent_id=request.agent_id,
query=request.query,
thinking_budget=request.thinking_budget,
max_tokens=request.max_tokens,
enable_trace=request.trace,
reranker=request.reranker,
fact_type='world'
)
# Convert trace to dict
trace_dict = trace.to_dict() if trace else None
return SearchResponse(
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/world_search: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
@app.post(
"/api/agent_search",
response_model=SearchResponse,
tags=["Search"],
summary="Search agent action facts",
description="Search only agent facts - memories about what the AI agent did, actions taken, and tasks performed"
)
async def api_agent_search(request: SearchRequest):
"""Search only agent facts (facts about what the agent did)."""
try:
# Run search with fact_type filter for 'agent'
results, trace = await app.state.memory.search_async(
agent_id=request.agent_id,
query=request.query,
thinking_budget=request.thinking_budget,
max_tokens=request.max_tokens,
enable_trace=request.trace,
reranker=request.reranker,
fact_type='agent'
)
# Convert trace to dict
trace_dict = trace.to_dict() if trace else None
return SearchResponse(
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/agent_search: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
@app.post(
"/api/opinion_search",
response_model=SearchResponse,
tags=["Search"],
summary="Search agent opinions",
description="Search only opinion facts - the agent's formed beliefs, perspectives, and viewpoints"
)
async def api_opinion_search(request: SearchRequest):
"""Search only opinion facts (agent's formed opinions and perspectives)."""
try:
# Run search with fact_type filter for 'opinion'
results, trace = await app.state.memory.search_async(
agent_id=request.agent_id,
query=request.query,
thinking_budget=request.thinking_budget,
max_tokens=request.max_tokens,
enable_trace=request.trace,
reranker=request.reranker,
fact_type='opinion'
)
# Convert trace to dict
trace_dict = trace.to_dict() if trace else None
return SearchResponse(
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/opinion_search: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
@app.post(
"/api/think",
response_model=ThinkResponse,
@ -562,6 +475,17 @@ def _register_routes(app: FastAPI):
agent_id
)
# Get pending operations count
pending_ops_result = await conn.fetchrow(
"""
SELECT COUNT(*) as count
FROM async_operations
WHERE agent_id = $1
""",
agent_id
)
pending_operations = pending_ops_result['count'] if pending_ops_result else 0
# Format results
nodes_by_type = {row['fact_type']: row['count'] for row in node_stats}
links_by_type = {row['link_type']: row['count'] for row in link_stats}
@ -574,7 +498,8 @@ def _register_routes(app: FastAPI):
"total_nodes": total_nodes,
"total_links": total_links,
"nodes_by_type": nodes_by_type,
"links_by_type": links_by_type
"links_by_type": links_by_type,
"pending_operations": pending_operations
}
except Exception as e:
@ -696,9 +621,28 @@ def _register_routes(app: FastAPI):
content_dict["context"] = item.context
contents.append(content_dict)
# Submit task to background queue
# Generate UUID for this operation
operation_id = uuid.uuid4()
# Insert operation record into database BEFORE scheduling task
pool = await app.state.memory._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"""
INSERT INTO async_operations (id, agent_id, task_type, items_count, document_id)
VALUES ($1, $2, $3, $4, $5)
""",
operation_id,
request.agent_id,
'batch_put',
len(contents),
request.document_id
)
# Submit task to background queue with operation_id
await app.state.memory._task_backend.submit_task({
'type': 'batch_put',
'operation_id': str(operation_id),
'agent_id': request.agent_id,
'contents': contents,
'document_id': request.document_id,
@ -706,7 +650,7 @@ def _register_routes(app: FastAPI):
'upsert': request.upsert
})
logging.info(f"Batch put task queued for agent_id={request.agent_id}, {len(contents)} items")
logging.info(f"Batch put task queued for agent_id={request.agent_id}, {len(contents)} items, operation_id={operation_id}")
return BatchPutAsyncResponse(
success=True,

View file

@ -178,9 +178,6 @@ class EntityResolver:
# Batch create new entities using multi-row VALUES
if entities_to_create:
import logging
logger = logging.getLogger(__name__)
create_start = time.time()
# Build multi-row VALUES statement
# VALUES ($1, $2, ...), ($N+1, $N+2, ...), ...
values_clauses = []
@ -212,7 +209,6 @@ class EntityResolver:
for i, (idx, entity_data) in enumerate(entities_to_create):
entity_ids[idx] = created_rows[i]['id']
logger.info(f" [6.2.2.X] Batch created {len(entities_to_create)} new entities in {time.time() - create_start:.3f}s")
return entity_ids

264
memora/remote_client.py Normal file
View file

@ -0,0 +1,264 @@
"""
Remote client for Memora API.
This module provides a client that connects to a remote Memora API server
instead of using the local TemporalSemanticMemory instance directly.
"""
import httpx
import asyncio
from typing import List, Dict, Any, Optional, Tuple
from datetime import datetime
class RemoteMemoryClient:
"""
HTTP client for Memora API that provides the same interface as TemporalSemanticMemory.
This allows benchmarks to connect to a remote Memora API server instead of
accessing the memory system directly.
"""
def __init__(self, base_url: str = "http://localhost:8000", timeout: float = 300.0):
"""
Initialize remote memory client.
Args:
base_url: Base URL of the Memora API server (default: http://localhost:8000)
timeout: Request timeout in seconds (default: 300s for large batch operations)
"""
self.base_url = base_url.rstrip('/')
self.timeout = timeout
self.client = httpx.AsyncClient(timeout=timeout)
async def initialize(self):
"""Initialize the client (no-op for remote client)."""
pass
async def close(self):
"""Close the HTTP client."""
await self.client.aclose()
async def put_batch_async(
self,
agent_id: str,
contents: List[Dict[str, Any]],
document_id: Optional[str] = None,
document_metadata: Optional[Dict[str, Any]] = None,
upsert: bool = False
) -> Dict[str, Any]:
"""
Store multiple memory items via API.
Args:
agent_id: Agent identifier
contents: List of content dicts with 'content', 'event_date', 'context' keys
document_id: Optional document identifier
document_metadata: Optional document metadata
upsert: Whether to upsert (update if exists)
Returns:
Result dict with success status
"""
# Convert contents to API format
items = []
for content in contents:
item = {
"content": content["content"]
}
if "event_date" in content and content["event_date"]:
# Convert datetime to ISO format string
event_date = content["event_date"]
if isinstance(event_date, datetime):
item["event_date"] = event_date.isoformat()
else:
item["event_date"] = event_date
if "context" in content and content["context"]:
item["context"] = content["context"]
items.append(item)
# Make API request
request_data = {
"agent_id": agent_id,
"items": items,
"upsert": upsert
}
if document_id:
request_data["document_id"] = document_id
if document_metadata:
request_data["document_metadata"] = document_metadata
response = await self.client.post(
f"{self.base_url}/api/memories/batch_async",
json=request_data
)
response.raise_for_status()
return response.json()
async def search_async(
self,
agent_id: str,
query: str,
thinking_budget: int = 100,
max_tokens: int = 4096,
enable_trace: bool = False,
reranker: str = "heuristic",
fact_type: Optional[str] = None
) -> Tuple[List[Dict[str, Any]], Optional[Dict[str, Any]]]:
"""
Search memories via API.
Args:
agent_id: Agent identifier
query: Search query
thinking_budget: Budget for spreading activation
max_tokens: Maximum tokens to retrieve
enable_trace: Whether to return trace information
reranker: Reranker type ("heuristic" or other)
fact_type: Optional fact type filter (world/agent/opinion)
Returns:
Tuple of (results, trace)
"""
request_data = {
"agent_id": agent_id,
"query": query,
"thinking_budget": thinking_budget,
"max_tokens": max_tokens,
"trace": enable_trace,
"reranker": reranker
}
if fact_type:
request_data["fact_type"] = fact_type
response = await self.client.post(
f"{self.base_url}/api/search",
json=request_data
)
response.raise_for_status()
result = response.json()
return result.get("results", []), result.get("trace")
async def think_async(
self,
agent_id: str,
query: str,
thinking_budget: int = 50
) -> Dict[str, Any]:
"""
Generate answer using think API.
Args:
agent_id: Agent identifier
query: Question to answer
thinking_budget: Budget for memory exploration
Returns:
Dict with 'text', 'based_on', and 'new_opinions' keys
"""
request_data = {
"agent_id": agent_id,
"query": query,
"thinking_budget": thinking_budget
}
response = await self.client.post(
f"{self.base_url}/api/think",
json=request_data
)
response.raise_for_status()
return response.json()
async def delete_agent(self, agent_id: str) -> Dict[str, Any]:
"""
Delete all data for an agent.
Note: This endpoint may not be available in the API.
For now, this is a no-op that returns success.
Args:
agent_id: Agent identifier
Returns:
Result dict
"""
# Note: delete_agent is not exposed in the API yet
# For benchmarks, we might need to manually clear data or use unique agent IDs
return {"success": True, "message": "Delete agent not supported via API"}
async def list_agents(self) -> List[str]:
"""
List all agents.
Returns:
List of agent IDs
"""
response = await self.client.get(f"{self.base_url}/api/agents")
response.raise_for_status()
result = response.json()
return result.get("agents", [])
async def get_agent_stats(self, agent_id: str) -> Dict[str, Any]:
"""
Get statistics for an agent.
Args:
agent_id: Agent identifier
Returns:
Dict with statistics including total_nodes, total_links, and pending_operations
"""
response = await self.client.get(f"{self.base_url}/api/stats/{agent_id}")
response.raise_for_status()
return response.json()
async def wait_for_backlog_completion(
self,
agent_id: str,
poll_interval: float = 1.0,
timeout: float = 300.0,
verbose: bool = True
) -> Dict[str, Any]:
"""
Poll agent stats until pending_operations is zero or timeout is reached.
Args:
agent_id: Agent identifier
poll_interval: Time to wait between polls in seconds (default: 1.0)
timeout: Maximum time to wait in seconds (default: 300)
verbose: Whether to print status updates
Returns:
Final stats dict
Raises:
TimeoutError: If pending_operations doesn't clear within timeout
"""
import time
start_time = time.time()
while True:
elapsed = time.time() - start_time
if elapsed > timeout:
raise TimeoutError(
f"Timeout waiting for pending operations to clear for agent '{agent_id}' "
f"after {timeout}s"
)
stats = await self.get_agent_stats(agent_id)
pending_operations = stats.get("pending_operations", 0)
if verbose:
print(
f"Agent '{agent_id}' pending operations: {pending_operations} "
f"(elapsed: {elapsed:.1f}s)"
)
if pending_operations == 0:
if verbose:
print(f"All operations completed for agent '{agent_id}' in {elapsed:.1f}s")
return stats
await asyncio.sleep(poll_interval)

View file

@ -224,17 +224,32 @@ class TemporalSemanticMemory(
Example: {'type': 'access_count_update', 'node_ids': [...]}
"""
task_type = task_dict.get('type')
operation_id = task_dict.get('operation_id')
if task_type == 'access_count_update':
await self._handle_access_count_update(task_dict)
elif task_type == 'reinforce_opinion':
await self._handle_reinforce_opinion(task_dict)
elif task_type == 'form_opinion':
await self._handle_form_opinion(task_dict)
elif task_type == 'batch_put':
await self._handle_batch_put(task_dict)
else:
logger.error(f"Unknown task type: {task_type}")
try:
if task_type == 'access_count_update':
await self._handle_access_count_update(task_dict)
elif task_type == 'reinforce_opinion':
await self._handle_reinforce_opinion(task_dict)
elif task_type == 'form_opinion':
await self._handle_form_opinion(task_dict)
elif task_type == 'batch_put':
await self._handle_batch_put(task_dict)
else:
logger.error(f"Unknown task type: {task_type}")
finally:
# Delete operation record if operation_id is present
if operation_id:
try:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"DELETE FROM async_operations WHERE id = $1",
uuid.UUID(operation_id)
)
logger.debug(f"Deleted async operation record: {operation_id}")
except Exception as e:
logger.error(f"Failed to delete async operation record {operation_id}: {e}")
async def initialize(self):
"""Initialize the connection pool and background workers."""

View file

@ -2,7 +2,7 @@
"openapi": "3.1.0",
"info": {
"title": "Agent Memory API",
"description": "\nA temporal-semantic memory system for AI agents that stores, retrieves, and reasons over memories.\n\n## Features\n\n* **Batch Memory Storage**: Store multiple memories efficiently with automatic fact extraction\n* **Semantic Search**: Find relevant memories using natural language queries\n* **Fact Type Filtering**: Search across world facts, agent actions, and opinions separately\n* **Think Endpoint**: Generate contextual answers based on agent identity and memories\n* **Graph Visualization**: Interactive memory graph visualization\n* **Document Tracking**: Track and manage memory documents with upsert support\n\n## Architecture\n\nThe system uses:\n- **Temporal Links**: Connect memories that are close in time\n- **Semantic Links**: Connect semantically similar memories\n- **Entity Links**: Connect memories that mention the same entities\n- **Spreading Activation**: Intelligent traversal for memory retrieval\n ",
"description": "\nA temporal-semantic memory system for AI agents that stores, retrieves, and reasons over memories.\n\n## Features\n\n* **Batch Memory Storage**: Store multiple memories efficiently with automatic fact extraction\n* **Semantic Search**: Find relevant memories using natural language queries\n* **Fact Type Filtering**: Search across world facts, agent actions, and opinions separately\n* **Think Endpoint**: Generate contextual answers based on agent identity and memories\n* **Graph Visualization**: Interactive memory graph visualization\n* **Document Tracking**: Track and manage memory documents with upsert support\n\n## Architecture\n\nThe system uses:\n- **Temporal Links**: Connect memories that are close in time\n- **Semantic Links**: Connect semantically similar memories\n- **Entity Links**: Connect memories that mention the same entities\n- **Spreading Activation**: Intelligent traversal for memory retrieval\n ",
"contact": {
"name": "Memory System"
},
@ -84,8 +84,8 @@
"tags": [
"Search"
],
"summary": "Search all memory types",
"description": "Search across all memory types (world, agent, opinion) using semantic similarity and spreading activation",
"summary": "Search memory",
"description": "Search memory using semantic similarity and spreading activation.\n\n The fact_type parameter is required and must be one of:\n - 'world': General knowledge about people, places, events, and things that happen\n - 'agent': Memories about what the AI agent did, actions taken, and tasks performed\n - 'opinion': The agent's formed beliefs, perspectives, and viewpoints",
"operationId": "api_search_api_search_post",
"requestBody": {
"content": {
@ -121,139 +121,13 @@
}
}
},
"/api/world_search": {
"post": {
"tags": [
"Search"
],
"summary": "Search world facts",
"description": "Search only world facts - general knowledge about people, places, events, and things that happen",
"operationId": "api_world_search_api_world_search_post",
"requestBody": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/SearchRequest"
}
}
},
"required": true
},
"responses": {
"200": {
"description": "Successful Response",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/SearchResponse"
}
}
}
},
"422": {
"description": "Validation Error",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
}
}
}
}
},
"/api/agent_search": {
"post": {
"tags": [
"Search"
],
"summary": "Search agent action facts",
"description": "Search only agent facts - memories about what the AI agent did, actions taken, and tasks performed",
"operationId": "api_agent_search_api_agent_search_post",
"requestBody": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/SearchRequest"
}
}
},
"required": true
},
"responses": {
"200": {
"description": "Successful Response",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/SearchResponse"
}
}
}
},
"422": {
"description": "Validation Error",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
}
}
}
}
},
"/api/opinion_search": {
"post": {
"tags": [
"Search"
],
"summary": "Search agent opinions",
"description": "Search only opinion facts - the agent's formed beliefs, perspectives, and viewpoints",
"operationId": "api_opinion_search_api_opinion_search_post",
"requestBody": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/SearchRequest"
}
}
},
"required": true
},
"responses": {
"200": {
"description": "Successful Response",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/SearchResponse"
}
}
}
},
"422": {
"description": "Validation Error",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
}
}
}
}
},
"/api/think": {
"post": {
"tags": [
"Reasoning"
],
"summary": "Think and generate answer",
"description": "Think and formulate an answer using agent identity, world facts, and opinions.\n\nThis endpoint:\n1. Retrieves agent facts (agent's identity)\n2. Retrieves world facts relevant to the query\n3. Retrieves existing opinions (agent's perspectives)\n4. Uses LLM to formulate a contextual answer\n5. Extracts and stores any new opinions formed\n6. Returns plain text answer, the facts used, and new opinions",
"description": "Think and formulate an answer using agent identity, world facts, and opinions.\n\n This endpoint:\n 1. Retrieves agent facts (agent's identity)\n 2. Retrieves world facts relevant to the query\n 3. Retrieves existing opinions (agent's perspectives)\n 4. Uses LLM to formulate a contextual answer\n 5. Extracts and stores any new opinions formed\n 6. Returns plain text answer, the facts used, and new opinions",
"operationId": "api_think_api_think_post",
"requestBody": {
"content": {
@ -311,13 +185,54 @@
}
}
},
"/api/stats/{agent_id}": {
"get": {
"tags": [
"Memory Statistics"
],
"summary": "Get memory statistics for an agent",
"description": "Get statistics about nodes and links for a specific agent",
"operationId": "api_stats_api_stats__agent_id__get",
"parameters": [
{
"name": "agent_id",
"in": "path",
"required": true,
"schema": {
"type": "string",
"title": "Agent Id"
}
}
],
"responses": {
"200": {
"description": "Successful Response",
"content": {
"application/json": {
"schema": {}
}
}
},
"422": {
"description": "Validation Error",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
}
}
}
}
},
"/api/memories/batch": {
"post": {
"tags": [
"Memory Storage"
],
"summary": "Store multiple memories",
"description": "Store multiple memory items in batch with automatic fact extraction.\n\nFeatures:\n- Efficient batch processing\n- Automatic fact extraction from natural language\n- Entity recognition and linking\n- Document tracking with optional upsert\n- Temporal and semantic linking\n\nThe system automatically:\n1. Extracts semantic facts from the content\n2. Generates embeddings\n3. Deduplicates similar facts\n4. Creates temporal, semantic, and entity links\n5. Tracks document metadata",
"description": "Store multiple memory items in batch with automatic fact extraction.\n\n Features:\n - Efficient batch processing\n - Automatic fact extraction from natural language\n - Entity recognition and linking\n - Document tracking with optional upsert\n - Temporal and semantic linking\n\n The system automatically:\n 1. Extracts semantic facts from the content\n 2. Generates embeddings\n 3. Deduplicates similar facts\n 4. Creates temporal, semantic, and entity links\n 5. Tracks document metadata",
"operationId": "api_batch_put_api_memories_batch_post",
"requestBody": {
"content": {
@ -353,11 +268,67 @@
}
}
},
"/api/locomo": {
"get": {
"summary": "Api Locomo",
"description": "Get Locomo benchmark results.",
"operationId": "api_locomo_api_locomo_get",
"/api/memories/batch_async": {
"post": {
"tags": [
"Memory Storage"
],
"summary": "Store multiple memories asynchronously",
"description": "Store multiple memory items in batch asynchronously using the task backend.\n\n This endpoint returns immediately after queuing the task, without waiting for completion.\n The actual processing happens in the background.\n\n Features:\n - Immediate response (non-blocking)\n - Background processing via task queue\n - Efficient batch processing\n - Automatic fact extraction from natural language\n - Entity recognition and linking\n - Document tracking with optional upsert\n - Temporal and semantic linking\n\n The system automatically:\n 1. Queues the batch put task\n 2. Returns immediately with success=True, queued=True\n 3. Processes in background: extracts facts, generates embeddings, creates links",
"operationId": "api_batch_put_async_api_memories_batch_async_post",
"requestBody": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/BatchPutRequest"
}
}
},
"required": true
},
"responses": {
"200": {
"description": "Successful Response",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/BatchPutAsyncResponse"
}
}
}
},
"422": {
"description": "Validation Error",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
}
}
}
}
},
"/api/memory/{unit_id}": {
"delete": {
"tags": [
"Memory Storage"
],
"summary": "Delete a memory unit",
"description": "Delete a single memory unit and all its associated links (temporal, semantic, and entity links)",
"operationId": "api_delete_memory_unit_api_memory__unit_id__delete",
"parameters": [
{
"name": "unit_id",
"in": "path",
"required": true,
"schema": {
"type": "string",
"title": "Unit Id"
}
}
],
"responses": {
"200": {
"description": "Successful Response",
@ -366,6 +337,16 @@
"schema": {}
}
}
},
"422": {
"description": "Validation Error",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
}
}
}
}
@ -397,6 +378,59 @@
]
}
},
"BatchPutAsyncResponse": {
"properties": {
"success": {
"type": "boolean",
"title": "Success"
},
"message": {
"type": "string",
"title": "Message"
},
"agent_id": {
"type": "string",
"title": "Agent Id"
},
"document_id": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Document Id"
},
"items_count": {
"type": "integer",
"title": "Items Count"
},
"queued": {
"type": "boolean",
"title": "Queued"
}
},
"type": "object",
"required": [
"success",
"message",
"agent_id",
"items_count",
"queued"
],
"title": "BatchPutAsyncResponse",
"description": "Response model for async batch put endpoint.",
"example": {
"agent_id": "user123",
"document_id": "conversation_123",
"items_count": 2,
"message": "Batch put task queued for background processing",
"queued": true,
"success": true
}
},
"BatchPutRequest": {
"properties": {
"agent_id": {
@ -526,12 +560,26 @@
},
"type": "array",
"title": "Edges"
},
"table_rows": {
"items": {
"additionalProperties": true,
"type": "object"
},
"type": "array",
"title": "Table Rows"
},
"total_units": {
"type": "integer",
"title": "Total Units"
}
},
"type": "object",
"required": [
"nodes",
"edges"
"edges",
"table_rows",
"total_units"
],
"title": "GraphDataResponse",
"description": "Response model for graph data endpoint.",
@ -555,7 +603,17 @@
"label": "Bob went hiking",
"type": "world"
}
]
],
"table_rows": [
{
"context": "Work info",
"date": "2024-01-15 10:30",
"entities": "Alice (PERSON), Google (ORGANIZATION)",
"id": "abc12345...",
"text": "Alice works at Google"
}
],
"total_units": 2
}
},
"HTTPValidationError": {
@ -613,12 +671,35 @@
"event_date": "2024-01-15T10:30:00Z"
}
},
"OpinionItem": {
"properties": {
"text": {
"type": "string",
"title": "Text"
},
"confidence": {
"type": "number",
"title": "Confidence"
}
},
"type": "object",
"required": [
"text",
"confidence"
],
"title": "OpinionItem",
"description": "Model for an opinion with confidence score."
},
"SearchRequest": {
"properties": {
"query": {
"type": "string",
"title": "Query"
},
"fact_type": {
"type": "string",
"title": "Fact Type"
},
"agent_id": {
"type": "string",
"title": "Agent Id",
@ -629,15 +710,15 @@
"title": "Thinking Budget",
"default": 100
},
"top_k": {
"max_tokens": {
"type": "integer",
"title": "Top K",
"default": 10
"title": "Max Tokens",
"default": 4096
},
"mmr_lambda": {
"type": "number",
"title": "Mmr Lambda",
"default": 0.5
"reranker": {
"type": "string",
"title": "Reranker",
"default": "heuristic"
},
"trace": {
"type": "boolean",
@ -647,16 +728,18 @@
},
"type": "object",
"required": [
"query"
"query",
"fact_type"
],
"title": "SearchRequest",
"description": "Request model for search endpoint.",
"example": {
"agent_id": "user123",
"mmr_lambda": 0.5,
"fact_type": "world",
"max_tokens": 4096,
"query": "What did Alice say about machine learning?",
"reranker": "heuristic",
"thinking_budget": 100,
"top_k": 10,
"trace": true
}
},
@ -719,11 +802,6 @@
"type": "integer",
"title": "Thinking Budget",
"default": 50
},
"top_k": {
"type": "integer",
"title": "Top K",
"default": 10
}
},
"type": "object",
@ -735,8 +813,7 @@
"example": {
"agent_id": "user123",
"query": "What do you think about artificial intelligence?",
"thinking_budget": 50,
"top_k": 10
"thinking_budget": 50
}
},
"ThinkResponse": {
@ -758,7 +835,7 @@
},
"new_opinions": {
"items": {
"type": "string"
"$ref": "#/components/schemas/OpinionItem"
},
"type": "array",
"title": "New Opinions",
@ -794,7 +871,10 @@
]
},
"new_opinions": [
"AI has great potential when used responsibly"
{
"confidence": 0.95,
"text": "AI has great potential when used responsibly"
}
],
"text": "Based on my understanding, AI is a transformative technology..."
}

View file

@ -28,6 +28,7 @@ dependencies = [
"pytest-timeout>=2.4.0",
"dateparser>=1.2.0",
"tiktoken>=0.12.0",
"httpx>=0.27.0",
]
[tool.hatch.build.targets.wheel]

View file

@ -795,6 +795,7 @@ dependencies = [
{ name = "dateparser" },
{ name = "fastapi", extra = ["standard"] },
{ name = "greenlet" },
{ name = "httpx" },
{ name = "langchain-text-splitters" },
{ name = "openai" },
{ name = "pgvector" },
@ -818,6 +819,7 @@ requires-dist = [
{ name = "dateparser", specifier = ">=1.2.0" },
{ name = "fastapi", extras = ["standard"], specifier = ">=0.120.3" },
{ name = "greenlet", specifier = ">=3.2.4" },
{ name = "httpx", specifier = ">=0.27.0" },
{ name = "langchain-text-splitters", specifier = ">=0.3.0" },
{ name = "openai", specifier = ">=1.0.0" },
{ name = "pgvector", specifier = ">=0.4.1" },