Memory

Memory System
    ├── Memory Types
    │   ├── Short-Term Memory
    │   │   ├── Conversation Context
    │   │   ├── Task Sequence State
    │   │   └── Session-Scoped
    │   │
    │   ├── Long-Term Memory
    │   │   ├── Persistent Storage
    │   │   ├── Cross-Session Knowledge
    │   │   └── Historical Data
    │   │
    │   └── Entity Memory
    │       ├── Named Entity Tracking
    │       ├── Relationship Graphs
    │       └── Entity Attributes
    │
    ├── Memory Storage
    │   ├── MemoryStorage (Interface)
    │   ├── In-Memory Store
    │   ├── Persistent Store
    │   └── Embedding Index
    │
    ├── Memory Configuration
    │   ├── Storage Limits
    │   ├── Persistence Settings
    │   ├── Embedding Configuration
    │   ├── Retrieval Thresholds
    │   └── Maintenance Rules
    │
    └── Memory Operations
        ├── Store
        ├── Retrieve (Semantic Search)
        ├── Update
        ├── Delete
        ├── Compress
        └── Prune

let conversationalAgent = Agent(
    role: "Customer Support Assistant",
    purpose: "Help customers with their inquiries",
    context: """
    Friendly support agent who remembers the conversation
    and can reference previous topics discussed.
    """,
    memory: true  // Enable short-term memory
)

// Usage scenario:
// User: "I need help with my order"
// Agent: "I'd be happy to help! Can you provide your order number?"
// User: "It's 12345"
// Agent: "Thank you! Let me look up order 12345 for you..."
//       [Agent remembers the order number from previous message]

let personalAssistant = Agent(
    role: "Personal Assistant",
    purpose: "Provide personalized help based on user history",
    context: """
    Remembers user preferences, past requests, and habits
    to provide increasingly personalized assistance.
    """,
    longTermMemory: true  // Enable persistent memory
)

// Usage scenario:
// Session 1:
// User: "I prefer emails in the morning at 8 AM"
// Agent: "Got it! I'll remember to send summaries at 8 AM"
//       [Stores preference in long-term memory]
//
// Session 2 (next day):
// Agent: "Good morning! Here's your 8 AM email summary"
//       [Retrieved preference from previous session]

let crmAgent = Agent(
    role: "Sales Assistant",
    purpose: "Manage customer relationships and track interactions",
    context: """
    Expert at remembering customer details, preferences,
    and interaction history to provide personalized service.
    """,
    entityMemory: true  // Enable entity tracking
)

// Usage scenario:
// User: "John Smith from Acme Corp called about the Enterprise plan"
// Agent: "I'll note that John Smith (Acme Corp) is interested in Enterprise"
//       [Creates entities: Person(John Smith), Organization(Acme Corp)]
//       [Links: John Smith WORKS_AT Acme Corp]
//       [Links: Acme Corp INTERESTED_IN Enterprise Plan]
//
// Later:
// User: "What did John want?"
// Agent: "John Smith from Acme Corp was interested in the Enterprise plan"
//       [Retrieved entity relationships]

let sophisticatedAgent = Agent(
    role: "Executive Assistant",
    purpose: "Comprehensive personal and professional assistance",
    context: """
    Advanced assistant with full memory capabilities:
    - Remembers conversations (short-term)
    - Learns preferences over time (long-term)
    - Tracks contacts and relationships (entity)
    """,
    memory: true,           // Short-term: conversation context
    longTermMemory: true,   // Long-term: persistent preferences
    entityMemory: true      // Entity: people and organizations
)

let basicConfig = MemoryConfiguration(
    maxMemoryItems: 100,
    persistencePath: "./memory",
    similarityThreshold: 0.75
)

let agent = Agent(
    role: "Assistant",
    purpose: "General assistance with memory",
    context: "Helpful assistant",
    memory: true,
    longTermMemory: true,
    memoryConfig: basicConfig
)

let highQualityConfig = MemoryConfiguration(
    maxMemoryItems: 200,
    persistencePath: "./memory/high-quality",
    embeddingModel: "text-embedding-3-large",  // Best quality
    similarityThreshold: 0.85,                 // Precise retrieval
    compressionEnabled: false,                 // Keep full detail
    pruneOldItems: false                       // Keep all items
)

