local setup and speed

This commit is contained in:
Nicolò Boschi 2025-11-05 15:35:07 +01:00
parent d8fdbb94ef
commit 39801d9f8b
50 changed files with 462024 additions and 50576 deletions

21
.env.dev Normal file
View file

@ -0,0 +1,21 @@
# Development/Production environment
# Database
DATABASE_URL=postgresql://postgres.goflmrvzwagridonyxxn:OjUtCVVtoV0nGPPP@aws-1-us-east-1.pooler.supabase.com:6543/postgres
# Disable tokenizers parallelism warning (happens with forked processes)
TOKENIZERS_PARALLELISM=false
# Main LLM Configuration (for memory operations: put/think/opinions)
# Choose one: "openai", "groq", or "ollama"
MEMORY_LLM_PROVIDER=groq
MEMORY_LLM_API_KEY=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
MEMORY_LLM_MODEL=openai/gpt-oss-120b
# MEMORY_LLM_BASE_URL=http://localhost:11434/v1 # For ollama or custom endpoints
# Judge LLM Configuration (for benchmark evaluation)
# If not set, falls back to main LLM configuration
JUDGE_LLM_PROVIDER=groq
JUDGE_LLM_API_KEY=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
JUDGE_LLM_MODEL=openai/gpt-oss-120b
# JUDGE_LLM_BASE_URL=https://api.custom.com/v1 # Optional custom endpoint

831
README.md
View file

@ -1,625 +1,410 @@
# Entity-Aware Memory System for AI Agents
# Memora - Entity-Aware Memory System for AI Agents
A proof-of-concept memory system that enables AI agents to store, retrieve, and connect memories using temporal, semantic, and entity-based relationships.
## Overview
This system implements a sophisticated graph-based memory architecture where memories are connected through three complementary networks:
1. **Temporal Network** - Memories linked by time proximity
2. **Semantic Network** - Memories linked by meaning similarity
3. **Entity Network** - Memories linked by shared entities (people, organizations, places)
The combination of these three networks enables powerful memory retrieval that goes beyond simple vector search, allowing agents to find relevant memories through multiple pathways.
A temporal-semantic-entity memory system that enables AI agents to store, retrieve, and reason over memories using graph-based spreading activation search.
## Architecture
### Triple Network Design
### Three Memory Networks
The system maintains three separate but interconnected memory networks:
1. **World Network** (`fact_type='world'`)
- General knowledge and facts about the world
- Information not specific to the agent's actions
- Example: "Alice works at Google", "Yosemite is in California"
**1. World Network** (`fact_type='world'`)
- General knowledge and facts about the world
- Information not specific to the agent's actions
- Example: "Alice works at Google", "Yosemite is in California"
2. **Agent Network** (`fact_type='agent'`)
- Facts about what the AI agent specifically did
- Agent's own actions and experiences
- Example: "The agent helped debug a Python script", "The agent recommended Yosemite"
**2. Agent Network** (`fact_type='agent'`)
- Facts about what the AI agent specifically did
- Agent's own actions and experiences
- Example: "The agent helped debug a Python script", "The agent recommended Yosemite"
3. **Opinion Network** (`fact_type='opinion'`)
- Agent's formed opinions and perspectives
- Automatically extracted during think operations
- Includes reasons and confidence scores (0.0-1.0)
- Immutable once formed (event_date = when opinion was formed)
- Example: "Python is better for data science than JavaScript (Reasons: has better libraries like pandas and numpy) [confidence: 0.85]"
**3. Opinion Network** (`fact_type='opinion'`)
- Agent's formed opinions and perspectives
- Automatically extracted during think operations
- Includes reasons and confidence scores (0.0-1.0)
- Immutable once formed (event_date = when opinion was formed)
- Example: "Python is better for data science than JavaScript (Reasons: has better libraries like pandas and numpy) [confidence: 0.85]"
All three networks share the same infrastructure (temporal/semantic/entity links) but can be searched independently or together. The **think** operation combines all three networks to formulate consistent, contextual answers while forming new opinions.
All three networks share the same infrastructure (temporal/semantic/entity links) but can be searched independently or together.
### Core Concepts
### Core Components
**Memory Units**: Individual sentence-level memories that are:
- Self-contained (pronouns resolved to actual referents)
- Self-contained (pronouns resolved to actual referents by LLM)
- Validated to have subject + verb (complete thoughts)
- Embedded as vectors for semantic similarity
- Embedded as 384-dim vectors using `BAAI/bge-small-en-v1.5`
- Timestamped for temporal relationships
- Linked to extracted entities
- Classified as either 'world' or 'agent' fact type
- Linked to extracted entities via spaCy NER
- Classified as 'world', 'agent', or 'opinion'
**Entity Resolution**: Named entities (PERSON, ORG, GPE, etc.) are:
**Entity Resolution**: Named entities (PERSON, ORG, PLACE, PRODUCT, CONCEPT, OTHER) are:
- Extracted using spaCy NER
- Disambiguated using a scoring algorithm
- Disambiguated using scoring algorithm (name similarity 50%, co-occurrence 30%, temporal proximity 20%)
- Tracked with canonical IDs across all memories
- Used to create strong connections between related memories
### Three Types of Memory Links
#### 1. Temporal Links (Time-Based)
**Purpose**: Connect memories that occurred close together in time
**How it works**:
- When storing a new memory, find all memories within a time window (default: 24 hours)
- Create weighted links based on temporal proximity
- Weight formula: `weight = max(0.3, 1.0 - (time_diff / window_size))`
**1. Temporal Links** (Time-Based)
- Connect memories within time window (default: 24 hours)
- Weight: `max(0.3, 1.0 - (time_diff / window_size))`
- Closer in time = stronger link
- Use case: "What happened recently?" or understanding sequences
**Visualization**: Cyan, dashed lines
**Use case**: "What happened recently?" or understanding sequences of events
#### 2. Semantic Links (Meaning-Based)
**Purpose**: Connect memories with similar content/meaning
**How it works**:
- Generate embeddings using local `bge-small-en-v1.5` model (384 dimensions)
- Store embeddings in PostgreSQL with pgvector extension
- When storing a new memory, find top-k similar memories using cosine similarity
- Create links only if similarity exceeds threshold (default: 0.7)
**2. Semantic Links** (Meaning-Based)
- Connect memories with similar embeddings
- Uses pgvector with HNSW index for fast nearest neighbor search
- Create links only if cosine similarity > threshold (default: 0.7)
- Weight = cosine similarity score
- Use case: "Tell me about hiking" retrieves all semantically related activities
**Visualization**: Pink, solid lines
**Technology**:
- **SentenceTransformers** - Local embedding model (BAAI/bge-small-en-v1.5)
- **pgvector** - PostgreSQL extension for vector operations
- **HNSW index** - Fast approximate nearest neighbor search
**Use case**: "Tell me about hiking" retrieves all semantically related outdoor activities
#### 3. Entity Links (Identity-Based)
**Purpose**: Connect ALL memories about the same person, organization, or place
**How it works**:
- Extract entities from text using spaCy NER
- Resolve entity identity using disambiguation algorithm:
- Name similarity (50% weight) - using SequenceMatcher
- Co-occurring entities (30% weight) - entities that appear together
- Temporal proximity (20% weight) - recent mentions more likely same entity
- If score > threshold (0.4 for PERSON with exact match, 0.6 otherwise): reuse existing entity
- If score < threshold: create new entity
- Link all memories mentioning the same entity with weight 1.0 (no decay)
**Visualization**: Gold, thick lines
**Technology**:
- **spaCy** (`en_core_web_sm`) - Named Entity Recognition
- **difflib.SequenceMatcher** - String similarity matching
**Use case**: "What does Alice do?" returns ALL memories about Alice (hiking, work at Google, Python project) even if semantically distant
**Critical advantage**: Solves the problem where "Alice loves hiking" wouldn't normally connect to "Alice works at Google" through semantic similarity alone.
**3. Entity Links** (Identity-Based)
- Connect ALL memories mentioning the same entity
- No decay over time (weight 1.0)
- Critical advantage: Solves the problem where "Alice loves hiking" wouldn't normally connect to "Alice works at Google" through semantic similarity alone
- Use case: "What does Alice do?" returns ALL memories about Alice
### Spreading Activation Search
The search algorithm explores the memory graph using spreading activation:
The search algorithm explores the memory graph using spreading activation with backpressure limiting (max 32 concurrent searches):
1. **Entry Points**: Find top-3 semantically similar memories to the query (vector search, similarity ≥ 0.5)
2. **Activation Spreading**: Start with activation = actual similarity score (0.5 to 1.0) at entry points
1. **Entry Points**: Find top-3 semantically similar memories (vector search, similarity ≥ 0.5)
2. **Activation Spreading**: Start with activation = actual similarity score at entry points
3. **Graph Traversal**: Follow links to neighbors, spreading activation with decay (0.8 factor)
4. **Thinking Budget**: Limit exploration to N units (controls computational cost)
5. **Dynamic Weighting**: Combine activation, semantic similarity, recency, and frequency:
```
final_weight = w_a × activation + w_s × semantic_similarity + w_r × recency + w_f × frequency
# Default weights (configurable via search parameters):
w_a = 0.30 # Activation weight
Default weights (configurable):
w_a = 0.30 # Activation weight (graph structure)
w_s = 0.30 # Semantic similarity weight
w_r = 0.25 # Recency weight
w_f = 0.15 # Frequency weight
semantic_similarity = cosine_similarity(query_embedding, memory_embedding)
recency = 1 / (1 + log(1 + days_since/365)) # Logarithmic decay with 1-year half-life
frequency = normalized to [0, 1] from log(access_count + 1) / log(10)
w_r = 0.25 # Recency weight (logarithmic decay, 1-year half-life)
w_f = 0.15 # Frequency weight (normalized access_count)
```
**Weight Tuning**: All weights are configurable via `search_async()` parameters, enabling benchmark experiments with different scoring strategies (e.g., emphasizing graph structure vs semantic similarity).
Recency uses logarithmic decay to provide meaningful differentiation over years:
- Today: 1.000 (100% weight)
- 1 week: 0.981 (barely any decay)
- 1 month: 0.927 (still very recent)
- 3 months: 0.819 (recent)
- 6 months: 0.714
- 1 year: 0.591 (half-life point)
- 2 years: 0.477 ✓
- 5 years: 0.358 ✓ (clearly different from 2 years!)
- 10 years: 0.294 ✓
This ensures old memories (2yr vs 5yr) have different weights, unlike exponential decay.
6. **Return Top-K**: Sort by final weight and return top results
6. **MMR Diversification**: Optional Maximal Marginal Relevance to balance relevance with diversity
7. **Return Top-K**: Sort by final weight and return top results
This approach ensures:
- Semantic relevance to query is always considered (default 30% weight)
- Graph structure influences results through activation (default 30% weight)
- Recently accessed memories get boosted (default 25% weight - recency bias)
- Frequently accessed memories get boosted (default 15% weight - importance signal)
### Search Tracing & Debugging
The system includes comprehensive search tracing to understand and debug the search process:
**Enable tracing**:
```python
results, trace = memory.search(
agent_id="agent_1",
query="Who works at Google?",
enable_trace=True # Returns detailed SearchTrace object
)
```
**Trace captures**:
- Every node visited with parent/child relationships
- All links explored (followed or pruned) with reasons
- Weight calculations broken down by component
- Entry points selected and their similarity scores
- Pruning decisions (already visited, activation too low, budget exhausted)
- Performance metrics for each search phase
**Export trace for visualization**:
```python
# Save trace as JSON for external visualization tools
trace_json = trace.to_json()
with open("trace.json", "w") as f:
f.write(trace_json)
```
**Use cases**:
- Understanding why certain memories were/weren't retrieved
- Debugging search behavior
- Analyzing link type effectiveness
- Performance profiling
- Building custom visualization layers
See `SEARCH_TRACE.md` for complete trace API documentation and `examples/trace_example.py` for a working demo.
### Self-Contained Memory Units
Every memory unit is self-contained through LLM fact extraction:
**Problem**: "She joined Google last year" - unclear who "she" is
**Solution**: LLM-based fact extraction that:
- Resolves pronouns to actual referents during extraction
- Makes facts readable without original context
- Includes all relevant details (WHO, WHAT, WHERE, WHEN, WHY, HOW)
- Processes facts in parallel for speed
**Result**: "Alice joined Google last year" - fully self-contained
**Technology**:
- LLM fact extraction with detailed prompts for pronoun resolution
- Structured output using Pydantic models
- Batch processing for efficiency
- Semantic relevance to query is always considered (30%)
- Graph structure influences results through activation (30%)
- Recently accessed memories get boosted (25%)
- Frequently accessed memories get boosted (15%)
### LLM-Based Fact Extraction
Raw content is processed through an LLM to extract meaningful facts before storage:
Raw content is processed through an LLM (Groq by default) to extract meaningful facts:
**Problem**: Raw text contains noise (greetings, filler words, reactions) that waste storage and reduce retrieval quality
**Solution**: LLM-based extraction with optimized prompting:
- Filters out social pleasantries and non-informative content
- Extracts only facts with substance (biographical, events, opinions, recommendations, descriptions, relationships)
- Filters out noise (greetings, filler words)
- Extracts only substantive facts (biographical, events, opinions, recommendations)
- Creates self-contained statements with subject+action+context
- Categorizes and attributes facts to speakers
**Technology**:
- **OpenAI-compatible API** - Supports Groq (default), OpenAI, and other providers
- **Structured output** - Uses Pydantic models for reliable fact extraction
- **Optimized prompting** - Concise prompts (~300 chars) emphasize dense output with no fluff
- **Automatic chunking** - Large documents (>120k chars) split at sentence boundaries
- **Fast sentence splitting** - Regex-based splitter (no heavy NLP models)
- **Progress tracking** - Logs chunk processing for transparency
**For large documents (e.g., podcast transcripts)**:
- Documents <120k chars: processed in one pass
- Documents >120k chars: automatically chunked at sentence boundaries
- Each chunk kept under ~30k tokens to avoid output token limits
- Facts aggregated across all chunks
- Resolves pronouns to actual referents
- Automatic chunking for large documents (>120k chars)
- Structured output using Pydantic models
- Retry logic for JSON validation failures
### Technology Stack
**Database**:
- PostgreSQL 15+ with extensions:
- `pgvector` - Vector similarity operations
- `uuid-ossp` - UUID generation
- PostgreSQL 15+ with `pgvector` and `uuid-ossp` extensions
**Python Libraries**:
- `psycopg2-binary` - PostgreSQL client
- `sentence-transformers` - Local embedding model (bge-small-en-v1.5)
- `torch` - Deep learning framework (for embeddings)
- `spacy` - NLP (NER, dependency parsing, tokenization)
- `langchain-text-splitters` - Intelligent text chunking
- `networkx` - Graph operations
- `pyvis` - Interactive HTML graph visualization
- `rich` - Terminal UI
- `asyncpg` - Async PostgreSQL client with connection pooling
- `sentence-transformers` - Local embedding model (BAAI/bge-small-en-v1.5)
- `openai` - LLM API client (supports Groq, OpenAI)
- `fastapi` - Web API framework
**Models**:
- BAAI/bge-small-en-v1.5 - Local embedding model (384 dimensions)
**Architecture Patterns**:
- Mixin pattern for code organization
- Connection pooling with backpressure (32 concurrent searches max)
- Background task management for opinion storage
- Cached LLM client for performance
## Quick Start
### Prerequisites
1. PostgreSQL 15+ with pgvector extension
2. Python 3.11+
### Setup
1. Install dependencies:
```bash
uv sync
```
2. Install spaCy model:
```bash
uv pip install https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.7.1/en_core_web_sm-3.7.1-py3-none-any.whl
```
3. Create database and run schema:
```bash
psql -U postgres -c "CREATE DATABASE memory_poc"
psql -U postgres -d memory_poc -f schema.sql
```
4. Configure environment:
```bash
cp .env.example .env
# Edit .env with your DATABASE_URL
```
### Run Tests
Run the full test suite:
```bash
uv run pytest tests/ -v
uv sync
```
Run specific test files:
2. Configure environment files:
Create `.env.local` for local development:
```bash
uv run pytest tests/test_memory_operations.py -v
uv run pytest tests/test_entity_linking.py -v
cat > .env.local << 'EOF'
# Database
DATABASE_URL=postgresql://memora:memora_dev@localhost:5432/memora
# LLM Provider: "openai", "groq", or "ollama"
LLM_PROVIDER=groq
# API Key (not needed for ollama)
LLM_API_KEY=your_api_key_here
# Optional: Custom base URL (for ollama or custom endpoints)
# LLM_BASE_URL=http://localhost:11434/v1
EOF
```
Run a single test:
Create `.env.dev` for dev/production environment:
```bash
uv run pytest tests/test_memory_operations.py::test_put_creates_memory_units -v
cat > .env.dev << 'EOF'
# Database
DATABASE_URL=postgresql://user:password@host:5432/memora
# LLM Provider: "openai", "groq", or "ollama"
LLM_PROVIDER=groq
# API Key (not needed for ollama)
LLM_API_KEY=your_api_key_here
# Optional: Custom base URL
# LLM_BASE_URL=https://api.custom-provider.com/v1
EOF
```
### Run Demo
### LLM Provider Configuration
The system supports multiple LLM providers with separate configuration for main operations and benchmark evaluation:
#### Main LLM (for memory operations)
**Groq** (default, fast inference):
```bash
LLM_PROVIDER=groq
LLM_API_KEY=your_groq_api_key
```
**OpenAI**:
```bash
LLM_PROVIDER=openai
LLM_API_KEY=your_openai_api_key
```
**Ollama** (local, no API key needed):
```bash
LLM_PROVIDER=ollama
LLM_BASE_URL=http://localhost:11434/v1 # Default, can be customized
```
#### Judge LLM (for benchmark evaluation)
Benchmarks can use a separate LLM for evaluation (e.g., using Groq for fast answer generation but OpenAI GPT-4 for accurate judging):
```bash
uv run python demos/demo_entity.py
# If not set, falls back to main LLM configuration
JUDGE_LLM_PROVIDER=openai
JUDGE_LLM_API_KEY=your_openai_api_key
# JUDGE_LLM_BASE_URL=https://api.custom.com/v1 # Optional
```
This will:
1. Clear previous demo data
2. Store sample memories about Alice, Bob, Google, Yosemite
3. Search for "What does Alice do?"
4. Show entity resolution results
5. Generate interactive HTML graph visualization
Open `memory_graph_interactive.html` in your browser to explore the memory graph!
## Using as a Library (Local Import)
You can import this project from another Poetry project using a local path dependency:
### 1. Add to your project's `pyproject.toml`:
```toml
[tool.poetry.dependencies]
memory-poc = {path = "../memory-poc", develop = true}
```
Or using poetry CLI:
**Example: Fast generation, accurate judging**:
```bash
poetry add ../memory-poc --editable
# Main LLM - Groq for speed
LLM_PROVIDER=groq
LLM_API_KEY=your_groq_key
# Judge LLM - OpenAI GPT-4 for accuracy
JUDGE_LLM_PROVIDER=openai
JUDGE_LLM_API_KEY=your_openai_key
```
### 2. Import the memory system:
```python
from memora import TemporalSemanticMemory
# Initialize memory
memory = TemporalSemanticMemory()
await memory.initialize()
# Use the memory system
await memory.put_batch_async(
agent_id="my_agent",
contents=["Alice works at Google", "Bob loves hiking"],
event_date=datetime.now(timezone.utc)
)
results, trace = await memory.search_async(
agent_id="my_agent",
query="Who works at Google?"
)
```
### 3. Import the FastAPI app:
```python
from memora.web import app, memory
# Use the FastAPI app in your own project
# You can mount it as a sub-application or run it directly
import uvicorn
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000)
```
### 4. Example: Extending the FastAPI app
```python
from fastapi import FastAPI
from memora.web import app as memory_app, memory
# Create your own app
my_app = FastAPI()
# Mount the memory app as a sub-application
my_app.mount("/memory", memory_app)
# Add your own endpoints that use the memory system
@my_app.post("/my-custom-endpoint")
async def my_endpoint():
# Use the shared memory instance
results, _ = await memory.search_async(
agent_id="my_agent",
query="some query"
)
return {"results": results}
if __name__ == "__main__":
import uvicorn
uvicorn.run(my_app, host="0.0.0.0", port=8000)
```
### Available Exports
**From `memory` package:**
- `TemporalSemanticMemory` - Main memory system class
- `SearchTrace`, `SearchTracer` - Search tracing utilities
- `QueryInfo`, `EntryPoint`, `NodeVisit`, etc. - Trace data structures
**From `memory.web` package:**
- `app` - FastAPI application instance
- `memory` - Shared TemporalSemanticMemory instance
### Web Server
To run the web interface:
### Local Development
```bash
# Development mode with auto-reload
uvicorn memora.web.server:app --reload --port 8000
# Start local PostgreSQL (with initialization)
./scripts/start-local-db.sh
# Production mode
uvicorn memora.web.server:app --host 0.0.0.0 --port 8000 --workers 4
# Start the server with local environment (default)
./scripts/start-server.sh --env local
# Start the server with dev environment
./scripts/start-server.sh --env dev
# Erase local database (stop + cleanup)
./scripts/erase-local-db.sh
```
Then open http://localhost:8000 in your browser to access the visualization interface.
The server will start at http://localhost:8080
## Project Structure
**API Endpoints**:
- `GET /` - Interactive visualization UI
- `POST /api/memories/batch` - Store memories
- `POST /api/search` - Search all networks
- `POST /api/world_search` - Search world facts only
- `POST /api/agent_search` - Search agent facts only
- `POST /api/opinion_search` - Search opinions only
- `POST /api/think` - Think and generate contextual answers
- `GET /api/graph` - Get graph data for visualization
- `GET /api/agents` - List all agents
```
memory-poc/
├── memory/ # Core memory system package
│ ├── temporal_semantic_memory.py # Main memory system class
│ ├── operations/ # Modular operation mixins
│ │ ├── embedding_operations.py # Embedding generation with process pool
│ │ ├── link_operations.py # Entity, temporal, semantic links
│ │ ├── batch_operations.py # Placeholder for future extraction
│ │ └── search_operations.py # Placeholder for future extraction
│ ├── entity_resolver.py # Entity extraction and disambiguation
│ ├── llm_client.py # LLM-based fact extraction
│ └── utils.py # Utility functions
├── demos/ # Demo scripts
│ └── demo_entity.py # Main entity-aware demo
├── visualizations/ # Visualization tools
│ └── interactive_graph.py # Interactive HTML graph (pyvis)
├── schema.sql # Database schema
├── pyproject.toml # Dependencies
└── README.md # This file
## API Examples (curl)
### Store Memories (PUT)
```bash
# Store memories for an agent
curl -X POST http://localhost:8080/api/memories/batch \
-H "Content-Type: application/json" \
-d '{
"agent_id": "alice_agent",
"items": [
{
"content": "Alice works at Google as a software engineer. She joined last year and focuses on machine learning infrastructure.",
"context": "career discussion",
"event_date": "2024-01-15T10:00:00Z"
},
{
"content": "Alice loves hiking in Yosemite National Park. She goes every weekend and has climbed Half Dome three times.",
"context": "hobby conversation"
}
],
"document_id": "conversation_001"
}'
```
The memory system uses a **mixin pattern** for code organization:
- `TemporalSemanticMemory` inherits from `EmbeddingOperationsMixin` and `LinkOperationsMixin`
- This reduced the main file from 1,720 lines to 1,420 lines (17% reduction)
- See `memory/operations/README.md` for detailed refactoring documentation
## Key Features
**Triple network architecture**: Separate networks for world knowledge, agent actions, and opinions
**Opinion formation**: Automatically extracts and stores opinions with confidence scores during thinking
**Three-layered linking**: Temporal + Semantic + Entity
**Entity disambiguation**: Resolves "Alice" across different contexts
**Self-contained units**: Pronouns resolved to actual referents
**Spreading activation**: Graph-aware search beyond vector similarity
**Think operation**: Combines all three networks for consistent, contextual answers
**Confidence scores**: Opinions include confidence ratings (0.0-1.0) based on supporting evidence
**Interactive visualization**: Explore memory graph in browser
**Recency & frequency weighting**: Recent and important memories boosted
**Linguistic validation**: Memory units verified to have subject + verb
**Modular architecture**: Mixin pattern with 17% code size reduction
## API Usage
### Store Memories
```python
from memora import TemporalSemanticMemory
memory = TemporalSemanticMemory()
memory.put(
agent_id="agent_1",
content="Alice works at Google as a software engineer. She joined last year.",
context="Career discussion",
event_date=datetime.now(timezone.utc)
)
Response:
```json
{
"success": true,
"message": "Successfully stored 2 memory items",
"agent_id": "alice_agent",
"document_id": "conversation_001",
"items_count": 2
}
```
### Search Memories
```python
# Basic search (trace disabled by default)
results, trace = memory.search(
agent_id="agent_1",
query="What does Alice do?",
thinking_budget=50, # How many units to explore
top_k=10 # Number of results to return
)
for result in results:
print(f"{result['text']} (weight: {result['weight']:.3f})")
# Search only world facts
results, trace = memory.search(
agent_id="agent_1",
query="What does Alice do?",
fact_type="world" # Only search world network
)
# Search only agent facts
results, trace = memory.search(
agent_id="agent_1",
query="What have I done?",
fact_type="agent" # Only search agent network
)
# Search only opinions
results, trace = memory.search(
agent_id="agent_1",
query="What do I think about Python?",
fact_type="opinion" # Only search opinion network
)
# Search with tracing for debugging
results, trace = memory.search(
agent_id="agent_1",
query="What does Alice do?",
thinking_budget=50,
top_k=10,
enable_trace=True # Returns detailed SearchTrace object
)
# Analyze trace
print(f"Nodes visited: {trace.summary.total_nodes_visited}")
print(f"Entry points: {len(trace.entry_points)}")
trace_json = trace.to_json() # Export for visualization
# Search with custom weight tuning
results, trace = memory.search(
agent_id="agent_1",
query="What does Alice do?",
thinking_budget=50,
top_k=10,
weight_activation=0.40, # Emphasize graph structure
weight_semantic=0.40, # Emphasize semantic similarity
weight_recency=0.10, # De-emphasize recency
weight_frequency=0.10 # De-emphasize frequency
)
```bash
# Search across all networks
curl -X POST http://localhost:8080/api/search \
-H "Content-Type: application/json" \
-d '{
"agent_id": "alice_agent",
"query": "What does Alice do?",
"thinking_budget": 100,
"top_k": 10,
"mmr_lambda": 0.5,
"trace": false
}'
```
### Think and Formulate Answers
The `think` operation combines all three networks to formulate consistent, contextual answers:
```python
result = await memory.think_async(
agent_id="agent_1",
query="What do you think about Python?",
thinking_budget=50,
top_k=10
)
print(result["text"]) # Plain text answer from LLM
# Access facts used to formulate the answer
for fact in result["based_on"]["world"]:
print(f"World: {fact['text']}")
for fact in result["based_on"]["agent"]:
print(f"Agent: {fact['text']}")
for fact in result["based_on"]["opinion"]:
print(f"Opinion: {fact['text']} (confidence: {fact.get('confidence_score', 'N/A')})")
# Check for newly formed opinions
for opinion in result["new_opinions"]:
print(f"New opinion formed: {opinion['text']} (confidence: {opinion['confidence']})")
Response:
```json
{
"results": [
{
"id": "550e8400-e29b-41d4-a716-446655440000",
"text": "Alice works at Google as a software engineer",
"context": "career discussion",
"event_date": "2024-01-15T10:00:00Z",
"weight": 0.95,
"fact_type": "world"
},
{
"id": "550e8400-e29b-41d4-a716-446655440001",
"text": "Alice joined Google last year",
"weight": 0.87,
"fact_type": "world"
}
],
"trace": null
}
```
The think operation:
1. Searches the agent network to understand the agent's identity and actions
2. Searches the world network for relevant general knowledge
3. Searches the opinion network for existing perspectives
4. Uses an LLM (Groq by default) to formulate a coherent answer, being consistent with existing opinions
5. Extracts any new opinions formed during thinking with confidence scores
6. Stores new opinions with the current timestamp and query context
7. Returns plain text response with supporting facts from all networks and new opinions
### Think and Generate Answer
## How It Works: Example
```bash
# Think operation: combines agent identity, world knowledge, and opinions
curl -X POST http://localhost:8080/api/think \
-H "Content-Type: application/json" \
-d '{
"agent_id": "alice_agent",
"query": "What do you know about Alice?",
"thinking_budget": 50,
"top_k": 10
}'
```
**Input memories**:
1. "Alice loves hiking in the mountains" (7 days ago)
2. "She goes hiking every weekend in Yosemite" (7 days ago)
3. "Alice works at Google as a software engineer" (3 days ago)
4. "She joined Google last year" (3 days ago)
Response:
```json
{
"text": "Alice is a software engineer at Google who joined last year. She specializes in machine learning infrastructure. In her free time, she's an avid hiker who frequents Yosemite National Park on weekends and has climbed Half Dome three times.",
"based_on": {
"world": [
{
"text": "Alice works at Google as a software engineer",
"weight": 0.95,
"id": "550e8400-e29b-41d4-a716-446655440000"
},
{
"text": "Alice loves hiking in Yosemite National Park",
"weight": 0.89,
"id": "550e8400-e29b-41d4-a716-446655440002"
}
],
"agent": [],
"opinion": []
},
"new_opinions": []
}
```
**Processing**:
1. ✅ Coreference resolution → "Alice goes hiking...", "Alice joined Google..."
2. ✅ Entity extraction → Identifies "Alice" (PERSON), "Google" (ORG), "Yosemite" (GPE)
3. ✅ Entity resolution → All "Alice" mentions = same person
4. ✅ Create links:
- Temporal: Memory 1 ↔ Memory 2 (same day)
- Semantic: "hiking" memories link together, "Google" memories link together
- Entity: ALL Alice memories strongly linked (weight 1.0)
## Running Benchmarks
**Query: "What does Alice do?"**
1. Vector search finds "Alice works at Google" as top entry point
2. Spreading activation follows entity links to find:
- "Alice joined Google..." (entity link: Alice)
- "Alice loves hiking..." (entity link: Alice)
- "Alice goes hiking..." (entity link: Alice)
3. Returns ALL Alice memories, properly ranked by relevance
The system includes two benchmarks for evaluating memory retrieval quality:
## Why This Architecture?
### LoComo Benchmark
**Problem with vector-only search**: "Alice loves hiking" and "Alice works at Google" are semantically distant - pure vector search might miss this connection.
Long-term Conversational Memory benchmark - evaluates multi-turn conversation understanding:
**Solution**: Entity links ensure memories about the same person/place/organization are strongly connected regardless of semantic distance.
```bash
# Run full benchmark with think API (uses local env by default)
./scripts/benchmarks/run-locomo.sh --use-think
**Result**: More human-like memory retrieval that understands identity and relationships.
# Run with dev environment
./scripts/benchmarks/run-locomo.sh --use-think --env dev
# Run with limits for quick testing
./scripts/benchmarks/run-locomo.sh --use-think --max-conversations 5 --max-questions 3
# Skip ingestion (use existing data)
./scripts/benchmarks/run-locomo.sh --use-think --skip-ingestion
```
### LongMemEval Benchmark
Long-term Memory Evaluation benchmark - tests memory retention and retrieval:
```bash
# Run full benchmark (uses local env by default)
./scripts/benchmarks/run-longmemeval.sh
# Run with dev environment
./scripts/benchmarks/run-longmemeval.sh --env dev
# Run with arguments (pass any args directly)
./scripts/benchmarks/run-longmemeval.sh --max-instances 10 --max-questions 5
# Skip ingestion
./scripts/benchmarks/run-longmemeval.sh --skip-ingestion
```
### Visualizer
View benchmark results in an interactive web interface:
```bash
# Start the visualizer server
./scripts/benchmarks/start-visualizer.sh
```
The visualizer will be available at http://localhost:8001
**Benchmark Results**: Results are saved to `benchmark_results.json` in each benchmark directory with metrics including accuracy, F1 score, and per-question performance.
## License

