File size: 3,734 Bytes
dd5890c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | import Foundation
#if canImport(PDFKit)
import PDFKit
#endif
/// A local-file retrieval tier: scans a folder for text/markdown/PDF, splits into chunks, embeds
/// each with BGE, and answers cosine queries. Used as an optional knowledge source for `lookup`
/// turns (Mac primarily, but works on iOS too). Nil-by-default in `Assistant` → zero behaviour change.
public final class LocalFileIndex {
public struct Chunk { public let text: String; let emb: [Float]; public let source: String }
private var chunks: [Chunk] = []
private let bge: Embedder
public private(set) var indexedFiles: Int = 0
public init(bge: Embedder) { self.bge = bge }
public var count: Int { chunks.count }
private static let exts: Set<String> = ["txt", "md", "markdown", "text", "pdf"]
/// Recursively index a directory. Replaces any prior index.
@discardableResult
public func build(from dir: URL, maxFiles: Int = 300, maxCharsPerFile: Int = 200_000) async -> Int {
chunks.removeAll(); indexedFiles = 0
let fm = FileManager.default
guard let en = fm.enumerator(at: dir, includingPropertiesForKeys: nil) else { return 0 }
let all = en.allObjects.compactMap { $0 as? URL }
var files = [URL]()
for url in all where Self.exts.contains(url.pathExtension.lowercased()) {
files.append(url); if files.count >= maxFiles { break }
}
for url in files {
guard let raw = extractText(url) else { continue }
let body = String(raw.prefix(maxCharsPerFile))
let name = url.lastPathComponent
for piece in chunkText(body) {
chunks.append(Chunk(text: piece, emb: await bge.encode(piece, isQuery: false), source: name))
}
indexedFiles += 1
}
return chunks.count
}
/// Top-`k` chunks by cosine ≥ `minSim`, as "filename: text".
public func search(_ query: String, k: Int = 3, minSim: Float = 0.5) async
-> [(text: String, source: String, sim: Float)] {
guard !chunks.isEmpty else { return [] }
let qv = await bge.encode(query, isQuery: true)
var scored = chunks.map { (c: $0, s: cosine($0.emb, qv)) }.filter { $0.s >= minSim }
scored.sort { $0.s > $1.s }
return scored.prefix(k).map { ($0.c.text, $0.c.source, $0.s) }
}
// MARK: - extraction / chunking
private func extractText(_ url: URL) -> String? {
if url.pathExtension.lowercased() == "pdf" {
#if canImport(PDFKit)
return PDFDocument(url: url)?.string
#else
return nil
#endif
}
return try? String(contentsOf: url, encoding: .utf8)
}
/// Pack paragraphs into ~`target`-char chunks; hard-split paragraphs longer than that.
private func chunkText(_ text: String, target: Int = 500) -> [String] {
let paras = text.components(separatedBy: "\n\n")
.map { $0.trimmingCharacters(in: .whitespacesAndNewlines) }
.filter { $0.count >= 20 }
var out = [String](); var cur = ""
func flush() { if !cur.isEmpty { out.append(cur); cur = "" } }
for p in paras {
if p.count > target {
flush()
var s = Substring(p)
while s.count > target {
let idx = s.index(s.startIndex, offsetBy: target)
out.append(String(s[..<idx])); s = s[idx...]
}
cur = String(s)
} else {
if cur.count + p.count + 1 > target { flush() }
cur = cur.isEmpty ? p : cur + "\n" + p
}
}
flush()
return out
}
}
|