PR 11: GraphRAG - Unified Service #12

/// Response from a GraphRAG query
struct GraphRAGResponse: Sendable {
    let answer: String
    let sourceDocuments: [String]
    let relatedEntities: [String]
    let confidence: Double
    let graphContext: GraphContext?
    
    struct GraphContext: Sendable {
        let nodesTraversed: Int
        let edgesFollowed: Int
        let maxDepthReached: Int
    }
}

/// Response from document indexing
struct IndexingResult: Sendable {
    let documentId: String
    let entitiesExtracted: Int
    let relationshipsInferred: Int
    let embeddingGenerated: Bool
    let processingTime: TimeInterval
}

enum GraphRAGError: LocalizedError {
    case indexNotInitialized
    case embeddingFailed(String)
    case noResultsFound
    case graphTraversalFailed(Error)
    case documentNotFound(String)
    case relationshipExtractionFailed(Error)
    
    var errorDescription: String? {
        switch self {
        case .indexNotInitialized:
            return "Vector index has not been initialized"
        case .embeddingFailed(let reason):
            return "Failed to generate embedding: \(reason)"
        case .noResultsFound:
            return "No relevant results found in knowledge graph"
        case .graphTraversalFailed(let error):
            return "Graph traversal failed: \(error.localizedDescription)"
        case .documentNotFound(let id):
            return "Document not found: \(id)"
        case .relationshipExtractionFailed(let error):
            return "Relationship extraction failed: \(error.localizedDescription)"
        }
    }
}

import GRDB
import FoundationModels