let premiumAgent = Agent(
    role: "Premium Assistant",
    purpose: "High-quality personalized assistance",
    context: "Premium service with excellent memory",
    memory: true,
    longTermMemory: true,
    memoryConfig: highQualityConfig
)

let efficientConfig = MemoryConfiguration(
    maxMemoryItems: 50,                        // Smaller limit
    persistencePath: "./memory/compact",
    embeddingModel: "text-embedding-3-small",  // Efficient model
    similarityThreshold: 0.7,                  // Broader retrieval
    compressionEnabled: true,                  // Auto-compress
    autoSummarize: true,                       // Summarize old items
    pruneOldItems: true                        // Remove old entries
)

let efficientAgent = Agent(
    role: "Efficient Assistant",
    purpose: "Memory-efficient assistance",
    context: "Optimized for resource constraints",
    memory: true,
    longTermMemory: true,
    memoryConfig: efficientConfig
)

let largeScaleConfig = MemoryConfiguration(
    maxMemoryItems: 1000,                      // Large capacity
    persistencePath: "./memory/large-scale",
    embeddingModel: "text-embedding-3-large",
    similarityThreshold: 0.80,
    compressionEnabled: true,                  // Manage size
    autoSummarize: true,                       // Keep summaries
    pruneOldItems: true                        // Auto-prune
)

let knowledgeAgent = Agent(
    role: "Knowledge Assistant",
    purpose: "Manage large knowledge bases",
    context: "Expert with extensive memory",
    memory: true,
    longTermMemory: true,
    entityMemory: true,
    memoryConfig: largeScaleConfig
)

// Agent with all memory types
let fullMemoryAgent = Agent(
    role: "Personal Assistant",
    purpose: "Comprehensive personal assistance",
    context: "Assistant with full memory capabilities",
    memory: true,           // Short-term memory
    longTermMemory: true,   // Persistent memory
    entityMemory: true,     // Entity tracking
    memoryConfig: memoryConfig
)

// Agent with selective memory
let selectiveAgent = Agent(
    role: "Conversational Agent",
    purpose: "Handle conversations",
    context: "Conversational interface",
    memory: true,           // Only short-term
    longTermMemory: false,  // No persistence
    entityMemory: false     // No entity tracking
)

// Agent with no memory
let statelessAgent = Agent(
    role: "Simple Responder",
    purpose: "Answer single questions",
    context: "Stateless question answering",
    memory: false,          // No memory overhead
    longTermMemory: false,
    entityMemory: false
)

let sharedMemoryOrbit = try await Orbit.create(
    name: "Collaborative Workflow",
    agents: [agent1, agent2, agent3],
    tasks: tasks,
    memory: true,           // Shared short-term memory
    longTermMemory: true,   // Shared persistent memory
    entityMemory: true,     // Shared entity tracking
    memoryConfig: memoryConfig
)

// Explicitly no memory
let noMemoryAgent = Agent(
    role: "Stateless Worker",
    purpose: "Process independent tasks",
    context: "Simple task processor",
    memory: false,
    longTermMemory: false,
    entityMemory: false
)

// Default is no memory (can omit parameters)
let defaultAgent = Agent(
    role: "Default Agent",
    purpose: "Basic task execution",
    context: "Simple agent"
    // memory defaults to false
)

public struct TaskExecutionContext: Sendable {
    /// Previous task outputs
    public var taskOutputs: [TaskOutput]

    /// Orbit-level inputs
    public var inputs: Metadata

    /// Shared memory storage
    public var memory: MemoryStorage?

    /// Knowledge base access
    public var knowledgeBase: KnowledgeBase?

    /// Available tools
    public var availableTools: [String]
}

public protocol MemoryStorage: Sendable {
    /// Store a memory item
    func store(
        key: String,
        value: String,
        metadata: [String: String]?
    ) async throws

    /// Retrieve memories by semantic search
    func retrieve(
        query: String,
        limit: Int,
        threshold: Double?
    ) async throws -> [MemoryItem]

    /// Get specific memory by key
    func get(key: String) async throws -> MemoryItem?

    /// Update existing memory
    func update(
        key: String,
        value: String,
        metadata: [String: String]?
    ) async throws

    /// Delete memory item
    func delete(key: String) async throws

    /// Clear all memories
    func clear() async throws

    /// Get memory count
    var count: Int { get async }