View file

@ -3,6 +3,8 @@ Alembic environment configuration for SQLAlchemy with pgvector.
Uses synchronous psycopg2 driver for migrations to avoid pgbouncer issues.
"""
import os
import sys
from pathlib import Path
from logging.config import fileConfig
from sqlalchemy import pool, engine_from_config
@ -12,10 +14,26 @@ from alembic import context
from dotenv import load_dotenv
# Import your models here
from memory.models import Base
from memora.models import Base
# Load environment variables
load_dotenv()
# Load environment variables based on DATABASE_URL env var or default to local
def load_env():
"""Load environment variables from .env.local or .env.dev"""
# Check if DATABASE_URL is already set (e.g., by CI/CD)
if os.getenv("DATABASE_URL"):
return
# Default to local environment
env_file = ".env.local"
if Path(env_file).exists():
load_dotenv(env_file)
else:
# Fallback to dev
env_file = ".env.dev"
if Path(env_file).exists():
load_dotenv(env_file)
load_env()
# this is the Alembic Config object, which provides
# access to the values within the .ini file in use.

View file

@ -1,219 +0,0 @@
# Benchmark Suite
This directory contains a common benchmark framework and benchmark-specific implementations for evaluating the memory system.
## Structure
```
benchmarks/
├── common/ # Common benchmark framework
│ ├── benchmark_runner.py # Main runner with all optimizations
│ └── __init__.py
├── locomo/ # LoComo benchmark
│ ├── locomo_benchmark.py # LoComo-specific implementations
│ ├── run_benchmark.py # Runner script
│ └── locomo10.json # Dataset (place here)
├── longmemeval/ # LongMemEval benchmark
│ ├── longmemeval_benchmark.py # LongMemEval-specific implementations
│ ├── run_benchmark.py # Runner script
│ └── longmemeval_s_cleaned.json # Dataset (auto-downloaded)
└── visualizer/ # Web-based benchmark visualizer
├── server.py # FastAPI server
├── serve.sh # Launch script
└── static/ # Frontend assets (HTML, CSS, JS)
```
## Common Framework
The common framework provides a unified interface with optimizations from the working LoComo implementation:
- **Batch ingestion** via `put_batch_async`
- **Parallel question processing** with rate limiting
- **Parallel LLM judging** with configurable semaphore
- **Progress tracking** with Rich
- **Comprehensive metrics** collection
## LoComo Benchmark
**Location**: `locomo/`
**Purpose**: Evaluate long-term conversational memory through Question Answering on multi-session conversations.
### Quick Start
1. **Run full benchmark** (10 conversations, ~2000 questions):
```bash
cd locomo
uv run python run_benchmark.py
```
2. **Run with think API** (integrated search + answer generation, skips separate search step):
```bash
cd locomo
uv run python run_benchmark.py --use-think
```
Note: Think mode uses the memory system's integrated `think_async()` API which performs its own retrieval and reasoning in a single call, making it more efficient than the traditional two-step approach.
3. **Run quick test** (1 conversation, 10 questions):
```bash
cd locomo
uv run python run_benchmark.py --max-conversations 1 --max-questions 10
```
4. **View results**:
- Detailed report: `locomo/RESULTS.md` or `locomo/results_table_think.md`
- Raw data: `locomo/benchmark_results.json` or `locomo/benchmark_results_think.json`
### Dataset
- **Source**: [Snap Research LoComo](https://github.com/snap-research/locomo)
- **File**: `locomo10.json` (10 conversations)
- **Size**: Each conversation has ~300 turns over ~35 sessions spanning several months
- **Tasks**: Question Answering with 3 reasoning types (single-hop, temporal, multi-hop)
### Methodology
1. **Ingest** each conversation turn-by-turn with timestamps
2. **Apply** coreference resolution and entity extraction
3. **Create** temporal, semantic, and entity links
4. **Answer** questions using spreading activation search
5. **Evaluate** using LLM-as-judge (GPT-4o-mini)
### Expected Performance
Based on published results:
- **Human**: ~95%
- **Letta (GPT-4o-mini)**: 74.0%
- **Mem0 Graph**: 68.5%
- **Our target**: 65-75% (competitive with state-of-the-art)
### Computational Cost
**Per conversation** (~300 turns):
- ~300 embedding API calls (ingestion)
- ~200 embedding API calls (queries)
- ~200 LLM API calls (answer generation)
- ~200 LLM API calls (judgment)
**Estimated runtime**: 2-5 minutes per conversation (API-dependent)
**Estimated cost**: $0.50-1.00 per conversation (OpenAI pricing)
## LongMemEval Benchmark
**Location**: `longmemeval/`
**Purpose**: Evaluate five core long-term interactive memory abilities: information extraction, multi-session reasoning, temporal reasoning, knowledge updates, and abstention.
### Quick Start
1. **Download dataset**:
```bash
cd longmemeval
curl -L "https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json" -o longmemeval_s_cleaned.json
```
2. **Run full benchmark** (500 questions):
```bash
cd longmemeval
uv run python run_benchmark.py
```
3. **Run quick test** (5 instances):
```bash
cd longmemeval
uv run python run_benchmark.py --max-instances 5
```
4. **View results**:
- Raw data: `longmemeval/benchmark_results.json`
### Dataset
- **Source**: [LongMemEval (ICLR 2025)](https://github.com/xiaowu0162/LongMemEval)
- **File**: `longmemeval_s_cleaned.json` (500 instances)
- **Size**: ~40 sessions per instance (~115k tokens)
- **Tasks**: 5 memory abilities across different question types
### Methodology
1. **Ingest** multi-session conversations with timestamps
2. **Apply** coreference resolution and entity extraction
3. **Create** temporal, semantic, and entity links
4. **Retrieve** relevant memories using spreading activation
5. **Generate** answers using GPT-4o-mini
6. **Evaluate** using GPT-4o as judge
### Expected Performance
Based on published results:
- **Human**: ~95%
- **Zep**: 75.2%
- **Letta (GPT-4o-mini)**: 74.0%
- **Mem0 Graph**: 68.5%
- **Our target**: 65-75% (competitive with state-of-the-art)
### Computational Cost
**Full benchmark** (500 instances):
- Embeddings: Free (local model)
- Answer generation: 500 × GPT-4o-mini calls
- Evaluation: 500 × GPT-4o calls
- **Estimated runtime**: 2-4 hours
- **Estimated cost**: $50-80 (OpenAI API)
## Benchmark Visualizer
**Location**: `visualizer/`
**Purpose**: Web-based interface for visualizing and analyzing benchmark results.
### Quick Start
1. **Start the visualizer**:
```bash
cd visualizer
./serve.sh
```
2. **Open browser**: http://localhost:8001
3. **Select benchmark**: Choose from:
- "LoComo (search)" - Traditional search + LLM
- "LoComo (think)" - Integrated think API
### Features
- Interactive visualization of benchmark results
- Category-wise performance breakdown (Multi-hop, Single-hop, Temporal, Open-domain)
- Filter by correctness (all/correct/incorrect answers)
- Detailed Q&A view with reasoning and retrieved memories
- Overall and per-item accuracy statistics
- Think mode displays fact types with color-coded borders:
- Green: World facts
- Orange: Agent facts
- Purple: Opinion facts
See `visualizer/README.md` for more details.
## Future Benchmarks
- **MemGPT Tasks**: Long-context question answering
- **Custom Temporal Reasoning**: Time-based memory retrieval
- **Entity-Centric Queries**: Testing entity link effectiveness
## Adding New Benchmarks
1. Create a new directory: `benchmarks/{benchmark_name}/`
2. Add dataset: `benchmarks/{benchmark_name}/data/`
3. Implement adapter: `benchmarks/{benchmark_name}/run_benchmark.py`
4. Document results: `benchmarks/{benchmark_name}/RESULTS.md`
## Results Summary
| Benchmark | Metric | Our System | Best Published | Status |
|-----------|--------|------------|----------------|--------|
| LoComo QA | Accuracy | {TBD}% | 74.0% (Letta) | In Progress |
| LongMemEval | Accuracy | {TBD}% | 75.2% (Zep) | Ready to Run |
*Last updated: 2025-10-30*

View file

@ -217,6 +217,7 @@ class BenchmarkRunner:
weight_semantic=weight_semantic,
weight_recency=weight_recency,
weight_frequency=weight_frequency,
fact_type="world"
)
if not results:

File diff suppressed because it is too large Load diff

View file

@ -16,13 +16,11 @@ import asyncio
import pydantic
from openai import AsyncOpenAI
import os
from dotenv import load_dotenv
load_dotenv()
# Import common framework
sys.path.insert(0, str(Path(__file__).parent.parent))
from common.benchmark_runner import BenchmarkDataset, LLMAnswerGenerator, LLMAnswerEvaluator
from memora.llm_wrapper import LLMConfig
class LoComoDataset(BenchmarkDataset):
"""LoComo dataset implementation."""
@ -113,19 +111,13 @@ class QuestionAnswer(pydantic.BaseModel):
class LoComoAnswerGenerator(LLMAnswerGenerator):
"""LoComo-specific answer generator using Groq."""
"""LoComo-specific answer generator using configurable LLM provider."""
def __init__(self):
"""Initialize with Groq client."""
groq_api_key = os.getenv('GROQ_API_KEY')
if not groq_api_key:
raise ValueError("GROQ_API_KEY environment variable not set")
base_url = os.getenv('GROQ_BASE_URL', 'https://api.groq.com/openai/v1')
self.client = AsyncOpenAI(
api_key=groq_api_key,
base_url=base_url
)
"""Initialize with LLM configuration for memory operations."""
self.llm_config = LLMConfig.for_memory()
self.client = self.llm_config.client
self.model = self.llm_config.model
async def generate_answer(
self,
@ -146,10 +138,9 @@ class LoComoAnswerGenerator(LLMAnswerGenerator):
context = json.dumps(context_parts)
# Use Groq to generate answer
# Use LLM to generate answer
try:
response = await self.client.beta.chat.completions.parse(
model="openai/gpt-oss-120b",
answer_obj = await self.llm_config.call(
messages=[
{
"role": "system",
@ -201,9 +192,9 @@ Answer:
"""
}
],
response_format=QuestionAnswer
response_format=QuestionAnswer,
scope="memory"
)
answer_obj = response.choices[0].message.parsed
return answer_obj.answer, answer_obj.reasoning, None
except Exception as e:
return f"Error generating answer: {str(e)}", "Error occurred during answer generation.", None
@ -260,7 +251,6 @@ class LoComoThinkAnswerGenerator(LLMAnswerGenerator):
query=question,
thinking_budget=self.thinking_budget,
top_k=self.top_k,
model="openai/gpt-oss-120b",
temperature=0.7,
max_tokens=1000
)
@ -329,19 +319,13 @@ class JudgeResponse(pydantic.BaseModel):
class LoComoAnswerEvaluator(LLMAnswerEvaluator):
"""LoComo-specific answer evaluator using Groq."""
"""LoComo-specific answer evaluator using configurable LLM provider."""
def __init__(self):
"""Initialize with Groq client."""
groq_api_key = os.getenv('GROQ_API_KEY')
if not groq_api_key:
raise ValueError("GROQ_API_KEY environment variable not set")
base_url = os.getenv('GROQ_BASE_URL', 'https://api.groq.com/openai/v1')
self.client = AsyncOpenAI(
api_key=groq_api_key,
base_url=base_url
)
"""Initialize with LLM configuration for judge/evaluator."""
self.llm_config = LLMConfig.for_judge()
self.client = self.llm_config.client
self.model = self.llm_config.model
async def judge_answer(
self,
@ -358,8 +342,7 @@ class LoComoAnswerEvaluator(LLMAnswerEvaluator):
"""
async with semaphore:
try:
response = await self.client.beta.chat.completions.parse(
model="openai/gpt-oss-120b",
judgement = await self.llm_config.call(
messages=[
{
"role": "system",
@ -368,9 +351,9 @@ class LoComoAnswerEvaluator(LLMAnswerEvaluator):
{
"role": "user",
"content": f"""
Your task is to label an answer to a question as CORRECT or WRONG. You williolw23 be given the following data:
(1) a question (posed by one user to another user),
(2) a gold (ground truth) answer,
Your task is to label an answer to a question as 'CORRECT' or 'WRONG'. You williolw23 be given the following data:
(1) a question (posed by one user to another user),
(2) a 'gold' (ground truth) answer,
(3) a generated answer
which you will score as CORRECT/WRONG.
@ -378,11 +361,11 @@ Your task is to label an answer to a question as CORRECT or WRONG. Y
The gold answer will usually be a concise and short answer that includes the referenced topic, for example:
Question: Do you remember what I got the last time I went to Hawaii?
Gold answer: A shell necklace
The generated answer might be much longer, but you should be generous with your grading - as long as it touches on the same topic as the gold answer, it should be counted as CORRECT.
The generated answer might be much longer, but you should be generous with your grading - as long as it touches on the same topic as the gold answer, it should be counted as CORRECT.
For time related questions, the gold answer will be a specific date, month, year, etc. The generated answer might be much longer or use relative time references (like "last Tuesday" or "next month"), but you should be generous with your grading - as long as it refers to the same date or time period as the gold answer, it should be counted as CORRECT. Even if the format differs (e.g., "May 7th" vs "7 May"), consider it CORRECT if it's the same date.
Now its time for the real question:
Now it's time for the real question:
Question: {question}
Gold answer: {correct_answer}
Generated answer: {predicted_answer}
@ -392,12 +375,12 @@ Your task is to label an answer to a question as CORRECT or WRONG. Y
"""
}
],
response_format=JudgeResponse,
scope="judge",
temperature=0,
max_tokens=4096,
response_format=JudgeResponse
max_tokens=4096
)
judgement = response.choices[0].message.parsed
return judgement.correct, judgement.reasoning
except Exception as e:

View file

@ -1,7 +1,16 @@
# LoComo Benchmark Results
**Overall Accuracy**: 65.33% (98/150)
**Overall Accuracy**: 49.15% (690/1404)
| Sample ID | Sessions | Questions | Correct | Accuracy | Multi-hop | Single-hop | Temporal | Open-domain |
|-----------|----------|-----------|---------|----------|-----------|------------|----------|-------------|
| conv-26 | 19 | 150 | 98 | 65.33% | N/A | N/A | N/A | N/A |
| conv-26 | -1 | 150 | 98 | 65.33% | N/A | N/A | N/A | N/A |
| conv-30 | -1 | 81 | 46 | 56.79% | N/A | N/A | N/A | N/A |
| conv-41 | -1 | 150 | 80 | 53.33% | N/A | N/A | N/A | N/A |
| conv-42 | -1 | 150 | 58 | 38.67% | N/A | N/A | N/A | N/A |
| conv-43 | -1 | 150 | 69 | 46.00% | N/A | N/A | N/A | N/A |
| conv-44 | -1 | 123 | 78 | 63.41% | N/A | N/A | N/A | N/A |
| conv-47 | -1 | 150 | 74 | 49.33% | N/A | N/A | N/A | N/A |
| conv-48 | -1 | 150 | 84 | 56.00% | N/A | N/A | N/A | N/A |
| conv-49 | -1 | 150 | 37 | 24.67% | N/A | N/A | N/A | N/A |
| conv-50 | -1 | 150 | 66 | 44.00% | N/A | N/A | N/A | N/A |

View file

@ -14,6 +14,7 @@ sys.path.insert(0, str(Path(__file__).parent.parent))
import asyncio
import argparse
import logging
from memora import TemporalSemanticMemory
from locomo_benchmark import LoComoDataset, LoComoAnswerGenerator, LoComoThinkAnswerGenerator, LoComoAnswerEvaluator
from common.benchmark_runner import BenchmarkRunner
@ -149,7 +150,6 @@ def generate_markdown_table(results: dict, use_think: bool = False):
if __name__ == "__main__":
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
parser = argparse.ArgumentParser(description='Run LoComo benchmark')

View file

@ -1,242 +0,0 @@
# LongMemEval Benchmark
**Purpose**: Evaluate long-term interactive memory capabilities across five core abilities.
## Overview
LongMemEval is a comprehensive benchmark that tests chat assistants on realistic long-term memory scenarios. The benchmark evaluates five core memory abilities:
1. **Information Extraction** - Retrieving specific facts from conversation history
2. **Multi-Session Reasoning** - Connecting information across multiple conversations
3. **Temporal Reasoning** - Understanding time-based relationships and changes
4. **Knowledge Updates** - Handling conflicting or updated information
5. **Abstention** - Recognizing when information is insufficient to answer
## Dataset
- **Source**: [LongMemEval (ICLR 2025)](https://github.com/xiaowu0162/LongMemEval)
- **File**: `longmemeval_s_cleaned.json`
- **Size**: 500 question-answer pairs
- **Context**: ~40 sessions per instance (~115k tokens)
- **Format**: Multi-turn conversations with timestamped sessions
### Dataset Structure
Each instance contains:
- `question_id`: Unique identifier
- `question_type`: Category (single-session, multi-session, temporal, knowledge-update, abstention)
- `question`: Query text
- `answer`: Expected answer
- `question_date`: Query timestamp
- `haystack_sessions`: List of conversation sessions with turns
- `answer_session_ids`: Evidence session identifiers
## Quick Start
### Prerequisites
1. Python 3.11+ with dependencies installed (`uv sync`)
2. PostgreSQL database configured
3. OpenAI API key set in environment
```bash
export OPENAI_API_KEY="your-api-key"
```
### Download Dataset
```bash
cd benchmarks/longmemeval
curl -L "https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json" -o longmemeval_s_cleaned.json
```
### Run Benchmark
**Full evaluation** (500 questions):
```bash
uv run python run_benchmark.py
```
**Quick test** (first 5 instances):
```bash
uv run python run_benchmark.py --max-instances 5
```
**Custom settings**:
```bash
uv run python run_benchmark.py \
--max-instances 10 \
--thinking-budget 100 \
--top-k 20 \
--output my_results.json
```
### View Results
Results are saved to `benchmark_results.json` and include:
- Per-question scores and predictions
- Performance breakdown by question type
- Retrieved memory units for debugging
- Evaluation explanations
## Methodology
### 1. Ingestion Phase
For each instance:
1. Parse all conversation sessions with timestamps
2. Process each turn (user and assistant messages)
3. Store in memory system with:
- Coreference resolution (pronouns → entities)
- Entity extraction and disambiguation
- Temporal, semantic, and entity link creation
### 2. Retrieval Phase
For each question:
1. Generate query embedding
2. Find entry points (top-3 similar memories)
3. Spread activation through memory graph
4. Apply recency and frequency weighting
5. Return top-k most relevant memory units
### 3. Answer Generation
1. Format retrieved memories as context
2. Generate answer using GPT-4o-mini
3. Enforce answering only from provided memories
4. Handle abstention cases appropriately
### 4. Evaluation
1. Compare predicted answer to gold answer
2. Use GPT-4o as judge for semantic equivalence
3. Binary scoring (1 = correct, 0 = incorrect)
4. Aggregate by question type
## Parameters
| Parameter | Default | Description |
|-----------|---------|-------------|
| `--max-instances` | 500 | Number of instances to evaluate |
| `--max-questions` | None | Limit questions per instance (for testing) |
| `--thinking-budget` | 100 | Exploration budget for spreading activation |
| `--top-k` | 20 | Number of memory units to retrieve |
| `--output` | `benchmark_results.json` | Output file path |
## Expected Performance
Based on published results:
| System | Accuracy |
|--------|----------|
| Human | ~95% |
| Zep | 75.2% |
| Letta (GPT-4o-mini) | 74.0% |
| Mem0 Graph | 68.5% |
| **Target** | **65-75%** |
## Performance by Question Type
Expected breakdown:
- **Single-session**: 70-80% (easiest - information in one session)
- **Multi-session**: 60-70% (requires connecting across sessions)
- **Temporal reasoning**: 60-70% (requires time-based reasoning)
- **Knowledge updates**: 50-65% (hardest - handling conflicting info)
- **Abstention**: 65-75% (recognizing insufficient information)
## Computational Cost
**Per instance** (~40 sessions, ~200 turns):
- Ingestion: ~200 embedding generations (local model, fast)
- Query: 1 embedding generation + graph search
- Answer: 1 GPT-4o-mini call (~200 tokens)
- Evaluation: 1 GPT-4o call (~150 tokens)
**Full benchmark** (500 instances):
- Runtime: 2-4 hours (depends on API rate limits)
- Cost: ~$50-80 (OpenAI API for answer generation + evaluation)
- Embeddings: Free (local model)
## Example Output
```
LongMemEval Benchmark Evaluation
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Overall Performance
┏━━━━━━━━━━━━━━━━┳━━━━━━━┓
┃ Metric ┃ Value ┃
┡━━━━━━━━━━━━━━━━╇━━━━━━━┩
│ Total │ 500 │
│ Correct │ 345 │
│ Incorrect │ 155 │
│ Accuracy │ 69.0% │
└────────────────┴───────┘
Performance by Question Type
┏━━━━━━━━━━━━━━━━━━━━┳━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━┓
┃ Question Type ┃ Total ┃ Correct ┃ Accuracy ┃
┡━━━━━━━━━━━━━━━━━━━━╇━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━┩
│ single-session │ 150 │ 115 │ 76.7% │
│ multi-session │ 120 │ 78 │ 65.0% │
│ temporal-reasoning │ 100 │ 65 │ 65.0% │
│ knowledge-update │ 80 │ 48 │ 60.0% │
│ abstention │ 50 │ 39 │ 78.0% │
└────────────────────┴───────┴─────────┴──────────┘
```
## Troubleshooting
### Dataset not found
```bash
cd benchmarks/longmemeval
curl -L "https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json" -o longmemeval_s_cleaned.json
```
### OpenAI API key error
```bash
export OPENAI_API_KEY="your-api-key"
```
### Memory ingestion slow
- This is expected for first-time entity resolution
- Subsequent queries are fast (graph search)
- Consider using `--max-instances` for quick testing
### Low accuracy
- Try increasing `--thinking-budget` (default: 100)
- Try increasing `--top-k` (default: 20)
- Check retrieved memories in results JSON for debugging
## Architecture Integration
This benchmark tests the full memory system architecture:
1. ✅ **Coreference Resolution**: Makes memories self-contained
2. ✅ **Entity Extraction**: Identifies people, organizations, places
3. ✅ **Entity Disambiguation**: Links mentions across sessions
4. ✅ **Temporal Links**: Connects memories by time proximity
5. ✅ **Semantic Links**: Connects memories by meaning
6. ✅ **Entity Links**: Connects memories by shared entities
7. ✅ **Spreading Activation**: Graph-aware retrieval
8. ✅ **Recency/Frequency Weighting**: Importance signals
## Citation
If you use this benchmark, please cite:
```bibtex
@inproceedings{wu2025longmemeval,
title={LongMemEval: Benchmarking Chat Assistants on Long-Term Interactive Memory},
author={Wu, Di and Wang, Hongwei and Liu, Wenhao and Wang, Jiaheng and Li, Zihan and Huang, Yiqin and Patel, Zelin and Liu, Yiheng and Meng, Bo and Pan, Sinong and others},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025}
}
```
## Related Benchmarks
- **LoComo**: Multi-session conversational QA
- **MemGPT Tasks**: Long-context question answering
- **Custom Temporal Reasoning**: Time-based memory retrieval

View file

@ -15,13 +15,11 @@ from typing import List, Dict, Any, Tuple, Optional
import asyncio
from openai import AsyncOpenAI
import os
from dotenv import load_dotenv
load_dotenv()
# Import common framework
sys.path.insert(0, str(Path(__file__).parent.parent))
from common.benchmark_runner import BenchmarkDataset, LLMAnswerGenerator, LLMAnswerEvaluator
from memora.llm_wrapper import LLMConfig
class LongMemEvalDataset(BenchmarkDataset):
@ -126,15 +124,13 @@ class LongMemEvalDataset(BenchmarkDataset):
class LongMemEvalAnswerGenerator(LLMAnswerGenerator):
"""LongMemEval-specific answer generator using OpenAI."""
"""LongMemEval-specific answer generator using configurable LLM provider."""
def __init__(self, model: str = "gpt-4o-mini"):
"""Initialize with OpenAI client."""
self.model = model
openai_api_key = os.getenv("OPENAI_API_KEY")
if not openai_api_key:
raise ValueError("OPENAI_API_KEY environment variable not set")
self.client = AsyncOpenAI(api_key=openai_api_key)
def __init__(self):
"""Initialize with LLM configuration for memory operations."""
self.llm_config = LLMConfig.for_memory()
self.client = self.llm_config.client
self.model = self.llm_config.model
async def generate_answer(
self,
@ -170,28 +166,25 @@ Instructions:
Answer:"""
try:
response = await self.client.chat.completions.create(
model=self.model,
answer = await self.llm_config.call(
messages=[{"role": "user", "content": prompt}],
scope="memory",
temperature=0.0,
max_tokens=300
)
answer = response.choices[0].message.content.strip()
return answer, "" # LongMemEval doesn't use reasoning
return answer.strip(), "" # LongMemEval doesn't use reasoning
except Exception as e:
return f"Error generating answer: {str(e)}", ""
class LongMemEvalAnswerEvaluator(LLMAnswerEvaluator):
"""LongMemEval-specific answer evaluator using OpenAI."""
"""LongMemEval-specific answer evaluator using configurable LLM provider."""
def __init__(self, model: str = "gpt-4o"):
"""Initialize with OpenAI client."""
self.model = model
openai_api_key = os.getenv("OPENAI_API_KEY")
if not openai_api_key:
raise ValueError("OPENAI_API_KEY environment variable not set")
self.client = AsyncOpenAI(api_key=openai_api_key)
def __init__(self):
"""Initialize with LLM configuration for judge/evaluator."""
self.llm_config = LLMConfig.for_judge()
self.client = self.llm_config.client
self.model = self.llm_config.model
async def judge_answer(
self,
@ -227,14 +220,14 @@ Score: [0 or 1]
Explanation: [brief explanation]"""
try:
response = await self.client.chat.completions.create(
model=self.model,
content = await self.llm_config.call(
messages=[{"role": "user", "content": prompt}],
scope="judge",
temperature=0.0,
max_tokens=200
)
content = response.choices[0].message.content.strip()
content = content.strip()
# Parse score and explanation
lines = content.split('\n')

View file

@ -25,6 +25,7 @@ sys.path.insert(0, str(Path(__file__).parent.parent))
import asyncio
import argparse
import subprocess
import logging
from rich.console import Console
from memora import TemporalSemanticMemory
from longmemeval_benchmark import LongMemEvalDataset, LongMemEvalAnswerGenerator, LongMemEvalAnswerEvaluator
@ -97,8 +98,8 @@ async def run_benchmark(
# Initialize components
dataset = LongMemEvalDataset()
answer_generator = LongMemEvalAnswerGenerator(model="gpt-4o-mini")
answer_evaluator = LongMemEvalAnswerEvaluator(model="gpt-4o")
answer_generator = LongMemEvalAnswerGenerator()
answer_evaluator = LongMemEvalAnswerEvaluator()
memory = TemporalSemanticMemory()
# Create benchmark runner
@ -172,6 +173,8 @@ def generate_type_report(results: dict):
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
parser = argparse.ArgumentParser(description="Run LongMemEval benchmark")
parser.add_argument(
"--max-instances",

View file

@ -1,62 +0,0 @@
# Benchmark Visualizer
A standalone web service for visualizing benchmark results. Currently supports the LoComo benchmark with plans to add more benchmarks in the future.
## Features
- Interactive web interface for viewing benchmark results
- Detailed breakdown by category (Multi-hop, Single-hop, Temporal, Open-domain)
- Filter options to view all, correct, or incorrect answers
- Expandable Q&A details with reasoning and retrieved memories
- Overall and per-item accuracy statistics
## Running the Visualizer
### Option 1: Using the serve script (recommended)
```bash
cd benchmarks/visualizer
./serve.sh
```
### Option 2: Using uvicorn directly
```bash
cd benchmarks/visualizer
uv run uvicorn server:app --reload --host 0.0.0.0 --port 8001
```
Then open your browser to: http://localhost:8001
## Usage
1. Select a benchmark from the dropdown:
- **LoComo (search)**: Traditional two-step approach (search → LLM answer generation)
- **LoComo (think)**: Integrated approach using think API (single call for retrieval + reasoning)
2. The visualization will automatically load and display:
- Overall accuracy statistics
- Category-wise performance breakdown
- Detailed results for each conversation
3. Use the filter controls to show all answers, only incorrect, or only correct answers
4. Expand individual conversations to see Q&A details, reasoning, and retrieved memories
## API Endpoints
- `GET /` - Main visualizer page
- `GET /api/locomo?mode={search|think}` - Returns LoComo benchmark results as JSON
- `mode=search` (default): Returns results from `benchmark_results.json`
- `mode=think`: Returns results from `benchmark_results_think.json`
## Requirements
- FastAPI
- Uvicorn
- Python 3.11+
The visualizer reads benchmark results from:
- `benchmarks/locomo/benchmark_results.json` for search mode
- `benchmarks/locomo/benchmark_results_think.json` for think mode
Make sure to run the benchmark first to generate results:
- Search mode: `cd benchmarks/locomo && uv run python run_benchmark.py`
- Think mode: `cd benchmarks/locomo && uv run python run_benchmark.py --use-think`

View file

@ -1,4 +0,0 @@
#!/bin/bash
# Start the Benchmark Visualizer server with hot reload
cd "$(dirname "$0")"
uv run uvicorn server:app --reload --host 0.0.0.0 --port 8001

6
docker/.gitignore vendored Normal file
View file

@ -0,0 +1,6 @@
# Environment files
.env
.env.local
# Docker volumes (data persistence)
postgres_data/

23
docker/docker-compose.yml Normal file
View file

@ -0,0 +1,23 @@
version: '3.8'
services:
postgres:
image: pgvector/pgvector:pg16
container_name: memora-postgres
environment:
POSTGRES_USER: memora
POSTGRES_PASSWORD: memora_dev
POSTGRES_DB: memora
ports:
- "5432:5432"
volumes:
- postgres_data:/var/lib/postgresql/data
- ./init-extensions.sql:/docker-entrypoint-initdb.d/01-init-extensions.sql
healthcheck:
test: ["CMD-SHELL", "pg_isready -U memora"]
interval: 5s
timeout: 5s
retries: 5
volumes:
postgres_data:

View file

@ -0,0 +1,2 @@
-- Enable pgvector extension
CREATE EXTENSION IF NOT EXISTS vector;

View file

@ -1,22 +0,0 @@
-- ============================================================================
-- DROP ALL TABLES FOR MEMORY POC DATABASE
-- ============================================================================
--
-- WARNING: This will completely remove all tables and data!
-- Use with caution, especially in production environments.
--
-- Usage:
-- psql -d your_database -f drop.sql
--
-- After running this, you'll need to recreate the schema:
-- psql -d your_database -f schema.sql
-- ============================================================================
-- Drop all tables in reverse dependency order
-- CASCADE ensures dependent objects are also dropped
DROP TABLE IF EXISTS memory_links CASCADE;
DROP TABLE IF EXISTS entity_cooccurrences CASCADE;
DROP TABLE IF EXISTS unit_entities CASCADE;
DROP TABLE IF EXISTS entities CASCADE;
DROP TABLE IF EXISTS memory_units CASCADE;

View file

@ -1,94 +0,0 @@
"""
Example of using memory-poc as a library in another project.
This demonstrates how to:
1. Import and use the memory system directly
2. Import and extend the FastAPI app
3. Mount the memory app as a sub-application
"""
import asyncio
from memora.web import app, memory
async def example_memory_usage():
"""Example of using the memory system directly."""
# Initialize memory system
await memory.initialize()
try:
# Store some memories
await memory.put_async(
agent_id="example_agent",
content="Example memory content",
context="test context"
)
# Search for memories
results, trace = await memory.search_async(
agent_id="example_agent",
query="example",
top_k=5
)
print(f"Found {len(results)} results")
for result in results:
print(f" - {result['text']} (score: {result['score']:.4f})")
# Use think functionality
think_result = await memory.think_async(
agent_id="example_agent",
query="What do you know?",
thinking_budget=50
)
print(f"\nThink result: {think_result['text']}")
if think_result.get('new_opinions'):
print(f"New opinions formed: {len(think_result['new_opinions'])}")
finally:
# Clean up
await memory.close()
def example_fastapi_extension():
"""Example of extending the FastAPI app with custom endpoints."""
from fastapi import FastAPI
# Option 1: Add endpoints directly to the imported app
@app.get("/api/custom")
async def custom_endpoint():
return {"message": "Custom endpoint added to memory-poc app"}
# Option 2: Mount as sub-application
main_app = FastAPI(title="My Application")
@main_app.get("/")
async def root():
return {"message": "My main application"}
# Mount the memory app at /memory
main_app.mount("/memory", app)
# Now you can access:
# - / -> your main app
# - /memory/ -> memory visualization
# - /memory/api/graph -> memory graph API
# - /memory/api/search -> memory search API
return main_app
if __name__ == "__main__":
# Example 1: Use memory system directly
print("=" * 80)
print("Example 1: Direct memory system usage")
print("=" * 80)
asyncio.run(example_memory_usage())
# Example 2: FastAPI extension
print("\n" + "=" * 80)
print("Example 2: FastAPI app extension")
print("=" * 80)
extended_app = example_fastapi_extension()
print("FastAPI app extended successfully")
print("To run: uvicorn library_usage_example:extended_app --reload")

View file

@ -1,46 +0,0 @@
"""
Example: Running many think operations in parallel with optimized connection pooling.
For 100 parallel think operations:
- Each think does 3 searches (world, agent, opinion)
- Each search acquires 1-3 connections briefly
- Total: ~300 concurrent connection requests
Solution: Increase pool_max_size to handle the concurrency.
"""
import asyncio
from memora import TemporalSemanticMemory
async def main():
# For 100 parallel think operations, use a larger pool
# Rule of thumb: pool_max_size >= (num_parallel_thinks * 3)
memory = TemporalSemanticMemory(
pool_min_size=10, # Keep some connections warm
pool_max_size=200 # Allow up to 200 concurrent connections
)
await memory.initialize()
# Example: Run 100 think operations in parallel
queries = [f"Query {i}" for i in range(100)]
tasks = [
memory.think_async(
agent_id="test_agent",
query=query,
thinking_budget=50,
top_k=10
)
for query in queries
]
# Run all thinks in parallel
results = await asyncio.gather(*tasks)
print(f"Completed {len(results)} think operations")
await memory.close()
if __name__ == "__main__":
asyncio.run(main())

View file

@ -1,7 +1,8 @@
"""
LLM client for fact extraction and other AI-powered operations.
Fact extraction from text using LLM.
Uses OpenAI-compatible API (works with Groq, OpenAI, etc.)
Extracts semantic facts, entities, and temporal information from text.
Uses the LLMConfig wrapper for all LLM calls.
"""
import os
import json
@ -91,46 +92,13 @@ def chunk_text(text: str, max_chars: int = 120000) -> List[str]:
return splitter.split_text(text)
def get_llm_client() -> AsyncOpenAI:
"""
Get configured async LLM client.
Supports:
- Groq (default): Set GROQ_API_KEY and optionally GROQ_BASE_URL
- OpenAI: Set OPENAI_API_KEY
Returns:
Configured AsyncOpenAI client
"""
# Check for Groq configuration first
groq_api_key = os.getenv('GROQ_API_KEY')
if groq_api_key:
base_url = os.getenv('GROQ_BASE_URL', 'https://api.groq.com/openai/v1')
return AsyncOpenAI(
api_key=groq_api_key,
base_url=base_url
)
# Fall back to OpenAI
openai_api_key = os.getenv('OPENAI_API_KEY')
if openai_api_key:
return AsyncOpenAI(api_key=openai_api_key)
raise ValueError(
"No LLM API key found. Set GROQ_API_KEY or OPENAI_API_KEY environment variable."
)
async def _extract_facts_from_chunk(
chunk: str,
chunk_index: int,
total_chunks: int,
event_date: datetime,
context: str,
model: str,
temperature: float,
max_tokens: int,
client: AsyncOpenAI
llm_config: 'LLMConfig'
) -> List[Dict[str, str]]:
"""
Extract facts from a single chunk (internal helper for parallel processing).
@ -374,8 +342,7 @@ Remember:
for attempt in range(max_retries):
try:
llm_call_start = time.time()
response = await client.beta.chat.completions.parse(
model=model,
extraction_response = await llm_config.call(
messages=[
{
"role": "system",
@ -386,16 +353,14 @@ Remember:
"content": prompt
}
],
temperature=temperature,
max_tokens=max_tokens,
response_format=FactExtractionResponse,
extra_body={"service_tier": "auto"},
scope="memory_extract_facts",
temperature=0.1,
max_tokens=65000,
extra_body={"service_tier": "auto"}
)
llm_call_time = time.time() - llm_call_start
# Extract the parsed response
extraction_response = response.choices[0].message.parsed
# Convert to dict format
chunk_facts = [fact.model_dump() for fact in extraction_response.facts]
@ -421,9 +386,7 @@ async def extract_facts_from_text(
text: str,
event_date: datetime,
context: str = "",
model: str = "openai/gpt-oss-120b",
temperature: float = 0.1,
max_tokens: int = 65000,
llm_config: Optional['LLMConfig'] = None,
chunk_size: int = 5000
) -> List[Dict[str, str]]:
"""
@ -436,9 +399,7 @@ async def extract_facts_from_text(
text: Input text (conversation, article, etc.)
event_date: Reference date for resolving relative times
context: Context about the conversation/document
model: LLM model to use
temperature: Sampling temperature (lower = more focused)
max_tokens: Maximum tokens in response
llm_config: LLM configuration to use (if None, uses default from environment)
chunk_size: Maximum characters per chunk
Returns:
@ -446,10 +407,16 @@ async def extract_facts_from_text(
"""
import time
import logging
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from .llm_wrapper import LLMConfig
logger = logging.getLogger(__name__)
client = get_llm_client()
if llm_config is None:
from .llm_wrapper import LLMConfig
llm_config = LLMConfig.for_memory()
# Chunk text if necessary
chunk_start = time.time()
@ -466,10 +433,7 @@ async def extract_facts_from_text(
total_chunks=len(chunks),
event_date=event_date,
context=context,
model=model,
temperature=temperature,
max_tokens=max_tokens,
client=client
llm_config=llm_config
)
for i, chunk in enumerate(chunks)
]

229
memora/llm_wrapper.py Normal file
View file

@ -0,0 +1,229 @@
"""
LLM wrapper for unified configuration across providers.
"""
import os
import time
import asyncio
from typing import Optional, Any, Dict, List
from openai import AsyncOpenAI, RateLimitError, APIError
import logging
logger = logging.getLogger(__name__)
# Disable httpx logging
logging.getLogger("httpx").setLevel(logging.WARNING)
class LLMConfig:
"""Configuration for an LLM provider."""
def __init__(
self,
provider: Optional[str] = None,
api_key: Optional[str] = None,
base_url: Optional[str] = None,
model: Optional[str] = None,
provider_env: str = "MEMORY_LLM_PROVIDER",
api_key_env: str = "MEMORY_LLM_API_KEY",
base_url_env: str = "MEMORY_LLM_BASE_URL",
model_env: str = "MEMORY_LLM_MODEL",
):
"""
Initialize LLM configuration.
Args:
provider: Provider name ("openai", "groq", "ollama"). If None, reads from provider_env.
api_key: API key. If None, reads from api_key_env.
base_url: Base URL. If None, reads from base_url_env.
model: Model name. If None, reads from model_env.
provider_env: Environment variable name for provider (default: "MEMORY_LLM_PROVIDER")
api_key_env: Environment variable name for API key (default: "MEMORY_LLM_API_KEY")
base_url_env: Environment variable name for base URL (default: "MEMORY_LLM_BASE_URL")
model_env: Environment variable name for model (default: "MEMORY_LLM_MODEL")
"""
self.provider = (provider or os.getenv(provider_env, "groq")).lower()
self.api_key = api_key or os.getenv(api_key_env)
self.base_url = base_url or os.getenv(base_url_env)
self.model = model or os.getenv(model_env, "openai/gpt-oss-120b")
# Validate provider
if self.provider not in ["openai", "groq", "ollama"]:
raise ValueError(
f"Invalid LLM provider: {self.provider}. Must be 'openai', 'groq', or 'ollama'."
)
# Set default base URLs
if not self.base_url:
if self.provider == "groq":
self.base_url = "https://api.groq.com/openai/v1"
elif self.provider == "ollama":
self.base_url = "http://localhost:11434/v1"
# Validate API key (not needed for ollama)
if self.provider != "ollama" and not self.api_key:
raise ValueError(
f"API key not found for {self.provider}. Set {api_key_env} environment variable."
)
# Create client
if self.provider == "ollama":
self.client = AsyncOpenAI(api_key="ollama", base_url=self.base_url)
elif self.base_url:
self.client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url)
else:
self.client = AsyncOpenAI(api_key=self.api_key)
logger.info(
f"Initialized LLM: provider={self.provider}, model={self.model}, base_url={self.base_url}"
)
async def call(
self,
messages: List[Dict[str, str]],
response_format: Optional[Any] = None,
scope: str = "memory",
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 60.0,
**kwargs
) -> Any:
"""
Make an LLM API call with consistent configuration and retry logic.
Args:
messages: List of message dicts with 'role' and 'content'
response_format: Optional Pydantic model for structured output
scope: Scope identifier (e.g., 'memory', 'judge') for future tracking
max_retries: Maximum number of retry attempts (default: 5)
initial_backoff: Initial backoff time in seconds (default: 1.0)
max_backoff: Maximum backoff time in seconds (default: 60.0)
**kwargs: Additional parameters to pass to the API (temperature, max_tokens, etc.)
Returns:
Parsed response if response_format is provided, otherwise the text content
Raises:
Exception: Re-raises any API errors after all retries are exhausted
"""
start_time = time.time()
call_params = {
"model": self.model,
"messages": messages,
**kwargs
}
last_exception = None
for attempt in range(max_retries + 1):
try:
if response_format is not None:
# Use structured output parsing and return .parsed
response = await self.client.beta.chat.completions.parse(
response_format=response_format,
**call_params
)
result = response.choices[0].message.parsed
else:
# Standard completion and return text content
response = await self.client.chat.completions.create(**call_params)
result = response.choices[0].message.content
# Log call details on success
duration = time.time() - start_time
usage = response.usage
logger.info(
f"model={self.provider}/{self.model}, "
f"input_tokens={usage.prompt_tokens}, output_tokens={usage.completion_tokens}, "
f"total_tokens={usage.total_tokens}, time={duration:.3f}s"
)
return result
except RateLimitError as e:
last_exception = e
if attempt < max_retries:
# Calculate exponential backoff with jitter
backoff = min(initial_backoff * (2 ** attempt), max_backoff)
# Add jitter (±20%)
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
sleep_time = backoff + jitter
logger.warning(
f"Rate limit error (429) on attempt {attempt + 1}/{max_retries + 1}. "
f"Retrying in {sleep_time:.2f}s... Error: {str(e)}"
)
await asyncio.sleep(sleep_time)
else:
logger.error(
f"Rate limit error (429) after {max_retries + 1} attempts. Giving up. Error: {str(e)}"
)
raise
except APIError as e:
last_exception = e
# Check if it's a retryable error (5xx server errors)
if hasattr(e, 'status_code') and 500 <= e.status_code < 600:
if attempt < max_retries:
backoff = min(initial_backoff * (2 ** attempt), max_backoff)
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
sleep_time = backoff + jitter
logger.warning(
f"API error ({e.status_code}) on attempt {attempt + 1}/{max_retries + 1}. "
f"Retrying in {sleep_time:.2f}s... Error: {str(e)}"
)
await asyncio.sleep(sleep_time)
else:
logger.error(
f"API error ({e.status_code}) after {max_retries + 1} attempts. Giving up. Error: {str(e)}"
)
raise
else:
# Non-retryable API error, raise immediately
logger.error(f"Non-retryable API error: {str(e)}")
raise
except Exception as e:
# Any other exception, log and raise immediately
logger.error(f"Unexpected error during LLM call: {type(e).__name__}: {str(e)}")
raise
# This should never be reached, but just in case
if last_exception:
raise last_exception
raise RuntimeError("LLM call failed after all retries with no exception captured")
@classmethod
def for_memory(cls) -> "LLMConfig":
"""Create configuration for memory operations."""
return cls(
provider_env="MEMORY_LLM_PROVIDER",
api_key_env="MEMORY_LLM_API_KEY",
base_url_env="MEMORY_LLM_BASE_URL",
model_env="MEMORY_LLM_MODEL",
)
@classmethod
def for_judge(cls) -> "LLMConfig":
"""
Create configuration for judge/evaluator operations.
Falls back to memory LLM config if judge-specific config not set.
"""
# Check if judge-specific config exists, otherwise fall back to memory config
judge_provider = os.getenv("JUDGE_LLM_PROVIDER", os.getenv("MEMORY_LLM_PROVIDER", "groq"))
judge_api_key = os.getenv("JUDGE_LLM_API_KEY", os.getenv("MEMORY_LLM_API_KEY"))
judge_base_url = os.getenv("JUDGE_LLM_BASE_URL", os.getenv("MEMORY_LLM_BASE_URL"))
judge_model = os.getenv("JUDGE_LLM_MODEL", os.getenv("MEMORY_LLM_MODEL", "openai/gpt-oss-120b"))
return cls(
provider=judge_provider,
api_key=judge_api_key,
base_url=judge_base_url,
model=judge_model,
provider_env="JUDGE_LLM_PROVIDER",
api_key_env="JUDGE_LLM_API_KEY",
base_url_env="JUDGE_LLM_BASE_URL",
model_env="JUDGE_LLM_MODEL",
)

View file

@ -1,112 +0,0 @@
# Memory Operations Modules
This directory contains specialized operation modules for the TemporalSemanticMemory class.
## Refactoring Results
✅ **Successfully Completed!**
**File Size Reduction:**
- Before: 2,065 lines (temporal_semantic_memory.py)
- After: 1,846 lines (temporal_semantic_memory.py)
- **Removed: 219 lines (11% reduction)**
**Modules Created:**
- `embedding_operations.py` - Embedding generation with process pool parallelism
- `link_operations.py` - Entity, temporal, and semantic link creation (300+ lines)
- `think_operations.py` - Think operations with opinion handling (230+ lines)
- `batch_operations.py` - Placeholder for future extraction
- `search_operations.py` - Placeholder for future extraction
## Architecture
The memory system now uses a **mixin pattern** for better code organization:
```python
class TemporalSemanticMemory(
EmbeddingOperationsMixin,
LinkOperationsMixin,
ThinkOperationsMixin,
):
"""
Advanced memory system using temporal and semantic linking.
Mixins provide:
- EmbeddingOperationsMixin: _generate_embedding, _generate_embeddings_batch
- LinkOperationsMixin: Entity, temporal, semantic link operations
- ThinkOperationsMixin: think_async, _extract_opinions_from_text
"""
# Core infrastructure and batch operations
pass
```
## What Was Extracted
### EmbeddingOperationsMixin (embedding_operations.py)
- `_generate_embedding()` - Single embedding generation
- `_generate_embeddings_batch()` - Parallel batch embedding generation
- Process pool worker functions for CPU parallelism
### LinkOperationsMixin (link_operations.py)
- `_extract_entities_batch_optimized()` - Entity resolution and linking
- `_create_temporal_links_batch_per_fact()` - Time-based connections
- `_create_semantic_links_batch()` - Meaning-based connections
- `_insert_entity_links_batch()` - Batch link insertion
### ThinkOperationsMixin (think_operations.py)
- `think_async()` - Formulate answers using agent, world, and opinion facts
- `_extract_opinions_from_text()` - Extract opinions from generated text with LLM
- Parallel fact retrieval with `asyncio.gather`
- Opinion formation and storage as background tasks
### Remaining in Main Class
- Database connection management (`__init__`, `_get_pool`, `close`)
- Batch storage operations (`put`, `put_async`, `put_batch_async`)
- Search operations (`search`, `search_async`, `_apply_mmr`)
- Document management (`get_document`, `delete_document`, `delete_agent`)
- Deduplication (`_find_duplicate_facts_batch`)
- Opinion evaluation (`_evaluate_opinion_update_async`)
## Benefits Achieved
1. ✅ **Better Organization** - Related methods grouped in focused modules
2. ✅ **Reduced Complexity** - Main file is 17% smaller
3. ✅ **Reusability** - Mixins can be composed and tested independently
4. ✅ **Maintainability** - Easier to find and modify specific operations
5. ✅ **All Tests Pass** - No breaking changes to public API
## Usage
The public API remains unchanged:
```python
from memory import TemporalSemanticMemory
memory = TemporalSemanticMemory()
# All operations work exactly as before
result = await memory.think_async(
agent_id="agent_1",
query="What have I done?"
)
results, trace = await memory.search_async(
agent_id="agent_1",
query="example query",
fact_type="world"
)
```
## Future Work (Optional)
The foundation is now in place for further extraction:
- Extract batch operations to `batch_operations.py`
- Extract search operations to `search_operations.py`
- Split large methods into smaller, focused functions
## Design Principles Followed
1. ✅ **Preserved batch mechanisms** - Performance maintained
2. ✅ **No breaking changes** - All tests pass
3. ✅ **Clear separation** - Each mixin has focused responsibility
4. ✅ **Gradual refactoring** - Can continue incrementally

View file

@ -1,18 +0,0 @@
"""
Search operations for memory retrieval using spreading activation.
NOTE: This is a placeholder for future refactoring.
The actual implementation is currently in temporal_semantic_memory.py
"""
class SearchOperationsMixin:
"""
Mixin class for search operations.
Methods to be extracted:
- search
- search_async
- _apply_mmr
"""
pass

View file

@ -20,9 +20,6 @@ class ThinkOperationsMixin:
query: str,
thinking_budget: int = 50,
top_k: int = 10,
model: str = "openai/gpt-oss-120b",
temperature: float = 0.7,
max_tokens: int = 1000,
) -> Dict[str, Any]:
"""
Think and formulate an answer using agent identity, world facts, and opinions.
@ -31,7 +28,7 @@ class ThinkOperationsMixin:
1. Retrieves agent facts (agent's identity and past actions)
2. Retrieves world facts (general knowledge)
3. Retrieves existing opinions (agent's formed perspectives)
4. Uses Groq LLM to formulate an answer
4. Uses LLM to formulate an answer
5. Extracts and stores any new opinions formed during thinking
6. Returns plain text answer and the facts used
@ -40,9 +37,6 @@ class ThinkOperationsMixin:
query: Question to answer
thinking_budget: Number of memory units to explore
top_k: Maximum facts to retrieve
model: LLM model to use (default: openai/gpt-oss-120b)
temperature: Sampling temperature
max_tokens: Maximum tokens in response
Returns:
Dict with:
@ -50,11 +44,9 @@ class ThinkOperationsMixin:
- based_on: Dict with 'world', 'agent', and 'opinion' fact lists
- new_opinions: List of newly formed opinions
"""
# Use cached LLM client
if self._llm_client is None:
raise ValueError("GROQ_API_KEY environment variable not set")
client = self._llm_client
# Use cached LLM config
if self._llm_config is None:
raise ValueError("Memory LLM API key not set. Set MEMORY_LLM_API_KEY environment variable.")
# Steps 1-3: Run all three searches in parallel
(agent_results, _), (world_results, _), (opinion_results, _) = await asyncio.gather(
@ -148,44 +140,27 @@ Provide a helpful, accurate answer based on the facts above. Be consistent with
If you form any new opinions while thinking about this question, state them clearly in your answer."""
response = await client.chat.completions.create(
model=model,
answer_text = await self._llm_config.call(
messages=[
{"role": "system", "content": "You are a helpful AI assistant. Always respond in plain text without markdown formatting. You can form and express opinions based on facts."},
{"role": "user", "content": prompt}
],
temperature=temperature,
max_tokens=max_tokens
scope="memory_think",
temperature=0.7,
max_tokens=1000
)
answer_text = response.choices[0].message.content.strip()
answer_text = answer_text.strip()
# Step 6: Extract new opinions from the answer
new_opinions = await self._extract_opinions_from_text(
client=client,
text=answer_text,
model=model
)
# Step 6: Extract and store new opinions asynchronously (fire and forget)
await self._task_backend.submit_task({
'type': 'form_opinion',
'agent_id': agent_id,
'answer_text': answer_text,
'query': query
})
# Step 7: Store new opinions (schedule as background tasks, don't wait)
if new_opinions:
current_time = datetime.now(timezone.utc)
for opinion_dict in new_opinions:
task = asyncio.create_task(
self.put_async(
agent_id=agent_id,
content=opinion_dict["text"],
context=f"formed during thinking about: {query}",
event_date=current_time,
fact_type_override='opinion',
confidence_score=opinion_dict["confidence"]
)
)
# Track task and auto-remove when done
self._background_tasks.add(task)
task.add_done_callback(self._background_tasks.discard)
# Step 8: Return response with facts split by type
# Step 7: Return response with facts split by type (don't wait for opinions)
return {
"text": answer_text,
"based_on": {
@ -193,22 +168,57 @@ If you form any new opinions while thinking about this question, state them clea
"agent": agent_results,
"opinion": opinion_results
},
"new_opinions": new_opinions
"new_opinions": [] # Opinions are being extracted asynchronously
}
async def _extract_and_store_opinions_async(
self,
agent_id: str,
answer_text: str,
query: str
):
"""
Background task to extract and store opinions from think response.
This runs asynchronously and does not block the think response.
Args:
agent_id: Agent identifier
answer_text: The generated answer text
query: The original query
"""
try:
# Extract opinions from the answer
new_opinions = await self._extract_opinions_from_text(text=answer_text, query=query)
# Store new opinions
if new_opinions:
current_time = datetime.now(timezone.utc)
for opinion_dict in new_opinions:
await self.put_async(
agent_id=agent_id,
content=opinion_dict["text"],
context=f"formed during thinking about: {query}",
event_date=current_time,
fact_type_override='opinion',
confidence_score=opinion_dict["confidence"]
)
logger.debug(f"[THINK] Extracted and stored {len(new_opinions)} new opinions")
except Exception as e:
logger.warning(f"[THINK] Failed to extract/store opinions: {str(e)}")
async def _extract_opinions_from_text(
self,
client,
text: str,
model: str
query: str
) -> List[Dict[str, Any]]:
"""
Extract opinions with reasons and confidence from text using LLM.
Args:
client: OpenAI client
text: Text to extract opinions from
model: LLM model to use
query: The original query that prompted this response
Returns:
List of dicts with keys: 'text' (opinion with reasons), 'confidence' (score 0-1)
@ -226,31 +236,41 @@ If you form any new opinions while thinking about this question, state them clea
description="List of opinions formed with their supporting reasons and confidence scores"
)
extraction_prompt = f"""Extract any opinions or perspectives that were formed in the following text.
An opinion is a judgment, viewpoint, or conclusion that goes beyond just stating facts.
extraction_prompt = f"""Extract any NEW opinions or perspectives that were formed while answering the following question.
TEXT:
ORIGINAL QUESTION:
{query}
ANSWER PROVIDED:
{text}
For each opinion found, provide:
1. The opinion itself
2. The reasons or facts that support it
3. A confidence score (0.0 to 1.0) indicating how confident the agent is in this opinion based on the available information
An opinion is a judgment, viewpoint, or conclusion that goes beyond just stating facts. It represents a formed perspective or belief.
If no clear opinions are expressed, return an empty list."""
IMPORTANT: Do NOT extract statements like:
- "I don't have enough information"
- "The facts don't contain information about X"
- "I cannot answer because..."
- Simple acknowledgments or meta-statements about the query itself
ONLY extract actual opinions, judgments, or perspectives about substantive topics.
For each opinion found, provide:
1. The opinion itself (what the agent believes or concludes)
2. The reasons or facts that support it
3. A confidence score (0.0 to 1.0) indicating how confident the agent is in this opinion
If no genuine opinions are expressed (e.g., the response just says "I don't know"), return an empty list."""
try:
response = await client.beta.chat.completions.parse(
model=model,
result = await self._llm_config.call(
messages=[
{"role": "system", "content": "You extract opinions and perspectives from text."},
{"role": "user", "content": extraction_prompt}
],
response_format=OpinionExtractionResponse
response_format=OpinionExtractionResponse,
scope="memory_extract_opinion"
)
result = response.choices[0].message.parsed
# Format opinions with reasons included in the text and confidence score
formatted_opinions = []
for op in result.opinions:

200
memora/task_backend.py Normal file
View file

@ -0,0 +1,200 @@
"""
Abstract task backend for running async tasks.
This provides an abstraction that can be adapted to different execution models:
- AsyncIO queue (default implementation)
- Pub/Sub architectures (future)
- Message brokers (future)
"""
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, Callable, Awaitable
import asyncio
import logging
logger = logging.getLogger(__name__)
class TaskBackend(ABC):
"""
Abstract base class for task execution backends.
Implementations must:
1. Store/publish task events (as serializable dicts)
2. Execute tasks through a provided executor callback
The backend treats tasks as pure dictionaries that can be serialized
and sent over the network. The executor (typically TemporalSemanticMemory.execute_task)
receives the dict and routes it to the appropriate handler.
"""
def __init__(self):
"""Initialize the task backend."""
self._executor: Optional[Callable[[Dict[str, Any]], Awaitable[None]]] = None
self._initialized = False
def set_executor(self, executor: Callable[[Dict[str, Any]], Awaitable[None]]):
"""
Set the executor callback for processing tasks.
Args:
executor: Async function that takes a task dict and executes it
"""
self._executor = executor
@abstractmethod
async def initialize(self):
"""
Initialize the backend (e.g., start workers, connect to broker).
"""
pass
@abstractmethod
async def submit_task(self, task_dict: Dict[str, Any]):
"""
Submit a task for execution.
Args:
task_dict: Task as a dictionary (must be serializable)
"""
pass
@abstractmethod
async def shutdown(self):
"""
Shutdown the backend gracefully (e.g., stop workers, close connections).
"""
pass
async def _execute_task(self, task_dict: Dict[str, Any]):
"""
Execute a task through the registered executor.
Args:
task_dict: Task dictionary to execute
"""
if self._executor is None:
task_type = task_dict.get('type', 'unknown')
logger.warning(f"No executor registered, skipping task {task_type}")
return
try:
await self._executor(task_dict)
except Exception as e:
task_type = task_dict.get('type', 'unknown')
logger.error(f"Error executing task {task_type}: {e}")
import traceback
traceback.print_exc()
class AsyncIOQueueBackend(TaskBackend):
"""
Task backend implementation using asyncio queues.
This is the default implementation that uses in-process asyncio queues
and a periodic consumer worker.
"""
def __init__(
self,
batch_size: int = 100,
batch_interval: float = 1.0
):
"""
Initialize AsyncIO queue backend.
Args:
batch_size: Maximum number of tasks to process in one batch
batch_interval: Maximum time (seconds) to wait before processing batch
"""
super().__init__()
self._queue: Optional[asyncio.Queue] = None
self._worker_task: Optional[asyncio.Task] = None
self._shutdown_event: Optional[asyncio.Event] = None
self._batch_size = batch_size
self._batch_interval = batch_interval
async def initialize(self):
"""Initialize the queue and start the worker."""
if self._initialized:
return
self._queue = asyncio.Queue()
self._shutdown_event = asyncio.Event()
self._worker_task = asyncio.create_task(self._worker())
self._initialized = True
logger.info("AsyncIOQueueBackend initialized")
async def submit_task(self, task_dict: Dict[str, Any]):
"""
Submit a task by putting it in the queue.
Args:
task_dict: Task dictionary to execute
"""
if not self._initialized:
await self.initialize()
await self._queue.put(task_dict)
task_type = task_dict.get('type', 'unknown')
task_id = task_dict.get('id')
logger.debug(f"Task submitted: {task_type} (id: {task_id})")
async def shutdown(self):
"""Shutdown the worker and drain the queue."""
if not self._initialized:
return
logger.info("Shutting down AsyncIOQueueBackend...")
# Signal shutdown
self._shutdown_event.set()
# Cancel worker
if self._worker_task is not None:
self._worker_task.cancel()
try:
await self._worker_task
except asyncio.CancelledError:
logger.debug("Worker task cancelled successfully")
self._initialized = False
logger.info("AsyncIOQueueBackend shutdown complete")
async def _worker(self):
"""
Background worker that processes tasks in batches.
Collects tasks for up to batch_interval seconds or batch_size items,
then processes them.
"""
while not self._shutdown_event.is_set():
try:
# Collect tasks for batching
tasks = []
deadline = asyncio.get_event_loop().time() + self._batch_interval
while len(tasks) < self._batch_size and asyncio.get_event_loop().time() < deadline:
try:
remaining_time = max(0.1, deadline - asyncio.get_event_loop().time())
task_dict = await asyncio.wait_for(
self._queue.get(),
timeout=remaining_time
)
tasks.append(task_dict)
except asyncio.TimeoutError:
break
# Process batch
if tasks:
logger.debug(f"Processing batch of {len(tasks)} tasks")
# Execute tasks concurrently
await asyncio.gather(
*[self._execute_task(task_dict) for task_dict in tasks],
return_exceptions=True
)
except asyncio.CancelledError:
break
except Exception as e:
logger.error(f"Worker error: {e}")
await asyncio.sleep(1) # Backoff on error

View file

@ -12,7 +12,6 @@ import os
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional, Tuple
import asyncpg
from dotenv import load_dotenv
import asyncio
from .embeddings import Embeddings, SentenceTransformersEmbeddings
import time
@ -27,6 +26,8 @@ from .utils import (
)
from .entity_resolver import EntityResolver
from .operations import EmbeddingOperationsMixin, LinkOperationsMixin, ThinkOperationsMixin
from .llm_wrapper import LLMConfig
from .task_backend import TaskBackend, AsyncIOQueueBackend
def utcnow():
@ -55,30 +56,43 @@ class TemporalSemanticMemory(
def __init__(
self,
db_url: Optional[str] = None,
memory_llm_provider: Optional[str] = None,
memory_llm_api_key: Optional[str] = None,
memory_llm_base_url: Optional[str] = None,
memory_llm_model: Optional[str] = None,
embeddings: Optional[Embeddings] = None,
embedding_model: Optional[str] = None,
pool_min_size: int = 5,
pool_max_size: int = 100,
task_backend: Optional[TaskBackend] = None,
):
"""
Initialize the temporal + semantic memory system.
Args:
db_url: PostgreSQL connection URL (postgresql://user:pass@host:port/dbname)
db_url: PostgreSQL connection URL (postgresql://user:pass@host:port/dbname).
If not provided, reads from DATABASE_URL environment variable.
memory_llm_provider: LLM provider for memory operations: "openai", "groq", or "ollama".
If not provided, reads from MEMORY_LLM_PROVIDER environment variable (default: "groq").
memory_llm_api_key: API key for the LLM provider.
If not provided, reads from MEMORY_LLM_API_KEY environment variable.
memory_llm_base_url: Base URL for the LLM API (for ollama or custom endpoints).
If not provided, reads from MEMORY_LLM_BASE_URL environment variable.
Defaults: groq="https://api.groq.com/openai/v1", ollama="http://localhost:11434/v1"
memory_llm_model: Model name to use for all memory operations (put/think/opinions).
If not provided, reads from MEMORY_LLM_MODEL environment variable (default: "openai/gpt-oss-120b").
embeddings: Embeddings implementation to use. If not provided, uses SentenceTransformersEmbeddings
embedding_model: (Deprecated) Name of the SentenceTransformer model to use. Use embeddings parameter instead.
pool_min_size: Minimum number of connections in the pool (default: 5)
pool_max_size: Maximum number of connections in the pool (default: 100)
Increase for parallel think/search operations (e.g., 200-300 for 100+ parallel thinks)
task_backend: Custom task backend for async task execution. If not provided, uses AsyncIOQueueBackend
"""
load_dotenv()
# Initialize PostgreSQL connection URL
self.db_url = db_url or os.getenv("DATABASE_URL")
if not self.db_url:
raise ValueError(
"Database URL not found. "
"Set DATABASE_URL environment variable."
"Provide db_url parameter or set DATABASE_URL environment variable."
)
# Connection pool (will be created in initialize())
@ -94,75 +108,77 @@ class TemporalSemanticMemory(
if embeddings is not None:
self.embeddings = embeddings
else:
# Default to SentenceTransformersEmbeddings
model_name = embedding_model or "BAAI/bge-small-en-v1.5"
self.embeddings = SentenceTransformersEmbeddings(model_name)
self.embeddings = SentenceTransformersEmbeddings("BAAI/bge-small-en-v1.5")
# Initialize LLM client (cached for reuse across operations)
from openai import AsyncOpenAI
groq_api_key = os.getenv("GROQ_API_KEY")
if groq_api_key:
self._llm_client = AsyncOpenAI(
api_key=groq_api_key,
base_url="https://api.groq.com/openai/v1"
)
else:
self._llm_client = None # Will be created on-demand if needed
# Initialize LLM configuration
self._llm_config = LLMConfig(
provider=memory_llm_provider,
api_key=memory_llm_api_key,
base_url=memory_llm_base_url,
model=memory_llm_model,
provider_env="MEMORY_LLM_PROVIDER",
api_key_env="MEMORY_LLM_API_KEY",
base_url_env="MEMORY_LLM_BASE_URL",
model_env="MEMORY_LLM_MODEL",
)
# Background queue for access count updates (to avoid blocking searches)
self._access_count_queue = asyncio.Queue()
self._access_count_worker_task = None
self._shutdown_event = asyncio.Event()
# Store client and model for convenience
self._llm_client = self._llm_config.client
self._llm_model = self._llm_config.model
# Track background opinion PUT tasks to ensure clean shutdown
self._background_tasks = set()
# Initialize task backend
self._task_backend = task_backend or AsyncIOQueueBackend(
batch_size=100,
batch_interval=1.0
)
# Backpressure mechanism: limit concurrent searches to prevent overwhelming the database
self._search_semaphore = asyncio.Semaphore(32)
async def _access_count_worker(self):
"""Background worker that processes access count updates in batches."""
pool = self._pool # Pool is guaranteed to exist when worker starts
async def _handle_access_count_update(self, task_dict: Dict[str, Any]):
"""
Handler for access count update tasks.
while not self._shutdown_event.is_set():
try:
# Collect updates for up to 1 second or 1000 items
updates = {}
deadline = asyncio.get_event_loop().time() + 1.0
Args:
task_dict: Dict with 'node_ids' key containing list of node IDs to update
"""
node_ids = task_dict.get('node_ids', [])
if not node_ids:
return
while len(updates) < 1000 and asyncio.get_event_loop().time() < deadline:
try:
# Wait for items with short timeout
remaining_time = max(0.1, deadline - asyncio.get_event_loop().time())
node_ids = await asyncio.wait_for(
self._access_count_queue.get(),
timeout=remaining_time
)
# Deduplicate by adding to set
for node_id in node_ids:
updates[node_id] = True
except asyncio.TimeoutError:
break
pool = await self._get_pool()
try:
# Convert string UUIDs to UUID type for faster matching
uuid_list = [uuid.UUID(nid) for nid in node_ids]
async with pool.acquire() as conn:
await conn.execute(
"UPDATE memory_units SET access_count = access_count + 1 WHERE id = ANY($1::uuid[])",
uuid_list
)
except Exception as e:
logger.error(f"Access count handler: Error updating access counts: {e}")
# Process batch if we have updates
if updates:
node_id_list = list(updates.keys())
try:
# Convert string UUIDs to UUID type for faster matching
uuid_list = [uuid.UUID(nid) for nid in node_id_list]
async with pool.acquire() as conn:
await conn.execute(
"UPDATE memory_units SET access_count = access_count + 1 WHERE id = ANY($1::uuid[])",
uuid_list
)
except Exception as e:
logger.error(f"Access count worker: Error updating access counts: {e}")
async def execute_task(self, task_dict: Dict[str, Any]):
"""
Execute a task by routing it to the appropriate handler.
except asyncio.CancelledError:
break
except Exception as e:
logger.error(f"Access count worker: Unexpected error: {e}")
await asyncio.sleep(1) # Backoff on error
This method is called by the task backend to execute tasks.
It receives a plain dict that can be serialized and sent over the network.
Args:
task_dict: Task dictionary with 'type' key and other payload data
Example: {'type': 'access_count_update', 'node_ids': [...]}
"""
task_type = task_dict.get('type')
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)
else:
logger.error(f"Unknown task type: {task_type}")
async def initialize(self):
"""Initialize the connection pool and background workers."""
@ -183,11 +199,12 @@ class TemporalSemanticMemory(
# Initialize entity resolver with pool
self.entity_resolver = EntityResolver(self._pool)
# Start access count worker
self._access_count_worker_task = asyncio.create_task(self._access_count_worker())
# Set executor for task backend and initialize
self._task_backend.set_executor(self.execute_task)
await self._task_backend.initialize()
self._initialized = True
logger.info("Memory system initialized (pool and workers started)")
logger.info("Memory system initialized (pool and task backend started)")
async def _get_pool(self) -> asyncpg.Pool:
"""Get the connection pool (must call initialize() first)."""
@ -195,37 +212,14 @@ class TemporalSemanticMemory(
await self.initialize()
return self._pool
async def wait_for_background_tasks(self):
"""Wait for all background tasks (e.g., opinion PUTs) to complete."""
if self._background_tasks:
await asyncio.gather(*self._background_tasks, return_exceptions=True)
async def close(self):
"""Close the connection pool and shutdown background workers."""
logger.info("close() started")
# Signal shutdown to worker
self._shutdown_event.set()
logger.info("shutdown event set")
# Wait for background opinion PUT tasks to complete
if self._background_tasks:
logger.debug(f"waiting for {len(self._background_tasks)} background tasks to complete")
await self.wait_for_background_tasks()
logger.debug("background tasks completed")
# Cancel and wait for worker task
if self._access_count_worker_task is not None:
logger.debug("cancelling worker task")
self._access_count_worker_task.cancel()
try:
logger.debug("waiting for worker task to finish")
await self._access_count_worker_task
logger.debug("worker task finished")
except asyncio.CancelledError:
logger.debug("worker task cancelled successfully")
else:
logger.debug("no worker task to cancel")
# Shutdown task backend
logger.debug("shutting down task backend")
await self._task_backend.shutdown()
logger.debug("task backend shutdown complete")
# Close pool
if self._pool is not None:
@ -236,6 +230,7 @@ class TemporalSemanticMemory(
else:
logger.debug("no pool to close")
self._initialized = False
logger.debug("close() completed")
async def _find_duplicate_facts_batch(
@ -466,14 +461,14 @@ class TemporalSemanticMemory(
# Step 1: Extract facts from ALL contents in parallel
step_start = time.time()
# Create tasks for parallel fact extraction
# Create tasks for parallel fact extraction using configured LLM
fact_extraction_tasks = []
for item in contents:
content = item["content"]
context = item.get("context", "")
event_date = item.get("event_date") or utcnow()
task = extract_facts(content, event_date, context)
task = extract_facts(content, event_date, context, llm_config=self._llm_config)
fact_extraction_tasks.append((task, event_date, context))
# Wait for all fact extractions to complete
@ -703,14 +698,13 @@ class TemporalSemanticMemory(
# Trigger opinion reinforcement in background (non-blocking)
# Only trigger if there are entities in the new units
if any(filtered_entities):
asyncio.create_task(
self._reinforce_opinions_async(
agent_id=agent_id,
created_unit_ids=created_unit_ids,
unit_texts=filtered_sentences,
unit_entities=filtered_entities
)
)
await self._task_backend.submit_task({
'type': 'reinforce_opinion',
'agent_id': agent_id,
'created_unit_ids': created_unit_ids,
'unit_texts': filtered_sentences,
'unit_entities': filtered_entities
})
logger.debug("[PUT_BATCH_ASYNC] Opinion reinforcement task queued in background")
return result_unit_ids
@ -725,6 +719,7 @@ class TemporalSemanticMemory(
self,
agent_id: str,
query: str,
fact_type: str,
thinking_budget: int = 50,
top_k: int = 10,
enable_trace: bool = False,
@ -733,7 +728,6 @@ class TemporalSemanticMemory(
weight_recency: float = 0.25,
weight_frequency: float = 0.15,
mmr_lambda: float = 0.5,
fact_type: Optional[str] = None,
) -> tuple[List[Dict[str, Any]], Optional[Any]]:
"""
Search memories using spreading activation (synchronous wrapper).
@ -744,6 +738,7 @@ class TemporalSemanticMemory(
Args:
agent_id: Agent ID to search for
query: Search query
fact_type: Required filter for fact type ('world', 'agent', or 'opinion')
thinking_budget: How many units to explore (computational budget)
top_k: Number of results to return
enable_trace: If True, returns detailed SearchTrace object
@ -752,21 +747,21 @@ class TemporalSemanticMemory(
weight_recency: Weight for recency component (default: 0.25)
weight_frequency: Weight for frequency component (default: 0.15)
mmr_lambda: Lambda for MMR diversification (0=max diversity, 1=no diversity, default: 0.5)
fact_type: Optional filter for fact type ('world' or 'agent')
Returns:
Tuple of (results, trace)
"""
# Run async version synchronously
return asyncio.run(self.search_async(
agent_id, query, thinking_budget, top_k, enable_trace,
weight_activation, weight_semantic, weight_recency, weight_frequency, mmr_lambda, fact_type
agent_id, query, fact_type, thinking_budget, top_k, enable_trace,
weight_activation, weight_semantic, weight_recency, weight_frequency, mmr_lambda
))
async def search_async(
self,
agent_id: str,
query: str,
fact_type: str,
thinking_budget: int = 50,
top_k: int = 10,
enable_trace: bool = False,
@ -775,7 +770,6 @@ class TemporalSemanticMemory(
weight_recency: float = 0.25,
weight_frequency: float = 0.15,
mmr_lambda: float = 0.5,
fact_type: Optional[str] = None,
max_neighbors_per_node: int = 20,
) -> tuple[List[Dict[str, Any]], Optional[Any]]:
"""
@ -835,36 +829,21 @@ class TemporalSemanticMemory(
if conn_acquire_time > 0.1: # Log if waiting > 100ms
log_buffer.append(f" [2.1] Waited {conn_acquire_time:.3f}s for connection (pool busy)")
# Build entry point query with optional fact_type filter
if fact_type:
entry_points = await conn.fetch(
"""
SELECT id, text, context, event_date, access_count, embedding,
1 - (embedding <=> $1::vector) AS similarity
FROM memory_units
WHERE agent_id = $2
AND embedding IS NOT NULL
AND fact_type = $3
AND (1 - (embedding <=> $1::vector)) >= 0.5
ORDER BY embedding <=> $1::vector
LIMIT 3
""",
query_embedding_str, agent_id, fact_type
)
else:
entry_points = await conn.fetch(
"""
SELECT id, text, context, event_date, access_count, embedding,
1 - (embedding <=> $1::vector) AS similarity
FROM memory_units
WHERE agent_id = $2
AND embedding IS NOT NULL
AND (1 - (embedding <=> $1::vector)) >= 0.5
ORDER BY embedding <=> $1::vector
LIMIT 3
""",
query_embedding_str, agent_id
)
# Find entry points using vector similarity
entry_points = await conn.fetch(
"""
SELECT id, text, context, event_date, access_count, embedding,
1 - (embedding <=> $1::vector) AS similarity
FROM memory_units
WHERE agent_id = $2
AND embedding IS NOT NULL
AND fact_type = $3
AND (1 - (embedding <=> $1::vector)) >= 0.5
ORDER BY embedding <=> $1::vector
LIMIT 3
""",
query_embedding_str, agent_id, fact_type
)
step_duration = time.time() - step_start
log_buffer.append(f" [2] Find entry points: {len(entry_points)} found in {step_duration:.3f}s")
@ -942,49 +921,29 @@ class TemporalSemanticMemory(
substep_start = time.time()
uuid_array = [uuid.UUID(nid) for nid in node_ids]
# Build neighbor query with optional fact_type filter
# Query neighbors for batch, limiting to top N per node
# OPTIMIZATION: Limit neighbors per node to reduce data transfer
# Dense graphs can have 100+ neighbors per node, but spreading activation
# only needs top-weighted neighbors. This reduces query from 9000→1000 rows.
# Configurable via max_neighbors_per_node parameter (default: 20)
if fact_type:
all_neighbors = await conn.fetch(
"""
SELECT * FROM (
SELECT ml.from_unit_id, ml.to_unit_id, ml.weight, ml.link_type, ml.entity_id,
mu.text, mu.context, mu.event_date, mu.access_count,
mu.id as neighbor_id,
ROW_NUMBER() OVER (PARTITION BY ml.from_unit_id ORDER BY ml.weight DESC) as rn
FROM memory_links ml
JOIN memory_units mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.weight >= 0.1
AND mu.fact_type = $2
) sub
WHERE rn <= $3
ORDER BY from_unit_id, weight DESC
""",
uuid_array, fact_type, max_neighbors_per_node
)
else:
all_neighbors = await conn.fetch(
"""
SELECT * FROM (
SELECT ml.from_unit_id, ml.to_unit_id, ml.weight, ml.link_type, ml.entity_id,
mu.text, mu.context, mu.event_date, mu.access_count,
mu.id as neighbor_id,
ROW_NUMBER() OVER (PARTITION BY ml.from_unit_id ORDER BY ml.weight DESC) as rn
FROM memory_links ml
JOIN memory_units mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.weight >= 0.1
) sub
WHERE rn <= $2
ORDER BY from_unit_id, weight DESC
""",
uuid_array, max_neighbors_per_node
)
all_neighbors = await conn.fetch(
"""
SELECT * FROM (
SELECT ml.from_unit_id, ml.to_unit_id, ml.weight, ml.link_type, ml.entity_id,
mu.text, mu.context, mu.event_date, mu.access_count,
mu.id as neighbor_id,
ROW_NUMBER() OVER (PARTITION BY ml.from_unit_id ORDER BY ml.weight DESC) as rn
FROM memory_links ml
JOIN memory_units mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.weight >= 0.1
AND mu.fact_type = $2
) sub
WHERE rn <= $3
ORDER BY from_unit_id, weight DESC
""",
uuid_array, fact_type, max_neighbors_per_node
)
neighbor_query_time = time.time() - substep_start
if neighbor_query_time > 1.0: # Log slow neighbor queries
log_buffer.append(f" [3.3.3] Slow NEIGHBOR query: {neighbor_query_time:.3f}s for {len(node_ids)} nodes → {len(all_neighbors)} neighbors")
@ -1190,7 +1149,10 @@ class TemporalSemanticMemory(
# Step 4: Queue access count updates (background worker will process them)
if visited_node_ids:
await self._access_count_queue.put(visited_node_ids)
await self._task_backend.submit_task({
'type': 'access_count_update',
'node_ids': visited_node_ids
})
log_buffer.append(f" [4] Queued access count updates for {len(visited_node_ids)} nodes")
# Step 5: Sort by final weight and apply MMR for diversity
@ -1427,6 +1389,36 @@ class TemporalSemanticMemory(
"memory_units_deleted": units_count if deleted else 0
}
async def delete_memory_unit(self, unit_id: str) -> Dict[str, Any]:
"""
Delete a single memory unit and all its associated links.
Due to CASCADE DELETE constraints, this will automatically delete:
- All links from this unit (memory_links where from_unit_id = unit_id)
- All links to this unit (memory_links where to_unit_id = unit_id)
- All entity associations (unit_entities where unit_id = unit_id)
Args:
unit_id: UUID of the memory unit to delete
Returns:
Dictionary with deletion result
"""
pool = await self._get_pool()
async with pool.acquire() as conn:
async with conn.transaction():
# Delete the memory unit (cascades to links and associations)
deleted = await conn.fetchval(
"DELETE FROM memory_units WHERE id = $1 RETURNING id",
unit_id
)
return {
"success": deleted is not None,
"unit_id": str(deleted) if deleted else None,
"message": "Memory unit and all its links deleted successfully" if deleted else "Memory unit not found"
}
async def delete_agent(self, agent_id: str) -> Dict[str, int]:
"""
Delete all data for a specific agent (multi-tenant cleanup).
@ -1649,23 +1641,19 @@ class TemporalSemanticMemory(
async def _evaluate_opinion_update_async(
self,
client,
opinion_text: str,
opinion_confidence: float,
new_event_text: str,
entity_name: str,
model: str = "openai/gpt-oss-120b",
) -> Optional[Dict[str, Any]]:
"""
Evaluate if an opinion should be updated based on a new event.
Args:
client: OpenAI client
opinion_text: Current opinion text (includes reasons)
opinion_confidence: Current confidence score (0.0-1.0)
new_event_text: Text of the new event
entity_name: Name of the entity this opinion is about
model: LLM model to use
Returns:
Dict with 'action' ('keep'|'update'), 'new_confidence', 'new_text' (if action=='update')
@ -1707,18 +1695,16 @@ Guidelines:
- Small changes in confidence are normal; large jumps should be rare"""
try:
response = await client.beta.chat.completions.parse(
model=model,
result = await self._llm_config.call(
messages=[
{"role": "system", "content": "You evaluate and update opinions based on new information."},
{"role": "user", "content": evaluation_prompt}
],
response_format=OpinionEvaluation,
scope="memory_evaluate_opinion",
temperature=0.3 # Lower temperature for more consistent evaluation
)
result = response.choices[0].message.parsed
# Only return updates if something actually changed
if result.action == 'keep' and abs(result.new_confidence - opinion_confidence) < 0.01:
return None
@ -1734,13 +1720,48 @@ Guidelines:
logger.warning(f"Failed to evaluate opinion update: {str(e)}")
return None
async def _handle_form_opinion(self, task_dict: Dict[str, Any]):
"""
Handler for form opinion tasks.
Args:
task_dict: Dict with keys: 'agent_id', 'answer_text', 'query'
"""
agent_id = task_dict['agent_id']
answer_text = task_dict['answer_text']
query = task_dict['query']
await self._extract_and_store_opinions_async(
agent_id=agent_id,
answer_text=answer_text,
query=query
)
async def _handle_reinforce_opinion(self, task_dict: Dict[str, Any]):
"""
Handler for reinforce opinion tasks.
Args:
task_dict: Dict with keys: 'agent_id', 'created_unit_ids', 'unit_texts', 'unit_entities'
"""
agent_id = task_dict['agent_id']
created_unit_ids = task_dict['created_unit_ids']
unit_texts = task_dict['unit_texts']
unit_entities = task_dict['unit_entities']
await self._reinforce_opinions_async(
agent_id=agent_id,
created_unit_ids=created_unit_ids,
unit_texts=unit_texts,
unit_entities=unit_entities
)
async def _reinforce_opinions_async(
self,
agent_id: str,
created_unit_ids: List[str],
unit_texts: List[str],
unit_entities: List[List[Dict[str, str]]],
model: str = "openai/gpt-oss-120b",
):
"""
Background task to reinforce opinions based on newly ingested events.
@ -1752,7 +1773,6 @@ Guidelines:
created_unit_ids: List of newly created memory unit IDs
unit_texts: Texts of the newly created units
unit_entities: Entities extracted from each unit
model: LLM model to use for evaluation
"""
try:
# Extract all unique entity names from the new units
@ -1790,11 +1810,10 @@ Guidelines:
logger.debug(f"[REINFORCE] Found {len(opinions)} opinions to potentially reinforce")
# Use cached LLM client
if self._llm_client is None:
logger.error("[REINFORCE] LLM client not available, skipping opinion reinforcement")
# Use cached LLM config
if self._llm_config is None:
logger.error("[REINFORCE] LLM config not available, skipping opinion reinforcement")
return
client = self._llm_client
# Evaluate each opinion against the new events
updates_to_apply = []
@ -1818,12 +1837,10 @@ Guidelines:
# Evaluate if opinion should be updated
evaluation = await self._evaluate_opinion_update_async(
client,
opinion_text,
opinion_confidence,
combined_events,
entity_name,
model
entity_name
)
if evaluation:

View file

@ -2,11 +2,15 @@
Utility functions for memory system.
"""
from datetime import datetime
from typing import List, Dict
from .llm_client import extract_facts_from_text
from typing import List, Dict, TYPE_CHECKING
if TYPE_CHECKING:
from .llm_wrapper import LLMConfig
from .fact_extraction import extract_facts_from_text
async def extract_facts(text: str, event_date: datetime, context: str = "") -> List[Dict[str, str]]:
async def extract_facts(text: str, event_date: datetime, context: str = "", llm_config: 'LLMConfig' = None) -> List[Dict[str, str]]:
"""
Extract semantic facts from text using LLM.
@ -20,6 +24,7 @@ async def extract_facts(text: str, event_date: datetime, context: str = "") -> L
text: Input text (conversation, article, etc.)
event_date: Reference date for resolving relative times
context: Context about the conversation/document
llm_config: LLM configuration to use
Returns:
List of fact dictionaries with keys: 'fact' (text) and 'date' (ISO string)
@ -30,7 +35,7 @@ async def extract_facts(text: str, event_date: datetime, context: str = "") -> L
if not text or not text.strip():
return []
fact_dicts = await extract_facts_from_text(text, event_date, context)
fact_dicts = await extract_facts_from_text(text, event_date, context, llm_config=llm_config)
if not fact_dicts:
raise Exception(f"LLM extracted 0 facts from text of length {len(text)}. This may indicate the text contains no meaningful information, or the LLM failed to extract facts.")

View file

@ -22,9 +22,11 @@ from memora.embeddings import Embeddings
import logging
load_dotenv()
logging.basicConfig(level=logging.INFO)
# Environment variables are loaded by the shell script that calls this module
# No need to load .env files here as they're sourced by start-server.sh
def create_app(embeddings: Optional[Embeddings] = None, db_url: Optional[str] = None) -> FastAPI:
"""
@ -105,6 +107,7 @@ class SearchRequest(BaseModel):
top_k: int = 10
mmr_lambda: float = 0.5
trace: bool = False
fact_type: Optional[str] = None
class Config:
json_schema_extra = {
@ -114,7 +117,8 @@ class SearchRequest(BaseModel):
"thinking_budget": 100,
"top_k": 10,
"mmr_lambda": 0.5,
"trace": True
"trace": True,
"fact_type": "world"
}
}
@ -325,8 +329,8 @@ def _register_routes(app: FastAPI):
"/api/search",
response_model=SearchResponse,
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. Optionally filter by fact_type (world, agent, opinion)"
)
async def api_search(request: SearchRequest):
"""Run a search and return results with trace."""
@ -338,7 +342,8 @@ def _register_routes(app: FastAPI):
thinking_budget=request.thinking_budget,
top_k=request.top_k,
enable_trace=request.trace,
mmr_lambda=request.mmr_lambda
mmr_lambda=request.mmr_lambda,
fact_type=request.fact_type
)
# Convert trace to dict
@ -585,6 +590,30 @@ def _register_routes(app: FastAPI):
raise HTTPException(status_code=500, detail=str(e))
@app.delete(
"/api/memory/{unit_id}",
tags=["Memory Storage"],
summary="Delete a memory unit",
description="Delete a single memory unit and all its associated links (temporal, semantic, and entity links)"
)
async def api_delete_memory_unit(unit_id: str):
"""Delete a memory unit and all its links."""
try:
result = await app.state.memory.delete_memory_unit(unit_id)
if not result["success"]:
raise HTTPException(status_code=404, detail=result["message"])
return result
except HTTPException:
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
print(f"Error in /api/memory/{unit_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
# Create default app instance

View file

@ -758,3 +758,20 @@ body {
border: 2px solid #333;
box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
/* Delete Button */
.delete-button {
padding: 4px 12px;
background: #ef5350;
color: white;
border: none;
border-radius: 4px;
cursor: pointer;
font-size: 12px;
font-weight: bold;
transition: background 0.2s;
}
.delete-button:hover {
background: #d32f2f;
}

View file

@ -326,6 +326,13 @@ function updateDataTable(factType, data) {
<td>${row.context}</td>
<td>${row.date}</td>
<td>${row.entities}</td>
<td>
<button onclick="deleteRecord('${factType}', '${row.id}')"
class="delete-button"
title="Delete this record and all its links">
Delete
</button>
</td>
</tr>
`).join('');
}
@ -351,6 +358,33 @@ function updateDataTable(factType, data) {
}
}
// Delete a record and all its links
window.deleteRecord = async function(factType, recordId) {
if (!confirm('Are you sure you want to delete this record and all its links? This action cannot be undone.')) {
return;
}
try {
const response = await fetch(`api/memory/${encodeURIComponent(recordId)}`, {
method: 'DELETE'
});
if (!response.ok) {
const error = await response.json();
throw new Error(error.detail || `HTTP ${response.status}`);
}
const result = await response.json();
alert(result.message || 'Record deleted successfully');
// Reload the table data
await loadDataView(factType);
} catch (e) {
console.error('Error deleting record:', e);
alert('Error deleting record: ' + e.message);
}
}
// Load data from API (old function - kept for backward compatibility)
async function loadGraphData() {
try {
@ -626,6 +660,7 @@ function addDebugPane() {
<option value="all">All Facts</option>
<option value="world">World Facts</option>
<option value="agent">Agent Facts</option>
<option value="opinion">Opinion Facts</option>
</select>
</div>
<div>
@ -801,29 +836,29 @@ window.runSearchInPane = async function(paneId) {
}
try {
// Determine endpoint based on search type
let endpoint = 'api/search';
if (searchType === 'world') {
endpoint = 'api/world_search';
} else if (searchType === 'agent') {
endpoint = 'api/agent_search';
// Prepare request body with optional fact_type
const requestBody = {
query: query,
agent_id: agentId,
thinking_budget: thinkingBudget,
top_k: topK,
mmr_lambda: mmrLambda,
trace: true
};
// Add fact_type if not 'all'
if (searchType !== 'all') {
requestBody.fact_type = searchType;
}
statusBar.innerHTML = '<span style="color: #ff9800;">🔄 Searching...</span>';
const response = await fetch(endpoint, {
const response = await fetch('api/search', {
method: 'POST',
headers: {
'Content-Type': 'application/json'
},
body: JSON.stringify({
query: query,
agent_id: agentId,
thinking_budget: thinkingBudget,
top_k: topK,
mmr_lambda: mmrLambda,
trace: true
})
body: JSON.stringify(requestBody)
});
const data = await response.json();

View file

@ -83,10 +83,10 @@
<div class="table-container">
<table class="memory-table">
<thead>
<tr><th>ID</th><th>Text</th><th>Context</th><th>Date</th><th>Entities</th></tr>
<tr><th>ID</th><th>Text</th><th>Context</th><th>Date</th><th>Entities</th><th>Actions</th></tr>
</thead>
<tbody id="world-table-body">
<tr><td colspan="5" class="empty-message">Click "Load World Facts" to view data</td></tr>
<tr><td colspan="6" class="empty-message">Click "Load World Facts" to view data</td></tr>
</tbody>
</table>
</div>
@ -143,10 +143,10 @@
<div class="table-container">
<table class="memory-table">
<thead>
<tr><th>ID</th><th>Text</th><th>Context</th><th>Date</th><th>Entities</th></tr>
<tr><th>ID</th><th>Text</th><th>Context</th><th>Date</th><th>Entities</th><th>Actions</th></tr>
</thead>
<tbody id="agent-table-body">
<tr><td colspan="5" class="empty-message">Click "Load Agent Facts" to view data</td></tr>
<tr><td colspan="6" class="empty-message">Click "Load Agent Facts" to view data</td></tr>
</tbody>
</table>
</div>
@ -203,10 +203,10 @@
<div class="table-container">
<table class="memory-table">
<thead>
<tr><th>ID</th><th>Text</th><th>Context</th><th>Date</th><th>Entities</th></tr>
<tr><th>ID</th><th>Text</th><th>Context</th><th>Date</th><th>Entities</th><th>Actions</th></tr>
</thead>
<tbody id="opinions-table-body">
<tr><td colspan="5" class="empty-message">Click "Load Opinions" to view data</td></tr>
<tr><td colspan="6" class="empty-message">Click "Load Opinions" to view data</td></tr>
</tbody>
</table>
</div>

View file

@ -37,4 +37,4 @@ log_cli_level = "INFO"
log_cli_format = "%(asctime)s %(levelname)s %(message)s"
log_cli_date_format = "%Y-%m-%d %H:%M:%S"
addopts = "--timeout 60 -p no:warnings"
asyncio_default_fixture_loop_scope = "session"
asyncio_default_fixture_loop_scope = "function"

View file

@ -1,145 +0,0 @@
-- Enable the pgvector extension and uuid extension
CREATE EXTENSION IF NOT EXISTS vector;
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
-- ============================================================================
-- TEMPORAL + SEMANTIC + ENTITY MEMORY ARCHITECTURE
-- ============================================================================
-- Documents: Source of memory units (for tracking, updates, and deletion)
CREATE TABLE IF NOT EXISTS documents (
id TEXT NOT NULL, -- User-provided document ID
agent_id TEXT NOT NULL,
PRIMARY KEY (id, agent_id),
original_text TEXT, -- Full original content (for context expansion)
content_hash TEXT, -- SHA256 hash for deduplication
metadata JSONB DEFAULT '{}'::jsonb, -- User-provided metadata
created_at TIMESTAMPTZ DEFAULT NOW(),
updated_at TIMESTAMPTZ DEFAULT NOW()
);
-- Memory Units: Individual sentence-level memories
CREATE TABLE IF NOT EXISTS memory_units (
id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
agent_id TEXT NOT NULL,
document_id TEXT, -- Link to source document
text TEXT NOT NULL,
embedding vector(384), -- bge-small-en-v1.5 dimension
context TEXT, -- What was happening when this memory was formed
event_date TIMESTAMPTZ NOT NULL, -- When the event occurred
fact_type TEXT NOT NULL DEFAULT 'world', -- 'world' (general facts), 'agent' (agent actions), or 'opinion' (agent opinions)
confidence_score FLOAT, -- Confidence score for opinions (0.0 to 1.0, only used for fact_type='opinion')
access_count INTEGER DEFAULT 0, -- For recency/frequency weighting
created_at TIMESTAMPTZ DEFAULT NOW(),
updated_at TIMESTAMPTZ DEFAULT NOW(),
CHECK (fact_type IN ('world', 'agent', 'opinion')),
CHECK (confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0))
);
-- Entities: Resolved entities (people, organizations, locations, etc.)
CREATE TABLE IF NOT EXISTS entities (
id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
canonical_name TEXT NOT NULL, -- "Alice", "Google", "San Francisco"
entity_type TEXT NOT NULL, -- PERSON, ORG, GPE, etc.
agent_id TEXT NOT NULL, -- Entities are scoped to agents
metadata JSONB DEFAULT '{}'::jsonb, -- Additional entity info
first_seen TIMESTAMPTZ DEFAULT NOW(),
last_seen TIMESTAMPTZ DEFAULT NOW(),
mention_count INTEGER DEFAULT 1
);
-- Unit-Entity associations: Which entities appear in which units
CREATE TABLE IF NOT EXISTS unit_entities (
unit_id UUID REFERENCES memory_units(id) ON DELETE CASCADE,
entity_id UUID REFERENCES entities(id) ON DELETE CASCADE,
PRIMARY KEY (unit_id, entity_id)
);
-- Entity Co-occurrences: Materialized cache of which entities appear together
-- This dramatically speeds up entity resolution by avoiding expensive joins
CREATE TABLE IF NOT EXISTS entity_cooccurrences (
entity_id_1 UUID REFERENCES entities(id) ON DELETE CASCADE,
entity_id_2 UUID REFERENCES entities(id) ON DELETE CASCADE,
cooccurrence_count INTEGER DEFAULT 1,
last_cooccurred TIMESTAMPTZ DEFAULT NOW(),
PRIMARY KEY (entity_id_1, entity_id_2),
CHECK (entity_id_1 < entity_id_2) -- Enforce ordering to avoid duplicates
);
-- Memory Links: Temporal, semantic, AND entity connections
CREATE TABLE IF NOT EXISTS memory_links (
from_unit_id UUID REFERENCES memory_units(id) ON DELETE CASCADE,
to_unit_id UUID REFERENCES memory_units(id) ON DELETE CASCADE,
link_type TEXT NOT NULL, -- 'temporal', 'semantic', or 'entity'
weight FLOAT NOT NULL DEFAULT 1.0, -- Link strength
entity_id UUID REFERENCES entities(id) ON DELETE CASCADE, -- Set for entity links
created_at TIMESTAMPTZ DEFAULT NOW()
);
-- Unique constraint to prevent duplicate links
CREATE UNIQUE INDEX IF NOT EXISTS idx_memory_links_unique
ON memory_links (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid));
-- ============================================================================
-- FOREIGN KEY CONSTRAINTS
-- ============================================================================
-- Add foreign key from memory_units to documents (composite key)
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_constraint WHERE conname = 'memory_units_document_fkey'
) THEN
ALTER TABLE memory_units
ADD CONSTRAINT memory_units_document_fkey
FOREIGN KEY (document_id, agent_id)
REFERENCES documents(id, agent_id)
ON DELETE CASCADE;
END IF;
END $$;
-- ============================================================================
-- INDEXES
-- ============================================================================
-- Document indexes
CREATE INDEX IF NOT EXISTS idx_documents_agent_id ON documents(agent_id);
CREATE INDEX IF NOT EXISTS idx_documents_content_hash ON documents(content_hash);
-- Memory unit indexes
CREATE INDEX IF NOT EXISTS idx_memory_units_agent_id ON memory_units(agent_id);
CREATE INDEX IF NOT EXISTS idx_memory_units_document_id ON memory_units(document_id);
CREATE INDEX IF NOT EXISTS idx_memory_units_event_date ON memory_units(event_date DESC);
CREATE INDEX IF NOT EXISTS idx_memory_units_agent_date ON memory_units(agent_id, event_date DESC);
CREATE INDEX IF NOT EXISTS idx_memory_units_access_count ON memory_units(access_count DESC);
CREATE INDEX IF NOT EXISTS idx_memory_units_fact_type ON memory_units(fact_type);
CREATE INDEX IF NOT EXISTS idx_memory_units_agent_fact_type ON memory_units(agent_id, fact_type);
-- Vector similarity index (HNSW for fast approximate nearest neighbor)
CREATE INDEX IF NOT EXISTS idx_memory_units_embedding ON memory_units
USING hnsw (embedding vector_cosine_ops);
-- Entity indexes
CREATE INDEX IF NOT EXISTS idx_entities_agent_id ON entities(agent_id);
CREATE INDEX IF NOT EXISTS idx_entities_canonical_name ON entities(canonical_name);
CREATE INDEX IF NOT EXISTS idx_entities_type ON entities(entity_type);
CREATE INDEX IF NOT EXISTS idx_entities_agent_name_type ON entities(agent_id, canonical_name, entity_type);
-- Unit-entity indexes
CREATE INDEX IF NOT EXISTS idx_unit_entities_unit ON unit_entities(unit_id);
CREATE INDEX IF NOT EXISTS idx_unit_entities_entity ON unit_entities(entity_id);
-- Entity co-occurrence indexes for fast lookups
CREATE INDEX IF NOT EXISTS idx_entity_cooccurrences_entity1 ON entity_cooccurrences(entity_id_1);
CREATE INDEX IF NOT EXISTS idx_entity_cooccurrences_entity2 ON entity_cooccurrences(entity_id_2);
CREATE INDEX IF NOT EXISTS idx_entity_cooccurrences_count ON entity_cooccurrences(cooccurrence_count DESC);
-- Link indexes for graph traversal
CREATE INDEX IF NOT EXISTS idx_memory_links_from ON memory_links(from_unit_id);
CREATE INDEX IF NOT EXISTS idx_memory_links_to ON memory_links(to_unit_id);
CREATE INDEX IF NOT EXISTS idx_memory_links_type ON memory_links(link_type);
CREATE INDEX IF NOT EXISTS idx_memory_links_entity ON memory_links(entity_id) WHERE entity_id IS NOT NULL;
-- Composite index for spreading activation neighbor queries (from_unit_id + weight filter)
CREATE INDEX IF NOT EXISTS idx_memory_links_from_weight ON memory_links(from_unit_id, weight DESC)
WHERE weight >= 0.1;

View file

@ -0,0 +1,43 @@
#!/bin/bash
set -e
cd "$(dirname "$0")/../.."
# Parse --env argument to source the right env file
ENV_MODE="local"
ARGS=()
while [[ $# -gt 0 ]]; do
case $1 in
--env)
ENV_MODE="$2"
if [[ "$ENV_MODE" != "local" && "$ENV_MODE" != "dev" ]]; then
echo "Error: --env must be 'local' or 'dev'"
exit 1
fi
shift 2
;;
*)
ARGS+=("$1")
shift
;;
esac
done
# Source environment file
ENV_FILE=".env.${ENV_MODE}"
if [ ! -f "$ENV_FILE" ]; then
echo "Error: Environment file $ENV_FILE not found"
exit 1
fi
echo "🚀 Starting LoComo Benchmark with '${ENV_MODE}' environment..."
echo "📄 Loading environment from $ENV_FILE"
echo ""
# Export all variables from env file
set -a
source "$ENV_FILE"
set +a
uv run python benchmarks/locomo/run_benchmark.py "${ARGS[@]}"

View file

@ -0,0 +1,43 @@
#!/bin/bash
set -e
cd "$(dirname "$0")/../.."
# Parse --env argument to source the right env file
ENV_MODE="local"
ARGS=()
while [[ $# -gt 0 ]]; do
case $1 in
--env)
ENV_MODE="$2"
if [[ "$ENV_MODE" != "local" && "$ENV_MODE" != "dev" ]]; then
echo "Error: --env must be 'local' or 'dev'"
exit 1
fi
shift 2
;;
*)
ARGS+=("$1")
shift
;;
esac
done
# Source environment file
ENV_FILE=".env.${ENV_MODE}"
if [ ! -f "$ENV_FILE" ]; then
echo "Error: Environment file $ENV_FILE not found"
exit 1
fi
echo "🚀 Starting LongMemEval Benchmark with '${ENV_MODE}' environment..."
echo "📄 Loading environment from $ENV_FILE"
echo ""
# Export all variables from env file
set -a
source "$ENV_FILE"
set +a
uv run python benchmarks/longmemeval/run_benchmark.py "${ARGS[@]}"

View file

@ -0,0 +1,13 @@
#!/bin/bash
set -e
cd "$(dirname "$0")/../.."
echo "🎨 Starting Benchmark Visualizer..."
echo ""
echo "Server will be available at: http://localhost:8001"
echo ""
cd benchmarks/visualizer
open http://localhost:8001
uv run uvicorn server:app --reload --host 0.0.0.0 --port 8001

15
scripts/erase-local-db.sh Executable file
View file

@ -0,0 +1,15 @@
#!/bin/bash
set -e
cd "$(dirname "$0")/.."
echo "🛑 Stopping and erasing local PostgreSQL..."
echo ""
# Stop and remove containers, networks, volumes
cd docker
docker-compose down -v
echo ""
echo "✅ Local database has been stopped and all data erased!"
echo ""

View file

@ -1,166 +0,0 @@
"""
Profile slow database queries to identify optimization opportunities.
Usage:
uv run python scripts/profile_queries.py
"""
import asyncio
import os
import asyncpg
from dotenv import load_dotenv
load_dotenv()
async def profile_entry_points_query():
"""Profile the vector similarity entry points query."""
db_url = os.getenv("DATABASE_URL")
conn = await asyncpg.connect(db_url)
# Generate a dummy embedding vector (384 dimensions for bge-small-en-v1.5)
dummy_embedding = str([0.1] * 384)
print("=" * 80)
print("PROFILING: Entry Points Query (Vector Similarity)")
print("=" * 80)
# Run EXPLAIN ANALYZE
explain = await conn.fetch("""
EXPLAIN (ANALYZE, BUFFERS, VERBOSE)
SELECT id, text, context, event_date, access_count, embedding,
1 - (embedding <=> $1::vector) AS similarity
FROM memory_units
WHERE agent_id = $2
AND embedding IS NOT NULL
AND (1 - (embedding <=> $1::vector)) >= 0.5
ORDER BY embedding <=> $1::vector
LIMIT 3
""", dummy_embedding, "test_agent")
for row in explain:
print(row[0])
await conn.close()
async def profile_neighbors_query(sample_node_ids):
"""Profile the neighbors JOIN query."""
db_url = os.getenv("DATABASE_URL")
conn = await asyncpg.connect(db_url)
print("\n" + "=" * 80)
print("PROFILING: Neighbors Query (Graph Traversal)")
print(f"Sample size: {len(sample_node_ids)} nodes")
print("=" * 80)
# Run EXPLAIN ANALYZE
explain = await conn.fetch("""
EXPLAIN (ANALYZE, BUFFERS, VERBOSE)
SELECT ml.from_unit_id, ml.to_unit_id, ml.weight, ml.link_type, ml.entity_id,
mu.text, mu.context, mu.event_date, mu.access_count,
mu.id as neighbor_id
FROM memory_links ml
JOIN memory_units mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.weight >= 0.1
ORDER BY ml.from_unit_id, ml.weight DESC
""", sample_node_ids)
for row in explain:
print(row[0])
await conn.close()
async def profile_embeddings_query(sample_node_ids):
"""Profile the batch embeddings fetch query."""
db_url = os.getenv("DATABASE_URL")
conn = await asyncpg.connect(db_url)
print("\n" + "=" * 80)
print("PROFILING: Embeddings Query (Batch Fetch)")
print(f"Sample size: {len(sample_node_ids)} nodes")
print("=" * 80)
# Run EXPLAIN ANALYZE
explain = await conn.fetch("""
EXPLAIN (ANALYZE, BUFFERS, VERBOSE)
SELECT id, embedding
FROM memory_units
WHERE id = ANY($1::uuid[])
""", sample_node_ids)
for row in explain:
print(row[0])
await conn.close()
async def get_sample_node_ids(batch_size=50):
"""Get sample node IDs for profiling."""
db_url = os.getenv("DATABASE_URL")
conn = await asyncpg.connect(db_url)
rows = await conn.fetch(f"""
SELECT id FROM memory_units
LIMIT {batch_size}
""")
await conn.close()
return [row['id'] for row in rows]
async def check_indexes():
"""Check what indexes exist."""
db_url = os.getenv("DATABASE_URL")
conn = await asyncpg.connect(db_url)
print("\n" + "=" * 80)
print("CURRENT INDEXES")
print("=" * 80)
indexes = await conn.fetch("""
SELECT
tablename,
indexname,
indexdef
FROM pg_indexes
WHERE schemaname = 'public'
AND tablename IN ('memory_units', 'memory_links')
ORDER BY tablename, indexname
""")
for idx in indexes:
print(f"\nTable: {idx['tablename']}")
print(f"Index: {idx['indexname']}")
print(f"Definition: {idx['indexdef']}")
await conn.close()
async def main():
print("Starting Query Profiling...")
# Check indexes first
await check_indexes()
# Get sample node IDs
sample_ids = await get_sample_node_ids(50)
if sample_ids:
print(f"\nGot {len(sample_ids)} sample node IDs for profiling")
# Profile each query type
await profile_entry_points_query()
await profile_neighbors_query(sample_ids)
await profile_embeddings_query(sample_ids)
else:
print("\nNo data in database - run ingestion first")
print("\n" + "=" * 80)
print("PROFILING COMPLETE")
print("=" * 80)
if __name__ == "__main__":
asyncio.run(main())

42
scripts/start-local-db.sh Executable file
View file

@ -0,0 +1,42 @@
#!/bin/bash
set -e
cd "$(dirname "$0")/.."
echo "🚀 Starting local PostgreSQL..."
echo ""
# Start docker compose
cd docker
docker-compose up -d
echo ""
echo "⏳ Waiting for PostgreSQL to be ready..."
until docker exec memora-postgres pg_isready -U memora > /dev/null 2>&1; do
sleep 1
done
echo ""
echo "✅ PostgreSQL is ready!"
echo ""
# Initialize database schema
cd ..
export DATABASE_URL="postgresql://memora:memora_dev@localhost:5432/memora"
echo "📊 Running database migrations..."
uv run alembic upgrade head
echo ""
echo "✅ Database initialized successfully!"
echo ""
echo "📊 Connection Info:"
echo " Host: localhost"
echo " Port: 5432"
echo " Database: memora"
echo " User: memora"
echo " Password: memora_dev"
echo ""
echo "🛑 To stop and clean up:"
echo " ./scripts/erase-local-db.sh"
echo ""

50
scripts/start-server.sh Executable file
View file

@ -0,0 +1,50 @@
#!/bin/bash
set -e
cd "$(dirname "$0")/.."
# Parse arguments
ENV_MODE="local"
while [[ $# -gt 0 ]]; do
case $1 in
--env)
ENV_MODE="$2"
if [[ "$ENV_MODE" != "local" && "$ENV_MODE" != "dev" ]]; then
echo "Error: --env must be 'local' or 'dev'"
exit 1
fi
shift 2
;;
*)
echo "Usage: $0 [--env local|dev]"
echo ""
echo "Options:"
echo " --env local Use local environment (default)"
echo " --env dev Use dev environment"
exit 1
;;
esac
done
# Source environment file
ENV_FILE=".env.${ENV_MODE}"
if [ ! -f "$ENV_FILE" ]; then
echo "Error: Environment file $ENV_FILE not found"
exit 1
fi
echo "🚀 Starting Memory API Server with '${ENV_MODE}' environment..."
echo "📄 Loading environment from $ENV_FILE"
echo ""
# Export all variables from env file
set -a
source "$ENV_FILE"
set +a
echo "Server will be available at: http://localhost:8080"
echo ""
open http://localhost:8080
uv run uvicorn memora.web.server:app --reload --host 0.0.0.0 --port 8080

View file

@ -1,3 +0,0 @@
#!/bin/bash
# Start the FastAPI server with hot reload
uv run uvicorn memora.web.server:app --reload --host 0.0.0.0 --port 8080

View file

@ -5,26 +5,61 @@ import pytest
import pytest_asyncio
import os
import asyncio
from pathlib import Path
from dotenv import load_dotenv
from memora import TemporalSemanticMemory
from memora.llm_wrapper import LLMConfig
import asyncpg
load_dotenv()
# Force local database URL for all tests
LOCAL_DB_URL = "postgresql://memora:memora_dev@localhost:5432/memora"
@pytest_asyncio.fixture(scope="session")
# Load environment variables from .env.local at the start of test session
def pytest_configure(config):
"""Load environment variables before running tests."""
env_file = Path(__file__).parent.parent / ".env.local"
if env_file.exists():
load_dotenv(env_file)
else:
print(f"Warning: {env_file} not found, tests may fail without proper configuration")
# Override DATABASE_URL to use local database
os.environ["DATABASE_URL"] = LOCAL_DB_URL
@pytest.fixture(scope="session")
def llm_config():
"""
Provide LLM configuration for tests.
This can be used by tests that need to call LLM directly without memory system.
"""
return LLMConfig.for_memory()
@pytest_asyncio.fixture(scope="function")
async def memory():
"""
Provide a shared memory system instance for all tests in the session.
This avoids reloading the embedding model for every test (saves 3+ seconds per test).
Provide a memory system instance for each test function.
Forces the use of local database URL.
Note: Using function scope to avoid event loop issues, but this means
the embedding model will be loaded for each test (adds ~3 seconds per test).
Tests should handle their own cleanup by calling memory.delete_agent(agent_id)
in their finally blocks. The fixture will attempt cleanup as a safeguard.
"""
mem = TemporalSemanticMemory()
mem = TemporalSemanticMemory(db_url=LOCAL_DB_URL)
await mem.initialize()
yield mem
# Ensure cleanup happens at end of session
# Attempt cleanup (tests should already have called close, but this is a safeguard)
try:
await mem.close()
if mem._pool and not mem._pool._closing:
await mem.close()
except Exception as e:
print(f"Warning: Error during memory cleanup: {e}")
# Ignore errors during fixture cleanup since test may have already closed
pass
@pytest_asyncio.fixture(scope="function")
@ -51,8 +86,9 @@ async def clean_agent(memory):
async def db_connection():
"""
Provide a database connection for direct DB queries in tests.
Uses the forced local database URL.
"""
conn = await asyncpg.connect(os.getenv('DATABASE_URL'), statement_cache_size=0)
conn = await asyncpg.connect(LOCAL_DB_URL, statement_cache_size=0)
yield conn
try:
await conn.close()

View file

@ -2,7 +2,7 @@
Test chunking functionality for large documents.
"""
import pytest
from memora.llm_client import chunk_text
from memora.fact_extraction import chunk_text
def test_chunk_text_small():

View file

@ -2,24 +2,16 @@
Tests for document tracking and upsert functionality.
"""
import logging
import os
import pytest
from datetime import datetime, timezone
from memora import TemporalSemanticMemory
@pytest.mark.asyncio
async def test_document_creation_and_retrieval():
async def test_document_creation_and_retrieval(memory):
"""Test that documents are created and can be retrieved."""
db_url = os.getenv("DATABASE_URL")
if not db_url:
pytest.skip("DATABASE_URL not set")
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
agent_id = f"test_doc_{datetime.now(timezone.utc).timestamp()}"
try:
agent_id = f"test_doc_{datetime.now(timezone.utc).timestamp()}"
document_id = "meeting-001"
# Store memory with document tracking
@ -42,21 +34,15 @@ async def test_document_creation_and_retrieval():
assert doc["unit_count"] > 0
finally:
await memory.close()
await memory.delete_agent(agent_id)
@pytest.mark.asyncio
async def test_document_upsert():
async def test_document_upsert(memory):
"""Test that upsert deletes old units and creates new ones."""
db_url = os.getenv("DATABASE_URL")
if not db_url:
pytest.skip("DATABASE_URL not set")
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
agent_id = f"test_upsert_{datetime.now(timezone.utc).timestamp()}"
try:
agent_id = f"test_upsert_{datetime.now(timezone.utc).timestamp()}"
document_id = "meeting-002"
# First version
@ -92,21 +78,15 @@ async def test_document_upsert():
assert set(units_v1).isdisjoint(set(units_v2))
finally:
await memory.close()
await memory.delete_agent(agent_id)
@pytest.mark.asyncio
async def test_document_deletion():
async def test_document_deletion(memory):
"""Test that deleting a document cascades to memory units."""
db_url = os.getenv("DATABASE_URL")
if not db_url:
pytest.skip("DATABASE_URL not set")
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
agent_id = f"test_delete_{datetime.now(timezone.utc).timestamp()}"
try:
agent_id = f"test_delete_{datetime.now(timezone.utc).timestamp()}"
document_id = "meeting-003"
# Create document
@ -132,22 +112,15 @@ async def test_document_deletion():
assert doc_after is None
finally:
await memory.close()
await memory.delete_agent(agent_id)
@pytest.mark.asyncio
async def test_memory_without_document():
async def test_memory_without_document(memory):
"""Test that memories can still be created without document tracking."""
db_url = os.getenv("DATABASE_URL")
if not db_url:
pytest.skip("DATABASE_URL not set")
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
agent_id = f"test_no_doc_{datetime.now(timezone.utc).timestamp()}"
try:
agent_id = f"test_no_doc_{datetime.now(timezone.utc).timestamp()}"
# Create memory without document_id (backward compatibility)
units = await memory.put_async(
agent_id=agent_id,
@ -158,4 +131,4 @@ async def test_memory_without_document():
assert len(units) > 0
finally:
await memory.close()
await memory.delete_agent(agent_id)

View file

@ -5,12 +5,10 @@ This test loads a long conversation (419 dialogues across 19 sessions),
ingests it into memory, and runs searches to measure performance.
"""
import logging
import os
import json
import pytest
from datetime import datetime, timezone
from pathlib import Path
from memora import TemporalSemanticMemory
# Configure logging to show performance metrics
@ -22,16 +20,12 @@ logging.basicConfig(
@pytest.mark.asyncio
@pytest.mark.timeout(300) # 5 minute timeout for performance test
async def test_batch_ingestion_single_call():
async def test_batch_ingestion_single_call(memory):
"""
Test ingesting entire conversation in ONE batch call.
This is the most efficient way - all sessions in one put_batch_async.
"""
db_url = os.getenv("DATABASE_URL")
if not db_url:
pytest.skip("DATABASE_URL not set")
# Load conversation fixture
fixture_path = Path(__file__).parent / "fixtures" / "locomo_conversation_sample.json"
with open(fixture_path) as f:
@ -42,10 +36,6 @@ async def test_batch_ingestion_single_call():
logging.info(f"BATCH INGESTION TEST: {sample_id}")
logging.info(f"{'='*80}")
# Initialize memory
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
agent_id = f"batch_test_{sample_id}_{datetime.now(timezone.utc).timestamp()}"
try:
@ -121,11 +111,3 @@ async def test_batch_ingestion_single_call():
# Cleanup
logging.info("\nCleaning up...")
await memory.delete_agent(agent_id)
await memory.close()
if __name__ == "__main__":
# Allow running directly for quick perf checks
import asyncio
logging.info("Running performance test...")
asyncio.run(test_batch_ingestion_single_call())

View file

@ -2,26 +2,17 @@
Test search tracing functionality.
"""
import pytest
import asyncio
import os
from memora.temporal_semantic_memory import TemporalSemanticMemory
from memora.search_trace import SearchTrace
from datetime import datetime, timezone
@pytest.mark.asyncio
async def test_search_with_trace():
async def test_search_with_trace(memory):
"""Test that search with enable_trace=True returns a valid SearchTrace."""
# Use test database
db_url = os.getenv("DATABASE_URL")
if not db_url:
pytest.skip("DATABASE_URL not set")
memory = TemporalSemanticMemory(db_url=db_url)
# Generate a unique agent ID for this test
agent_id = f"test_trace_{datetime.now(timezone.utc).timestamp()}"
try:
# Generate a unique agent ID for this test
agent_id = f"test_trace_{datetime.now(timezone.utc).timestamp()}"
# Store some test memories
await memory.put_async(
@ -131,24 +122,17 @@ async def test_search_with_trace():
print(f" - Results returned: {trace.summary.results_returned}")
print(f" - Duration: {trace.summary.total_duration_seconds:.3f}s")
finally:
# Cleanup
await memory.delete_agent(agent_id)
finally:
await memory.close()
@pytest.mark.asyncio
async def test_search_without_trace():
async def test_search_without_trace(memory):
"""Test that search with enable_trace=False returns None for trace."""
db_url = os.getenv("DATABASE_URL")
if not db_url:
pytest.skip("DATABASE_URL not set")
memory = TemporalSemanticMemory(db_url=db_url)
agent_id = f"test_no_trace_{datetime.now(timezone.utc).timestamp()}"
try:
agent_id = f"test_no_trace_{datetime.now(timezone.utc).timestamp()}"
# Store a test memory
await memory.put_async(
@ -172,14 +156,6 @@ async def test_search_without_trace():
print("\n✓ Search without trace test passed!")
finally:
# Cleanup
await memory.delete_agent(agent_id)
finally:
await memory.close()
if __name__ == "__main__":
# Run tests directly
asyncio.run(test_search_with_trace())
asyncio.run(test_search_without_trace())

View file

@ -3,11 +3,11 @@ Test temporal extraction and per-fact dating.
"""
import pytest
from datetime import datetime, timezone, timedelta
from memora.llm_client import extract_facts_from_text
from memora.fact_extraction import extract_facts_from_text
@pytest.mark.asyncio
async def test_extract_facts_with_relative_dates():
async def test_extract_facts_with_relative_dates(memory):
"""Test that relative dates are converted to absolute dates."""
reference_date = datetime(2024, 3, 20, 14, 0, 0, tzinfo=timezone.utc)
@ -18,7 +18,7 @@ async def test_extract_facts_with_relative_dates():
This morning I had coffee with Alice.
"""
facts = await extract_facts_from_text(text, reference_date, "Personal diary")
facts = await extract_facts_from_text(text, reference_date, "Personal diary", llm_config=memory._llm_config)
print(f"\nExtracted {len(facts)} facts:")
for fact in facts:
@ -44,14 +44,14 @@ async def test_extract_facts_with_relative_dates():
@pytest.mark.asyncio
async def test_extract_facts_with_no_temporal_info():
async def test_extract_facts_with_no_temporal_info(memory):
"""Test that facts without temporal info use the reference date."""
reference_date = datetime(2024, 3, 20, 14, 0, 0, tzinfo=timezone.utc)
text = "Alice works at Google. She loves Python programming."
facts = await extract_facts_from_text(text, reference_date, "General info")
facts = await extract_facts_from_text(text, reference_date, "General info", llm_config=memory._llm_config)
print(f"\nExtracted {len(facts)} facts:")
for fact in facts:
@ -66,7 +66,7 @@ async def test_extract_facts_with_no_temporal_info():
@pytest.mark.asyncio
async def test_extract_facts_with_absolute_dates():
async def test_extract_facts_with_absolute_dates(memory):
"""Test that absolute dates in text are preserved."""
reference_date = datetime(2024, 3, 20, 14, 0, 0, tzinfo=timezone.utc)
@ -76,7 +76,7 @@ async def test_extract_facts_with_absolute_dates():
Bob will start his vacation on April 1st.
"""
facts = await extract_facts_from_text(text, reference_date, "Calendar events")
facts = await extract_facts_from_text(text, reference_date, "Calendar events", llm_config=memory._llm_config)
print(f"\nExtracted {len(facts)} facts:")
for fact in facts:

View file

@ -2,28 +2,20 @@
Test think function for opinion generation and consistency.
"""
import pytest
import os
from datetime import datetime, timezone
from memora import TemporalSemanticMemory
@pytest.mark.asyncio
async def test_think_opinion_consistency():
async def test_think_opinion_consistency(memory):
"""
Test that think function:
1. Generates an opinion
2. Stores the opinion in the database
3. Returns consistent response on subsequent calls with the same query
"""
db_url = os.getenv("DATABASE_URL")
if not db_url:
pytest.skip("DATABASE_URL not set")
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
agent_id = f"test_think_{datetime.now(timezone.utc).timestamp()}"
try:
agent_id = f"test_think_{datetime.now(timezone.utc).timestamp()}"
# Store some initial facts to give context for opinion formation
await memory.put_async(
@ -137,41 +129,29 @@ async def test_think_opinion_consistency():
await memory.delete_agent(agent_id)
except Exception as e:
print(f"Warning: Error during cleanup: {e}")
await memory.close()
@pytest.mark.asyncio
async def test_think_without_prior_context():
async def test_think_without_prior_context(memory):
"""
Test that think function handles queries when there's no relevant context.
"""
db_url = os.getenv("DATABASE_URL")
if not db_url:
pytest.skip("DATABASE_URL not set")
agent_id = f"test_think_no_context_{datetime.now(timezone.utc).timestamp()}"
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
# Call think without storing any prior facts
result = await memory.think_async(
agent_id=agent_id,
query="What is the capital of France?",
thinking_budget=20,
top_k=5
)
try:
agent_id = f"test_think_no_context_{datetime.now(timezone.utc).timestamp()}"
print(f"\n=== Think Without Context ===")
print(f"Answer: {result['text']}")
# Call think without storing any prior facts
result = await memory.think_async(
agent_id=agent_id,
query="What is the capital of France?",
thinking_budget=20,
top_k=5
)
print(f"\n=== Think Without Context ===")
print(f"Answer: {result['text']}")
# Should still return an answer (even if it says it doesn't have enough info)
assert result['text'], "Should return some answer"
assert 'based_on' in result, "Should return based_on structure"
finally:
await memory.close()
# Should still return an answer (even if it says it doesn't have enough info)
assert result['text'], "Should return some answer"
assert 'based_on' in result, "Should return based_on structure"
if __name__ == "__main__":

View file

@ -1,53 +0,0 @@
#!/usr/bin/env python3
"""Update existing memory_units to have fact_type='world' if NULL."""
import asyncio
import asyncpg
import os
from dotenv import load_dotenv
load_dotenv()
async def main():
conn = await asyncpg.connect(os.getenv('DATABASE_URL'))
# Check if fact_type column exists
result = await conn.fetchrow("""
SELECT column_name
FROM information_schema.columns
WHERE table_name = 'memory_units' AND column_name = 'fact_type'
""")
if not result:
print("fact_type column does not exist! Run migrations first:")
print(" uv run alembic upgrade head")
await conn.close()
return
print("fact_type column exists ✓")
# Update existing rows
count = await conn.fetchval("""
UPDATE memory_units
SET fact_type = 'world'
WHERE fact_type IS NULL
RETURNING (SELECT COUNT(*) FROM memory_units WHERE fact_type IS NULL)
""")
print(f"Updated {count} rows with fact_type='world'")
# Show distribution
distribution = await conn.fetch("""
SELECT fact_type, COUNT(*) as count
FROM memory_units
GROUP BY fact_type
ORDER BY count DESC
""")
print("\nFact type distribution:")
for row in distribution:
print(f" {row['fact_type']}: {row['count']}")
await conn.close()
if __name__ == "__main__":
asyncio.run(main())