PR 9: GraphRAG - Knowledge Graph Storage (GRDB) #10

import GRDB

/// A node in the knowledge graph
struct GraphNode: Codable, FetchableRecord, PersistableRecord, Sendable {
    var id: Int64?
    var externalId: String      // UUID or document ID
    var type: String            // "entity", "document", "category"
    var name: String            // Human-readable name
    var properties: Data?       // JSON-encoded additional properties
    var embedding: Data?        // Float array as binary data
    var createdAt: Date
    var updatedAt: Date
    
    static let databaseTableName = "graph_nodes"
    
    // MARK: - Convenience
    
    enum NodeType: String, CaseIterable {
        case entity = "entity"
        case document = "document"
        case category = "category"
    }
    
    var nodeType: NodeType? {
        NodeType(rawValue: type)
    }
    
    var embeddingVector: [Float]? {
        get {
            guard let data = embedding else { return nil }
            return data.withUnsafeBytes { ptr in
                Array(ptr.bindMemory(to: Float.self))
            }
        }
        set {
            guard let vector = newValue else {
                embedding = nil
                return
            }
            embedding = vector.withUnsafeBytes { Data($0) }
        }
    }
    
    /// Decode properties as a specific type
    func decodeProperties<T: Decodable>(as type: T.Type) -> T? {
        guard let data = properties else { return nil }
        return try? JSONDecoder().decode(type, from: data)
    }
    
    /// Encode properties from a value
    mutating func encodeProperties<T: Encodable>(_ value: T) {
        properties = try? JSONEncoder().encode(value)
    }
}

// MARK: - Associations

extension GraphNode {
    static let outgoingEdges = hasMany(GraphEdge.self, using: GraphEdge.sourceForeignKey)
    static let incomingEdges = hasMany(GraphEdge.self, using: GraphEdge.targetForeignKey)
    
    var outgoingEdges: QueryInterfaceRequest<GraphEdge> {
        request(for: Self.outgoingEdges)
    }
    
    var incomingEdges: QueryInterfaceRequest<GraphEdge> {
        request(for: Self.incomingEdges)
    }
}

import GRDB

/// An edge (relationship) in the knowledge graph
struct GraphEdge: Codable, FetchableRecord, PersistableRecord, Sendable {
    var id: Int64?
    var sourceNodeId: Int64
    var targetNodeId: Int64
    var relationshipType: String    // "mentions", "categorized_as", "similar_to"
    var weight: Double              // Relationship strength 0.0-1.0
    var properties: Data?           // JSON-encoded additional properties
    var createdAt: Date
    
    static let databaseTableName = "graph_edges"
    
    // MARK: - Foreign Keys
    
    static let sourceForeignKey = ForeignKey(["sourceNodeId"])
    static let targetForeignKey = ForeignKey(["targetNodeId"])
    
    // MARK: - Relationship Types
    
    enum RelationshipType: String, CaseIterable {
        case mentions = "mentions"
        case categorizedAs = "categorized_as"
        case similarTo = "similar_to"
        case relatedTo = "related_to"
        case worksFor = "works_for"
        case locatedIn = "located_in"
        case authoredBy = "authored_by"
    }
    
    var relationship: RelationshipType? {
        RelationshipType(rawValue: relationshipType)
    }
}

// MARK: - Associations

extension GraphEdge {
    static let sourceNode = belongsTo(GraphNode.self, using: sourceForeignKey)
    static let targetNode = belongsTo(GraphNode.self, using: targetForeignKey)
    
    var sourceNode: QueryInterfaceRequest<GraphNode> {
        request(for: Self.sourceNode)
    }
    
    var targetNode: QueryInterfaceRequest<GraphNode> {
        request(for: Self.targetNode)
    }
}

import GRDB

