# Memory Management System

This package provides a sophisticated memory management system for RAG (Retrieval-Augmented Generation) agents, inspired by the [CoALA paper](https://arxiv.org/abs/2309.02427) 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

```go
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

```go
// 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

```go
// 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

```go
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

```bash
# 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:

```bash
# 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

```go
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.