Deployment
Automated deployment update
4b1daed
|
Raw
History Blame Contribute Delete
8.78 kB

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.