actor KnowledgeGraphRepository {
    private let database: DatabaseQueue
    
    init(database: DatabaseQueue) throws {
        self.database = database
        try createTables()
    }
    
    // MARK: - Schema
    
    private func createTables() throws {
        try database.write { db in
            // Nodes table
            try db.create(table: "graph_nodes", ifNotExists: true) { t in
                t.autoIncrementedPrimaryKey("id")
                t.column("externalId", .text).notNull().unique()
                t.column("type", .text).notNull().indexed()
                t.column("name", .text).notNull()
                t.column("properties", .blob)
                t.column("embedding", .blob)
                t.column("createdAt", .datetime).notNull()
                t.column("updatedAt", .datetime).notNull()
            }
            
            // Edges table
            try db.create(table: "graph_edges", ifNotExists: true) { t in
                t.autoIncrementedPrimaryKey("id")
                t.column("sourceNodeId", .integer).notNull()
                    .references("graph_nodes", onDelete: .cascade)
                t.column("targetNodeId", .integer).notNull()
                    .references("graph_nodes", onDelete: .cascade)
                t.column("relationshipType", .text).notNull()
                t.column("weight", .double).notNull()
                t.column("properties", .blob)
                t.column("createdAt", .datetime).notNull()
            }
            
            // Indices for fast lookups
            try db.create(index: "idx_edges_source", on: "graph_edges", 
                         columns: ["sourceNodeId"], ifNotExists: true)
            try db.create(index: "idx_edges_target", on: "graph_edges", 
                         columns: ["targetNodeId"], ifNotExists: true)
            try db.create(index: "idx_edges_type", on: "graph_edges", 
                         columns: ["relationshipType"], ifNotExists: true)
            try db.create(index: "idx_nodes_name", on: "graph_nodes", 
                         columns: ["name"], ifNotExists: true)
        }
    }
    
    // MARK: - Node Operations
    
    func addNode(_ node: GraphNode) throws -> GraphNode {
        try database.write { db in
            var mutableNode = node
            mutableNode.createdAt = Date()
            mutableNode.updatedAt = Date()
            try mutableNode.insert(db)
            return mutableNode
        }
    }
    
    func updateNode(_ node: GraphNode) throws -> GraphNode {
        try database.write { db in
            var mutableNode = node
            mutableNode.updatedAt = Date()
            try mutableNode.update(db)
            return mutableNode
        }
    }
    
    func findNode(byId id: Int64) throws -> GraphNode? {
        try database.read { db in
            try GraphNode.fetchOne(db, key: id)
        }
    }
    
    func findNode(byExternalId externalId: String) throws -> GraphNode? {
        try database.read { db in
            try GraphNode.filter(Column("externalId") == externalId).fetchOne(db)
        }
    }
    
    func findNodes(byType type: String) throws -> [GraphNode] {
        try database.read { db in
            try GraphNode.filter(Column("type") == type).fetchAll(db)
        }
    }
    
    func findNodes(matching name: String) throws -> [GraphNode] {
        try database.read { db in
            try GraphNode.filter(Column("name").like("%\(name)%")).fetchAll(db)
        }
    }
    
    func deleteNode(id: Int64) throws {
        try database.write { db in
            try GraphNode.deleteOne(db, key: id)
        }
    }
    
    // MARK: - Edge Operations
    
    func addEdge(_ edge: GraphEdge) throws -> GraphEdge {
        try database.write { db in
            var mutableEdge = edge
            mutableEdge.createdAt = Date()
            try mutableEdge.insert(db)
            return mutableEdge
        }
    }
    
    func findEdges(from nodeId: Int64) throws -> [GraphEdge] {
        try database.read { db in
            try GraphEdge.filter(Column("sourceNodeId") == nodeId).fetchAll(db)
        }
    }
    
    func findEdges(to nodeId: Int64) throws -> [GraphEdge] {
        try database.read { db in
            try GraphEdge.filter(Column("targetNodeId") == nodeId).fetchAll(db)
        }
    }
    
    func findEdges(between sourceId: Int64, and targetId: Int64) throws -> [GraphEdge] {
        try database.read { db in
            try GraphEdge
                .filter(Column("sourceNodeId") == sourceId)
                .filter(Column("targetNodeId") == targetId)
                .fetchAll(db)
        }
    }
    
    func findEdges(ofType type: String) throws -> [GraphEdge] {
        try database.read { db in
            try GraphEdge.filter(Column("relationshipType") == type).fetchAll(db)
        }
    }
    
    // MARK: - Graph Traversal
    
    /// Breadth-first traversal from a starting node
    func traverse(
        from nodeId: Int64,
        depth: Int = 2,
        relationshipTypes: [String]? = nil,
        direction: TraversalDirection = .outgoing
    ) throws -> [GraphNode] {
        try database.read { db in
            var visited = Set<Int64>()
            var result: [GraphNode] = []
            var queue: [(Int64, Int)] = [(nodeId, 0)]
            
            while !queue.isEmpty {
                let (currentId, currentDepth) = queue.removeFirst()
                
                if visited.contains(currentId) || currentDepth > depth {
                    continue
                }
                visited.insert(currentId)
                
                if let node = try GraphNode.fetchOne(db, key: currentId) {
                    result.append(node)
                }
                
                // Get connected edges based on direction
                let edges: [GraphEdge]
                switch direction {
                case .outgoing:
                    edges = try GraphEdge.filter(Column("sourceNodeId") == currentId).fetchAll(db)
                case .incoming:
                    edges = try GraphEdge.filter(Column("targetNodeId") == currentId).fetchAll(db)
                case .both:
                    let outgoing = try GraphEdge.filter(Column("sourceNodeId") == currentId).fetchAll(db)
                    let incoming = try GraphEdge.filter(Column("targetNodeId") == currentId).fetchAll(db)
                    edges = outgoing + incoming
                }
                
                // Filter by relationship type if specified
                let filteredEdges = relationshipTypes.map { types in
                    edges.filter { types.contains($0.relationshipType) }
                } ?? edges
                
                for edge in filteredEdges {
                    let nextId = direction == .incoming ? edge.sourceNodeId : edge.targetNodeId
                    queue.append((nextId, currentDepth + 1))
                }
            }
            
            return result
        }
    }
    
    enum TraversalDirection {
        case outgoing
        case incoming
        case both
    }
    
    /// Find shortest path between two nodes
    func shortestPath(from sourceId: Int64, to targetId: Int64, maxDepth: Int = 5) throws -> [GraphNode]? {
        try database.read { db in
            var visited = Set<Int64>()
            var queue: [(Int64, [Int64])] = [(sourceId, [sourceId])]
            
            while !queue.isEmpty {
                let (currentId, path) = queue.removeFirst()
                
                if currentId == targetId {
                    // Found path, fetch nodes
                    return try path.compactMap { try GraphNode.fetchOne(db, key: $0) }
                }
                
                if visited.contains(currentId) || path.count > maxDepth {
                    continue
                }
                visited.insert(currentId)
                
                let edges = try GraphEdge.filter(Column("sourceNodeId") == currentId).fetchAll(db)
                for edge in edges {
                    if !visited.contains(edge.targetNodeId) {
                        queue.append((edge.targetNodeId, path + [edge.targetNodeId]))
                    }
                }
            }
            
            return nil
        }
    }
    
    // MARK: - Statistics
    
    func nodeCount() throws -> Int {
        try database.read { db in
            try GraphNode.fetchCount(db)
        }
    }
    
    func edgeCount() throws -> Int {
        try database.read { db in
            try GraphEdge.fetchCount(db)
        }
    }
}

