Knowledge Base

Knowledge Base RAG Pipeline
    ├── 1. Document Ingestion
    │   ├── Load documents from file paths
    │   ├── Parse content (PDF, MD, JSON, TXT)
    │   ├── Split into chunks (with overlap)
    │   └── Extract metadata
    │
    ├── 2. Embedding & Indexing
    │   ├── Generate embeddings for each chunk
    │   ├── Create vector index
    │   ├── Store in vector database
    │   └── Build metadata index
    │
    ├── 3. Query Processing
    │   ├── Agent generates query from task
    │   ├── Embed query using same model
    │   ├── Vector similarity search
    │   └── Rank results by relevance
    │
    ├── 4. Context Retrieval
    │   ├── Retrieve top-k relevant chunks
    │   ├── Apply similarity threshold
    │   ├── Rerank if needed
    │   └── Format as context
    │
    └── 5. Response Generation
        ├── Inject retrieved context into prompt
        ├── LLM generates response
        ├── Ground response in source documents
        └── Return with citations (optional)

OrbitAI Execution Flow with Knowledge Base
    ├── User Request
    │   └── "What are the agent capabilities?"
    │
    ├── Orbit Orchestrator
    │   └── Routes to appropriate agent
    │
    ├── Agent Execution
    │   ├── Analyzes task requirements
    │   ├── Formulates knowledge query
    │   └── Triggers knowledge retrieval
    │
    ├── Knowledge Base Query
    │   ├── Embeds query
    │   ├── Searches vector index
    │   ├── Retrieves relevant chunks
    │   └── Returns context
    │
    ├── LLM Processing
    │   ├── Context: Retrieved knowledge
    │   ├── Task: User request
    │   ├── Agent: Role and purpose
    │   └── Generates response
    │
    └── Response
        └── "Agents have tools, memory, and knowledge..."

import OrbitAI

let agent = Agent(
    role: "Documentation Assistant",
    purpose: "Answer questions about product documentation",
    context: "Expert assistant with access to all product docs",
    knowledgeSources: [
        "./docs/user-guide.pdf",
        "./docs/api-reference.md",
        "./docs/faq.txt"
    ]
)

let orbit = try await Orbit.create(
    name: "Documentation System",
    agents: [searchAgent, summaryAgent, answerAgent],
    tasks: [searchTask, summaryTask, answerTask],
    process: .sequential,
    knowledgeSources: [
        // Shared knowledge for all agents
        "./docs/product-manual.pdf",
        "./docs/api-reference.md",
        "./data/product-catalog.json"
    ]
)

// Create orbit
let orbit = try await Orbit.create(
    name: "Adaptive System",
    agents: [agent],
    tasks: [task],
    knowledgeSources: [
        "./initial-knowledge.pdf"
    ]
)

// Add new knowledge source at runtime
await orbit.addKnowledgeSource("./new-document.pdf")
await orbit.addKnowledgeSource("./updated-policy.md")

// Knowledge base automatically updates
// Agents can now access new documents
let result = try await orbit.run()

// Poor organization
knowledgeSources: [
    "./doc1.pdf",
    "./doc2.pdf",
    "./doc3.pdf",
    "./doc4.pdf",
    "./doc5.pdf"  // What are these?
]

// Good organization
knowledgeSources: [
    // Product documentation
    "./knowledge/products/user-guide.pdf",
    "./knowledge/products/api-reference.md",

    // Company policies
    "./knowledge/policies/employee-handbook.pdf",
    "./knowledge/policies/security-policy.md",

    // Support resources
    "./knowledge/support/faq.txt",
    "./knowledge/support/troubleshooting.md",

    // Data resources
    "./knowledge/data/product-catalog.json",
    "./knowledge/data/pricing.json"
]