    /// Persist to storage
    func persist() async throws

    /// Load from storage
    func load() async throws
}

public struct MemoryItem: Codable, Sendable {
    public let key: String
    public let value: String
    public let metadata: [String: String]
    public let timestamp: Date
    public let embedding: [Double]?
    public var accessCount: Int
    public var lastAccessed: Date
}

// Memory is automatically:
// 1. Stored when agents generate information
// 2. Retrieved when relevant to current task
// 3. Updated when information changes
// 4. Persisted when orbit completes
// 5. Pruned when limits are reached

let orbit = try await Orbit.create(
    name: "Automated Memory",
    agents: [agent],
    tasks: [task],
    memory: true,
    longTermMemory: true,
    memoryConfig: MemoryConfiguration(
        maxMemoryItems: 100,
        compressionEnabled: true,
        pruneOldItems: true  // Automatic management
    )
)

let output = try await orbit.run()
// Memory automatically persisted on completion

import OrbitAI

// Custom task with manual memory access
let customTask = ORTask(
    description: "Store and retrieve specific information",
    expectedOutput: "Memory operation results",
    customHandler: { context in
        guard let memory = context.memory else {
            return "Memory not available"
        }

        // Store information
        try await memory.store(
            key: "user_preference",
            value: "Morning emails at 8 AM",
            metadata: ["category": "preferences", "priority": "high"]
        )

        // Retrieve by semantic search
        let results = try await memory.retrieve(
            query: "email preferences",
            limit: 5,
            threshold: 0.75
        )

        // Get specific item
        if let preference = try await memory.get(key: "user_preference") {
            print("Found preference: \(preference.value)")
        }

        // Update existing memory
        try await memory.update(
            key: "user_preference",
            value: "Morning emails at 7 AM",  // Updated time
            metadata: ["category": "preferences", "priority": "high"]
        )

        return "Memory operations completed"
    }
)

Memory Lifecycle in Orbit Execution
    ├── 1. Initialization
    │   ├── Load existing memories (if long-term enabled)
    │   └── Initialize in-memory storage
    │
    ├── 2. During Execution
    │   ├── Agent generates output → Stored in memory
    │   ├── Agent needs context → Retrieved from memory
    │   ├── Information updates → Memory updated
    │   └── Memory limit reached → Automatic pruning
    │
    ├── 3. Retrieval Process
    │   ├── Query generated from task context
    │   ├── Query embedded using embedding model
    │   ├── Similarity search against memory embeddings
    │   ├── Results filtered by threshold
    │   └── Top-k results returned to agent
    │
    └── 4. Completion
        ├── Memory compression (if enabled)
        ├── Memory persistence (if long-term enabled)
        └── Memory cleanup (short-term cleared)

import OrbitAI

final class CustomMemoryStorage: MemoryStorage {
    private var storage: [String: MemoryItem] = [:]
    private let database: Database  // Custom database

    var count: Int {
        get async { storage.count }
    }

    func store(
        key: String,
        value: String,
        metadata: [String: String]?
    ) async throws {
        let item = MemoryItem(
            key: key,
            value: value,
            metadata: metadata ?? [:],
            timestamp: Date(),
            embedding: try await generateEmbedding(for: value),
            accessCount: 0,
            lastAccessed: Date()
        )

        storage[key] = item

        // Store in custom database
        try await database.insert(item)
    }

    func retrieve(
        query: String,
        limit: Int,
        threshold: Double?
    ) async throws -> [MemoryItem] {
        let queryEmbedding = try await generateEmbedding(for: query)
        let threshold = threshold ?? 0.7

        // Semantic search
        let scored = storage.values.compactMap { item -> (MemoryItem, Double)? in
            guard let embedding = item.embedding else { return nil }
            let similarity = cosineSimilarity(queryEmbedding, embedding)
            return similarity >= threshold ? (item, similarity) : nil
        }

        // Sort by similarity and take top-k
        return scored
            .sorted { $0.1 > $1.1 }
            .prefix(limit)
            .map { $0.0 }
    }

    func get(key: String) async throws -> MemoryItem? {
        storage[key]
    }

    func update(
        key: String,
        value: String,
        metadata: [String: String]?
    ) async throws {
        guard var item = storage[key] else {
            throw MemoryError.itemNotFound
        }

        item.value = value
        item.metadata = metadata ?? item.metadata
        item.lastAccessed = Date()

        storage[key] = item
        try await database.update(item)
    }

