9. Embeddings, Semantic Search, dan RAG di Device

// embeddings: Semantic similarity dengan NLEmbedding
import NaturalLanguage

let embedding = NLEmbedding.sentenceEmbedding(for: .english)!
let v1 = embedding.vector(for: "I love pizza")!
let v2 = embedding.vector(for: "Pizza is my favorite food")!
let v3 = embedding.vector(for: "The weather is sunny")!

func cosineSimilarity(_ a: [Double], _ b: [Double]) -> Double {
    let dot = zip(a, b).map(*).reduce(0, +)
    let normA = sqrt(a.map { $0 * $0 }.reduce(0, +))
    let normB = sqrt(b.map { $0 * $0 }.reduce(0, +))
    return dot / (normA * normB)
}

print(cosineSimilarity(v1, v2))  // ~0.85 (mirip)
print(cosineSimilarity(v1, v3))  // ~0.15 (beda)

1. Indexing (saat user tambah dokumen):
   - Chunk dokumen ke potongan 200-500 token
   - Generate embedding per chunk
   - Simpan ke local DB (SwiftData / SQLite)

2. Query (saat user tanya):
   - Generate embedding query
   - Cari chunk dengan similarity tertinggi (top-K)
   - Kirim chunk + query ke LLM (Foundation Models atau cloud)
   - Tampilkan jawaban dengan citation

// rag: Simple on-device vector store
@Model
final class DocumentChunk {
    var documentID: UUID
    var content: String
    var embedding: [Double]
    var order: Int

    init(documentID: UUID, content: String, embedding: [Double], order: Int) {
        self.documentID = documentID
        self.content = content
        self.embedding = embedding
        self.order = order
    }
}

actor VectorStore {
    private let context: ModelContext
    private let embedder: NLEmbedding

    init(context: ModelContext) {
        self.context = context
        self.embedder = NLEmbedding.sentenceEmbedding(for: .english)!
    }

    func index(documentID: UUID, content: String) async throws {
        let chunks = chunkText(content, maxTokens: 300)
        for (i, chunk) in chunks.enumerated() {
            guard let vector = embedder.vector(for: chunk) else { continue }
            let docChunk = DocumentChunk(
                documentID: documentID,
                content: chunk,
                embedding: vector,
                order: i
            )
            context.insert(docChunk)
        }
        try context.save()
    }

    func search(query: String, topK: Int = 5) async throws -> [DocumentChunk] {
        guard let queryVec = embedder.vector(for: query) else { return [] }
        let all = try context.fetch(FetchDescriptor<DocumentChunk>())
        let scored = all.map { chunk in
            (chunk, cosineSimilarity(chunk.embedding, queryVec))
        }
        return scored.sorted { $0.1 > $1.1 }.prefix(topK).map { $0.0 }
    }

    private func chunkText(_ text: String, maxTokens: Int) -> [String] {
        // Simple paragraph splitter; production pakai tokenizer proper
        text.components(separatedBy: "\n\n").filter { !$0.isEmpty }
    }
}

// rag: Compose RAG dengan Foundation Models
import FoundationModels

func answer(query: String, store: VectorStore) async throws -> String {
    let relevantChunks = try await store.search(query: query, topK: 3)
    let context = relevantChunks.map(\.content).joined(separator: "\n\n")

    let prompt = """
    Berdasarkan konteks berikut:

    \(context)

    Jawab pertanyaan: \(query)

    Kalau jawaban tidak ada di konteks, katakan tidak tahu.
    """

    let session = LanguageModelSession()
    let response = try await session.respond(to: prompt)
    return response.content
}