actor GraphRAGEnhancer {
    // MARK: - Dependencies
    
    private let entityExtractor: NativeEntityExtractor
    private let embeddingService: AppleNLEmbeddingService
    private let vectorStore: any VectorStore
    private let graphRepository: KnowledgeGraphRepository
    private let relationshipExtractor: AppleIntelligenceRelationshipExtractor?
    
    // MARK: - Configuration
    
    private let embeddingDimension = 512
    private let defaultSearchK = 10
    private let maxTraversalDepth = 3
    
    // MARK: - Initialization
    
    init(database: DatabaseQueue) async throws {
        self.entityExtractor = NativeEntityExtractor()
        self.embeddingService = AppleNLEmbeddingService()
        self.vectorStore = VectorStoreFactory.create(dimension: embeddingDimension)
        self.graphRepository = try KnowledgeGraphRepository(database: database)
        
        if #available(macOS 26.0, *) {
            self.relationshipExtractor = AppleIntelligenceRelationshipExtractor()
        } else {
            self.relationshipExtractor = nil
        }
    }
    
    // MARK: - Document Indexing
    
    /// Index a document into the knowledge graph
    func indexDocument(
        id: String,
        content: String,
        metadata: [String: Any] = [:]
    ) async throws -> IndexingResult {
        let startTime = Date()
        
        // 1. Extract entities using native NLTagger (fast)
        let entities = await entityExtractor.extractAll(from: content)
        
        // 2. Generate document embedding
        let embedding = await embeddingService.generateEmbedding(for: content)
        
        // 3. Create document node in graph
        var documentNode = GraphNode(
            externalId: id,
            type: GraphNode.NodeType.document.rawValue,
            name: metadata["filename"] as? String ?? id,
            properties: try? JSONSerialization.data(withJSONObject: metadata),
            createdAt: Date(),
            updatedAt: Date()
        )
        documentNode.embeddingVector = embedding
        documentNode = try graphRepository.addNode(documentNode)
        
        // 4. Add embedding to vector store
        var embeddingGenerated = false
        if let emb = embedding {
            try await vectorStore.add(id: id, vector: emb)
            embeddingGenerated = true
        }
        
        // 5. Create entity nodes and link to document
        for entity in entities {
            let entityId = "\(entity.type.rawValue):\(entity.text)"
            
            var entityNode: GraphNode
            if let existing = try graphRepository.findNode(byExternalId: entityId) {
                entityNode = existing
            } else {
                entityNode = GraphNode(
                    externalId: entityId,
                    type: GraphNode.NodeType.entity.rawValue,
                    name: entity.text,
                    properties: try? JSONEncoder().encode(["entityType": entity.type.rawValue]),
                    createdAt: Date(),
                    updatedAt: Date()
                )
                entityNode = try graphRepository.addNode(entityNode)
            }
            
            // Create edge: document -> mentions -> entity
            let edge = GraphEdge(
                sourceNodeId: documentNode.id!,
                targetNodeId: entityNode.id!,
                relationshipType: GraphEdge.RelationshipType.mentions.rawValue,
                weight: entity.confidence,
                createdAt: Date()
            )
            _ = try graphRepository.addEdge(edge)
        }
        
        // 6. Extract relationships using Apple Intelligence
        var relationshipsInferred = 0
        if let extractor = relationshipExtractor {
            do {
                let relationships = try await extractor.extractRelationships(
                    from: content,
                    entities: entities
                )
                
                // Create relationship edges
                for rel in relationships {
                    if let sourceNode = try graphRepository.findNode(
                        byExternalId: "PersonalName:\(rel.sourceEntity)"
                    ) ?? graphRepository.findNode(byExternalId: "OrganizationName:\(rel.sourceEntity)"),
                       let targetNode = try graphRepository.findNode(
                        byExternalId: "PersonalName:\(rel.targetEntity)"
                    ) ?? graphRepository.findNode(byExternalId: "OrganizationName:\(rel.targetEntity)") {
                        
                        let edge = GraphEdge(
                            sourceNodeId: sourceNode.id!,
                            targetNodeId: targetNode.id!,
                            relationshipType: rel.relationshipType.rawValue,
                            weight: rel.confidence,
                            createdAt: Date()
                        )
                        _ = try graphRepository.addEdge(edge)
                        relationshipsInferred += 1
                    }
                }
            } catch {
                NSLog("⚠️ [GraphRAG] Relationship extraction failed: %@", error.localizedDescription)
            }
        }
        
        let processingTime = Date().timeIntervalSince(startTime)
        
        return IndexingResult(
            documentId: id,
            entitiesExtracted: entities.count,
            relationshipsInferred: relationshipsInferred,
            embeddingGenerated: embeddingGenerated,
            processingTime: processingTime
        )
    }
    
    // MARK: - Querying
    
    /// Query the knowledge graph
    func query(_ question: String, k: Int? = nil) async throws -> GraphRAGResponse {
        let searchK = k ?? defaultSearchK
        
        // 1. Generate query embedding
        guard let queryEmbedding = await embeddingService.generateEmbedding(for: question) else {
            throw GraphRAGError.embeddingFailed("Could not generate embedding for query")
        }
        
        // 2. Find similar documents using vector search
        let similarDocs = try await vectorStore.search(query: queryEmbedding, k: searchK)
        
        guard !similarDocs.isEmpty else {
            throw GraphRAGError.noResultsFound
        }
        
        // 3. Get related entities through graph traversal
        var relatedEntities: [GraphNode] = []
        var graphContext = GraphRAGResponse.GraphContext(
            nodesTraversed: 0,
            edgesFollowed: 0,
            maxDepthReached: 0
        )
        
        for (docId, _) in similarDocs {
            if let docNode = try graphRepository.findNode(byExternalId: docId) {
                let connected = try graphRepository.traverse(
                    from: docNode.id!,
                    depth: maxTraversalDepth,
                    direction: .both
                )
                relatedEntities.append(contentsOf: connected)
            }
        }
        
        // 4. Build context for synthesis
        let context = buildContext(
            documents: similarDocs,
            entities: relatedEntities
        )
        
        // 5. Synthesize answer using Apple Intelligence
        let answer: String
        if #available(macOS 26.0, *) {
            answer = try await synthesizeAnswer(question: question, context: context)
        } else {
            answer = "Based on \(similarDocs.count) relevant documents: \(context.prefix(500))..."
        }
        
        // 6. Calculate confidence
        let confidence = calculateConfidence(similarDocs)
        
        return GraphRAGResponse(
            answer: answer,
            sourceDocuments: similarDocs.map { $0.id },
            relatedEntities: Array(Set(relatedEntities.map { $0.name })).prefix(20).map { $0 },
            confidence: confidence,
            graphContext: graphContext
        )
    }
    
    /// Find documents similar to a given document
    func findSimilar(documentId: String, k: Int = 5) async throws -> [(id: String, similarity: Double)] {
        guard let docNode = try graphRepository.findNode(byExternalId: documentId),
              let embedding = docNode.embeddingVector else {
            throw GraphRAGError.documentNotFound(documentId)
        }
        
        let results = try await vectorStore.search(query: embedding, k: k + 1)
        
        // Filter out the query document itself and convert distance to similarity
        return results
            .filter { $0.id != documentId }
            .prefix(k)
            .map { ($0.id, Double(1.0 / (1.0 + $0.distance))) }
    }
    
    // MARK: - Category Enhancement
    
    /// Enhance category suggestions using graph context
    func enhanceCategorySuggestion(
        for signature: FileSignature,
        baseCategorization: CategorizationResult
    ) async throws -> CategorizationResult {
        // Find similar documents
        guard let embedding = await embeddingService.generateEmbedding(
            for: signature.textContent ?? signature.url.lastPathComponent
        ) else {
            return baseCategorization
        }
        
        let similar = try await vectorStore.search(query: embedding, k: 5)
        
        // Check if similar documents have consistent categories
        // (This would integrate with the category storage)
        
        // For now, just boost confidence if we found similar docs
        if !similar.isEmpty {
            return CategorizationResult(
                categoryPath: baseCategorization.categoryPath,
                confidence: min(1.0, baseCategorization.confidence + 0.1),
                rationale: baseCategorization.rationale + " (supported by \(similar.count) similar documents)",
                extractedKeywords: baseCategorization.extractedKeywords,
                provider: baseCategorization.provider
            )
        }
        
        return baseCategorization
    }
    
    // MARK: - Private Helpers
    
    private func buildContext(
        documents: [(id: String, distance: Float)],
        entities: [GraphNode]
    ) -> String {
        var context = "**Related documents (by similarity):**\n"
        for (id, distance) in documents.prefix(5) {
            let similarity = 1.0 / (1.0 + distance)
            context += "- \(id) (similarity: \(String(format: "%.2f", similarity)))\n"
        }
        
        context += "\n**Related entities from knowledge graph:**\n"
        let uniqueEntities = Array(Set(entities.map { $0.name })).prefix(15)
        for entity in uniqueEntities {
            context += "- \(entity)\n"
        }
        
        return context
    }
    
    @available(macOS 26.0, *)
    private func synthesizeAnswer(question: String, context: String) async throws -> String {
        let session = LanguageModelSession()
        
        let prompt = """
        Based on the following context from the knowledge graph, answer the question.
        
        **Context:**
        \(context)
        
        **Question:** \(question)
        
        Provide a concise, accurate answer based only on the available context.
        If you cannot answer from the context, say so.
        """
        
        return try await session.respond(to: prompt)
    }
    
    private func calculateConfidence(_ results: [(id: String, distance: Float)]) -> Double {
        guard !results.isEmpty else { return 0 }
        let avgDistance = results.map { $0.distance }.reduce(0, +) / Float(results.count)
        return Double(max(0, min(1.0, 1.0 / (1.0 + avgDistance))))
    }
    
    // MARK: - Maintenance
    
    /// Rebuild vector index from graph
    func rebuildVectorIndex() async throws {
        try await vectorStore.clear()
        
        let documents = try graphRepository.findNodes(byType: "document")
        for doc in documents {
            if let embedding = doc.embeddingVector {
                try await vectorStore.add(id: doc.externalId, vector: embedding)
            }
        }
    }
    
    /// Get statistics
    func statistics() async throws -> GraphRAGStatistics {
        return GraphRAGStatistics(
            nodeCount: try graphRepository.nodeCount(),
            edgeCount: try graphRepository.edgeCount(),
            vectorCount: await vectorStore.count
        )
    }
}