    func delete(key: String) async throws {
        storage.removeValue(forKey: key)
        try await database.delete(key: key)
    }

    func clear() async throws {
        storage.removeAll()
        try await database.deleteAll()
    }

    func persist() async throws {
        // Custom persistence logic
        try await database.saveAll(Array(storage.values))
    }

    func load() async throws {
        // Custom loading logic
        let items = try await database.loadAll()
        storage = Dictionary(uniqueKeysWithValues: items.map { ($0.key, $0) })
    }

    private func generateEmbedding(for text: String) async throws -> [Double] {
        // Call embedding API
        // Return embedding vector
    }

    private func cosineSimilarity(_ a: [Double], _ b: [Double]) -> Double {
        // Calculate cosine similarity
    }
}

import OrbitAI
import Pinecone  // Example vector database

final class VectorMemoryStorage: MemoryStorage {
    private let pinecone: PineconeClient
    private let indexName: String
    private let embeddingModel: EmbeddingModel

    init(
        apiKey: String,
        indexName: String,
        embeddingModel: EmbeddingModel
    ) {
        self.pinecone = PineconeClient(apiKey: apiKey)
        self.indexName = indexName
        self.embeddingModel = embeddingModel
    }

    func store(
        key: String,
        value: String,
        metadata: [String: String]?
    ) async throws {
        // Generate embedding
        let embedding = try await embeddingModel.embed(text: value)

        // Store in Pinecone
        try await pinecone.upsert(
            index: indexName,
            vectors: [
                Vector(
                    id: key,
                    values: embedding,
                    metadata: metadata ?? [:]
                )
            ]
        )
    }

    func retrieve(
        query: String,
        limit: Int,
        threshold: Double?
    ) async throws -> [MemoryItem] {
        // Generate query embedding
        let queryEmbedding = try await embeddingModel.embed(text: query)

        // Query Pinecone
        let results = try await pinecone.query(
            index: indexName,
            vector: queryEmbedding,
            topK: limit,
            includeMetadata: true
        )

        // Convert to MemoryItems
        return results.matches.compactMap { match in
            guard let threshold = threshold,
                  match.score >= threshold else { return nil }

            return MemoryItem(
                key: match.id,
                value: match.metadata["value"] as? String ?? "",
                metadata: match.metadata as? [String: String] ?? [:],
                timestamp: Date(),
                embedding: match.values,
                accessCount: 0,
                lastAccessed: Date()
            )
        }
    }

    // Implement other methods...
}

final class CachedMemoryStorage: MemoryStorage {
    private let underlying: MemoryStorage
    private var cache: [String: (item: MemoryItem, expiry: Date)] = [:]
    private let cacheExpiry: TimeInterval = 300  // 5 minutes

    init(underlying: MemoryStorage) {
        self.underlying = underlying
    }

    func get(key: String) async throws -> MemoryItem? {
        // Check cache first
        if let cached = cache[key],
           cached.expiry > Date() {
            return cached.item
        }

        // Fetch from underlying storage
        if let item = try await underlying.get(key: key) {
            // Update cache
            cache[key] = (item, Date().addingTimeInterval(cacheExpiry))
            return item
        }

        return nil
    }

    func store(
        key: String,
        value: String,
        metadata: [String: String]?
    ) async throws {
        // Store in underlying
        try await underlying.store(key: key, value: value, metadata: metadata)

        // Invalidate cache
        cache.removeValue(forKey: key)
    }

    func retrieve(
        query: String,
        limit: Int,
        threshold: Double?
    ) async throws -> [MemoryItem] {
        // For semantic search, bypass cache
        return try await underlying.retrieve(
            query: query,
            limit: limit,
            threshold: threshold
        )
    }

    // Implement other methods with cache invalidation...
}

final class CompressingMemoryStorage: MemoryStorage {
    private let underlying: MemoryStorage
    private let llm: LLMProvider
    private let compressionThreshold: Int = 50

    init(underlying: MemoryStorage, llm: LLMProvider) {
        self.underlying = underlying
        self.llm = llm
    }