let hybridAgent = Agent(
    role: "Adaptive Assistant",
    purpose: "Provide personalized help using both knowledge and memory",
    context: """
    Intelligent assistant that:
    - Uses knowledge base for factual information
    - Uses memory for user preferences and history
    - Combines both for personalized, accurate responses
    """,
    // Memory for user interactions
    memory: true,
    longTermMemory: true,
    // Knowledge base for factual information
    knowledgeSources: [
        "./knowledge/product-docs.pdf",
        "./knowledge/company-info.md"
    ]
)

User Query: "What are the features of Product X?"
    ↓
Agent Processing:
    ├── Memory Check: "User previously asked about Product X pricing"
    ├── Knowledge Query: "Product X features from documentation"
    └── Synthesis: Personalized response combining both
    ↓
Response: "Product X has features A, B, C (from knowledge base).
           Based on your previous interest in pricing (from memory),
           you might also want to know that it's available at..."

let actionableAgent = Agent(
    role: "Executive Assistant",
    purpose: "Answer questions and take actions",
    context: "Assistant with knowledge and capabilities to act",

    // Knowledge for information
    knowledgeSources: [
        "./calendar-policies.md",
        "./company-contacts.json"
    ],

    // Tools for actions
    tools: [
        "apple_calendar",  // Create calendar events
        "send_email",      // Send emails
        "web_search"       // Search for info not in KB
    ]
)

let task = ORTask(
    description: """
    Check the company contacts in the knowledge base for John's email,
    then send him an email about the meeting using the send_email tool,
    and create a calendar event using the calendar tool based on
    the meeting policies in the knowledge base.
    """,
    expectedOutput: "Confirmation of email sent and event created"
)

let customTask = ORTask(
    description: "Custom knowledge retrieval task",
    expectedOutput: "Processed knowledge results",
    customHandler: { context in
        guard let knowledgeBase = context.knowledgeBase else {
            return "Knowledge base not available"
        }

        // Query knowledge base
        let results = try await knowledgeBase.query(
            query: "product specifications",
            limit: 5,
            threshold: 0.75
        )

        // Process results
        var output = "Found \(results.count) relevant documents:\n\n"

        for (index, result) in results.enumerated() {
            output += "\(index + 1). \(result.metadata.filename)\n"
            output += "   Relevance: \(result.score)\n"
            output += "   Content: \(result.content.prefix(100))...\n\n"
        }

        return output
    }
)

final class KnowledgeBaseDebugger {
    let knowledgeBase: KnowledgeBase

    init(knowledgeBase: KnowledgeBase) {
        self.knowledgeBase = knowledgeBase
    }

    func printStatistics() async throws {
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
        print("Knowledge Base Statistics")
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")

        let stats = try await knowledgeBase.getStatistics()

        print("Documents: \(stats.documentCount)")
        print("Chunks: \(stats.chunkCount)")
        print("Total size: \(stats.totalSizeBytes / 1024 / 1024) MB")
        print("Avg chunks per doc: \(stats.chunkCount / max(stats.documentCount, 1))")
        print("Index size: \(stats.indexSizeBytes / 1024 / 1024) MB")
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    }