struct GraphRAGStatistics: Sendable {
    let nodeCount: Int
    let edgeCount: Int
    let vectorCount: Int
}

// In SortAIPipeline or AppState
class SortAIPipeline {
    private var graphRAG: GraphRAGEnhancer?
    
    func initializeGraphRAG() async throws {
        self.graphRAG = try await GraphRAGEnhancer(database: database)
    }
    
    func processFile(_ url: URL) async throws {
        // ... existing processing ...
        
        // Index document in GraphRAG
        if let graphRAG = graphRAG, 
           let content = signature.textContent {
            let result = try await graphRAG.indexDocument(
                id: signature.url.path,
                content: content,
                metadata: ["filename": url.lastPathComponent]
            )
            NSLog("📊 [GraphRAG] Indexed: \(result.entitiesExtracted) entities, \(result.relationshipsInferred) relationships")
        }
    }
    
    func enhanceCategorization(
        _ result: CategorizationResult,
        for signature: FileSignature
    ) async throws -> CategorizationResult {
        guard let graphRAG = graphRAG else { return result }
        return try await graphRAG.enhanceCategorySuggestion(for: signature, baseCategorization: result)
    }
}

func testDocumentIndexing() async throws {
    let enhancer = try await GraphRAGEnhancer(database: makeTestDatabase())
    
    let result = try await enhancer.indexDocument(
        id: "test-doc-1",
        content: "Apple CEO Tim Cook announced new products in Cupertino.",
        metadata: ["filename": "announcement.txt"]
    )
    
    XCTAssertGreaterThan(result.entitiesExtracted, 0)
    XCTAssertTrue(result.embeddingGenerated)
}

func testSimilarDocumentSearch() async throws {
    let enhancer = try await GraphRAGEnhancer(database: makeTestDatabase())
    
    // Index two similar documents
    _ = try await enhancer.indexDocument(id: "doc1", content: "Financial quarterly report Q4 2025")
    _ = try await enhancer.indexDocument(id: "doc2", content: "Quarterly financial earnings Q4 2025")
    _ = try await enhancer.indexDocument(id: "doc3", content: "Cat playing with yarn")
    
    let similar = try await enhancer.findSimilar(documentId: "doc1", k: 2)
    
    XCTAssertEqual(similar.first?.id, "doc2", "doc2 should be most similar to doc1")
}

func testQuery() async throws {
    let enhancer = try await GraphRAGEnhancer(database: makeTestDatabase())
    
    _ = try await enhancer.indexDocument(
        id: "report",
        content: "Apple Inc reported record revenue of $124B in Q4 2025."
    )
    
    let response = try await enhancer.query("What was Apples revenue?")
    
    XCTAssertFalse(response.answer.isEmpty)
    XCTAssertTrue(response.sourceDocuments.contains("report"))
}