Spaces:
Build error
Memory Management System
This package provides a sophisticated memory management system for RAG (Retrieval-Augmented Generation) agents, inspired by the CoALA paper on cognitive architectures for language agents.
Architecture
┌─────────────────────────────────────────────────────────────────────┐
│ Agent Orchestrator │
├─────────────────────────────────────────────────────────────────────┤
│ Memory Integration Layer │
│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │
│ │ Working │ │ Temporal │ │ Context │ │ GDPR │ │
│ │ Memory │ │ Context │ │ Window │ │ Manager │ │
│ └─────────────┘ └─────────────┘ └─────────────┘ └─────────────┘ │
├─────────────────────────────────────────────────────────────────────┤
│ Memory Orchestrator │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ Query Analysis → Concurrent Retrieval → Ranking → Context │ │
│ └─────────────────────────────────────────────────────────────┘ │
├─────────────────────────────────────────────────────────────────────┤
│ Memory Backend │
│ ┌─────────────────────┐ ┌─────────────────────┐ │
│ │ Local Backend │ OR │ Rust Memory Client │ │
│ │ (Development) │ │ (Production) │ │
│ └─────────────────────┘ └─────────────────────┘ │
└─────────────────────────────────────────────────────────────────────┘
Memory Types
| Type | Description | TTL | Use Case |
|---|---|---|---|
| Episodic | Specific events/interactions | 7 days | "What did we discuss yesterday?" |
| Semantic | General knowledge/facts | 1 year | "What is the user's preferred language?" |
| Procedural | Learned patterns/skills | 90 days | "User prefers concise answers" |
| Temporal | Time-aware context | 30 days | "What happened last week?" |
Quick Start
Basic Usage
package main
import (
"context"
"github.com/AmaniQuery/amaniquery/internal/memory"
)
func main() {
// Create configuration
config := memory.DefaultMemoryConfig()
// Create integration (embedder and LLM client required for full functionality)
integration, err := memory.NewAgentMemoryIntegration(config, embedder, llmClient)
if err != nil {
panic(err)
}
defer integration.Stop()
// Start background workers
integration.Start()
ctx := context.Background()
// Store a conversation turn
err = integration.StoreConversationTurn(ctx, "user-123", "session-456",
"What's the weather like?",
"I don't have access to real-time weather data.")
// Get context for a query
memCtx, err := integration.GetContextForQuery(ctx, "user-123", "session-456",
"Tell me more about the weather forecast",
4096) // max tokens
// Use the context in your LLM prompt
prompt := buildPrompt(memCtx.ContextWindow, userQuery)
}
Using the Memory Worker
// Create and start the memory worker
worker := memory.NewMemoryWorker(integration, 4) // 4 worker goroutines
worker.Start()
defer worker.Stop()
// Async store (fire and forget)
worker.AsyncStoreEntry(ctx, &memory.MemoryEntry{
Type: memory.SemanticMemory,
Content: "User mentioned they live in Kenya",
UserID: "user-123",
SessionID: "session-456",
})
// Sync store with result
result, err := worker.SubmitWorkWithResult(ctx, memory.MemoryWorkItem{
Type: memory.WorkStoreEntry,
Data: entry,
})
GDPR Compliance
// Export user data (Right to Data Portability)
export, err := integration.ExportUserData(ctx, "user-123", "admin@example.com")
jsonData, _ := json.MarshalIndent(export, "", " ")
// Delete user data (Right to be Forgotten)
err = integration.DeleteUserData(ctx, "user-123", "admin@example.com")
Components
types.go
Core type definitions including MemoryEntry, MemoryQuery, MemoryManager interface, and configuration structures.
working.go
Session-specific working memory with automatic pruning based on size limits.
temporal.go
Time-aware context tracking with exponential decay scoring for recency-weighted retrieval.
context_window.go
Smart context window management with multiple formatting styles (Markdown, XML, JSON) and dynamic sizing.
orchestrator.go
Core memory orchestrator handling concurrent retrieval, LLM-based query analysis, and memory consolidation.
consolidation.go
Background workers for automatic memory consolidation, TTL cleanup, and retention policy enforcement.
gdpr.go
GDPR-compliant data management including export (Article 20), deletion (Article 17), and audit logging.
local_backend.go
In-memory backend for development and testing with full MemoryManager interface implementation.
rust_client.go
High-performance client for the Rust memory service with binary protocol support and automatic fallback.
integration.go
High-level integration API connecting memory with the agent orchestrator.
worker.go
Background worker for async memory operations, compatible with common worker pool patterns.
Configuration
config := &memory.MemoryConfig{
// Rust service (optional, for production)
RustServiceEnabled: true,
RustServiceAddress: "localhost",
RustServicePort: 9091,
ConnectionPoolSize: 10,
RequestTimeout: 5 * time.Second,
// Working memory
MaxWorkingMemorySize: 1024 * 1024, // 1MB
// TTL defaults
DefaultTTL: 24 * time.Hour,
// Consolidation
Consolidation: memory.ConsolidationConfig{
TurnThreshold: 50, // Consolidate after 50 turns
TimeThreshold: 30 * time.Minute,
EpisodicRetention: 7 * 24 * time.Hour,
SemanticRetention: 365 * 24 * time.Hour,
ProceduralRetention: 90 * 24 * time.Hour,
},
}
Testing
# Run all tests
go test ./internal/memory/...
# Run with verbose output
go test -v ./internal/memory/...
# Run benchmarks
go test -bench=. ./internal/memory/...
Rust Memory Service
For production deployments, the Rust memory service provides sub-millisecond latency:
# Start the service stack
docker-compose -f deployments/docker-compose.memory.yml up -d
# The Go client will automatically connect to the Rust service
# If unavailable, it falls back to the local backend
See rust-memory-service/README.md for more details.
Metrics
metrics := integration.GetMetrics()
// Available metrics:
// - TotalQueries
// - TotalStores
// - AvgRetrievalMs
// - ConsolidationRuns
// - SessionsProcessed
// - TTLDeletions
// - ActiveWorkingMemories
Performance Characteristics
| Operation | Local Backend | Rust Service |
|---|---|---|
| Store | < 1ms | < 0.5ms |
| Retrieve | < 5ms | < 1ms |
| Context Build | < 10ms | < 5ms |
License
Part of the AmaniQuery project.