    func listDocuments() async throws {
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
        print("Knowledge Base Documents")
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")

        let docs = try await knowledgeBase.listDocuments()

        for (index, doc) in docs.enumerated() {
            print("\n[\(index + 1)] \(doc.filename)")
            print("   Path: \(doc.path)")
            print("   Format: \(doc.format)")
            print("   Size: \(doc.sizeBytes / 1024) KB")
            print("   Chunks: \(doc.chunkCount)")
            print("   Indexed: \(doc.indexedDate)")
        }

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

    func testQuery(_ query: String) async throws {
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
        print("Testing Query: \"\(query)\"")
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")

        // Test various thresholds
        for threshold in [0.6, 0.7, 0.75, 0.8, 0.9] {
            let results = try await knowledgeBase.query(
                query: query,
                limit: 10,
                threshold: threshold
            )

            print("\nThreshold \(threshold): \(results.count) results")

            for (i, result) in results.prefix(3).enumerated() {
                print("  \(i+1). Score: \(String(format: "%.3f", result.score))")
                print("     Source: \(result.metadata.filename)")
                print("     Content: \(result.content.prefix(80))...")
            }
        }

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

    func validateSources(_ sources: [String]) {
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
        print("Validating Knowledge Sources")
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")

        var valid = 0
        var invalid = 0

        for source in sources {
            let exists = FileManager.default.fileExists(atPath: source)
            let readable = FileManager.default.isReadableFile(atPath: source)

            if exists && readable {
                print("✓ \(source)")
                valid += 1
            } else {
                print("✗ \(source)")
                if !exists {
                    print("  Error: File not found")
                } else if !readable {
                    print("  Error: Not readable")
                }
                invalid += 1
            }
        }

        print("\nSummary: \(valid) valid, \(invalid) invalid")
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    }
}

// Usage
let debugger = KnowledgeBaseDebugger(knowledgeBase: orbit.knowledgeBase!)
try await debugger.printStatistics()
try await debugger.listDocuments()
try await debugger.testQuery("product features")
debugger.validateSources(knowledgeSources)

let knowledgeConfig = KnowledgeConfiguration(
    // Chunking parameters
    chunkSize: 512,              // Characters per chunk
    chunkOverlap: 50,            // Overlap between chunks

    // Retrieval parameters
    retrievalLimit: 10,          // Max chunks to retrieve
    similarityThreshold: 0.75,   // Minimum similarity score

    // Embedding configuration
    embeddingModel: "text-embedding-ada-002",
    embeddingDimensions: 1536,

    // Processing options
    preprocessText: true,        // Clean text before embedding
    extractMetadata: true,       // Parse document metadata

    // Caching
    enableCache: true,
    cacheExpiry: 3600,          // Cache TTL in seconds

    // Vector storage
    vectorStore: .inMemory,     // or .persistent, .pinecone, etc.
    persistencePath: "./kb-vectors"
)

let agent = Agent(
    role: "Advanced Knowledge Agent",
    knowledgeSources: ["./docs/"],
    knowledgeConfig: knowledgeConfig
)

// Configure metadata extraction
let config = KnowledgeConfiguration(
    extractMetadata: true,
    metadataFields: [
        .title,
        .author,
        .createdDate,
        .modifiedDate,
        .category,
        .tags
    ]
)

// Query with metadata filters
let results = try await knowledgeBase.query(
    query: "product specifications",
    filters: [
        .category("products"),
        .dateRange(from: startDate, to: endDate),
        .tags(["version-2.0", "approved"])
    ]
)

// Example: Multi-tenant knowledge base
let customerResults = try await knowledgeBase.query(
    query: "pricing information",
    filters: [
        .metadata("customer_id", equals: customerId),
        .metadata("access_level", greaterThan: 2)
    ]
)

let config = KnowledgeConfiguration(
    retrievalLimit: 20,          // Initial broad retrieval
    rerankingEnabled: true,
    rerankingStrategy: .crossEncoder,
    rerankingLimit: 5,           // Final results after reranking
    crossEncoderModel: "cross-encoder/ms-marco-MiniLM-L-12-v2"
)

import Foundation

final class KnowledgeBaseBenchmark {
    let knowledgeBase: KnowledgeBase

    func runBenchmark() async throws -> BenchmarkResults {
        var results = BenchmarkResults()

        // Test queries
        let testQueries = [
            "product features",
            "pricing information",
            "technical specifications",
            "return policy",
            "customer support"
        ]

        // Warm up
        for query in testQueries {
            _ = try await knowledgeBase.query(query: query, limit: 5)
        }

        // Benchmark retrieval speed
        var retrievalTimes: [TimeInterval] = []

        for query in testQueries {
            let start = Date()
            let queryResults = try await knowledgeBase.query(
                query: query,
                limit: 5,
                threshold: 0.75
            )
            let duration = Date().timeIntervalSince(start)

            retrievalTimes.append(duration)
            results.queryResults[query] = queryResults.count
        }

        results.avgRetrievalTime = retrievalTimes.reduce(0, +) / Double(retrievalTimes.count)
        results.minRetrievalTime = retrievalTimes.min() ?? 0
        results.maxRetrievalTime = retrievalTimes.max() ?? 0

        // Benchmark embedding generation
        let embeddingStart = Date()
        _ = try await knowledgeBase.generateEmbedding(for: "test query")
        results.embeddingTime = Date().timeIntervalSince(embeddingStart)

        // Memory usage
        results.memoryUsage = getMemoryUsage()

        // Cache hit rate
        results.cacheHitRate = knowledgeBase.getCacheHitRate()

        return results
    }

    func printResults(_ results: BenchmarkResults) {
        print("""
        ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        Knowledge Base Benchmark Results
        ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

        Retrieval Performance:
          Average: \(String(format: "%.3f", results.avgRetrievalTime))s
          Min: \(String(format: "%.3f", results.minRetrievalTime))s
          Max: \(String(format: "%.3f", results.maxRetrievalTime))s

        Embedding Generation: \(String(format: "%.3f", results.embeddingTime))s
        Memory Usage: \(results.memoryUsage) MB
        Cache Hit Rate: \(String(format: "%.1f", results.cacheHitRate * 100))%

        Query Results:
        """)

        for (query, count) in results.queryResults {
            print("  \(query): \(count) results")
        }

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

    private func getMemoryUsage() -> Int {
        var info = mach_task_basic_info()
        var count = mach_msg_type_number_t(MemoryLayout<mach_task_basic_info>.size)/4

        let kerr: kern_return_t = withUnsafeMutablePointer(to: &info) {
            $0.withMemoryRebound(to: integer_t.self, capacity: 1) {
                task_info(mach_task_self_, task_flavor_t(MACH_TASK_BASIC_INFO), $0, &count)
            }
        }

        guard kerr == KERN_SUCCESS else { return 0 }
        return Int(info.resident_size) / 1024 / 1024
    }
}

struct BenchmarkResults {
    var avgRetrievalTime: TimeInterval = 0
    var minRetrievalTime: TimeInterval = 0
    var maxRetrievalTime: TimeInterval = 0
    var embeddingTime: TimeInterval = 0
    var memoryUsage: Int = 0
    var cacheHitRate: Double = 0
    var queryResults: [String: Int] = [:]
}

// Usage
let benchmark = KnowledgeBaseBenchmark(knowledgeBase: kb)
let results = try await benchmark.runBenchmark()
benchmark.printResults(results)

final class SecureKnowledgeBase: KnowledgeBase {
    private let sensitivePatterns: [String: NSRegularExpression] = [
        "ssn": try! NSRegularExpression(pattern: #"\d{3}-\d{2}-\d{4}"#),
        "credit_card": try! NSRegularExpression(pattern: #"\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}"#),
        "email": try! NSRegularExpression(pattern: #"[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}"#, options: .caseInsensitive),
        "phone": try! NSRegularExpression(pattern: #"\+?\d{10,}"#)
    ]

    override func processDocument(_ content: String) async throws -> ProcessedDocument {
        // Redact sensitive information
        var sanitized = content

        for (type, pattern) in sensitivePatterns {
            let range = NSRange(sanitized.startIndex..., in: sanitized)
            sanitized = pattern.stringByReplacingMatches(
                in: sanitized,
                range: range,
                withTemplate: "[\(type.uppercased())]"
            )
        }

        return try await super.processDocument(sanitized)
    }

    func setCustomPattern(name: String, pattern: String) throws {
        let regex = try NSRegularExpression(pattern: pattern)
        sensitivePatterns[name] = regex
    }
}

// Usage
let secureKB = SecureKnowledgeBase()
try secureKB.setCustomPattern(
    name: "employee_id",
    pattern: #"EMP-\d{6}"#
)

final class AccessControlledKnowledgeBase: KnowledgeBase {
    private var documentPermissions: [String: Set<String>] = [:]

    func setDocumentPermissions(
        documentPath: String,
        roles: Set<String>
    ) {
        documentPermissions[documentPath] = roles
    }

    override func query(
        query: String,
        userRole: String,
        limit: Int,
        threshold: Double?
    ) async throws -> [KnowledgeResult] {
        // Get all results
        let allResults = try await super.query(
            query: query,
            limit: limit * 2,  // Get more to filter
            threshold: threshold
        )

        // Filter by permissions
        let filtered = allResults.filter { result in
            guard let requiredRoles = documentPermissions[result.sourcePath] else {
                return true  // No restrictions
            }
            return requiredRoles.contains(userRole)
        }

        return Array(filtered.prefix(limit))
    }
}

// Usage
let acKB = AccessControlledKnowledgeBase()

// Set permissions
acKB.setDocumentPermissions(
    documentPath: "./confidential/executive-plan.pdf",
    roles: ["executive", "board"]
)

acKB.setDocumentPermissions(
    documentPath: "./public/user-guide.pdf",
    roles: ["employee", "customer", "public"]
)

// Query with role
let results = try await acKB.query(
    query: "company strategy",
    userRole: currentUser.role,
    limit: 5
)

import CryptoKit

final class EncryptedKnowledgeBase: KnowledgeBase {
    private let encryptionKey: SymmetricKey

    init(encryptionKey: SymmetricKey) {
        self.encryptionKey = encryptionKey
        super.init()
    }

    override func persistVectors(
        vectors: [Vector],
        path: String
    ) async throws {
        // Serialize vectors
        let data = try JSONEncoder().encode(vectors)

        // Encrypt
        let sealedBox = try AES.GCM.seal(
            data,
            using: encryptionKey
        )

        // Write encrypted data
        try sealedBox.combined?.write(to: URL(fileURLWithPath: path))
    }

    override func loadVectors(from path: String) async throws -> [Vector] {
        // Read encrypted data
        let encryptedData = try Data(contentsOf: URL(fileURLWithPath: path))

        // Decrypt
        let sealedBox = try AES.GCM.SealedBox(combined: encryptedData)
        let decryptedData = try AES.GCM.open(sealedBox, using: encryptionKey)

        // Deserialize
        return try JSONDecoder().decode([Vector].self, from: decryptedData)
    }
}

// Usage
let encryptionKey = SymmetricKey(size: .bits256)
let encryptedKB = EncryptedKnowledgeBase(encryptionKey: encryptionKey)

final class AuditedKnowledgeBase: KnowledgeBase {
    private let auditLogger: AuditLogger

    override func query(
        query: String,
        userId: String,
        limit: Int,
        threshold: Double?
    ) async throws -> [KnowledgeResult] {
        // Log query
        await auditLogger.log(
            event: .knowledgeQuery,
            userId: userId,
            details: [
                "query": query,
                "limit": "\(limit)",
                "threshold": "\(threshold ?? 0.7)"
            ]
        )

        let results = try await super.query(
            query: query,
            limit: limit,
            threshold: threshold
        )

        // Log results
        await auditLogger.log(
            event: .knowledgeResults,
            userId: userId,
            details: [
                "query": query,
                "resultCount": "\(results.count)",
                "sources": results.map { $0.sourcePath }.joined(separator: ", ")
            ]
        )

        return results
    }
}

// Audit Logger
actor AuditLogger {
    private var logs: [AuditLog] = []

    func log(event: AuditEvent, userId: String, details: [String: String]) {
        let log = AuditLog(
            timestamp: Date(),
            event: event,
            userId: userId,
            details: details
        )
        logs.append(log)

        // Persist to secure storage
        Task {
            try await persistLog(log)
        }
    }

    private func persistLog(_ log: AuditLog) async throws {
        // Write to secure audit log
    }
}

struct AuditLog {
    let timestamp: Date
    let event: AuditEvent
    let userId: String
    let details: [String: String]
}

enum AuditEvent {
    case knowledgeQuery
    case knowledgeResults
    case documentAdded
    case documentRemoved
}