PR 10: GraphRAG - Relationship Inference (Apple Intelligence)

[gilmanb1]

Overview

Implement relationship inference using Apple Intelligence. This component uses the LLM to identify relationships between extracted entities, bridging the entity extraction and knowledge graph components.

Dependencies

• PR 2 (PR 2: Apple Intelligence Provider Implementation #3): Apple Intelligence Provider Implementation
• PR 6 (PR 6: GraphRAG - Native Entity Extraction #7): GraphRAG - Native Entity Extraction (for entity types)

Files to Create

File	Action	Description
Sources/SortAI/Core/GraphRAG/AppleIntelligenceRelationshipExtractor.swift	Create	Relationship inference
Sources/SortAI/Core/GraphRAG/InferredRelationship.swift	Create	Relationship types

Implementation Details

1. InferredRelationship Struct

/// A relationship inferred between entities
struct InferredRelationship: Sendable, Codable {
    let sourceEntity: String
    let targetEntity: String
    let relationshipType: RelationshipKind
    let confidence: Double
    let context: String?
    
    enum RelationshipKind: String, Codable, CaseIterable {
        case worksFor = "works_for"
        case locatedIn = "located_in"
        case relatedTo = "related_to"
        case mentions = "mentions"
        case authoredBy = "authored_by"
        case partOf = "part_of"
        case owns = "owns"
        case collaboratesWith = "collaborates_with"
        
        var displayName: String {
            switch self {
            case .worksFor: return "works for"
            case .locatedIn: return "located in"
            case .relatedTo: return "related to"
            case .mentions: return "mentions"
            case .authoredBy: return "authored by"
            case .partOf: return "part of"
            case .owns: return "owns"
            case .collaboratesWith: return "collaborates with"
            }
        }
    }
}

2. @generable Types for Structured Output

import FoundationModels

@Generable
struct RelationshipExtractionResponse {
    @Guide(description: "List of relationships identified in the text")
    var relationships: [ExtractedRelationshipItem]
}

@Generable
struct ExtractedRelationshipItem {
    @Guide(description: "The source entity name")
    var source: String
    
    @Guide(description: "The target entity name")
    var target: String
    
    @Guide(description: "Relationship type: works_for, located_in, related_to, mentions, authored_by, part_of, owns, collaborates_with")
    var relationshipType: String
    
    @Guide(description: "Confidence from 0.0 to 1.0")
    var confidence: Double
    
    @Guide(description: "Brief context for why this relationship was inferred")
    var context: String?
}

3. AppleIntelligenceRelationshipExtractor

import FoundationModels

@available(macOS 26.0, *)
actor AppleIntelligenceRelationshipExtractor {
    private var session: LanguageModelSession?
    private let maxTextLength = 3000
    
    // MARK: - Extraction
    
    /// Extract relationships from text given pre-extracted entities
    func extractRelationships(
        from text: String,
        entities: [ExtractedEntity]
    ) async throws -> [InferredRelationship] {
        guard !entities.isEmpty else { return [] }
        
        let session = try await getOrCreateSession()
        
        // Build entity list for prompt
        let entityList = entities
            .map { "\($0.text) (\($0.type.displayName))" }
            .joined(separator: ", ")
        
        let truncatedText = String(text.prefix(maxTextLength))
        
        let prompt = """
        Given the following text and list of entities, identify relationships between them.
        
        **Entities found:**
        \(entityList)
        
        **Text:**
        \(truncatedText)
        
        **Instructions:**
        - Identify explicit and implicit relationships between the entities
        - Use these relationship types: works_for, located_in, related_to, mentions, authored_by, part_of, owns, collaborates_with
        - Assign confidence based on how explicitly the relationship is stated
        - Only include relationships you can justify from the text
        """
        
        let response: RelationshipExtractionResponse = try await session.respond(
            to: prompt,
            generating: RelationshipExtractionResponse.self
        )
        
        return response.relationships.compactMap { item in
            guard let kind = InferredRelationship.RelationshipKind(rawValue: item.relationshipType) else {
                return nil
            }
            
            return InferredRelationship(
                sourceEntity: item.source,
                targetEntity: item.target,
                relationshipType: kind,
                confidence: max(0, min(1, item.confidence)),
                context: item.context
            )
        }
    }
    
    /// Extract relationships without pre-extracted entities (LLM does both)
    func extractRelationshipsAndEntities(
        from text: String
    ) async throws -> (entities: [ExtractedEntity], relationships: [InferredRelationship]) {
        let session = try await getOrCreateSession()
        
        let truncatedText = String(text.prefix(maxTextLength))
        
        let prompt = """
        Analyze this text to:
        1. Extract named entities (people, organizations, locations, dates)
        2. Identify relationships between those entities
        
        **Text:**
        \(truncatedText)
        
        **Instructions:**
        - First identify all named entities
        - Then identify relationships using: works_for, located_in, related_to, mentions, authored_by, part_of, owns, collaborates_with
        """
        
        // For combined extraction, we could use a different @Generable type
        // or make two sequential calls
        
        let entityResponse: EntityExtractionResponse = try await session.respond(
            to: "Extract entities from: \(truncatedText)",
            generating: EntityExtractionResponse.self
        )
        
        let entities = entityResponse.entities.map { item in
            ExtractedEntity(
                text: item.text,
                type: EntityType(rawValue: item.type) ?? .keyword,
                confidence: 0.8
            )
        }
        
        let relationships = try await extractRelationships(from: text, entities: entities)
        
        return (entities, relationships)
    }
    
    // MARK: - Batch Processing
    
    /// Process multiple documents efficiently
    func extractRelationshipsBatch(
        documents: [(text: String, entities: [ExtractedEntity])]
    ) async throws -> [[InferredRelationship]] {
        var results: [[InferredRelationship]] = []
        
        for (text, entities) in documents {
            let relationships = try await extractRelationships(from: text, entities: entities)
            results.append(relationships)
        }
        
        return results
    }
    
    // MARK: - Session Management
    
    private func getOrCreateSession() async throws -> LanguageModelSession {
        if let session = session {
            return session
        }
        let newSession = LanguageModelSession()
        self.session = newSession
        return newSession
    }
    
    /// Reset session (useful for clearing context)
    func resetSession() {
        session = nil
    }
}

// MARK: - Fallback for Pre-macOS 26

/// Stub extractor for systems without Apple Intelligence
final class RelationshipExtractorUnavailable: Sendable {
    func extractRelationships(
        from text: String,
        entities: [ExtractedEntity]
    ) async throws -> [InferredRelationship] {
        // Return basic co-occurrence relationships
        var relationships: [InferredRelationship] = []
        
        for i in 0..<entities.count {
            for j in (i+1)..<entities.count {
                relationships.append(InferredRelationship(
                    sourceEntity: entities[i].text,
                    targetEntity: entities[j].text,
                    relationshipType: .relatedTo,
                    confidence: 0.3,
                    context: "Co-occurrence in document"
                ))
            }
        }
        
        return relationships
    }
}

Integration Example

// Usage in GraphRAG pipeline
func processDocument(_ text: String) async throws {
    // 1. Extract entities using NLTagger (fast)
    let entities = await entityExtractor.extractAll(from: text)
    
    // 2. Infer relationships using Apple Intelligence
    let relationships: [InferredRelationship]
    if #available(macOS 26.0, *) {
        let extractor = AppleIntelligenceRelationshipExtractor()
        relationships = try await extractor.extractRelationships(from: text, entities: entities)
    } else {
        let fallback = RelationshipExtractorUnavailable()
        relationships = try await fallback.extractRelationships(from: text, entities: entities)
    }
    
    // 3. Store in knowledge graph
    for relationship in relationships {
        // Create/find nodes and edges
        try await graphRepository.addRelationship(relationship)
    }
}

Performance Notes

Based on prototype testing:

• Relationship inference: ~1.8s per document
• Batch processing benefits from session reuse
• Context window limit: ~3000 chars recommended

Acceptance Criteria

• Extracts relationships from text given entities
• Uses @generable for type-safe structured output
• Maps relationship types correctly
• Provides fallback for pre-macOS 26
• Supports batch processing
• Returns confidence scores
• Includes context/justification
• Session is reused for performance

Testing

@available(macOS 26.0, *)
func testRelationshipExtraction() async throws {
    let extractor = AppleIntelligenceRelationshipExtractor()
    
    let text = """
    Tim Cook, CEO of Apple, announced new products at Apple Park in Cupertino.
    The event was attended by executives from Microsoft and Google.
    """
    
    let entities = [
        ExtractedEntity(text: "Tim Cook", type: .person),
        ExtractedEntity(text: "Apple", type: .organization),
        ExtractedEntity(text: "Apple Park", type: .location),
        ExtractedEntity(text: "Cupertino", type: .location),
        ExtractedEntity(text: "Microsoft", type: .organization),
        ExtractedEntity(text: "Google", type: .organization)
    ]
    
    let relationships = try await extractor.extractRelationships(from: text, entities: entities)
    
    XCTAssertFalse(relationships.isEmpty)
    
    // Should find "Tim Cook works_for Apple"
    let worksFor = relationships.first { 
        $0.sourceEntity == "Tim Cook" && 
        $0.relationshipType == .worksFor 
    }
    XCTAssertNotNil(worksFor)
}

func testFallbackExtraction() async throws {
    let fallback = RelationshipExtractorUnavailable()
    
    let entities = [
        ExtractedEntity(text: "Apple", type: .organization),
        ExtractedEntity(text: "Microsoft", type: .organization)
    ]
    
    let relationships = try await fallback.extractRelationships(from: "text", entities: entities)
    
    // Should create co-occurrence relationship
    XCTAssertEqual(relationships.count, 1)
    XCTAssertEqual(relationships[0].relationshipType, .relatedTo)
    XCTAssertEqual(relationships[0].confidence, 0.3)
}

Estimated Size

~150 lines of code

Risk Assessment

Medium - Depends on Apple Intelligence quality for relationship extraction. Mitigation: fallback to co-occurrence-based relationships.

[gilmanb1]

Dependencies: Blocked by #3, #7

[gilmanb1]

✅ Implementation Complete

This feature was already implemented as part of the LLM_Research.md implementation. The code exists in the codebase and all 233 tests pass.

A minor build fix was applied in commit 4341283 to fix the Apple Intelligence availability check API.