    func store(
        key: String,
        value: String,
        metadata: [String: String]?
    ) async throws {
        var finalValue = value

        // Compress long values
        if value.count > compressionThreshold {
            let prompt = """
            Summarize the following information concisely while retaining key details:

            \(value)
            """

            let summary = try await llm.generateResponse(
                for: LLMRequest(messages: [.user(prompt)])
            )

            finalValue = summary.content

            // Mark as compressed
            var newMetadata = metadata ?? [:]
            newMetadata["compressed"] = "true"
            newMetadata["original_length"] = "\(value.count)"

            try await underlying.store(
                key: key,
                value: finalValue,
                metadata: newMetadata
            )
        } else {
            try await underlying.store(
                key: key,
                value: finalValue,
                metadata: metadata
            )
        }
    }

    // Delegate other methods to underlying storage
}

// 1. Secure persistence paths
let config = MemoryConfiguration(
    persistencePath: "./memory/encrypted"  // Use encrypted storage
)

// 2. Implement data retention policies
let retentionConfig = MemoryConfiguration(
    maxMemoryItems: 100,
    pruneOldItems: true,  // Auto-delete old data
    // Custom: Delete after 30 days
)

// 3. Filter sensitive data
final class SecureMemoryStorage: MemoryStorage {
    func store(
        key: String,
        value: String,
        metadata: [String: String]?
    ) async throws {
        // Filter sensitive information
        let filtered = filterSensitiveData(value)

        // Encrypt before storage
        let encrypted = try encrypt(filtered)

        try await underlying.store(
            key: key,
            value: encrypted,
            metadata: metadata
        )
    }

    private func filterSensitiveData(_ text: String) -> String {
        var filtered = text

        // Remove credit card numbers
        filtered = filtered.replacingOccurrences(
            of: #"\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}"#,
            with: "[REDACTED]",
            options: .regularExpression
        )

        // Remove email addresses
        filtered = filtered.replacingOccurrences(
            of: #"[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}"#,
            with: "[EMAIL]",
            options: [.regularExpression, .caseInsensitive]
        )

        // Remove SSN
        filtered = filtered.replacingOccurrences(
            of: #"\d{3}-\d{2}-\d{4}"#,
            with: "[SSN]",
            options: .regularExpression
        )

        return filtered
    }
}

// 4. Implement access controls
final class AccessControlledMemory: MemoryStorage {
    func retrieve(
        query: String,
        limit: Int,
        threshold: Double?
    ) async throws -> [MemoryItem] {
        // Check permissions
        guard await hasReadPermission() else {
            throw MemoryError.accessDenied
        }

        return try await underlying.retrieve(
            query: query,
            limit: limit,
            threshold: threshold
        )
    }
}

final class MemoryDebugger {
    let memory: MemoryStorage

    init(memory: MemoryStorage) {
        self.memory = memory
    }

    func printAllMemories() async throws {
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
        print("Memory Contents")
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")

        let count = await memory.count
        print("Total items: \(count)")

        // Retrieve all memories (use broad threshold)
        let allMemories = try await memory.retrieve(
            query: "",
            limit: count,
            threshold: 0.0
        )

        for (index, item) in allMemories.enumerated() {
            print("\n[\(index + 1)] \(item.key)")
            print("   Value: \(item.value)")
            print("   Timestamp: \(item.timestamp)")
            print("   Access count: \(item.accessCount)")
            print("   Last accessed: \(item.lastAccessed)")
            print("   Metadata: \(item.metadata)")
            if let embedding = item.embedding {
                print("   Embedding: \(embedding.count) dimensions")
            }
        }

        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    }

    func testRetrieval(query: String) async throws {
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
        print("Testing Retrieval: \"\(query)\"")
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")

        for threshold in [0.5, 0.6, 0.7, 0.8, 0.9] {
            let results = try await memory.retrieve(
                query: query,
                limit: 10,
                threshold: threshold
            )

            print("\nThreshold \(threshold): \(results.count) results")
            for (index, item) in results.prefix(3).enumerated() {
                print("  \(index + 1). \(item.key): \(item.value.prefix(50))...")
            }
        }

        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    }

    func statistics() async -> MemoryStatistics {
        // Calculate and return statistics
    }
}

// Usage
let debugger = MemoryDebugger(memory: orbit.memory!)
try await debugger.printAllMemories()
try await debugger.testRetrieval(query: "user preferences")