func testNodeCreation() async throws {
    let repo = try KnowledgeGraphRepository(database: makeTestDatabase())
    
    var node = GraphNode(
        externalId: "test-1",
        type: "entity",
        name: "Test Entity",
        createdAt: Date(),
        updatedAt: Date()
    )
    
    node = try await repo.addNode(node)
    
    XCTAssertNotNil(node.id)
}

func testEdgeCreation() async throws {
    let repo = try KnowledgeGraphRepository(database: makeTestDatabase())
    
    // Create two nodes
    var node1 = GraphNode(externalId: "node-1", type: "entity", name: "Node 1", createdAt: Date(), updatedAt: Date())
    var node2 = GraphNode(externalId: "node-2", type: "entity", name: "Node 2", createdAt: Date(), updatedAt: Date())
    
    node1 = try await repo.addNode(node1)
    node2 = try await repo.addNode(node2)
    
    // Create edge
    var edge = GraphEdge(
        sourceNodeId: node1.id!,
        targetNodeId: node2.id!,
        relationshipType: "related_to",
        weight: 0.8,
        createdAt: Date()
    )
    
    edge = try await repo.addEdge(edge)
    
    XCTAssertNotNil(edge.id)
}

func testGraphTraversal() async throws {
    // Create a small graph: A -> B -> C
    let repo = try KnowledgeGraphRepository(database: makeTestDatabase())
    
    var nodeA = GraphNode(externalId: "A", type: "entity", name: "A", createdAt: Date(), updatedAt: Date())
    var nodeB = GraphNode(externalId: "B", type: "entity", name: "B", createdAt: Date(), updatedAt: Date())
    var nodeC = GraphNode(externalId: "C", type: "entity", name: "C", createdAt: Date(), updatedAt: Date())
    
    nodeA = try await repo.addNode(nodeA)
    nodeB = try await repo.addNode(nodeB)
    nodeC = try await repo.addNode(nodeC)
    
    _ = try await repo.addEdge(GraphEdge(sourceNodeId: nodeA.id!, targetNodeId: nodeB.id!, relationshipType: "related_to", weight: 1.0, createdAt: Date()))
    _ = try await repo.addEdge(GraphEdge(sourceNodeId: nodeB.id!, targetNodeId: nodeC.id!, relationshipType: "related_to", weight: 1.0, createdAt: Date()))
    
    // Traverse from A with depth 2
    let reachable = try await repo.traverse(from: nodeA.id!, depth: 2)
    
    XCTAssertEqual(reachable.count, 3) // A, B, C
}

func testShortestPath() async throws {
    // Create graph with multiple paths
    let repo = try KnowledgeGraphRepository(database: makeTestDatabase())
    
    // A -> B -> C (length 2)
    // A -> C (length 1)
    
    // Setup nodes and edges...
    
    let path = try await repo.shortestPath(from: nodeA.id!, to: nodeC.id!)
    
    XCTAssertEqual(path?.count, 2) // Direct path A -> C
}