diff --git a/Sources/DigUpKit/IndexStore.swift b/Sources/DigUpKit/IndexStore.swift index 7a9b8b82..c2c10c29 100644 --- a/Sources/DigUpKit/IndexStore.swift +++ b/Sources/DigUpKit/IndexStore.swift @@ -74,6 +74,25 @@ public final class IndexStore { file = directory.appendingPathComponent(name) db = try Database(path: file.path) try migrate() + try ensureGraphTables() + } + + /// Ensures the GraphRAG `edges` table exists without bumping `PRAGMA user_version`, so prebuilt `v0.7.0` helpers + /// and newer GraphRAG-enabled helpers can share `index.sqlite` and `code.sqlite` seamlessly. + public func ensureGraphTables() throws { + try db.execute(""" + CREATE TABLE IF NOT EXISTS edges( + id INTEGER PRIMARY KEY, + source_file INTEGER NOT NULL REFERENCES files(id) ON DELETE CASCADE, + source_segment INTEGER NOT NULL REFERENCES segments(id) ON DELETE CASCADE, + entity_type TEXT NOT NULL, + entity_key TEXT NOT NULL, + weight REAL NOT NULL DEFAULT 1.0 + ); + CREATE INDEX IF NOT EXISTS edges_key ON edges(entity_key); + CREATE INDEX IF NOT EXISTS edges_file ON edges(source_file); + CREATE INDEX IF NOT EXISTS edges_segment ON edges(source_segment); + """) } /// Bytes on disk: the database and its write-ahead log. @@ -434,7 +453,8 @@ public final class IndexStore { // MARK: Segments - /// Adds a segment; `text` goes into keyword search (normalized, see `SearchText`), `vector` into meaning search. + /// Adds a segment; `text` goes into keyword search (normalized, see `SearchText`), `vector` into meaning search, + /// and deterministic entities/identifiers go into `edges` for lightweight SQLite GraphRAG multi-hop traversal. public func addSegment(file: Int64, kind: SegmentKind, modality: Modality, loc: Double? = nil, locEnd: Double? = nil, excerpt: String? = nil, text: String? = nil, vector: [Float]? = nil) throws { try db.transaction { // the segment and its keyword text together @@ -442,11 +462,52 @@ public final class IndexStore { [.int(file), .text(kind.rawValue), .text(vector == nil ? Modality.none.rawValue : modality.rawValue), .optional(loc), .optional(locEnd), .optional(excerpt), vector.map { .blob(Vectors.half($0)) } ?? .null]) + let segmentID = db.lastInsertID if let text, !text.isEmpty { try db.run("INSERT INTO fts(rowid, name, body, original) VALUES(?, '', ?, ?)", - [.int(db.lastInsertID)] + Self.keywordText(text)) + [.int(segmentID)] + Self.keywordText(text)) + for entity in Self.extractGraphEntities(from: text) { + try db.run("INSERT INTO edges(source_file, source_segment, entity_type, entity_key, weight) VALUES(?, ?, ?, ?, ?)", + [.int(file), .int(segmentID), .text(entity.type), .text(entity.key), .double(entity.weight)]) + } + } + } + } + + /// Zero-LLM-overhead deterministic entity extractor for SQLite GraphRAG (`symbol`, `id_code`, `file_ref`). + public static func extractGraphEntities(from text: String) -> [(type: String, key: String, weight: Double)] { + var results: [(type: String, key: String, weight: Double)] = [] + var seen = Set() + let tokens = text.components(separatedBy: CharacterSet.alphanumerics.union(CharacterSet(charactersIn: "_-./")).inverted) + for raw in tokens where raw.count >= 4 && raw.count <= 80 { + let trimmed = raw.trimmingCharacters(in: CharacterSet(charactersIn: ".-_/")) + guard trimmed.count >= 4 else { continue } + let lower = trimmed.lowercased() + // 1. File references (e.g., auth.ts, settings.yaml, report.pdf) + if trimmed.contains("."), let ext = trimmed.split(separator: ".").last, + ["py", "ts", "tsx", "js", "rs", "swift", "md", "pdf", "json", "yaml", "yml", "toml", "sql", "sh"].contains(ext.lowercased()) { + let key = (trimmed as NSString).lastPathComponent.lowercased() + if seen.insert("file_ref:\(key)").inserted { + results.append(("file_ref", key, 2.2)) + } + } + // 2. CamelCase or snake_case symbols (e.g., refreshAccessToken, quantize_row_q8_0) + let hasUpper = trimmed.rangeOfCharacter(from: .uppercaseLetters) != nil + let hasLower = trimmed.rangeOfCharacter(from: .lowercaseLetters) != nil + let hasDigit = trimmed.rangeOfCharacter(from: .decimalDigits) != nil + if (hasUpper && hasLower && trimmed.count >= 6) || (trimmed.contains("_") && hasLower) { + if seen.insert("symbol:\(lower)").inserted { + results.append(("symbol", lower, 2.0)) + } + } else if hasDigit && (hasUpper || hasLower) && (trimmed.contains("-") || trimmed.contains("_") || hasUpper) { + // 3. Exact ID codes (invoice numbers, ticket IDs, error codes, CVEs) + if seen.insert("id_code:\(lower)").inserted { + results.append(("id_code", lower, 2.5)) + } } + if results.count >= 32 { break } } + return results } /// `body` and `original` for keyword search: the normalized text, and the text as written if that differs. @@ -457,13 +518,18 @@ public final class IndexStore { private func addNameRow(file: Int64, path: String) throws { try db.run("INSERT INTO segments(file, kind, modality) VALUES(?, 'name', 'none')", [.int(file)]) + let segID = db.lastInsertID try db.run("INSERT INTO fts(rowid, name, body) VALUES(?, ?, ?)", - [.int(db.lastInsertID), .text(SearchText.normalized(Self.nameText(path))), + [.int(segID), .text(SearchText.normalized(Self.nameText(path))), .text(SearchText.normalized(Self.folderText(path)))]) + let fileName = Self.nameText(path).lowercased() + try db.run("INSERT INTO edges(source_file, source_segment, entity_type, entity_key, weight) VALUES(?, ?, 'file_ref', ?, 2.5)", + [.int(file), .int(segID), .text(fileName)]) } private func deleteSegments(of file: Int64, keepingName: Bool) throws { let filter = keepingName ? "AND kind != 'name'" : "" + try db.run("DELETE FROM edges WHERE source_segment IN (SELECT id FROM segments WHERE file = ? \(filter))", [.int(file)]) try db.run("DELETE FROM fts WHERE rowid IN (SELECT id FROM segments WHERE file = ? \(filter))", [.int(file)]) try db.run("DELETE FROM segments WHERE file = ? \(filter)", [.int(file)]) } diff --git a/Sources/DigUpKit/ModelFiles.swift b/Sources/DigUpKit/ModelFiles.swift index dbd72053..cd81fcae 100644 --- a/Sources/DigUpKit/ModelFiles.swift +++ b/Sources/DigUpKit/ModelFiles.swift @@ -17,10 +17,22 @@ public enum ModelFiles { public let name: String public let bytes: Int64 public let sha256: String + public let repository: String + public let revision: String + + public init(name: String, bytes: Int64, sha256: String, + repository: String = ModelFiles.repository, + revision: String = ModelFiles.revision) { + self.name = name + self.bytes = bytes + self.sha256 = sha256 + self.repository = repository + self.revision = revision + } /// Hugging Face's address for this file at the pinned revision (it redirects to a CDN that takes HTTP ranges). public var url: URL { - URL(string: "https://huggingface.co/\(ModelFiles.repository)/resolve/\(ModelFiles.revision)/\(name)")! + URL(string: "https://huggingface.co/\(repository)/resolve/\(revision)/\(name)")! } } @@ -30,8 +42,28 @@ public enum ModelFiles { /// The image and audio encoders (555 MB). public static let media = File(name: "mmproj-embeddinggemma-2-Q8_0.gguf", bytes: 554_821_024, sha256: "c4a8a52691ecef40618438928bdf9e68379b854e24166f292592353db0aab64f") + + /// Primary local generative SLM for 8-core / 12GB RAM PC: Qwen2.5-3B-Instruct Q4_K_M (~1.93 GB). + public static let chatDefault = File( + name: "qwen2.5-3b-instruct-q4_k_m.gguf", + bytes: 2_104_932_128, + sha256: "6cf6efbb4dc67880f9b8fb2f00ca2bf0209fa0d9ed99ab3a241ff3814820b270", + repository: "Qwen/Qwen2.5-3B-Instruct-GGUF", + revision: "main" + ) + + /// Ultra-fast fallback SLM for CPU-constrained multitasking: Qwen2.5-1.5B-Instruct Q4_K_M (~1.12 GB). + public static let chatTurbo = File( + name: "qwen2.5-1.5b-instruct-q4_k_m.gguf", + bytes: 1_117_320_576, + sha256: "6ae5c20f9b4d05f46166fa4ca6838d16e54c57eb26ae3b4ca4b208f7da3bc60b", + repository: "Qwen/Qwen2.5-1.5B-Instruct-GGUF", + revision: "main" + ) + /// In download order. public static let all = [text, media] + public static let chatModels = [chatDefault, chatTurbo] public static let totalBytes = all.reduce(0) { $0 + $1.bytes } /// Where the app keeps them (`AppPaths`). diff --git a/Sources/DigUpKit/Searcher.swift b/Sources/DigUpKit/Searcher.swift index 4592cbe5..56598aef 100644 --- a/Sources/DigUpKit/Searcher.swift +++ b/Sources/DigUpKit/Searcher.swift @@ -435,4 +435,92 @@ public final class Searcher { }?.value }) } + + // MARK: Lightweight SQLite GraphRAG Multi-Hop Expansion + + public struct GraphContextPiece: Sendable { + public let path: String + public let kind: String + public let segmentID: Int64 + public let loc: Double? + public let locEnd: Double? + public let hop: Int + public let linkedVia: String? + public let text: String + } + + /// Expands initial hybrid search hits (`seeds`) across the SQLite `edges` graph (1 to 2 hops) to retrieve + /// connected code definitions, referenced files, or cross-document identifiers without LLM indexing overhead. + public func retrieveGraphContext(query: String, seeds: [SearchHit], hops: Int = 1, + maxSeeds: Int = 5, maxNeighbors: Int = 5) throws -> [GraphContextPiece] { + var pieces: [GraphContextPiece] = [] + var visitedSegments = Set() + var seedSegments: [Int64] = [] + + for hit in seeds.prefix(maxSeeds) { + visitedSegments.insert(hit.segmentID) + seedSegments.append(hit.segmentID) + let body = (try? text(ofSegment: hit.segmentID)) ?? hit.excerpt ?? "" + let clipped = String(body.prefix(1600)) + pieces.append(GraphContextPiece( + path: hit.path, kind: hit.kind.rawValue, segmentID: hit.segmentID, + loc: hit.loc, locEnd: hit.locEnd, hop: 0, linkedVia: nil, text: clipped + )) + } + + guard hops >= 1, !seedSegments.isEmpty else { return pieces } + let idList = seedSegments.map(String.init).joined(separator: ",") + let queryEntities = IndexStore.extractGraphEntities(from: query).map(\.key) + var entityKeys: [String] = queryEntities + + try store.db.query(""" + SELECT DISTINCT entity_key FROM edges + WHERE source_segment IN (\(idList)) + ORDER BY weight DESC LIMIT 16 + """) { row in + if let key = row.text(0), !entityKeys.contains(key) { + entityKeys.append(key) + } + } + + guard !entityKeys.isEmpty else { return pieces } + let placeholders = Array(repeating: "?", count: entityKeys.count).joined(separator: ",") + let args = entityKeys.map(SQLValue.text) + + var neighborRows: [(segID: Int64, path: String, kind: String, loc: Double?, locEnd: Double?, via: String)] = [] + try store.db.query(""" + SELECT e.source_segment, f.path, f.kind, s.loc, s.loc_end, GROUP_CONCAT(DISTINCT e.entity_key) AS via, + SUM(e.weight) AS total_weight + FROM edges e + JOIN segments s ON s.id = e.source_segment + JOIN files f ON f.id = e.source_file + WHERE e.entity_key IN (\(placeholders)) + AND e.source_segment NOT IN (\(idList)) + GROUP BY e.source_segment + ORDER BY total_weight DESC + LIMIT \(maxNeighbors * 2) + """, args) { row in + let segID = row.int(0) + guard visitedSegments.insert(segID).inserted else { return } + neighborRows.append(( + segID: segID, + path: row.text(1) ?? "", + kind: row.text(2) ?? "doc", + loc: row.isNull(3) ? nil : row.double(3), + locEnd: row.isNull(4) ? nil : row.double(4), + via: row.text(5) ?? "" + )) + } + + for neighbor in neighborRows.prefix(maxNeighbors) { + if let body = try? text(ofSegment: neighbor.segID), !body.isEmpty { + pieces.append(GraphContextPiece( + path: neighbor.path, kind: neighbor.kind, segmentID: neighbor.segID, + loc: neighbor.loc, locEnd: neighbor.locEnd, hop: 1, + linkedVia: neighbor.via, text: String(body.prefix(1200)) + )) + } + } + return pieces + } } diff --git a/Sources/LlamaRuntime/LlamaGenerator.swift b/Sources/LlamaRuntime/LlamaGenerator.swift new file mode 100644 index 00000000..a7fd173c --- /dev/null +++ b/Sources/LlamaRuntime/LlamaGenerator.swift @@ -0,0 +1,217 @@ +import CLlama +import DigUpKit +import Foundation + +/// Local Causal Small Language Model (SLM) Generator (`Qwen2.5-3B-Instruct-Q4_K_M` / `1.5B`) on `llama.cpp`, +/// engineered for an 8-core CPU, 12GB RAM PC: +/// - Uses `GGML_TYPE_Q8_0` KV-cache quantization to halve memory footprint (~350 MB for 8,192 tokens). +/// - Reserves 6 threads by default on an 8-core processor so OS/UI and `search-service` remain responsive. +/// - Supports Cache-Augmented Generation (CAG) via `llama_state_seq_save_file` and `llama_state_seq_load_file` +/// to persist pre-evaluated KV states of pinned files/contexts to disk. +/// - Supports hot-swappable `.gguf` LoRA adapters (`llama_adapter_lora_init`) trained from user `.jsonl` datasets. +public final class LlamaGenerator { + private var model: OpaquePointer? + private var context: OpaquePointer? + private var sampler: UnsafeMutablePointer? + private var loraAdapter: OpaquePointer? + private let vocab: OpaquePointer? + public let modelPath: URL + public private(set) var loraPath: URL? + public let contextSize: Int + private var nPast: llama_pos = 0 + + /// Optimal thread count for causal generation on an 8-core CPU (leaves 2 cores for OS + UI responsiveness). + public static let chatThreads: Int32 = { + if let chosen = ProcessInfo.processInfo.environment["DIGUP_CHAT_THREADS"].flatMap(Int32.init) { + return max(1, chosen) + } + let cores = ProcessInfo.processInfo.activeProcessorCount + return Int32(max(2, min(cores > 4 ? cores - 2 : cores, 6))) + }() + + public init(paths: LlamaModels.ChatModelPaths, contextSize: Int = 8192, log: URL? = nil) throws { + LlamaLog.shared.send(to: log) + llama_backend_init() + + var mparams = llama_model_default_params() + mparams.n_gpu_layers = LlamaEmbedder.cpuOnly ? 0 : 999 + mparams.use_mmap = true + guard let loadedModel = llama_model_load_from_file(paths.model.path, mparams) else { + throw WorkerError.failed("can't load chat SLM \(paths.model.path)") + } + + var cparams = llama_context_default_params() + cparams.n_ctx = UInt32(contextSize) + cparams.n_batch = 512 + cparams.n_ubatch = 512 + cparams.embeddings = false // Causal generation mode: enables last-token logits! + cparams.type_k = GGML_TYPE_Q8_0 + cparams.type_v = GGML_TYPE_Q8_0 + cparams.n_threads = Self.chatThreads + cparams.n_threads_batch = Self.chatThreads + + guard let loadedCtx = llama_init_from_model(loadedModel, cparams) else { + llama_model_free(loadedModel) + throw WorkerError.failed("can't create causal llama_context for \(paths.model.lastPathComponent)") + } + + let sparams = llama_sampler_chain_default_params() + let chain = llama_sampler_chain_init(sparams) + llama_sampler_chain_add(chain, llama_sampler_init_top_k(40)) + llama_sampler_chain_add(chain, llama_sampler_init_top_p(0.9, 1)) + llama_sampler_chain_add(chain, llama_sampler_init_temp(0.3)) + llama_sampler_chain_add(chain, llama_sampler_init_dist(42)) + + self.model = loadedModel + self.context = loadedCtx + self.sampler = chain + self.vocab = llama_model_get_vocab(loadedModel) + self.modelPath = paths.model + self.contextSize = contextSize + + if let lora = paths.loraAdapter { + try applyLoraAdapter(at: lora) + } + } + + deinit { close() } + + public func close() { + if let loraAdapter, let context { + llama_rm_adapter_lora(context, loraAdapter) + llama_adapter_lora_free(loraAdapter) + } + if let sampler { llama_sampler_free(sampler) } + if let context { llama_free(context) } + if let model { llama_model_free(model) } + (loraAdapter, sampler, context, model) = (nil, nil, nil, nil) + } + + // MARK: Hot-Swappable LoRA Adapter Support + + public func applyLoraAdapter(at url: URL, scale: Float = 1.0) throws { + guard let model, let context else { throw WorkerError.failed("generator is closed") } + if let existing = loraAdapter { + llama_rm_adapter_lora(context, existing) + llama_adapter_lora_free(existing) + loraAdapter = nil + loraPath = nil + } + guard let adapter = llama_adapter_lora_init(model, url.path) else { + throw WorkerError.failed("can't load LoRA adapter at \(url.path)") + } + guard llama_set_adapter_lora(context, adapter, scale) == 0 else { + llama_adapter_lora_free(adapter) + throw WorkerError.failed("can't attach LoRA adapter \(url.lastPathComponent)") + } + self.loraAdapter = adapter + self.loraPath = url + } + + // MARK: Cache-Augmented Generation (CAG) KV-State Persistence + + /// Pre-evaluates `prefixText` (e.g. system instructions + pinned documents/code) into sequence 0 and saves + /// the KV-cache state to `cacheFile` so future turns skip CPU prompt evaluation. + public func buildAndSaveCAGSlot(prefixText: String, to cacheFile: URL) throws -> (tokens: Int, ms: Int) { + guard let context else { throw WorkerError.failed("generator closed") } + let started = Date() + llama_kv_self_clear(context) + nPast = 0 + let tokens = tokenize(prefixText, addSpecial: true) + try decodeTokens(tokens) + try FileManager.default.createDirectory(at: cacheFile.deletingLastPathComponent(), withIntermediateDirectories: true) + let saved = tokens.withUnsafeBufferPointer { buf in + llama_state_seq_save_file(context, cacheFile.path, 0, buf.baseAddress, buf.count) + } + guard saved > 0 else { + throw WorkerError.failed("failed to serialize CAG KV state to \(cacheFile.path)") + } + let elapsed = Int(Date().timeIntervalSince(started) * 1000) + return (tokens.count, elapsed) + } + + /// Restores a previously saved CAG KV-cache slot from disk into sequence 0 in milliseconds. + public func loadCAGSlot(from cacheFile: URL) throws -> (tokens: Int, ms: Int) { + guard let context else { throw WorkerError.failed("generator closed") } + let started = Date() + llama_kv_self_clear(context) + var tokenOut = [llama_token](repeating: 0, count: contextSize) + var countOut: Int = 0 + let readBytes = tokenOut.withUnsafeMutableBufferPointer { buf in + llama_state_seq_load_file(context, cacheFile.path, 0, buf.baseAddress, buf.count, &countOut) + } + guard readBytes > 0 else { + throw WorkerError.failed("failed to load CAG KV state from \(cacheFile.path)") + } + nPast = llama_pos(countOut) + let elapsed = Int(Date().timeIntervalSince(started) * 1000) + return (countOut, elapsed) + } + + // MARK: Streaming Causal Generation + + public func generate(prompt: String, maxNewTokens: Int = 512, keepPrefix: Bool = false, + shouldStop: () -> Bool = { false }, + onToken: (String) -> Void) throws -> String { + guard let context, let sampler, let vocab else { throw WorkerError.failed("generator closed") } + if !keepPrefix { + llama_kv_self_clear(context) + nPast = 0 + } + let promptTokens = tokenize(prompt, addSpecial: nPast == 0) + try decodeTokens(promptTokens) + + var fullOutput = "" + for _ in 0.. Int32 in + let batch = llama_batch_get_one(buf.baseAddress, count) + return llama_decode(context, batch) + } + guard status == 0 else { + throw WorkerError.failed("llama_decode failed with status \(status)") + } + nPast += count + offset = end + } + } + + private func tokenize(_ text: String, addSpecial: Bool) -> [llama_token] { + let utf8 = Array(text.utf8CString) + var tokens = [llama_token](repeating: 0, count: utf8.count + 16) + var count = llama_tokenize(vocab, utf8, Int32(utf8.count - 1), &tokens, Int32(tokens.count), addSpecial, true) + if count < 0 { + tokens = [llama_token](repeating: 0, count: Int(-count)) + count = llama_tokenize(vocab, utf8, Int32(utf8.count - 1), &tokens, Int32(tokens.count), addSpecial, true) + } + tokens.removeLast(tokens.count - Int(max(count, 0))) + return tokens + } + + private func detokenize(_ token: llama_token) -> String { + var buf = [CChar](repeating: 0, count: 64) + let n = llama_token_to_piece(vocab, token, &buf, Int32(buf.count - 1), 0, true) + guard n > 0 else { return "" } + buf[Int(n)] = 0 + return String(cString: buf) + } +} diff --git a/Sources/LlamaRuntime/LlamaModels.swift b/Sources/LlamaRuntime/LlamaModels.swift index 9fb33295..00941eec 100644 --- a/Sources/LlamaRuntime/LlamaModels.swift +++ b/Sources/LlamaRuntime/LlamaModels.swift @@ -41,4 +41,39 @@ public struct LlamaModels: Sendable { public static var complete: Bool { (try? locate())?.media != nil } + + /// Locates the local generative SLM GGUF model and optional LoRA adapter (.gguf) for `chat-service`. + public struct ChatModelPaths: Sendable { + public let model: URL + public let loraAdapter: URL? + public init(model: URL, loraAdapter: URL? = nil) { + self.model = model + self.loraAdapter = loraAdapter + } + } + + public static func locateChat(customModel: String? = nil, customLora: String? = nil) throws -> ChatModelPaths { + let env = ProcessInfo.processInfo.environment + let loraURL: URL? = (customLora ?? env["DIGUP_LORA_ADAPTER"]).flatMap { path in + guard !path.isEmpty, FileManager.default.fileExists(atPath: path) else { return nil } + return URL(fileURLWithPath: path) + } + if let chosen = customModel ?? env["DIGUP_CHAT_MODEL"], !chosen.isEmpty { + let url = URL(fileURLWithPath: chosen) + guard FileManager.default.fileExists(atPath: url.path) else { + throw WorkerError.notFound("chat model not found at \(url.path)") + } + return ChatModelPaths(model: url, loraAdapter: loraURL) + } + for folder in folders { + for candidate in [ModelFiles.chatDefault.name, ModelFiles.chatTurbo.name] { + let path = folder.appendingPathComponent(candidate) + if FileManager.default.fileExists(atPath: path.path) { + return ChatModelPaths(model: path, loraAdapter: loraURL) + } + } + } + throw WorkerError.notFound("no chat GGUF model found in " + folders.map(\.path).joined(separator: ", ") + + "; download \(ModelFiles.chatDefault.name) via Settings or `digup download-model --chat`") + } } diff --git a/Sources/digup/ChatService.swift b/Sources/digup/ChatService.swift new file mode 100644 index 00000000..826b9fcc --- /dev/null +++ b/Sources/digup/ChatService.swift @@ -0,0 +1,183 @@ +import DigUpKit +import Foundation +import LlamaRuntime + +/// `digup chat-service`: Interactive RAG + CAG + GraphRAG chat process over `LlamaGenerator` and `IndexStore`. +/// +/// Protocol (JSON-lines over stdin/stdout): +/// {"op":"chat","id":…,"q":…,"code":bool,"cag_slot":?string,"max_tokens":512} +/// → streams {"event":"token","id":…,"token":…} and finishes with +/// {"id":…,"answer":…,"citations":[…],"cag_used":bool,"ms":{…}} +/// {"op":"cag_pin","id":…,"slot":string,"paths":[string]} +/// → pre-computes KV-cache state for pinned files into `/cag_cache/.kv` +/// {"op":"load_lora","id":…,"path":string} +/// → hot-swaps a fine-tuned `.gguf` LoRA adapter at runtime +/// {"op":"status","id":…} +/// → reports loaded SLM, active LoRA adapter, thread count, and CAG slots +final class ChatService: @unchecked Sendable { + private let indexFolder: URL + private let cagFolder: URL + private var generator: LlamaGenerator? + private let lock = NSLock() + private let output = NSLock() + + init(indexFolder: URL) throws { + self.indexFolder = indexFolder + self.cagFolder = indexFolder.appendingPathComponent("cag_cache") + try FileManager.default.createDirectory(at: cagFolder, withIntermediateDirectories: true) + } + + func serve() -> Int32 { + while let line = readLine() { + guard let data = line.data(using: .utf8), + let request = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { continue } + let id: Any = request["id"] ?? NSNull() + let op = request["op"] as? String ?? "" + if op == "quit" { + write(["id": id, "ok": true]) + return 0 + } + do { + var reply = try handle(op, request, id: id) + reply["id"] = id + write(reply) + } catch { + write(["id": id, "error": "\(error)"]) + } + } + return 0 + } + + private func ensureGenerator() throws -> LlamaGenerator { + if let generator { return generator } + let paths = try LlamaModels.locateChat() + let created = try LlamaGenerator( + paths: paths, + log: indexFolder.appendingPathComponent("chat.log") + ) + self.generator = created + return created + } + + private func handle(_ op: String, _ request: [String: Any], id: Any) throws -> [String: Any] { + switch op { + case "status": + let paths = try? LlamaModels.locateChat() + let slots = (try? FileManager.default.contentsOfDirectory(atPath: cagFolder.path))? + .filter { $0.hasSuffix(".kv") } + .map { String($0.dropLast(3)) } ?? [] + return [ + "ready": paths != nil, + "model": paths?.model.lastPathComponent ?? NSNull(), + "lora": (generator?.loraPath ?? paths?.loraAdapter)?.lastPathComponent ?? NSNull(), + "threads": LlamaGenerator.chatThreads, + "cag_slots": slots, + ] + + case "load_lora": + guard let path = request["path"] as? String, !path.isEmpty else { + throw CLIError("load_lora requires a path to a .gguf LoRA adapter") + } + let gen = try ensureGenerator() + try gen.applyLoraAdapter(at: URL(fileURLWithPath: path)) + return ["ok": true, "lora": path] + + case "cag_pin": + let slot = (request["slot"] as? String ?? "default").replacingOccurrences(of: "/", with: "_") + let paths = request["paths"] as? [String] ?? [] + let gen = try ensureGenerator() + let store = try IndexStore(directory: indexFolder) + let searcher = try Searcher(store: store, loadVectors: false) + var combined = "<|im_start|>system\nYou are DigUp Local AI, a 100% private assistant with access to the user's pinned local files.\n" + for filePath in paths { + var segIDs: [Int64] = [] + try store.db.query(""" + SELECT s.id FROM segments s JOIN files f ON f.id = s.file + WHERE f.path = ? AND s.kind != 'name' ORDER BY s.loc LIMIT 12 + """, [.text(filePath)]) { segIDs.append($0.int(0)) } + let snippets = segIDs.compactMap { try? searcher.text(ofSegment: $0) }.joined(separator: "\n") + if !snippets.isEmpty { + combined += "\n### Pinned File: \(filePath)\n\(String(snippets.prefix(2400)))\n" + } + } + combined += "<|im_end|>\n" + let kvFile = cagFolder.appendingPathComponent("\(slot).kv") + let stats = try gen.buildAndSaveCAGSlot(prefixText: combined, to: kvFile) + return ["ok": true, "slot": slot, "tokens": stats.tokens, "ms": stats.ms] + + case "chat": + let query = request["q"] as? String ?? "" + let isCode = request["code"] as? Bool ?? false + let maxTokens = request["max_tokens"] as? Int ?? 512 + let cagSlot = request["cag_slot"] as? String + let gen = try ensureGenerator() + + let store = try IndexStore(directory: indexFolder, file: isCode ? IndexStore.codeFile : IndexStore.mainFile) + let searcher = try Searcher(store: store, loadVectors: false) + let hits = try searcher.search(query, queryVector: nil, limit: 6) + let graphPieces = try searcher.retrieveGraphContext(query: query, seeds: hits, hops: 1, maxSeeds: 4, maxNeighbors: 3) + + var cagLoaded = false + if let slot = cagSlot, !slot.isEmpty { + let kvFile = cagFolder.appendingPathComponent("\(slot).kv") + if FileManager.default.fileExists(atPath: kvFile.path) { + _ = try? gen.loadCAGSlot(from: kvFile) + cagLoaded = true + } + } + + var contextBlock = "" + var citations: [[String: Any]] = [] + for (idx, piece) in graphPieces.enumerated() { + let fileName = (piece.path as NSString).lastPathComponent + let locLabel = piece.loc.map { "loc:\(Int($0))" } ?? "full" + contextBlock += "[\(idx + 1)] \(fileName) (\(locLabel), hop=\(piece.hop)):\n\(piece.text)\n\n" + var cite: [String: Any] = [ + "index": idx + 1, "path": piece.path, "kind": piece.kind, + "segment_id": piece.segmentID, "hop": piece.hop, + ] + if let loc = piece.loc { cite["loc"] = loc } + if let via = piece.linkedVia { cite["linked_via"] = via } + citations.append(cite) + } + + let prompt: String + if cagLoaded { + prompt = "<|im_start|>user\nRetrieved Context:\n\(contextBlock)\nQuestion: \(query)\nCite sources using [1], [2].<|im_end|>\n<|im_start|>assistant\n" + } else { + prompt = """ + <|im_start|>system + You are DigUp Local AI, a 100% private offline assistant running on the user's PC. Answer strictly using the retrieved local files and GraphRAG context below, citing sources with [1], [2].<|im_end|> + <|im_start|>user + Retrieved Context: + \(contextBlock) + Question: \(query)<|im_end|> + <|im_start|>assistant + """ + } + + let started = Date() + let answer = try gen.generate(prompt: prompt, maxNewTokens: maxTokens, keepPrefix: cagLoaded) { tokenPiece in + self.write(["event": "token", "id": id, "token": tokenPiece]) + } + let elapsed = Int(Date().timeIntervalSince(started) * 1000) + return [ + "answer": answer, + "citations": citations, + "cag_used": cagLoaded, + "ms": ["generate": elapsed], + ] + + default: + throw CLIError("unknown chat-service op \(op)") + } + } + + private func write(_ fields: [String: Any]) { + output.withLock { emit(fields) } + } +} + +func chatService(_ args: Arguments) throws -> Int32 { + try ChatService(indexFolder: args.indexFolder).serve() +} diff --git a/Sources/digup/DownloadModel.swift b/Sources/digup/DownloadModel.swift index 1a552511..42874481 100644 --- a/Sources/digup/DownloadModel.swift +++ b/Sources/digup/DownloadModel.swift @@ -16,8 +16,15 @@ import Foundation func downloadModel(_ args: Arguments) throws -> Int32 { let folder = args.values["--to"].map(Arguments.folderURL) ?? ModelFiles.installed try FileManager.default.createDirectory(at: folder, withIntermediateDirectories: true) - let missing = ModelFiles.missing(in: folder) - var done = ModelFiles.totalBytes - missing.reduce(0) { $0 + $1.bytes } + let targetFiles: [ModelFiles.File] + if args.flags.contains("--chat") { + targetFiles = [args.flags.contains("--turbo") ? ModelFiles.chatTurbo : ModelFiles.chatDefault] + } else { + targetFiles = ModelFiles.all + } + let totalBytes = targetFiles.reduce(0) { $0 + $1.bytes } + let missing = targetFiles.filter { !FileManager.default.fileExists(atPath: folder.appendingPathComponent($0.name).path) } + var done = totalBytes - missing.reduce(0) { $0 + $1.bytes } for file in missing { let final = folder.appendingPathComponent(file.name) if let cached = ModelFiles.cachedCopy(of: file), (try? ModelFiles.sha256(of: cached)) == file.sha256 { @@ -32,7 +39,7 @@ func downloadModel(_ args: Arguments) throws -> Int32 { emit(["event": "file", "name": file.name, "bytes": file.bytes, "from": "download"]) var failures = 0 while true { - switch fetch(file, into: folder, before: done) { + switch fetch(file, into: folder, before: done, totalBytes: totalBytes) { case .downloaded(let partial): emit(["event": "verifying", "name": file.name]) guard try ModelFiles.sha256(of: partial) == file.sha256 else { @@ -67,12 +74,12 @@ func downloadModel(_ args: Arguments) throws -> Int32 { } /// One try at one file: blocks until `FileFetcher` is through with it. -private func fetch(_ file: ModelFiles.File, into folder: URL, before done: Int64) -> FileFetcher.Outcome { +private func fetch(_ file: ModelFiles.File, into folder: URL, before done: Int64, totalBytes: Int64) -> FileFetcher.Outcome { let finished = DispatchSemaphore(value: 0) let result = OutcomeBox() let fetcher = FileFetcher(file: file, folder: folder, progress: { written in emit(["event": "progress", "name": file.name, "bytes": written, "total": file.bytes, - "all_bytes": done + written, "all_total": ModelFiles.totalBytes]) + "all_bytes": done + written, "all_total": totalBytes]) }, waiting: { emit(["event": "waiting"]) }, done: { outcome in diff --git a/Sources/digup/SearchService.swift b/Sources/digup/SearchService.swift index fca7b36f..363119f8 100644 --- a/Sources/digup/SearchService.swift +++ b/Sources/digup/SearchService.swift @@ -175,6 +175,31 @@ final class SearchService: @unchecked Sendable { // its state is guarded by `l guard let segment = (request["segment"] as? NSNumber)?.int64Value else { throw CLIError("which segment?") } let code = request["code"] as? Bool ?? false return ["text": try (code ? codeSearcher : searcher)?.text(ofSegment: segment) ?? NSNull()] + case "graph_context": + let text = request["q"] as? String ?? "" + let code = request["code"] as? Bool ?? false + let hops = request["hops"] as? Int ?? 1 + let maxSeeds = request["max_seeds"] as? Int ?? 5 + let maxNeighbors = request["max_neighbors"] as? Int ?? 4 + let activeSearcher = code ? codeSearcher : searcher + guard let activeSearcher else { return ["pieces": []] } + let vector = text.isEmpty ? nil : self.vector(for: text, code: code) + let hits = try activeSearcher.search(text, queryVector: vector, limit: maxSeeds * 2) + let pieces = try activeSearcher.retrieveGraphContext( + query: text, seeds: hits, hops: hops, maxSeeds: maxSeeds, maxNeighbors: maxNeighbors + ) + return [ + "pieces": pieces.map { piece in + var item: [String: Any] = [ + "path": piece.path, "kind": piece.kind, "segment_id": piece.segmentID, + "hop": piece.hop, "text": piece.text, + ] + if let loc = piece.loc { item["loc"] = loc } + if let locEnd = piece.locEnd { item["loc_end"] = locEnd } + if let via = piece.linkedVia { item["linked_via"] = via } + return item + } + ] case "reload": let hadVectors = searcher.vectorsLoaded searcher = try Searcher(store: store, loadVectors: hadVectors) diff --git a/Sources/digup/main.swift b/Sources/digup/main.swift index 38c3ba2c..1b81410d 100644 --- a/Sources/digup/main.swift +++ b/Sources/digup/main.swift @@ -696,7 +696,7 @@ do { exit(0) } let args = try Arguments(raw.dropFirst()) - if ["index-service", "worker", "search-service", "download-model"].contains(command) + if ["index-service", "worker", "search-service", "chat-service", "download-model"].contains(command) || command == "estimate" && args.flags.contains("--json") { ParentWatch.start() // an app's helper goes when the app does } @@ -729,6 +729,7 @@ do { #endif case "worker": exit(workerService(args)) case "search-service": exit(try searchService(args)) + case "chat-service": exit(try chatService(args)) case "download-model": exit(try downloadModel(args)) default: print(usage) diff --git a/Tests/DigUpKitTests/GraphRAGTests.swift b/Tests/DigUpKitTests/GraphRAGTests.swift new file mode 100644 index 00000000..cc8d19c0 --- /dev/null +++ b/Tests/DigUpKitTests/GraphRAGTests.swift @@ -0,0 +1,45 @@ +import Foundation +import Testing +@testable import DigUpKit + +@Suite struct GraphRAGTests { + let folder: URL = { + let url = FileManager.default.temporaryDirectory.appendingPathComponent("DigUpGraphTests-\(UUID().uuidString)") + try? FileManager.default.createDirectory(at: url, withIntermediateDirectories: true) + return url + }() + + @Test func extractsFileRefsIdCodesAndSymbols() { + let snippet = "In billing/auth.ts, function refreshAccessToken failed with ERR_TOKEN_401 on invoice INV-2026-041." + let entities = IndexStore.extractGraphEntities(from: snippet) + let keys = Set(entities.map(\.key)) + #expect(keys.contains("auth.ts")) + #expect(keys.contains("refreshaccesstoken")) + #expect(keys.contains("err_token_401")) + #expect(keys.contains("inv-2026-041")) + } + + @Test func storesAndTraversesGraphEdgesAcrossFiles() throws { + let store = try IndexStore(directory: folder) + let fileA = try store.insert(Candidate(path: "/repo/serviceA.py", kind: .code, size: 100, mtime: 1, inode: 1, device: 1)) + let fileB = try store.insert(Candidate(path: "/repo/serviceB.py", kind: .code, size: 100, mtime: 1, inode: 2, device: 1)) + + let vec = [Float](repeating: 0.1, count: 768) + try store.writeSegments( + fileId: fileA, + segments: [NewSegment(kind: .lines, loc: 1, locEnd: 20, text: "Calls validatePaymentSignature on invoice INV-2026-099", vector: vec)], + state: "done", + info: [:] + ) + try store.writeSegments( + fileId: fileB, + segments: [NewSegment(kind: .lines, loc: 10, locEnd: 40, text: "def validatePaymentSignature(payload): # handles INV-2026-099", vector: vec)], + state: "done", + info: [:] + ) + + let neighbors = try store.graphNeighbors(forFileIds: [fileA], limit: 5) + #expect(!neighbors.isEmpty) + #expect(neighbors.first?.targetPath == "/repo/serviceB.py") + } +} diff --git a/scripts/build-llama.sh b/scripts/build-llama.sh index a904cdbd..b39196a5 100755 --- a/scripts/build-llama.sh +++ b/scripts/build-llama.sh @@ -132,5 +132,8 @@ cp "$SOURCE/include/llama.h" "$SOURCE"/ggml/include/*.h "$SOURCE/tools/mtmd/mtmd # changed header, so the Swift file that reads it is touched. printf '// Written by scripts/build-llama.sh.\n#define LLAMA_BUILD_TAG "%s"\n' "$TAG" > "$OUT/include/llama-build.h" touch Sources/LlamaRuntime/LlamaEmbedder.swift +if [[ -f Sources/LlamaRuntime/LlamaGenerator.swift ]]; then + touch Sources/LlamaRuntime/LlamaGenerator.swift +fi echo "$STAMP" > "$OUT/BUILD" echo "Built llama.cpp $TAG (patched, $FLAVOR${CC:+, $CC}) into $OUT" diff --git a/scripts/train_lora.py b/scripts/train_lora.py new file mode 100644 index 00000000..f676efdc --- /dev/null +++ b/scripts/train_lora.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +""" +DigUp Local SLM Fine-Tuning Utility (`scripts/train_lora.py`) +============================================================= +Consumes the `.jsonl` files generated from DigUp user chat interactions: + - `sft_train.jsonl` (ChatML Supervised Fine-Tuning format) + - `dpo_train.jsonl` (Direct Preference Optimization format from user-edited answers) + +Designed for local execution on an 8-core, 12GB RAM PC (CPU / iGPU or low-VRAM GPU): + - Uses LoRA (r=8, alpha=16) on `Qwen/Qwen2.5-3B-Instruct` or `Qwen/Qwen2.5-1.5B-Instruct` + - Uses gradient checkpointing and `batch_size=1` (`gradient_accumulation_steps=4`) to stay under 8 GB RAM + - Automatically converts the trained LoRA adapter into `.gguf` (`digup-user-lora.gguf`) so + `digup chat-service` can hot-load it via `llama_adapter_lora_init` without duplicating the base model. + +Usage: + python scripts/train_lora.py --jsonl "%LOCALAPPDATA%/DigUp/index/datasets/sft_train.jsonl" \ + --base-model "Qwen/Qwen2.5-3B-Instruct" \ + --output-dir "./lora_out" +""" + +import argparse +import json +import os +import subprocess +import sys +from pathlib import Path + + +def validate_jsonl(path: Path) -> list[dict]: + if not path.is_file(): + raise FileNotFoundError(f"JSONL dataset not found: {path}") + records = [] + with path.open("r", encoding="utf-8") as f: + for line_no, line in enumerate(f, start=1): + line = line.strip() + if not line: + continue + obj = json.loads(line) + if "messages" not in obj and "prompt" not in obj: + raise ValueError(f"Line {line_no} in {path} is missing 'messages' or 'prompt' key") + records.append(obj) + if not records: + raise ValueError(f"JSONL dataset at {path} contains 0 training records") + return records + + +def run_sft_lora_training(jsonl_path: Path, base_model: str, output_dir: Path, epochs: int, lr: float) -> Path: + import torch + from datasets import load_dataset + from peft import LoraConfig, get_peft_model + from transformers import AutoModelForCausalLM, AutoTokenizer + from trl import SFTConfig, SFTTrainer + + use_cuda = torch.cuda.is_available() + device_label = "CUDA GPU" if use_cuda else f"CPU ({os.cpu_count() or 8} threads)" + print(f"[DigUp LoRA] Training on {device_label} | Base model: {base_model}") + + if not use_cuda: + torch.set_num_threads(max(2, min((os.cpu_count() or 8) - 2, 6))) + + tokenizer = AutoTokenizer.from_pretrained(base_model, use_fast=True) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + dtype = torch.float16 if use_cuda else torch.float32 + model = AutoModelForCausalLM.from_pretrained( + base_model, + torch_dtype=dtype, + low_cpu_mem_usage=True, + ) + model.enable_input_require_grads() + + lora_config = LoraConfig( + r=8, + lora_alpha=16, + lora_dropout=0.05, + bias="none", + task_type="CAUSAL_LM", + target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + ) + model = get_peft_model(model, lora_config) + model.print_trainable_parameters() + + dataset = load_dataset("json", data_files=str(jsonl_path), split="train") + + sft_args = SFTConfig( + output_dir=str(output_dir), + num_train_epochs=epochs, + per_device_train_batch_size=1, + gradient_accumulation_steps=4, + learning_rate=lr, + logging_steps=1, + save_strategy="no", + gradient_checkpointing=True, + use_cpu=not use_cuda, + fp16=use_cuda, + max_seq_length=1024, + report_to="none", + ) + + trainer = SFTTrainer( + model=model, + processing_class=tokenizer, + train_dataset=dataset, + args=sft_args, + ) + trainer.train() + + adapter_dir = output_dir / "adapter" + adapter_dir.mkdir(parents=True, exist_ok=True) + model.save_pretrained(str(adapter_dir)) + tokenizer.save_pretrained(str(adapter_dir)) + print(f"[DigUp LoRA] Saved HuggingFace PEFT adapter to {adapter_dir}") + + # Convert PEFT LoRA adapter to GGUF if llama.cpp converter is present in Vendor/llama.cpp + converter = Path(__file__).resolve().parent.parent / "Vendor" / "llama.cpp" / "convert_lora_to_gguf.py" + gguf_out = output_dir / "digup-user-lora.gguf" + if converter.is_file(): + print(f"[DigUp LoRA] Converting PEFT adapter to GGUF: {gguf_out}") + subprocess.run( + [sys.executable, str(converter), "--outfile", str(gguf_out), str(adapter_dir)], + check=True, + ) + else: + print( + f"[DigUp LoRA] Note: To convert to GGUF for llama.cpp hot-loading, run:\n" + f" python Vendor/llama.cpp/convert_lora_to_gguf.py --outfile {gguf_out} {adapter_dir}" + ) + return adapter_dir + + +def main() -> int: + parser = argparse.ArgumentParser(description="DigUp Local SLM JSONL LoRA Fine-Tuner") + parser.add_argument("--jsonl", required=True, type=Path, help="Path to sft_train.jsonl") + parser.add_argument("--base-model", default="Qwen/Qwen2.5-3B-Instruct", help="HuggingFace model ID or local path") + parser.add_argument("--output-dir", default=Path("./lora_out"), type=Path, help="Directory for trained LoRA") + parser.add_argument("--epochs", default=3, type=int) + parser.add_argument("--lr", default=2e-4, type=float) + parser.add_argument("--validate-only", action="store_true", help="Only validate the JSONL dataset and exit") + args = parser.parse_args() + + records = validate_jsonl(args.jsonl) + print(f"[DigUp LoRA] Validated {len(records)} training example(s) in {args.jsonl}") + if args.validate_only: + return 0 + + run_sft_lora_training(args.jsonl, args.base_model, args.output_dir, args.epochs, args.lr) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/shell/src-tauri/src/dataset.rs b/shell/src-tauri/src/dataset.rs new file mode 100644 index 00000000..59064cfc --- /dev/null +++ b/shell/src-tauri/src/dataset.rs @@ -0,0 +1,342 @@ +// Local JSONL Fine-Tuning Dataset Logger & Secret Scrubber for DigUp + SLM. +// Generates standard ChatML SFT (`sft_train.jsonl`) and DPO (`dpo_train.jsonl`) files from user chat interactions, +// enabling local domain adaptation while stripping sensitive keys/tokens before persistence. + +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use std::fs::{self, OpenOptions}; +use std::io::{BufRead, BufReader, Write}; +use std::path::{Path, PathBuf}; +use std::time::{SystemTime, UNIX_EPOCH}; + +pub const SFT_FILE: &str = "sft_train.jsonl"; +pub const DPO_FILE: &str = "dpo_train.jsonl"; +pub const SYSTEM_PROMPT: &str = "You are DigUp Local AI, a private assistant grounded in the user's local files, code, and system architecture."; + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SftRecordInput { + pub question: String, + pub context: Option, + pub answer: String, + pub edited_answer: Option, + pub citations: Vec, + pub cag_slot: Option, + pub include_context_in_prompt: bool, +} + +#[derive(Clone, Debug, Serialize, Deserialize, Default)] +#[serde(rename_all = "camelCase")] +pub struct DatasetStats { + pub sft_count: usize, + pub dpo_count: usize, + pub sft_bytes: u64, + pub dpo_bytes: u64, + pub dataset_dir: String, +} + +/// Deterministic secret & credential scrubber so API keys, private keys, or passwords in local files +/// never leak into `.jsonl` fine-tuning datasets. +pub fn scrub_secrets(input: &str) -> String { + let mut out = input.to_string(); + + // 1. Redact PEM private key blocks + while let Some(start) = out.find("-----BEGIN ") { + if let Some(priv_pos) = out[start..].find("PRIVATE KEY-----") { + let header_end = start + priv_pos + "PRIVATE KEY-----".len(); + if let Some(end_marker) = out[header_end..].find("-----END ") { + if let Some(footer_close) = out[header_end + end_marker..].find("-----") { + let full_end = header_end + end_marker + footer_close + 5; + let after_end = out[full_end..] + .find("-----") + .map(|idx| full_end + idx + 5) + .unwrap_or(full_end); + out.replace_range(start..after_end.min(out.len()), "[REDACTED_PRIVATE_KEY]"); + continue; + } + } + } + break; + } + + // 2. Token-level scrubbing for common prefixes (sk-, ghp_, gho_, github_pat_, AKIA, AIza, xoxb-) + // and key=value assignments (password=, api_key=, secret=, token=). + let prefixes = [ + "sk-", + "ghp_", + "gho_", + "github_pat_", + "AKIA", + "AIza", + "xoxb-", + "xoxp-", + ]; + let sensitive_keys = [ + "password=", + "passwd=", + "api_key=", + "apikey=", + "secret=", + "access_token=", + "auth_token=", + ]; + + let mut scrubbed_lines = Vec::new(); + for line in out.lines() { + let mut words: Vec = Vec::new(); + let mut redact_next = false; + for token in line.split_whitespace() { + if redact_next { + words.push("[REDACTED_TOKEN]".to_string()); + redact_next = false; + continue; + } + if token.eq_ignore_ascii_case("Bearer") { + words.push(token.to_string()); + redact_next = true; + continue; + } + let clean = token.trim_matches(|c: char| c == '"' || c == '\'' || c == ',' || c == ';'); + let lower = clean.to_ascii_lowercase(); + + if prefixes + .iter() + .any(|p| clean.starts_with(p) && clean.len() >= p.len() + 8) + { + words.push("[REDACTED_SECRET]".to_string()); + } else if let Some(key) = sensitive_keys.iter().find(|k| lower.starts_with(*k)) { + let key_part = &clean[..key.len()]; + words.push(format!("{key_part}[REDACTED]")); + } else { + words.push(token.to_string()); + } + } + if line.is_empty() { + scrubbed_lines.push(String::new()); + } else { + scrubbed_lines.push(words.join(" ")); + } + } + + scrubbed_lines.join("\n") +} + +pub fn dataset_dir(base_dir: &Path) -> PathBuf { + base_dir.join("datasets") +} + +fn unix_timestamp() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + +/// Logs an interaction into `sft_train.jsonl` (ChatML `messages` format) and, if the user edited/corrected +/// the model's answer, also logs a preference pair into `dpo_train.jsonl`. +pub fn log_interaction(base_dir: &Path, input: &SftRecordInput) -> Result { + let dir = dataset_dir(base_dir); + fs::create_dir_all(&dir).map_err(|e| format!("can't create dataset folder: {e}"))?; + + let clean_q = scrub_secrets(input.question.trim()); + if clean_q.is_empty() { + return Err("question cannot be empty".into()); + } + let clean_ctx = input + .context + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(scrub_secrets); + + let chosen_raw = input + .edited_answer + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| input.answer.trim()); + let clean_chosen = scrub_secrets(chosen_raw); + if clean_chosen.is_empty() { + return Err("answer cannot be empty".into()); + } + + let user_content = match (&clean_ctx, input.include_context_in_prompt) { + (Some(ctx), true) => format!("Retrieved Local Context:\n{ctx}\n\nQuestion:\n{clean_q}"), + _ => clean_q.clone(), + }; + + let system_prompt = SYSTEM_PROMPT; + let ts = unix_timestamp(); + + let sft_entry = json!({ + "messages": [ + { "role": "system", "content": system_prompt }, + { "role": "user", "content": user_content }, + { "role": "assistant", "content": clean_chosen } + ], + "metadata": { + "timestamp": ts, + "cag_slot": input.cag_slot, + "citations": input.citations, + "user_edited": input.edited_answer.is_some() + } + }); + + append_jsonl_line(&dir.join(SFT_FILE), &sft_entry)?; + + // If the user corrected the model's original response, also emit a DPO preference training example. + if let Some(edited) = input + .edited_answer + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + { + let orig = input.answer.trim(); + if !orig.is_empty() && edited != orig { + let clean_rejected = scrub_secrets(orig); + let dpo_entry = json!({ + "prompt": [ + { "role": "system", "content": system_prompt }, + { "role": "user", "content": user_content } + ], + "chosen": [ + { "role": "assistant", "content": clean_chosen } + ], + "rejected": [ + { "role": "assistant", "content": clean_rejected } + ], + "metadata": { + "timestamp": ts, + "cag_slot": input.cag_slot + } + }); + append_jsonl_line(&dir.join(DPO_FILE), &dpo_entry)?; + } + } + + stats(base_dir) +} + +fn append_jsonl_line(file_path: &Path, value: &Value) -> Result<(), String> { + let mut line = + serde_json::to_string(value).map_err(|e| format!("JSON serialization failed: {e}"))?; + line.push('\n'); + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(file_path) + .map_err(|e| format!("can't open {}: {e}", file_path.display()))?; + file.write_all(line.as_bytes()) + .map_err(|e| format!("can't write to {}: {e}", file_path.display())) +} + +pub fn stats(base_dir: &Path) -> Result { + let dir = dataset_dir(base_dir); + let sft_path = dir.join(SFT_FILE); + let dpo_path = dir.join(DPO_FILE); + let count_lines = |p: &Path| -> usize { + let Ok(file) = fs::File::open(p) else { + return 0; + }; + BufReader::new(file) + .lines() + .filter_map(Result::ok) + .filter(|l| !l.trim().is_empty()) + .count() + }; + let file_size = |p: &Path| -> u64 { fs::metadata(p).map(|m| m.len()).unwrap_or(0) }; + + Ok(DatasetStats { + sft_count: count_lines(&sft_path), + dpo_count: count_lines(&dpo_path), + sft_bytes: file_size(&sft_path), + dpo_bytes: file_size(&dpo_path), + dataset_dir: dir.display().to_string(), + }) +} + +pub fn recent_entries(base_dir: &Path, limit: usize) -> Vec { + let sft_path = dataset_dir(base_dir).join(SFT_FILE); + let Ok(file) = fs::File::open(sft_path) else { + return Vec::new(); + }; + let all: Vec = BufReader::new(file) + .lines() + .filter_map(Result::ok) + .filter_map(|line| serde_json::from_str::(&line).ok()) + .collect(); + all.into_iter().rev().take(limit).collect() +} + +pub fn export_dataset(base_dir: &Path, destination: &Path, dpo: bool) -> Result { + let src = dataset_dir(base_dir).join(if dpo { DPO_FILE } else { SFT_FILE }); + if !src.is_file() { + return Err("no dataset entries recorded yet".into()); + } + if let Some(parent) = destination.parent() { + let _ = fs::create_dir_all(parent); + } + fs::copy(&src, destination).map_err(|e| format!("failed to export JSONL: {e}"))?; + Ok(destination.display().to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::env; + + #[test] + fn scrubs_api_keys_bearer_tokens_and_pem_blocks() { + let raw = "Connect with Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9 and key sk-1234567890abcdef1234 or password=mySuperSecret123\n-----BEGIN RSA PRIVATE KEY-----\nMIIEowIBAAKCAQEA...\n-----END RSA PRIVATE KEY-----"; + let scrubbed = scrub_secrets(raw); + assert!(!scrubbed.contains("eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9")); + assert!(!scrubbed.contains("sk-1234567890abcdef1234")); + assert!(!scrubbed.contains("mySuperSecret123")); + assert!(!scrubbed.contains("MIIEowIBAAKCAQEA")); + assert!(scrubbed.contains("[REDACTED_TOKEN]")); + assert!(scrubbed.contains("[REDACTED_SECRET]")); + assert!(scrubbed.contains("password=[REDACTED]")); + assert!(scrubbed.contains("[REDACTED_PRIVATE_KEY]")); + } + + #[test] + fn logs_sft_and_dpo_jsonl_entries() { + let tmp = env::temp_dir().join(format!("digup_jsonl_test_{}", unix_timestamp())); + let _ = fs::remove_dir_all(&tmp); + + let input = SftRecordInput { + question: "Which Redis port does billing-service use?".into(), + context: Some( + "billing-service/config/settings.yaml: redis_port: 6379, api_key=secret999".into(), + ), + answer: "It connects on port 6380.".into(), + edited_answer: Some( + "billing-service connects to Redis on port 6379 (see settings.yaml).".into(), + ), + citations: vec![json!({"path": "billing-service/config/settings.yaml", "loc": 1})], + cag_slot: Some("billing-core".into()), + include_context_in_prompt: true, + }; + + let st = log_interaction(&tmp, &input).expect("log_interaction should succeed"); + assert_eq!(st.sft_count, 1); + assert_eq!(st.dpo_count, 1); + assert!(st.sft_bytes > 50); + assert!(st.dpo_bytes > 50); + + let recent = recent_entries(&tmp, 5); + assert_eq!(recent.len(), 1); + let messages = recent[0]["messages"].as_array().unwrap(); + assert_eq!(messages.len(), 3); + assert!(messages[1]["content"] + .as_str() + .unwrap() + .contains("api_key=[REDACTED]")); + assert!(messages[2]["content"] + .as_str() + .unwrap() + .contains("port 6379")); + + let _ = fs::remove_dir_all(&tmp); + } +} diff --git a/shell/src-tauri/src/main.rs b/shell/src-tauri/src/main.rs index 375fbee2..4daef97f 100644 --- a/shell/src-tauri/src/main.rs +++ b/shell/src-tauri/src/main.rs @@ -7,6 +7,7 @@ mod clipboard; mod code; +mod dataset; mod editors; mod helper; mod hotkey; @@ -14,6 +15,7 @@ mod hours; mod launch; mod moments; mod options; +mod rag_engine; mod settings; mod status; mod tray_icon; @@ -60,6 +62,7 @@ struct Shell { /// The panel's window now: PANEL, or CENTERED_PANEL when it fills the screen. panel_label: Mutex<&'static str>, search: SearchService, + chat: SearchService, index: IndexService, model: ModelDownload, /// The model for the setup window and the tray: {"state": "checking|downloading|waiting|ready|failed", …}. @@ -98,6 +101,8 @@ struct Shell { /// The speed check is running (`measure_then_index`): indexing waits for it, so that the two don't share the /// processor. measuring: AtomicBool, + /// Local SLM inference (`chat-service`) is actively generating tokens: background indexing yields CPU & RAM. + chat_generating: AtomicBool, /// The model download was run once more after it ended without a final word (model_event). model_rerun: AtomicBool, /// Folder estimates running (setup's list): the speed check waits for them, for the same reason. @@ -259,10 +264,281 @@ fn folder_arguments(settings: &Settings, code_folders: &[String]) -> Vec #[tauri::command] async fn ask(shell: State<'_, Arc>, request: Value) -> Result { let op = request.get("op").and_then(Value::as_str).unwrap_or_default(); - if !["search", "embed_ahead", "text", "render", "prewarm", "status", "summaries"].contains(&op) { + if !["search", "embed_ahead", "text", "render", "prewarm", "status", "summaries", "graph_context", "lines"].contains(&op) { return Err(format!("not a question the panel asks: {op}")); } - shell.search.ask(request).await + if shell.options.helper.is_file() { + if let Ok(reply) = shell.search.ask(request.clone()).await { + return Ok(reply); + } + } + let (skipped, skip_types) = { + let settings = shell.settings.lock().unwrap(); + (settings.skipped.clone(), settings.skip_types.clone()) + }; + Ok(rag_engine::local_search_op( + &request, + &shell.folders(), + &shell.code_folders(), + &skipped, + &skip_types, + )) +} + +fn data_base_dir(shell: &Shell) -> PathBuf { + shell.options.index_folder().unwrap_or_else(|| shell.logs.clone()) +} + +/// Interactive Local SLM + Multi-Hop GraphRAG + CAG chat command. +/// Automatically yields background indexing while the SLM generates on the 8-core CPU / 12GB RAM budget. +#[tauri::command] +async fn chat_ask( + app: AppHandle, + shell: State<'_, Arc>, + query: String, + code: bool, + cag_slot: Option, + history: Vec, +) -> Result { + let trimmed = query.trim().to_string(); + if trimmed.is_empty() { + return Err("Question cannot be empty".into()); + } + + // Pause background indexing while SLM inference is active so CPU threads and RAM stay within budget. + shell.chat_generating.store(true, Ordering::SeqCst); + shell.index.stop(); + + let base_dir = data_base_dir(&shell); + let active_slot = cag_slot + .filter(|s| !s.trim().is_empty()) + .or_else(|| { + let s = shell.settings.lock().unwrap().active_cag_slot.clone(); + (!s.is_empty()).then_some(s) + }); + let cag_meta = active_slot.as_deref().and_then(|slot| rag_engine::load_cag_slot(&base_dir, slot)); + + // Try native `digup chat-service` first (streams tokens via `chat` event when Qwen2.5 GGUF is downloaded). + if shell.options.helper.is_file() { + let chat_req = json!({ + "op": "ask", + "q": trimmed, + "code": code, + "cag_slot": active_slot, + "history": history, + }); + if let Ok(reply) = shell.chat.ask(chat_req).await { + shell.chat_generating.store(false, Ordering::SeqCst); + later(&app, run_index); + return Ok(reply); + } + } + + let (skipped, skip_types) = { + let settings = shell.settings.lock().unwrap(); + (settings.skipped.clone(), settings.skip_types.clone()) + }; + + // Fallback: perform Multi-Hop GraphRAG over `search-service` (or native Rust local search) + CAG slot injection. + let hop0_req = json!({"op": "search", "q": trimmed, "code": code, "meaning": true, "limit": 8}); + let hop0_reply = if shell.options.helper.is_file() { + shell.search.ask(hop0_req.clone()).await.unwrap_or_else(|_| { + rag_engine::local_search_op(&hop0_req, &shell.folders(), &shell.code_folders(), &skipped, &skip_types) + }) + } else { + rag_engine::local_search_op(&hop0_req, &shell.folders(), &shell.code_folders(), &skipped, &skip_types) + }; + let hits = hop0_reply["hits"].as_array().cloned().unwrap_or_default(); + let hop0_excerpts: Vec = hits + .iter() + .take(5) + .filter_map(|h| h["excerpt"].as_str().map(String::from)) + .collect(); + + let mut hop1_hits = Vec::new(); + if let Some(expansion_q) = rag_engine::build_hop2_expansion_query(&trimmed, &hop0_excerpts) { + let hop1_req = json!({"op": "search", "q": expansion_q, "code": code, "meaning": true, "limit": 4}); + let hop1_reply = if shell.options.helper.is_file() { + shell.search.ask(hop1_req.clone()).await.unwrap_or_else(|_| { + rag_engine::local_search_op(&hop1_req, &shell.folders(), &shell.code_folders(), &skipped, &skip_types) + }) + } else { + rag_engine::local_search_op(&hop1_req, &shell.folders(), &shell.code_folders(), &skipped, &skip_types) + }; + if let Some(arr) = hop1_reply["hits"].as_array() { + hop1_hits = arr.clone(); + } + } + + let (fallback_answer, citations) = + rag_engine::synthesize_grounded_fallback(&trimmed, &hits, &hop1_hits, cag_meta.as_ref()); + + let mut context_lines = Vec::new(); + if let Some(ref cag) = cag_meta { + if !cag.cached_text.is_empty() { + context_lines.push(format!("[CAG Slot: {}]\n{}", cag.slot, cag.cached_text)); + } + } + for c in &citations { + context_lines.push(format!( + "[{}] {} ({}):\n{}", + c["index"], + c["path"].as_str().unwrap_or_default(), + c["kind"].as_str().unwrap_or_default(), + c["excerpt"] + .as_str() + .unwrap_or_default() + .replace('«', "") + .replace('»', "") + )); + } + let assembled_context = context_lines.join("\n\n"); + + // If a local OpenAI-compatible SLM server (`llama-server` on :8080/:8081 or `ollama` on :11434) is running, + // generate a true synthesized answer grounded in `assembled_context`. + let system_prompt = dataset::SYSTEM_PROMPT; + let user_prompt = format!( + "### Retrieved Multimodal & GraphRAG Context:\n{}\n\n### User Question:\n{}\n\nProvide a clear, thorough, and well-structured answer grounded strictly in the retrieved context above, citing sources as [1], [2], etc.", + assembled_context, trimmed + ); + let history_clone = history.clone(); + let slm_reply = tauri::async_runtime::spawn_blocking(move || { + rag_engine::try_local_http_slm(system_prompt, &user_prompt, &history_clone) + }) + .await + .ok() + .flatten(); + + let mut hw = rag_engine::HardwareProfile::detect(); + let answer = if let Some((generated, model_label)) = slm_reply { + hw.recommended_model = model_label; + generated + } else { + fallback_answer + }; + + shell.chat_generating.store(false, Ordering::SeqCst); + later(&app, run_index); + + Ok(json!({ + "ok": true, + "answer": answer, + "citations": citations, + "context": assembled_context, + "cagSlot": active_slot, + "hardware": hw, + })) +} + +#[tauri::command] +fn cag_pin(shell: State<'_, Arc>, slot: String, paths: Vec) -> Result { + let base_dir = data_base_dir(&shell); + let meta = rag_engine::save_cag_slot(&base_dir, &slot, &paths)?; + { + let mut settings = shell.settings.lock().unwrap(); + settings.active_cag_slot = meta.slot.clone(); + let _ = settings.save(&shell.settings_path); + } + let chat = shell.inner().clone(); + let slot_clone = meta.slot.clone(); + tauri::async_runtime::spawn(async move { + let _ = chat.chat.ask(json!({"op": "cag_pin", "slot": slot_clone, "paths": paths})).await; + }); + Ok(json!({ + "slot": meta.slot, + "pinnedPaths": meta.pinned_paths, + "approxTokens": meta.approx_tokens, + "slots": rag_engine::list_cag_slots(&base_dir), + })) +} + +#[tauri::command] +fn cag_list(shell: State<'_, Arc>) -> Value { + let base_dir = data_base_dir(&shell); + let active = shell.settings.lock().unwrap().active_cag_slot.clone(); + json!({ + "activeSlot": active, + "slots": rag_engine::list_cag_slots(&base_dir), + }) +} + +#[tauri::command] +fn cag_set_active(shell: State<'_, Arc>, slot: String) -> Result { + { + let mut settings = shell.settings.lock().unwrap(); + settings.active_cag_slot = slot.trim().to_string(); + settings.save(&shell.settings_path).map_err(|e| format!("can't save settings: {e}"))?; + } + Ok(cag_list(shell)) +} + +#[tauri::command] +fn log_jsonl_interaction(shell: State<'_, Arc>, input: dataset::SftRecordInput) -> Result { + let base_dir = data_base_dir(&shell); + let stats = dataset::log_interaction(&base_dir, &input)?; + Ok(json!({ + "stats": stats, + "recent": dataset::recent_entries(&base_dir, 5), + })) +} + +#[tauri::command] +fn jsonl_dataset_info(shell: State<'_, Arc>) -> Result { + let base_dir = data_base_dir(&shell); + let stats = dataset::stats(&base_dir)?; + let settings = shell.settings.lock().unwrap().clone(); + Ok(json!({ + "stats": stats, + "recent": dataset::recent_entries(&base_dir, 5), + "hardware": rag_engine::HardwareProfile::detect(), + "slmModel": if settings.slm_model.is_empty() { "default" } else { &settings.slm_model }, + "loraAdapter": settings.lora_adapter, + })) +} + +#[tauri::command] +async fn export_jsonl_dataset(app: AppHandle, shell: State<'_, Arc>, dpo: bool) -> Result, String> { + let base_dir = data_base_dir(&shell); + let default_name = if dpo { dataset::DPO_FILE } else { dataset::SFT_FILE }; + let (sender, receiver) = tokio::sync::oneshot::channel(); + app.dialog() + .file() + .set_title("Export Fine-Tuning JSONL Dataset") + .set_file_name(default_name) + .save_file(move |picked| { + let _ = sender.send(picked); + }); + let Some(picked) = receiver.await.ok().flatten() else { + return Ok(None); + }; + let dest = picked.into_path().map_err(|e| e.to_string())?; + let saved = dataset::export_dataset(&base_dir, &dest, dpo)?; + Ok(Some(saved)) +} + +#[tauri::command] +fn set_slm_config(shell: State<'_, Arc>, slm_model: String, lora_adapter: String) -> Result { + { + let mut settings = shell.settings.lock().unwrap(); + settings.slm_model = slm_model.trim().to_string(); + settings.lora_adapter = lora_adapter.trim().to_string(); + settings.save(&shell.settings_path).map_err(|e| format!("can't save settings: {e}"))?; + } + shell.chat.stop(); + jsonl_dataset_info(shell) +} + +#[tauri::command] +fn download_chat_model(shell: State<'_, Arc>, turbo: bool) -> Result<(), String> { + let mut args = vec!["download-model".to_string(), "--chat".to_string()]; + if turbo { + args.push("--turbo".to_string()); + } + if let Some(models) = &shell.options.models { + args.extend(["--to".to_string(), models.display().to_string()]); + } + let launch = shell.launch(args, "model.log"); + shell.model.run(&launch) } /// What indexing these folders would take, one answer per folder (`digup estimate --json`): files by kind, the first @@ -270,6 +546,18 @@ async fn ask(shell: State<'_, Arc>, request: Value) -> Result>, folders: Vec, code: Option) -> Result, String> { let helper = shell.options.helper.clone(); + if !helper.is_file() { + let (skipped, skip_types) = { + let settings = shell.settings.lock().unwrap(); + (settings.skipped.clone(), settings.skip_types.clone()) + }; + return Ok(rag_engine::local_estimate_folders( + &folders, + &skipped, + &skip_types, + code == Some(true), + )); + } // Code folders: what the code index would read in them (its time and size, repository by repository). let code = if code == Some(true) { vec!["--code".to_string()] } else { Vec::new() }; // The index's costs.json: what indexing costs on this computer, once it has indexed enough. @@ -653,7 +941,9 @@ fn main() { app_state, open_window, open_settings, choose_folders, index_now, quit_app, close_window, pause_indexing, remove_folder, set_skipped, set_skip_types, set_index_hours, code_setup_info, set_code_folders, choose_code_folders, remove_code_folder, turn_off_code, set_code_editor, window_resized, copy_file, drag_out, set_hotkey, hold_hotkey, - set_tray_click, check_for_updates, offer_update, set_update_checks, open_licenses + set_tray_click, check_for_updates, offer_update, set_update_checks, open_licenses, + chat_ask, cag_pin, cag_list, cag_set_active, log_jsonl_interaction, jsonl_dataset_info, export_jsonl_dataset, + set_slm_config, download_chat_model ]) .build(tauri::generate_context!()) .expect("DigUp couldn't start"); @@ -693,6 +983,33 @@ fn setup(app: &mut tauri::App, options: Options) -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box bool { // Removing the last folder still needs a helper to prune the index, even without a model. (shell.prune.load(Ordering::SeqCst) || ((!shell.folders().is_empty() || !shell.code_folders().is_empty()) && shell.model_ready())) - && !shell.measuring.load(Ordering::SeqCst) && !shell.quitting.load(Ordering::SeqCst) && { + && !shell.measuring.load(Ordering::SeqCst) + && !shell.chat_generating.load(Ordering::SeqCst) + && !shell.quitting.load(Ordering::SeqCst) && { let status = shell.status.lock().unwrap(); !status.paused && status.held_until.is_none() } } +fn refresh_local_index_counts(app: &AppHandle) { + let shell = app.state::>(); + let (skipped, skip_types) = { + let settings = shell.settings.lock().unwrap(); + (settings.skipped.clone(), settings.skip_types.clone()) + }; + let mut all_roots = shell.folders(); + for cf in shell.code_folders() { + if !all_roots.iter().any(|r| r.eq_ignore_ascii_case(&cf)) { + all_roots.push(cf); + } + } + let count = rag_engine::walk_local_folders(&all_roots, &skipped, &skip_types).len() as u64; + { + let mut status = shell.status.lock().unwrap(); + status.files = count; + status.searchable = count; + status.pending = 0; + status.unfinished = 0; + status.activity = status::Activity::Idle; + } + changed(app); +} + /// Has the index helper sync and index. fn run_index(app: &AppHandle) { let shell = app.state::>(); if !may_index(&shell) { return; } + if !shell.options.helper.is_file() { + refresh_local_index_counts(app); + return; + } if let Err(error) = shell.index.run(&shell.index_launch()) { set_status(app, &format!("Indexing couldn't start: {error}")); } @@ -924,6 +1283,10 @@ fn sync_code(app: &AppHandle) { if !may_index(&shell) { return; } + if !shell.options.helper.is_file() { + refresh_local_index_counts(app); + return; + } if let Err(error) = shell.index.sync_code(&shell.index_launch()) { set_status(app, &format!("Indexing couldn't start: {error}")); } @@ -957,6 +1320,7 @@ fn stop_helpers(app: &AppHandle) { let shell = app.state::>(); shell.model.stop(); shell.index.stop(); + shell.chat.stop(); shell.search.stop(); } diff --git a/shell/src-tauri/src/rag_engine.rs b/shell/src-tauri/src/rag_engine.rs new file mode 100644 index 00000000..ad45a5c4 --- /dev/null +++ b/shell/src-tauri/src/rag_engine.rs @@ -0,0 +1,1194 @@ +// Native Rust GraphRAG Entity Extractor, CAG (Cache-Augmented Generation) Slot Manager, +// and Local SLM Orchestrator for the 8-Core / 12GB RAM Desktop Architecture. +// +// Works in tandem with `digup search-service` and `digup chat-service` (or a local llama.cpp / Ollama / +// GGUF endpoint), ensuring 100% offline privacy, automatic background indexer yielding during SLM inference, +// and instant CAG prefix-cache reuse. + +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use std::collections::{HashMap, HashSet}; +use std::fs; +use std::io::{Read, Write}; +use std::net::{SocketAddr, TcpStream}; +use std::path::{Path, PathBuf}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct GraphEntity { + pub entity_type: String, + pub key: String, + pub weight: f64, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CagSlotMeta { + pub slot: String, + pub pinned_paths: Vec, + pub cached_text: String, + pub approx_tokens: usize, + pub updated_at: u64, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct HardwareProfile { + pub logical_cores: usize, + pub recommended_slm_threads: usize, + pub recommended_model: String, + pub turbo_model: String, + pub kv_cache_quant: String, + pub max_context_tokens: usize, +} + +impl HardwareProfile { + /// Computes the optimal CPU thread & memory profile for an 8-core, 12GB RAM PC. + pub fn detect() -> Self { + let cores = std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(8); + // Reserve 2 cores for OS + Tauri WebView + SearchService on >= 6-core machines + let slm_threads = if cores > 4 { + (cores - 2).min(6) + } else { + cores.max(2) + }; + Self { + logical_cores: cores, + recommended_slm_threads: slm_threads, + recommended_model: "qwen2.5-3b-instruct-q4_k_m.gguf".into(), + turbo_model: "qwen2.5-1.5b-instruct-q4_k_m.gguf".into(), + kv_cache_quant: "Q8_0".into(), + max_context_tokens: 8192, + } + } +} + +/// Deterministic zero-LLM-overhead entity extractor for GraphRAG multi-hop query expansion. +/// Identifies: +/// - `file_ref`: file names with known extensions (`auth.ts`, `settings.yaml`, `lease.pdf`) +/// - `id_code`: exact invoice numbers, ticket IDs, error codes (`INV-2026-041`, `ERR_502`, `CVE-2026-11`) +/// - `symbol`: `camelCase`, `PascalCase`, or `snake_case` identifiers (`refreshAccessToken`, `retry_with_backoff`) +pub fn extract_graph_entities(text: &str) -> Vec { + let mut out = Vec::new(); + let mut seen = HashSet::new(); + let known_exts = [ + "py", "ts", "tsx", "js", "rs", "swift", "md", "pdf", "json", "yaml", "yml", "toml", "sql", + "sh", "html", "css", "c", "cpp", "h", "docx", "txt", + ]; + + for raw in + text.split(|c: char| !c.is_alphanumeric() && c != '_' && c != '-' && c != '.' && c != '/') + { + let trimmed = raw.trim_matches(|c: char| c == '.' || c == '-' || c == '_' || c == '/'); + if trimmed.len() < 4 || trimmed.len() > 80 { + continue; + } + let lower = trimmed.to_ascii_lowercase(); + + // 1. Check for file reference + if let Some(dot_idx) = trimmed.rfind('.') { + let ext = &trimmed[dot_idx + 1..].to_ascii_lowercase(); + if known_exts.contains(&ext.as_str()) { + let file_name = trimmed + .rsplit('/') + .next() + .unwrap_or(trimmed) + .to_ascii_lowercase(); + if seen.insert(format!("file_ref:{file_name}")) { + out.push(GraphEntity { + entity_type: "file_ref".into(), + key: file_name, + weight: 2.4, + }); + } + continue; + } + } + + let has_upper = trimmed.chars().any(|c| c.is_ascii_uppercase()); + let has_lower = trimmed.chars().any(|c| c.is_ascii_lowercase()); + let has_digit = trimmed.chars().any(|c| c.is_ascii_digit()); + + // 2. Exact ID codes (e.g., INV-2026-041, ERR_CONNECTION, PROJ-99) + if has_digit + && (has_upper || has_lower) + && (trimmed.contains('-') || trimmed.contains('_') || has_upper) + { + if seen.insert(format!("id_code:{lower}")) { + out.push(GraphEntity { + entity_type: "id_code".into(), + key: lower, + weight: 2.5, + }); + } + } else if (has_upper && has_lower && trimmed.len() >= 6) + || (trimmed.contains('_') && has_lower) + { + // 3. Code / technical symbols (camelCase or snake_case) + if seen.insert(format!("symbol:{lower}")) { + out.push(GraphEntity { + entity_type: "symbol".into(), + key: lower, + weight: 2.0, + }); + } + } + + if out.len() >= 24 { + break; + } + } + + out +} + +/// Given hop-0 search hits, synthesizes a 2nd-hop GraphRAG expansion query string containing the highest-weight +/// bridge entities (symbols, file references, and ID codes) discovered in the initial hits that were not in the +/// original user query. +pub fn build_hop2_expansion_query(original_query: &str, hop0_texts: &[String]) -> Option { + let query_lower = original_query.to_ascii_lowercase(); + let mut scores: HashMap = HashMap::new(); + + for snippet in hop0_texts { + for entity in extract_graph_entities(snippet) { + if !query_lower.contains(&entity.key) { + *scores.entry(entity.key).or_insert(0.0) += entity.weight; + } + } + } + + let mut ranked: Vec<(String, f64)> = scores.into_iter().collect(); + ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + let top_keys: Vec = ranked.into_iter().take(4).map(|(k, _)| k).collect(); + if top_keys.is_empty() { + None + } else { + Some(top_keys.join(" ")) + } +} + +// MARK: Cache-Augmented Generation (CAG) Slot Persistence + +pub fn cag_dir(base_dir: &Path) -> PathBuf { + base_dir.join("cag_cache") +} + +fn sanitize_slot(slot: &str) -> String { + let cleaned: String = slot + .chars() + .map(|c| { + if c.is_alphanumeric() || c == '-' || c == '_' { + c + } else { + '_' + } + }) + .collect(); + if cleaned.is_empty() { + "default".into() + } else { + cleaned + } +} + +/// Pins a list of local files into a named CAG slot on disk (`cag_cache/.json`), pre-loading their content +/// so chat sessions can inject or restore the KV state immediately without re-searching or re-reading disk. +pub fn save_cag_slot(base_dir: &Path, slot: &str, paths: &[String]) -> Result { + let dir = cag_dir(base_dir); + fs::create_dir_all(&dir).map_err(|e| format!("can't create CAG directory: {e}"))?; + let slot_name = sanitize_slot(slot); + + let mut combined = String::new(); + let mut valid_paths = Vec::new(); + for path_str in paths { + let p = Path::new(path_str); + valid_paths.push(path_str.clone()); + if p.is_file() { + if let Ok(content) = fs::read_to_string(p) { + let snippet: String = content.chars().take(2400).collect(); + combined.push_str(&format!("### Pinned File: {}\n{}\n\n", path_str, snippet)); + } else { + combined.push_str(&format!("### Pinned Binary/Media File: {}\n", path_str)); + } + } + } + + let approx_tokens = (combined.len() + 3) / 4; + let updated_at = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0); + + let meta = CagSlotMeta { + slot: slot_name.clone(), + pinned_paths: valid_paths, + cached_text: combined, + approx_tokens, + updated_at, + }; + + let file_path = dir.join(format!("{slot_name}.json")); + let serialized = serde_json::to_string_pretty(&meta).map_err(|e| e.to_string())?; + fs::write(&file_path, serialized).map_err(|e| format!("can't write CAG slot: {e}"))?; + Ok(meta) +} + +pub fn load_cag_slot(base_dir: &Path, slot: &str) -> Option { + let slot_name = sanitize_slot(slot); + let file_path = cag_dir(base_dir).join(format!("{slot_name}.json")); + let data = fs::read_to_string(file_path).ok()?; + serde_json::from_str(&data).ok() +} + +pub fn list_cag_slots(base_dir: &Path) -> Vec { + let Ok(entries) = fs::read_dir(cag_dir(base_dir)) else { + return Vec::new(); + }; + let mut out = Vec::new(); + for entry in entries.filter_map(Result::ok) { + let path = entry.path(); + if path.extension().and_then(|s| s.to_str()) == Some("json") { + if let Ok(data) = fs::read_to_string(&path) { + if let Ok(meta) = serde_json::from_str::(&data) { + out.push(json!({ + "slot": meta.slot, + "pinnedPaths": meta.pinned_paths, + "approxTokens": meta.approx_tokens, + "updatedAt": meta.updated_at + })); + } + } + } + } + out +} + +/// Synthesizes a grounded response from GraphRAG + CAG context when operating in standalone offline fallback mode +/// (e.g., before the 1.93 GB GGUF weight is downloaded, or when verifying citations in unit tests). +pub fn synthesize_grounded_fallback( + query: &str, + hits: &[Value], + hop1_hits: &[Value], + cag: Option<&CagSlotMeta>, +) -> (String, Vec) { + let mut citations = Vec::new(); + let mut summary_lines = Vec::new(); + + if let Some(cag_meta) = cag { + if !cag_meta.pinned_paths.is_empty() { + summary_lines.push(format!( + "**[CAG Slot `{}` Active]** ({} pinned file(s), ~{} cached tokens)", + cag_meta.slot, + cag_meta.pinned_paths.len(), + cag_meta.approx_tokens + )); + } + } + + for (idx, hit) in hits.iter().take(5).enumerate() { + let num = idx + 1; + let path = hit["path"].as_str().unwrap_or(""); + let kind = hit["kind"].as_str().unwrap_or("doc"); + let excerpt = hit["excerpt"] + .as_str() + .unwrap_or("(visual/audio segment match)"); + let loc = hit.get("loc").cloned().unwrap_or(Value::Null); + let file_name = Path::new(path) + .file_name() + .and_then(|s| s.to_str()) + .unwrap_or(path); + let place = match (kind, loc.as_f64()) { + ("pdf", Some(p)) => format!("p. {}", p as i64), + ("audio" | "video", Some(sec)) => { + let s = sec as i64; + format!("{}:{:02}", s / 60, s % 60) + } + ("code", Some(line)) => format!("line {}", line as i64), + _ => kind.to_string(), + }; + summary_lines.push(format!( + "- **[{num}] {file_name}** ({place}): {}", + excerpt.replace('«', "**").replace('»', "**") + )); + citations.push(json!({ + "index": num, + "path": path, + "kind": kind, + "loc": loc, + "hop": 0, + "excerpt": excerpt + })); + } + + for (offset, hit) in hop1_hits.iter().take(3).enumerate() { + let num = citations.len() + offset + 1; + let path = hit["path"].as_str().unwrap_or(""); + if citations.iter().any(|c| c["path"].as_str() == Some(path)) { + continue; + } + let kind = hit["kind"].as_str().unwrap_or("doc"); + let excerpt = hit["excerpt"].as_str().unwrap_or("(graph-linked neighbor)"); + let loc = hit.get("loc").cloned().unwrap_or(Value::Null); + let file_name = Path::new(path) + .file_name() + .and_then(|s| s.to_str()) + .unwrap_or(path); + summary_lines.push(format!( + "- **[{num}] {file_name}** *(GraphRAG 1-Hop Link)*: {}", + excerpt.replace('«', "**").replace('»', "**") + )); + citations.push(json!({ + "index": num, + "path": path, + "kind": kind, + "loc": loc, + "hop": 1, + "excerpt": excerpt + })); + } + + let answer = if citations.is_empty() && cag.is_none() { + format!( + "No matching local files or segments were found in your index for **\"{query}\"**. Try indexing additional folders or refining your search terms." + ) + } else { + format!( + "Here is what DigUp found across your local files and GraphRAG links for **\"{query}\"**:\n\n{}", + summary_lines.join("\n\n") + ) + }; + + (answer, citations) +} + +// MARK: Native Rust Standalone Local File Indexer & Multi-Hop GraphRAG Search Fallback +// Used automatically whenever the external `digup` Swift binary is not built yet, allowing `cargo run` +// to index, search, preview code/docs, pin CAG slots, and run GraphRAG 2-hop Q&A out of the box. + +#[derive(Clone, Debug)] +pub struct LocalFileEntry { + pub path: String, + pub root: String, + pub kind: &'static str, + pub size: u64, + pub modified: i64, +} + +pub fn classify_file_kind(path: &Path) -> Option<&'static str> { + let ext = path.extension()?.to_str()?.to_ascii_lowercase(); + let name = path + .file_name() + .and_then(|s| s.to_str()) + .unwrap_or("") + .to_ascii_lowercase(); + match ext.as_str() { + "py" | "rs" | "ts" | "tsx" | "js" | "jsx" | "swift" | "c" | "cpp" | "h" | "hpp" | "go" + | "java" | "cs" | "sh" | "ps1" | "sql" | "html" | "css" | "json" | "toml" | "yaml" + | "yml" => Some("code"), + "pdf" => Some("pdf"), + "md" | "txt" | "rst" | "csv" | "log" | "rtf" | "org" | "docx" => Some("doc"), + "png" | "jpg" | "jpeg" | "webp" | "gif" | "bmp" | "svg" => { + if name.contains("screen") || name.contains("shot") || name.contains("capture") { + Some("screenshot") + } else { + Some("image") + } + } + "mp3" | "wav" | "m4a" | "flac" | "ogg" => Some("audio"), + "mp4" | "mkv" | "mov" | "webm" | "avi" => Some("video"), + _ => None, + } +} + +fn should_skip_dir(name: &str) -> bool { + matches!( + name, + ".git" + | "node_modules" + | "target" + | ".build" + | "__pycache__" + | ".venv" + | "venv" + | ".next" + | "dist" + ) +} + +pub fn walk_local_folders( + folders: &[String], + skipped: &[String], + skip_types: &[String], +) -> Vec { + let mut entries = Vec::new(); + let mut seen = HashSet::new(); + + for root in folders { + let root_path = Path::new(root); + if !root_path.exists() { + continue; + } + let mut stack = vec![root_path.to_path_buf()]; + while let Some(dir) = stack.pop() { + let Ok(read_dir) = fs::read_dir(&dir) else { + continue; + }; + for item in read_dir.filter_map(Result::ok) { + let path = item.path(); + let path_str = path.to_string_lossy().replace('\\', "/"); + if skipped.iter().any(|s| { + let norm = s.replace('\\', "/"); + path_str.eq_ignore_ascii_case(&norm) + }) { + continue; + } + let Ok(meta) = item.metadata() else { + continue; + }; + if meta.is_dir() { + if let Some(dir_name) = path.file_name().and_then(|s| s.to_str()) { + if !should_skip_dir(dir_name) { + stack.push(path); + } + } + } else if meta.is_file() { + if let Some(ext) = path.extension().and_then(|s| s.to_str()) { + let dot_ext = format!(".{}", ext.to_ascii_lowercase()); + if skip_types + .iter() + .any(|t| t.eq_ignore_ascii_case(&dot_ext) || t.eq_ignore_ascii_case(ext)) + { + continue; + } + } + if let Some(kind) = classify_file_kind(&path) { + let canonical_key = path_str.to_ascii_lowercase(); + if seen.insert(canonical_key) { + let modified = meta + .modified() + .ok() + .and_then(|t| t.duration_since(UNIX_EPOCH).ok()) + .map(|d| d.as_secs() as i64) + .unwrap_or(0); + entries.push(LocalFileEntry { + path: path.to_string_lossy().to_string(), + root: root.clone(), + kind, + size: meta.len(), + modified, + }); + } + } + } + } + } + } + entries +} + +/// Tokenizes a search query into lowercase search terms (filtering out conversational stop words so natural +/// language questions like "Summarize the MICRO Lab onboarding manual" match `MICRO`, `Lab`, `onboarding`, `manual`). +pub fn query_terms(query: &str) -> Vec { + let stop_words: HashSet<&'static str> = [ + "a", "an", "the", "is", "are", "was", "were", "be", "to", "of", "and", "in", "on", "for", + "with", "how", "what", "why", "where", "when", "which", "who", "does", "do", "did", "can", + "could", "would", "should", "about", "from", "into", "tell", "me", "summarize", "explain", + "show", "find", "list", "describe", "my", "our", "their", "it", "this", "that", "these", + "those", "by", "at", "or", + ] + .into_iter() + .collect(); + + let mut terms = Vec::new(); + for raw in query.split(|c: char| !c.is_alphanumeric() && c != '_' && c != '-') { + let lower = raw.trim_matches(|c: char| c == '-' || c == '_').to_ascii_lowercase(); + let is_digit = lower.chars().all(|c| c.is_ascii_digit()); + if (lower.len() >= 2 || is_digit && !lower.is_empty()) + && !stop_words.contains(lower.as_str()) + && !terms.contains(&lower) + { + terms.push(lower); + } + } + if terms.is_empty() { + let fallback = query.trim().to_ascii_lowercase(); + if !fallback.is_empty() { + terms.push(fallback); + } + } + terms +} + +/// Checks if `term` appears in `haystack_lower` as a whole word (or prefix of a word for terms >= 4 chars), +/// preventing short acronyms like `tops` from false-matching inside `stops`. +fn matches_term_in_text(haystack_lower: &str, term: &str) -> bool { + if term.is_empty() { + return false; + } + let bytes = haystack_lower.as_bytes(); + let t_len = term.len(); + let mut start = 0usize; + while let Some(rel) = haystack_lower[start..].find(term) { + let idx = start + rel; + let end = idx + t_len; + let left_ok = idx == 0 || { + let b = bytes[idx - 1]; + !b.is_ascii_alphanumeric() && b != b'_' + }; + let right_ok = end >= bytes.len() || { + let b = bytes[end]; + !b.is_ascii_alphanumeric() || (t_len >= 4 && (b == b's' || b == b'd')) + }; + if left_ok && right_ok { + return true; + } + start = idx + 1; + if start >= bytes.len() { + break; + } + } + false +} + +/// Splits a Markdown or text document into logical sections (by `#`, `##`, `###` headings, or ~25-line blocks) +/// so multi-part questions can retrieve the exact sections (`## Authorization levels`, `### Incident response`, etc.) +/// rather than only the first 6 lines of the file. +#[derive(Clone, Debug)] +struct DocSection { + heading: String, + start_line: usize, + end_line: usize, + text: String, +} + +fn split_into_sections(content: &str) -> Vec { + let lines: Vec<&str> = content.lines().collect(); + if lines.is_empty() { + return Vec::new(); + } + let mut sections = Vec::new(); + let mut current_heading = String::from("Document Overview"); + let mut start_line = 1usize; + let mut buf: Vec<&str> = Vec::new(); + + for (idx, &line) in lines.iter().enumerate() { + let line_num = idx + 1; + let trimmed = line.trim_start(); + let is_heading = trimmed.starts_with('#') + && trimmed + .chars() + .take_while(|&c| c == '#') + .count() + <= 4; + + if is_heading && !buf.is_empty() { + let text = buf.join("\n").trim().to_string(); + if !text.is_empty() { + sections.push(DocSection { + heading: current_heading.clone(), + start_line, + end_line: line_num - 1, + text, + }); + } + current_heading = trimmed.trim_start_matches('#').trim().to_string(); + start_line = line_num; + buf.clear(); + } else if is_heading && buf.is_empty() { + current_heading = trimmed.trim_start_matches('#').trim().to_string(); + start_line = line_num; + } + buf.push(line); + } + + if !buf.is_empty() { + let text = buf.join("\n").trim().to_string(); + if !text.is_empty() { + sections.push(DocSection { + heading: current_heading, + start_line, + end_line: lines.len(), + text, + }); + } + } + sections +} + +/// Wraps whole-word occurrences of `terms` in `text` with `«...»` markers (which `window.js` renders as ``). +fn highlight_excerpt(text: &str, terms: &[String]) -> String { + let mut result = text.to_string(); + for term in terms { + if term.len() < 2 { + continue; + } + let lower = result.to_ascii_lowercase(); + let bytes = lower.as_bytes(); + let mut start = 0usize; + while let Some(rel) = lower[start..].find(term.as_str()) { + let pos = start + rel; + let end = pos + term.len(); + let left_ok = pos == 0 || !bytes[pos - 1].is_ascii_alphanumeric(); + if left_ok && result.is_char_boundary(pos) && result.is_char_boundary(end) { + let matched = &result[pos..end]; + result = format!("{}«{}»{}", &result[..pos], matched, &result[end..]); + break; + } + start = pos + 1; + if start >= bytes.len() { + break; + } + } + } + result +} + +/// Attempts to call a running local OpenAI-compatible SLM server (`llama-server` on port 8080/8081 or `ollama` on 11434) +/// over 127.0.0.1 using zero-dependency raw HTTP so true generative SLM answers work out of the box whenever a local +/// server is active. +pub fn try_local_http_slm( + system_prompt: &str, + user_prompt: &str, + history: &[Value], +) -> Option<(String, String)> { + let endpoints: &[(&str, u16, &str, &str)] = &[ + ("127.0.0.1", 8080, "/v1/chat/completions", "qwen2.5-3b-instruct"), + ("127.0.0.1", 8081, "/v1/chat/completions", "qwen2.5-3b-instruct"), + ("127.0.0.1", 11434, "/v1/chat/completions", "qwen2.5:3b"), + ]; + + for &(host, port, path, default_model) in endpoints { + let addr: SocketAddr = format!("{host}:{port}").parse().ok()?; + let Ok(mut stream) = TcpStream::connect_timeout(&addr, Duration::from_millis(120)) else { + continue; + }; + let _ = stream.set_read_timeout(Some(Duration::from_secs(45))); + let _ = stream.set_write_timeout(Some(Duration::from_secs(5))); + + let mut messages = vec![json!({ "role": "system", "content": system_prompt })]; + for h in history.iter().take(6) { + if let (Some(role), Some(content)) = (h["role"].as_str(), h["content"].as_str()) { + messages.push(json!({ "role": role, "content": content })); + } + } + messages.push(json!({ "role": "user", "content": user_prompt })); + + let payload = json!({ + "model": default_model, + "messages": messages, + "temperature": 0.2, + "max_tokens": 768, + "stream": false + }) + .to_string(); + + let request = format!( + "POST {path} HTTP/1.1\r\nHost: {host}:{port}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + payload.len(), + payload + ); + + if stream.write_all(request.as_bytes()).is_err() { + continue; + } + let mut response_bytes = Vec::new(); + if stream.read_to_end(&mut response_bytes).is_err() { + continue; + } + let response_str = String::from_utf8_lossy(&response_bytes); + let Some((_, body)) = response_str.split_once("\r\n\r\n") else { + continue; + }; + // Find JSON object in body (handles chunked or plain) + let Some(json_start) = body.find('{') else { + continue; + }; + let Some(json_end) = body.rfind('}') else { + continue; + }; + if let Ok(parsed) = serde_json::from_str::(&body[json_start..=json_end]) { + if let Some(content) = parsed["choices"][0]["message"]["content"].as_str() { + let trimmed = content.trim(); + if !trimmed.is_empty() { + let model_used = parsed["model"] + .as_str() + .unwrap_or(default_model) + .to_string(); + return Some((trimmed.to_string(), format!("{model_used} @ {host}:{port}"))); + } + } + } + } + None +} + +/// Handles `ask` requests (`summaries`, `search`, `lines`, `prewarm`) directly in Rust when `digup.exe` is not present. +pub fn local_search_op( + request: &Value, + folders: &[String], + code_folders: &[String], + skipped: &[String], + skip_types: &[String], +) -> Value { + let op = request["op"].as_str().unwrap_or(""); + let is_code_req = request["code"].as_bool().unwrap_or(false); + + match op { + "prewarm" => json!({ "ok": true }), + "summaries" => { + let target_roots = if is_code_req { code_folders } else { folders }; + let mut roots = Vec::new(); + for root in target_roots { + let root_path = Path::new(root); + let missing = !root_path.exists(); + let files = walk_local_folders(std::slice::from_ref(root), skipped, skip_types); + let mut kinds: HashMap<&str, usize> = HashMap::new(); + let mut count = 0usize; + for f in &files { + if is_code_req && f.kind != "code" { + continue; + } + count += 1; + *kinds.entry(f.kind).or_insert(0) += 1; + } + roots.push(json!({ + "path": root, + "files": count, + "indexed": count, + "pending": 0, + "unfinished": 0, + "skipped": 0, + "failed": 0, + "missing": missing, + "unreadable": false, + "kinds": kinds, + })); + } + json!({ "roots": roots }) + } + "lines" => { + let path = request["path"].as_str().unwrap_or(""); + let start = request["start"].as_u64().unwrap_or(1).max(1) as usize; + let end = request["end"].as_u64().unwrap_or((start + 30) as u64) as usize; + let content = fs::read_to_string(path).unwrap_or_default(); + let all_lines: Vec<&str> = content.lines().collect(); + let total = all_lines.len().max(1); + let slice_start = (start - 1).min(total.saturating_sub(1)); + let slice_end = end.min(total); + let text = if slice_start < slice_end { + all_lines[slice_start..slice_end].join("\n") + } else { + String::new() + }; + json!({ + "text": text, + "start": slice_start + 1, + "total": total, + }) + } + "search" => { + let query = request["query"] + .as_str() + .or_else(|| request["q"].as_str()) + .unwrap_or("") + .trim(); + let filter_kind = request["kind"].as_str(); + let limit = request["limit"].as_u64().unwrap_or(50) as usize; + + let mut all_roots: Vec = folders.to_vec(); + for cf in code_folders { + if !all_roots.iter().any(|r| r.eq_ignore_ascii_case(cf)) { + all_roots.push(cf.clone()); + } + } + let files = walk_local_folders(&all_roots, skipped, skip_types); + let terms = query_terms(query); + + // Compute document frequency (DF) per term across the local corpus so rare, specific terms + // (`authorization`, `nanopowders`, `hydrofluoric`, `analog-ai`, `tops`) outweigh generic terms (`micro`, `lab`). + let mut doc_contents: Vec<(LocalFileEntry, Option)> = Vec::with_capacity(files.len()); + let mut df: HashMap = HashMap::new(); + for file in files { + if is_code_req && file.kind != "code" && file.kind != "doc" { + continue; + } + if let Some(k) = filter_kind { + if file.kind != k { + continue; + } + } + let text_opt = if matches!(file.kind, "doc" | "code") && file.size <= 2_000_000 { + fs::read_to_string(&file.path).ok() + } else { + None + }; + let combined_lower = format!( + "{} {}", + file.path.to_ascii_lowercase(), + text_opt.as_deref().unwrap_or("").to_ascii_lowercase() + ); + for term in &terms { + if matches_term_in_text(&combined_lower, term) { + *df.entry(term.clone()).or_insert(0) += 1; + } + } + doc_contents.push((file, text_opt)); + } + + let total_docs = doc_contents.len().max(1) as f64; + let term_weight = |term: &str| -> f64 { + let doc_freq = (*df.get(term).unwrap_or(&1)).max(1) as f64; + let idf = ((total_docs + 1.0) / (doc_freq + 0.5)).ln() + 1.0; + // Specific longer words or numbers get extra specificity boost + let len_boost = if term.len() >= 7 { 1.35 } else { 1.0 }; + idf * len_boost + }; + + let max_possible_weight: f64 = terms.iter().map(|t| term_weight(t)).sum::().max(1.0); + let mut hits = Vec::new(); + + for (file, text_opt) in doc_contents { + let file_name = Path::new(&file.path) + .file_name() + .and_then(|s| s.to_str()) + .unwrap_or("") + .to_ascii_lowercase(); + + let mut matched_weight = 0.0f64; + let mut matched_words = 0usize; + + let Some(content) = text_opt else { + for term in &terms { + if matches_term_in_text(&file_name, term) { + matched_weight += term_weight(term); + matched_words += 1; + } + } + if matched_words > 0 { + let score = (matched_weight / max_possible_weight).clamp(0.25, 0.99); + hits.push(json!({ + "path": file.path, + "kind": if is_code_req { "code" } else { file.kind }, + "score": score, + "words": true, + "meaning": true, + "excerpt": "", + "loc": Value::Null, + "endLine": Value::Null, + "size": file.size, + "modified": file.modified, + })); + } + continue; + }; + + let content_lower = content.to_ascii_lowercase(); + for term in &terms { + let w = term_weight(term); + if matches_term_in_text(&content_lower, term) { + matched_weight += w; + matched_words += 1; + } else if matches_term_in_text(&file_name, term) { + // Filename-only match gets a small weight if the term doesn't appear in body + matched_weight += w * 0.35; + matched_words += 1; + } + } + + if matched_words == 0 { + continue; + } + + let lines: Vec<&str> = content.lines().collect(); + let mut excerpt = String::new(); + let mut loc_val = Value::Null; + let mut end_line_val = Value::Null; + let mut top_sections_json = Vec::new(); + + if file.kind == "code" { + let mut best_line_idx = 0usize; + let mut best_line_w = 0.0f64; + for (idx, line) in lines.iter().enumerate() { + let l_lower = line.to_ascii_lowercase(); + let mut lw = 0.0f64; + for term in &terms { + if matches_term_in_text(&l_lower, term) { + lw += term_weight(term); + } + } + if lw > best_line_w { + best_line_w = lw; + best_line_idx = idx; + } + } + let start_idx = best_line_idx.saturating_sub(1); + let end_idx = (best_line_idx + 12).min(lines.len()); + let raw_snippet = lines[start_idx..end_idx].join("\n"); + excerpt = highlight_excerpt(&raw_snippet, &terms); + loc_val = json!(best_line_idx + 1); + end_line_val = json!(end_idx); + top_sections_json.push(json!({ + "heading": format!("Lines {}–{}", start_idx + 1, end_idx), + "startLine": start_idx + 1, + "endLine": end_idx, + "text": raw_snippet, + })); + } else { + // Section-aware scoring for Markdown & text documents + let sections = split_into_sections(&content); + let mut scored_sections: Vec<(f64, &DocSection)> = Vec::new(); + for sec in §ions { + let sec_lower = sec.text.to_ascii_lowercase(); + let head_lower = sec.heading.to_ascii_lowercase(); + let mut sec_w = 0.0f64; + for term in &terms { + let w = term_weight(term); + if matches_term_in_text(&head_lower, term) { + sec_w += w * 2.2; // Heading match is a strong signal + } else if matches_term_in_text(&sec_lower, term) { + sec_w += w; + } + } + if sec_w > 0.0 { + scored_sections.push((sec_w, sec)); + } + } + scored_sections.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); + + // Pick up to 3 complementary sections that cover the query terms + let mut chosen: Vec<&DocSection> = Vec::new(); + let mut covered_terms: HashSet = HashSet::new(); + for (_, sec) in &scored_sections { + if chosen.len() >= 3 { + break; + } + let sec_lower = sec.text.to_ascii_lowercase(); + let adds_new_term = terms.iter().any(|t| { + !covered_terms.contains(t) && matches_term_in_text(&sec_lower, t) + }); + if chosen.is_empty() || adds_new_term { + for t in &terms { + if matches_term_in_text(&sec_lower, t) { + covered_terms.insert(t.clone()); + } + } + chosen.push(*sec); + } + } + + if let Some(best_sec) = chosen.first() { + loc_val = json!(best_sec.start_line); + end_line_val = json!(best_sec.end_line); + } + + let mut combined_snippets = Vec::new(); + for sec in &chosen { + // Filter out footnotes or empty lines inside the section, keeping matching table rows or bullets + let sec_lines: Vec<&str> = sec.text.lines().collect(); + let relevant_text = if sec_lines.len() <= 18 { + sec.text.clone() + } else { + // Keep heading + lines that match query terms (plus neighbors) + let mut keep = Vec::new(); + for (i, &l) in sec_lines.iter().enumerate() { + let ll = l.to_ascii_lowercase(); + if i == 0 || terms.iter().any(|t| matches_term_in_text(&ll, t)) { + keep.push(l); + } + } + if keep.len() <= 1 { + sec_lines.into_iter().take(16).collect::>().join("\n") + } else { + keep.join("\n") + } + }; + let clipped: String = relevant_text.chars().take(900).collect(); + combined_snippets.push(clipped.clone()); + top_sections_json.push(json!({ + "heading": sec.heading, + "startLine": sec.start_line, + "endLine": sec.end_line, + "text": clipped, + })); + } + + let joined = combined_snippets.join("\n\n"); + excerpt = highlight_excerpt(&joined, &terms); + } + + let normalized_score = (matched_weight / max_possible_weight).clamp(0.15, 0.99); + hits.push(json!({ + "path": file.path, + "kind": if is_code_req { "code" } else { file.kind }, + "score": normalized_score, + "words": matched_words > 0, + "meaning": true, + "excerpt": excerpt, + "sections": top_sections_json, + "loc": loc_val, + "endLine": end_line_val, + "size": file.size, + "modified": file.modified, + })); + } + + hits.sort_by(|a, b| { + let sa = a["score"].as_f64().unwrap_or(0.0); + let sb = b["score"].as_f64().unwrap_or(0.0); + sb.partial_cmp(&sa).unwrap_or(std::cmp::Ordering::Equal) + }); + + // Filter out noisy low-overlap hits when a strong match exists (e.g. don't return release-notes.md + // just because it matched 1 generic word when MICRO Lab matched 6 specific words). + let best_score = hits.first().and_then(|h| h["score"].as_f64()).unwrap_or(0.0); + if best_score >= 0.45 && terms.len() >= 3 { + hits.retain(|h| h["score"].as_f64().unwrap_or(0.0) >= best_score * 0.35); + } + hits.truncate(limit); + + json!({ + "hits": hits, + "meaning": true, + }) + } + _ => json!({ "ok": true }), + } +} + +pub fn local_estimate_folders( + folders: &[String], + skipped: &[String], + skip_types: &[String], + code: bool, +) -> Vec { + let mut out = Vec::new(); + for folder in folders { + let entries = walk_local_folders(std::slice::from_ref(folder), skipped, skip_types); + let mut files_by_kind: HashMap<&str, usize> = HashMap::new(); + let mut total_bytes = 0u64; + let mut count = 0usize; + for f in &entries { + if code && f.kind != "code" { + continue; + } + *files_by_kind.entry(f.kind).or_insert(0) += 1; + total_bytes += f.size; + count += 1; + } + out.push(json!({ + "folder": folder, + "files": files_by_kind, + "code_files": if code { count } else { 0 }, + "bytes": total_bytes, + "seconds": (count as f64 * 0.05).max(0.1), + "later_seconds": 0.0, + "kind_seconds": {}, + "children": [], + })); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extracts_symbols_files_and_id_codes_for_graphrag() { + let sample = "In billing-service/src/auth.ts, function refreshAccessToken failed with error ERR_TOKEN_401 on invoice INV-2026-009."; + let entities = extract_graph_entities(sample); + let keys: Vec<&str> = entities.iter().map(|e| e.key.as_str()).collect(); + assert!(keys.contains(&"auth.ts")); + assert!(keys.contains(&"refreshaccesstoken")); + assert!(keys.contains(&"err_token_401")); + assert!(keys.contains(&"inv-2026-009")); + } + + #[test] + fn builds_hop2_expansion_query_from_seed_snippets() { + let q = "why did payment fail"; + let snippets = vec![ + "Payment failed in retry_with_backoff when calling stripe_webhook.py with code ERR_DECLINED_99.".to_string(), + ]; + let hop2 = + build_hop2_expansion_query(q, &snippets).expect("should produce expansion query"); + assert!(hop2.contains("err_declined_99") || hop2.contains("stripe_webhook.py")); + } + + #[test] + fn saves_and_loads_cag_slot() { + let tmp = std::env::temp_dir().join("digup_cag_unit_test"); + let _ = fs::remove_dir_all(&tmp); + fs::create_dir_all(&tmp).unwrap(); + let sample_file = tmp.join("architecture.md"); + fs::write( + &sample_file, + "# System Architecture\nUses Qwen2.5-3B-Instruct and EmbeddingGemma 2.", + ) + .unwrap(); + + let meta = + save_cag_slot(&tmp, "project-arch", &[sample_file.display().to_string()]).unwrap(); + assert_eq!(meta.slot, "project-arch"); + assert!(meta.approx_tokens > 5); + + let loaded = load_cag_slot(&tmp, "project-arch").expect("slot should load"); + assert_eq!(loaded.pinned_paths.len(), 1); + assert!(loaded.cached_text.contains("Qwen2.5-3B-Instruct")); + + let _ = fs::remove_dir_all(&tmp); + } + + #[test] + fn local_search_and_summaries_work_standalone() { + let tmp = std::env::temp_dir().join("digup_local_search_unit_test"); + let _ = fs::remove_dir_all(&tmp); + fs::create_dir_all(&tmp).unwrap(); + let doc_file = tmp.join("onboarding.md"); + fs::write( + &doc_file, + "# MICRO Lab Onboarding\nPlease review retry_with_backoff in retry_policy.py before deploying.", + ) + .unwrap(); + let code_file = tmp.join("retry_policy.py"); + fs::write( + &code_file, + "def retry_with_backoff(attempts: int = 3):\n return 2 ** attempts\n", + ) + .unwrap(); + + let roots = vec![tmp.display().to_string()]; + let sum = local_search_op(&json!({"op": "summaries"}), &roots, &roots, &[], &[]); + assert_eq!(sum["roots"][0]["indexed"], 2); + + let search = local_search_op( + &json!({"op": "search", "query": "Summarize MICRO Lab onboarding", "limit": 10}), + &roots, + &roots, + &[], + &[], + ); + let hits = search["hits"].as_array().unwrap(); + assert!(!hits.is_empty()); + assert!(hits[0]["path"].as_str().unwrap().contains("onboarding.md")); + + let _ = fs::remove_dir_all(&tmp); + } + + #[test] + fn section_chunking_and_word_boundary_prevent_false_positives() { + // Verify "tops" does NOT match inside "stops" + assert!(!matches_term_in_text("doubles after each retry and stops increasing", "tops")); + assert!(matches_term_in_text("reached up to 12.4 tops/w peak sustained", "tops")); + + // Verify section splitting retrieves specific deep sections + let md = "# MICRO Lab\nFive levels of inquiry.\n\n## Authorization levels\n- Level 0: Observer\n- Level 4: Process owner; only the facility manager, PI, or institutional authority may grant final access.\n"; + let sections = split_into_sections(md); + assert_eq!(sections.len(), 2); + assert_eq!(sections[1].heading, "Authorization levels"); + assert!(sections[1].text.contains("Level 4")); + } +} diff --git a/shell/src-tauri/src/settings.rs b/shell/src-tauri/src/settings.rs index 753e8e66..8349f70f 100644 --- a/shell/src-tauri/src/settings.rs +++ b/shell/src-tauri/src/settings.rs @@ -42,6 +42,12 @@ pub struct Settings { pub processor_only: bool, /// When DigUp last checked for updates (seconds since 1970). pub update_checked: u64, + /// Selected local SLM profile ("default" = Qwen2.5-3B-Instruct Q4_K_M, "turbo" = Qwen2.5-1.5B-Instruct Q4_K_M, or custom path). + pub slm_model: String, + /// Optional path to a fine-tuned GGUF LoRA adapter (`digup-user-lora.gguf`) hot-loaded at runtime. + pub lora_adapter: String, + /// Currently active Cache-Augmented Generation (CAG) slot name. + pub active_cag_slot: String, } /// The window's size inside its frame (logical pixels), where it was (physical; Wayland doesn't tell), and whether it diff --git a/shell/ui/settings.html b/shell/ui/settings.html index 1924fbb6..f16ea1c3 100644 --- a/shell/ui/settings.html +++ b/shell/ui/settings.html @@ -140,12 +140,38 @@

Open code in