@pheem49/mint 1.2.4 → 1.3.0

package/src/AI_Brain/Gemini_API.js

@@ -2,11 +2,13 @@ const { GoogleGenAI } = require('@google/genai');
 const { readChatHistory, writeChatHistory, clearChatHistory } = require('../System/chat_history_manager');
 const { readConfig } = require('../System/config_manager');
 const pluginManager = require('../Plugins/plugin_manager');
+const mcpManager = require('../Plugins/mcp_manager');
 
 let ai = null;
 let activeApiKey = '';
 const initialEnvKey = (process.env.GEMINI_API_KEY || '').trim();
-const DEFAULT_GEMINI_MODEL = 'gemini-2.5-flash'; // Optimized model
+const axios = require('axios');
+const DEFAULT_GEMINI_MODEL = 'gemini-2.5-flash';
 
 function decodeUnicode(str) {
   if (!str) return '';
@@ -53,12 +55,24 @@ Always respond exactly with valid JSON containing NO MARKDOWN FORMATTING (do not
 {
   "response": "Your conversational reply here (Matches user language).",
   "action": {
-    "type": "none" | "open_url" | "open_app" | "search" | "web_automation" | "create_folder" | "open_file" | "delete_file" | "clipboard_write" | "system_info" | "plugin" | "learn_file" | "system_automation",
+    "type": "none" | "open_url" | "open_app" | "search" | "web_automation" | "create_folder" | "open_file" | "open_folder" | "delete_file" | "clipboard_write" | "system_info" | "plugin" | "learn_file" | "learn_folder" | "system_automation" | "mcp_tool" | "mouse_click" | "mouse_move" | "type_text" | "key_tap",
+
     "pluginName": "only if type is plugin",
-    "target": "target string based on type or plugin instruction"
+    "server": "only if type is mcp_tool (server name)",
+    "target": "target string based on type (tool name if mcp_tool, text to type if type_text, key name if key_tap)",
+    "x": 0-1000, // required for mouse_click and mouse_move
+    "y": 0-1000, // required for mouse_click and mouse_move
+    "button": 1 | 2 | 3, // optional for mouse_click, 1=left, 2=middle, 3=right
+    "args": { "param": "value" } // only if type is mcp_tool
   }
 }
 
+COORDINATE SYSTEM:
+- When analyzing an image, use a coordinate system from 0 to 1000.
+- (0, 0) is the Top-Left corner.
+- (1000, 1000) is the Bottom-Right corner.
+- To click an element, estimate its center point and provide x and y.
+
 Examples:
 Input: "Hi, what is your name?"
 Output: { "response": "Hello! My name is Mint, your personal AI assistant. How can I help you today?", "action": { "type": "none", "target": "" } }
@@ -122,7 +136,19 @@ const MAX_HISTORY_MESSAGES = 20; // Keep only the last 20 messages (approx 10 tu
 function createChat(history = []) {
   // Load plugins and get dynamic description for the prompt
   pluginManager.loadPlugins();
-  const dynamicPrompt = systemInstruction + pluginManager.getPromptDescriptions();
+  // Inject MCP Tools
+  const mcpTools = mcpManager.getAllTools();
+  let mcpPrompt = "\n\nAVAILABLE MCP TOOLS (Model Context Protocol):\n";
+  if (mcpTools.length > 0) {
+      mcpTools.forEach(tool => {
+          mcpPrompt += `- Server: ${tool.serverName}, Tool: ${tool.name}\n  Desc: ${tool.description}\n  Args: ${JSON.stringify(tool.inputSchema.properties)}\n`;
+      });
+      mcpPrompt += "\nTo use these tools, use action type 'mcp_tool', specify the 'server' name, set 'target' to the tool name, and provide 'args'.\n";
+  } else {
+      mcpPrompt += "No MCP tools currently connected.\n";
+  }
+
+  const dynamicPrompt = systemInstruction + pluginManager.getPromptDescriptions() + mcpPrompt;
 
   // Truncate history and strip custom fields like 'timestamp' before passing to SDK
   const cleanedHistory = (history || []).map(msg => ({
@@ -151,7 +177,18 @@ resolveApiKey();
 initAiClient();
 createChat(readChatHistory());
 
-const { searchKnowledge } = require('./knowledge_base');
+function shouldUseKnowledgeSearch(message) {
+  const text = (message || '').trim().toLowerCase();
+  if (!text) return false;
+
+  const knowledgeHints = [
+    'readme', 'docs', 'documentation', 'manual', 'guide', 'knowledge', 'rag',
+    'search local', 'search files', 'learn file', 'project files', 'source code',
+    'ไฟล์', 'เอกสาร', 'คู่มือ', 'ค้นหาในเครื่อง', 'ค้นหาไฟล์', 'ข้อมูลในเครื่อง', 'โค้ดโปรเจค'
+  ];
+
+  return knowledgeHints.some(hint => text.includes(hint));
+}
 
 async function handleChat(message, base64Image = null, base64Audio = null) {
   try {
@@ -175,7 +212,8 @@ async function handleChat(message, base64Image = null, base64Audio = null) {
     let finalMessage = message;
     
     // Inject Local RAG Context
-    if (message && message.trim().length > 0) {
+    if (message && message.trim().length > 0 && shouldUseKnowledgeSearch(message)) {
+        const { searchKnowledge } = require('./knowledge_base');
         const retrievedDocs = await searchKnowledge(message);
         if (retrievedDocs && retrievedDocs.length > 0) {
             let contextString = `\n\n[LOCAL KNOWLEDGE BASE - USE THIS CONTEXT TO ANSWER]\n`;
@@ -187,9 +225,17 @@ async function handleChat(message, base64Image = null, base64Audio = null) {
     }
 
     if (provider === 'ollama') {
-        const axios = require('axios');
-        return await handleOllamaChat(finalMessage, base64Image, base64Audio, config, axios);
+        return await handleOllamaChat(finalMessage, base64Image, base64Audio, config);
+    }
+    
+    if (provider === 'anthropic') {
+        return await handleAnthropicChat(finalMessage, base64Image, config);
     }
+    
+    if (provider === 'openai') {
+        return await handleOpenAIChat(finalMessage, base64Image, config);
+    }
+
 
     const desiredModel = resolveGeminiModel();
     if (!chat || activeModel !== desiredModel) {
@@ -288,7 +334,127 @@ async function handleChat(message, base64Image = null, base64Audio = null) {
   }
 }
 
-async function handleOllamaChat(finalMessage, base64Image, base64Audio, config, axios) {
+async function handleAnthropicChat(finalMessage, base64Image, config) {
+    const history = readChatHistory() || [];
+    const apiKey = config.anthropicApiKey || process.env.ANTHROPIC_API_KEY;
+    if (!apiKey) return { response: "กรุณาใส่ Anthropic API Key ในการตั้งค่าก่อนนะคะ", action: { type: "none" } };
+
+    const mcpTools = mcpManager.getAllTools();
+    let mcpPrompt = "\n\nAVAILABLE MCP TOOLS:\n";
+    mcpTools.forEach(tool => {
+        mcpPrompt += `- Server: ${tool.serverName}, Tool: ${tool.name}\n  Desc: ${tool.description}\n  Args: ${JSON.stringify(tool.inputSchema.properties)}\n`;
+    });
+
+    const systemPrompt = systemInstruction + pluginManager.getPromptDescriptions() + mcpPrompt;
+    
+    const messages = [];
+    for (const msg of history.slice(-MAX_HISTORY_MESSAGES)) {
+        const role = msg.role === 'model' ? 'assistant' : 'user';
+        let text = Array.isArray(msg.parts) ? msg.parts.map(p => p.text || '').join('\n') : '';
+        if (text) messages.push({ role, content: text });
+    }
+
+    const content = [];
+    if (base64Image) {
+        const base64Data = base64Image.replace(/^data:image\/\w+;base64,/, '');
+        const mimeType = base64Image.match(/^data:(image\/\w+);base64,/)[1];
+        content.push({
+            type: "image",
+            source: { type: "base64", media_type: mimeType, data: base64Data }
+        });
+    }
+    content.push({ type: "text", text: finalMessage || "Analyze this." });
+    messages.push({ role: "user", content });
+
+    const response = await axios.post('https://api.anthropic.com/v1/messages', {
+        model: config.anthropicModel || 'claude-3-5-sonnet-latest',
+        max_tokens: 4096,
+        system: systemPrompt,
+        messages: messages
+    }, {
+        headers: {
+            'x-api-key': apiKey,
+            'anthropic-version': '2023-06-01',
+            'content-type': 'application/json'
+        }
+    });
+
+    const outputText = response.data.content[0].text;
+    history.push({ role: 'user', parts: [{ text: finalMessage }] });
+    history.push({ role: 'model', parts: [{ text: outputText }] });
+    writeChatHistory(history.slice(-MAX_HISTORY_MESSAGES));
+
+    return parseAiResponse(outputText);
+}
+
+async function handleOpenAIChat(finalMessage, base64Image, config) {
+    const history = readChatHistory() || [];
+    const apiKey = config.openaiApiKey || process.env.OPENAI_API_KEY;
+    if (!apiKey) return { response: "กรุณาใส่ OpenAI API Key ในการตั้งค่าก่อนนะคะ", action: { type: "none" } };
+
+    const mcpTools = mcpManager.getAllTools();
+    let mcpPrompt = "\n\nAVAILABLE MCP TOOLS:\n";
+    mcpTools.forEach(tool => {
+        mcpPrompt += `- Server: ${tool.serverName}, Tool: ${tool.name}\n  Desc: ${tool.description}\n  Args: ${JSON.stringify(tool.inputSchema.properties)}\n`;
+    });
+
+    const systemPrompt = systemInstruction + pluginManager.getPromptDescriptions() + mcpPrompt;
+    
+    const messages = [{ role: "system", content: systemPrompt }];
+    for (const msg of history.slice(-MAX_HISTORY_MESSAGES)) {
+        const role = msg.role === 'model' ? 'assistant' : 'user';
+        let text = Array.isArray(msg.parts) ? msg.parts.map(p => p.text || '').join('\n') : '';
+        if (text) messages.push({ role, content: text });
+    }
+
+    const content = [{ type: "text", text: finalMessage || "Analyze this." }];
+    if (base64Image) {
+        content.push({
+            type: "image_url",
+            image_url: { url: base64Image }
+        });
+    }
+    messages.push({ role: "user", content });
+
+    const response = await axios.post('https://api.openai.com/v1/chat/completions', {
+        model: config.openaiModel || 'gpt-4o',
+        messages: messages,
+        response_format: { type: "json_object" }
+    }, {
+        headers: {
+            'Authorization': `Bearer ${apiKey}`,
+            'Content-Type': 'application/json'
+        }
+    });
+
+    const outputText = response.data.choices[0].message.content;
+    history.push({ role: 'user', parts: [{ text: finalMessage }] });
+    history.push({ role: 'model', parts: [{ text: outputText }] });
+    writeChatHistory(history.slice(-MAX_HISTORY_MESSAGES));
+
+    return parseAiResponse(outputText);
+}
+
+function parseAiResponse(outputText) {
+    let parsedResult;
+    try {
+        parsedResult = JSON.parse(outputText);
+    } catch (e) {
+        const jsonMatch = outputText.match(/```json\n([\s\S]*?)\n```/) || outputText.match(/\{[\s\S]*\}/);
+        if (jsonMatch) {
+            parsedResult = JSON.parse(jsonMatch[jsonMatch.length > 1 ? 1 : 0]);
+        } else {
+            parsedResult = { response: outputText, action: { type: "none", target: "" } };
+        }
+    }
+    if (parsedResult && typeof parsedResult.response === 'string') {
+        parsedResult.response = decodeUnicode(parsedResult.response);
+    }
+    parsedResult.timestamp = new Date().toISOString();
+    return parsedResult;
+}
+
+async function handleOllamaChat(finalMessage, base64Image, base64Audio, config) {
     const history = readChatHistory() || [];
     pluginManager.loadPlugins();
     

package/src/AI_Brain/knowledge_base.js

@@ -1,15 +1,14 @@
 const fs = require('fs');
 const path = require('path');
 const os = require('os');
+const crypto = require('crypto');
 const { GoogleGenAI } = require('@google/genai');
 const pdf = require('pdf-parse');
 const mammoth = require('mammoth');
 const xlsx = require('xlsx');
-const axios = require('axios');
-const cheerio = require('cheerio');
 const { readConfig } = require('../System/config_manager');
 
-// Handle electron dependency safely for benchmarks/tests
+// Handle electron dependency safely
 let app;
 try {
     const electron = require('electron');
@@ -20,7 +19,7 @@ try {
 
 let ai = null;
 let activeApiKey = '';
-const initialEnvKey = (process.env.GEMINI_API_KEY || '').trim();
+let DatabaseSync = null;
 
 function resolveApiKey() {
     let settingsKey = '';
@@ -30,53 +29,63 @@ function resolveApiKey() {
     } catch (e) {
         settingsKey = '';
     }
-
-    const envKey = initialEnvKey;
-    const selectedKey = settingsKey || envKey || '';
-
-    if (selectedKey !== (process.env.GEMINI_API_KEY || '')) {
-        process.env.GEMINI_API_KEY = selectedKey;
-    }
-
+    const selectedKey = settingsKey || process.env.GEMINI_API_KEY || '';
     activeApiKey = selectedKey;
     return selectedKey;
 }
 
 function getAiClient() {
-    const prevKey = activeApiKey;
-    const nextKey = resolveApiKey();
-    if (!ai || nextKey !== prevKey) {
-        ai = new GoogleGenAI({ apiKey: nextKey });
+    const key = resolveApiKey();
+    if (!ai || activeApiKey !== key) {
+        ai = new GoogleGenAI({ apiKey: key });
     }
     return ai;
 }
 
 function getDbPath() {
+    const fileName = 'mint-knowledge.sqlite';
     if (app && app.getPath) {
-        return path.join(app.getPath('userData'), 'mint-knowledge.json');
+        return path.join(app.getPath('userData'), fileName);
     }
-    // Use global .mint directory for CLI/Benchmarking
     const mintDir = path.join(os.homedir(), '.mint');
-    if (!fs.existsSync(mintDir)) {
-        fs.mkdirSync(mintDir, { recursive: true });
-    }
-    return path.join(mintDir, 'mint-knowledge.json');
+    if (!fs.existsSync(mintDir)) fs.mkdirSync(mintDir, { recursive: true });
+    return path.join(mintDir, fileName);
 }
 
-function loadDb() {
-    try {
-        const p = getDbPath();
-        if (fs.existsSync(p)) {
-            return JSON.parse(fs.readFileSync(p, 'utf8'));
-        }
-    } catch (err) {
-        console.error('[KnowledgeBase] Load Error:', err);
+function getDatabaseSync() {
+    if (!DatabaseSync) {
+        ({ DatabaseSync } = require('node:sqlite'));
     }
-    return { documents: [] };
+    return DatabaseSync;
 }
 
-function saveDb(db) {
-    fs.writeFileSync(getDbPath(), JSON.stringify(db, null, 2));
+// Initialize Database
+let dbInstance = null;
+function getDb() {
+    if (dbInstance) return dbInstance;
+    const dbPath = getDbPath();
+    const Database = getDatabaseSync();
+    dbInstance = new Database(dbPath);
+
+    // Create Tables
+    dbInstance.exec(`
+        CREATE TABLE IF NOT EXISTS sources (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            path TEXT UNIQUE,
+            name TEXT,
+            hash TEXT,
+            last_indexed DATETIME DEFAULT CURRENT_TIMESTAMP
+        );
+        CREATE TABLE IF NOT EXISTS chunks (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            source_id INTEGER,
+            text TEXT,
+            embedding BLOB,
+            FOREIGN KEY(source_id) REFERENCES sources(id) ON DELETE CASCADE
+        );
+        CREATE INDEX IF NOT EXISTS idx_chunks_source ON chunks(source_id);
+    `);
+    return dbInstance;
 }
 
 async function generateEmbedding(text) {
@@ -85,138 +94,203 @@ async function generateEmbedding(text) {
         model: 'gemini-embedding-001',
         contents: text,
     });
-    // The google/genai package returns an array of embeddings
     return response.embeddings[0].values;
 }
 
+
 function cosineSimilarity(vecA, vecB) {
-    let dotProduct = 0.0;
-    let normA = 0.0;
-    let normB = 0.0;
+    let dotProduct = 0, normA = 0, normB = 0;
     for (let i = 0; i < vecA.length; i++) {
         dotProduct += vecA[i] * vecB[i];
         normA += vecA[i] * vecA[i];
         normB += vecB[i] * vecB[i];
     }
-    if (normA === 0 || normB === 0) return 0;
     return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
 }
 
+function getFileHash(filePath) {
+    const content = fs.readFileSync(filePath);
+    return crypto.createHash('md5').update(content).digest('hex');
+}
+
 function chunkText(text, maxChars = 1000, overlap = 200) {
     const chunks = [];
     let current = 0;
-    const step = maxChars - overlap;
     while (current < text.length) {
         chunks.push(text.slice(current, current + maxChars));
-        current += step;
+        current += (maxChars - overlap);
+        if (current >= text.length) break;
     }
     return chunks;
 }
 
-/**
- * Reads a local file or URL, chunks its text, generates embeddings, and saves to knowledge base.
- */
-async function indexFile(resourcePath) {
+async function indexFile(filePath) {
     try {
-        if (!resourcePath || resourcePath.trim() === '') return "ไม่พบข้อมูล กรุณาระบุ Path หรือ URL ค่ะ";
-        
+        if (!fs.existsSync(filePath)) return `ไม่พบไฟล์: ${filePath}`;
+        const stats = fs.statSync(filePath);
+        if (stats.isDirectory()) return await indexFolder(filePath);
+        if (stats.size > 10 * 1024 * 1024) return `ไฟล์ใหญ่เกินไป (> 10MB): ${filePath}`;
+
+        const hash = getFileHash(filePath);
+        const db = getDb();
+
+        // Check if already indexed and unchanged
+        const checkStmt = db.prepare("SELECT id, hash FROM sources WHERE path = ?");
+        const existing = checkStmt.get(filePath);
+
+        if (existing && existing.hash === hash) {
+            return `⏩ ${path.basename(filePath)} ไม่มีการเปลี่ยนแปลง (ข้ามการอ่าน)`;
+        }
+
+        console.log(`[RAG] Indexing ${filePath}...`);
         let content = '';
-        let sourceName = '';
-        let resourceId = '';
-
-        // Handle Web URLs
-        if (resourcePath.startsWith('http://') || resourcePath.startsWith('https://')) {
-            sourceName = resourcePath;
-            resourceId = resourcePath;
-            try {
-                const response = await axios.get(resourcePath);
-                const $ = cheerio.load(response.data);
-                $('script, style, noscript, nav, footer, header').remove();
-                content = $('body').text().replace(/\s+/g, ' ').trim();
-            } catch (e) {
-                return `ไม่สามารถดึงข้อมูลจากเว็บไซต์ได้ค่ะ: ${e.message}`;
+        const ext = path.extname(filePath).toLowerCase();
+
+        // Extraction logic
+        if (ext === '.pdf') {
+            const data = await pdf(fs.readFileSync(filePath));
+            content = data.text;
+        } else if (ext === '.docx') {
+            const res = await mammoth.extractRawText({ path: filePath });
+            content = res.value;
+        } else if (ext === '.xlsx') {
+            const wb = xlsx.readFile(filePath);
+            content = wb.SheetNames.map(n => xlsx.utils.sheet_to_csv(wb.Sheets[n])).join('\n');
+        } else {
+            content = fs.readFileSync(filePath, 'utf8');
+        }
+
+        if (!content.trim()) return `⚠️ ไฟล์ไม่มีข้อความ: ${filePath}`;
+
+        // Begin transaction
+        db.exec("BEGIN TRANSACTION");
+        try {
+            if (existing) {
+                db.prepare("DELETE FROM chunks WHERE source_id = ?").run(existing.id);
+                db.prepare("UPDATE sources SET hash = ?, last_indexed = CURRENT_TIMESTAMP WHERE id = ?").run(hash, existing.id);
+            } else {
+                db.prepare("INSERT INTO sources (path, name, hash) VALUES (?, ?, ?)").run(filePath, path.basename(filePath), hash);
             }
-        } 
-        // Handle Local Files
-        else {
-            const filePath = resourcePath;
-            if (!fs.existsSync(filePath)) return `ไม่พบไฟล์: ${filePath}`;
             
-            const stats = fs.statSync(filePath);
-            if (stats.size > 5 * 1024 * 1024) return `ขนาดไฟล์ใหญ่เกินไป (> 5MB): ${filePath}`; 
+            const sourceId = existing ? existing.id : db.prepare("SELECT last_insert_rowid() as id").get().id;
+            const chunks = chunkText(content);
             
-            sourceName = path.basename(filePath);
-            resourceId = filePath;
-            const ext = path.extname(filePath).toLowerCase();
-
-            if (ext === '.pdf') {
-                const dataBuffer = fs.readFileSync(filePath);
-                const data = await pdf(dataBuffer);
-                content = data.text;
-            } else if (ext === '.docx') {
-                const result = await mammoth.extractRawText({path: filePath});
-                content = result.value;
-            } else if (ext === '.xlsx') {
-                const workbook = xlsx.readFile(filePath);
-                content = '';
-                for (const sheetName of workbook.SheetNames) {
-                    const sheet = workbook.Sheets[sheetName];
-                    const csv = xlsx.utils.sheet_to_csv(sheet);
-                    content += `\n--- Sheet: ${sheetName} ---\n` + csv;
-                }
-            } else {
-                content = fs.readFileSync(filePath, 'utf8');
+            const insertChunk = db.prepare("INSERT INTO chunks (source_id, text, embedding) VALUES (?, ?, ?)");
+            for (const chunk of chunks) {
+                const embedding = await generateEmbedding(chunk);
+                const embeddingBlob = Buffer.from(new Float32Array(embedding).buffer);
+                insertChunk.run(sourceId, chunk, embeddingBlob);
             }
+            db.exec("COMMIT");
+            return `✅ Successfully indexed ${path.basename(filePath)} (${chunks.length} chunks)`;
+        } catch (e) {
+            db.exec("ROLLBACK");
+            throw e;
         }
-
-        if (!content || content.trim().length === 0) return `ข้อมูลว่างเปล่าหรือไม่มีข้อความ: ${resourcePath}`;
-        
-        const chunks = chunkText(content);
-        const db = loadDb();
-        
-        for (let i = 0; i < chunks.length; i++) {
-            const embedding = await generateEmbedding(chunks[i]);
-            db.documents.push({
-                id: `${resourceId}#${i}-${Date.now()}`,
-                source: sourceName,
-                path: resourcePath,
-                text: chunks[i],
-                embedding
-            });
-        }
-        
-        saveDb(db);
-        return `✅ เรียนรู้ข้อมูลจาก ${sourceName} เรียบร้อยแล้ว (แบ่งเป็น ${chunks.length} ส่วน)`;
     } catch (err) {
-        console.error('[KnowledgeBase] Indexing error:', err);
-        return `❌ เกิดข้อผิดพลาดในการเรียนรู้ไฟล์: ${err.message}`;
+        console.error('[RAG] Error:', err);
+        return `❌ Failed to index: ${err.message}`;
     }
 }
 
 /**
- * Searches the local knowledge base for relevant chunks.
+ * Recursively gets all files in a directory asynchronously
  */
+async function getAllFiles(dirPath, arrayOfFiles = []) {
+    const files = await fs.promises.readdir(dirPath, { withFileTypes: true });
+    
+    for (const file of files) {
+        const fullPath = path.join(dirPath, file.name);
+        if (file.isDirectory()) {
+            await getAllFiles(fullPath, arrayOfFiles);
+        } else {
+            arrayOfFiles.push(fullPath);
+        }
+    }
+    return arrayOfFiles;
+}
+
+async function indexFolder(folderPath) {
+    console.log(`[RAG] Indexing folder: ${folderPath}`);
+    const files = await getAllFiles(folderPath);
+    console.log(`[RAG] Found ${files.length} files to check.`);
+    
+    // Process in small batches to avoid blocking
+    const BATCH_SIZE = 5;
+    let indexedCount = 0;
+    let skippedCount = 0;
+
+    for (let i = 0; i < files.length; i += BATCH_SIZE) {
+        const batch = files.slice(i, i + BATCH_SIZE);
+        await Promise.all(batch.map(async (file) => {
+            const res = await indexFile(file);
+            if (res && res.startsWith('✅')) indexedCount++;
+            else skippedCount++;
+        }));
+    }
+    
+    console.log(`[RAG] Indexing complete. ${indexedCount} new/updated, ${skippedCount} skipped.`);
+    return `📂 Folder indexing complete: ${indexedCount} learned, ${skippedCount} skipped.`;
+}
+
 async function searchKnowledge(query, topK = 3) {
-    const db = loadDb();
-    if (!db.documents || db.documents.length === 0) return null;
+    const startTime = Date.now();
+    const db = getDb();
+    const MAX_CHUNKS_TO_SEARCH = 2000; // Limit search to keep it fast
     
+    const countRes = db.prepare("SELECT COUNT(*) as count FROM chunks").get();
+    if (!countRes || countRes.count === 0) return null;
+
     try {
         const queryVector = await generateEmbedding(query);
-        const results = db.documents.map(doc => ({
-            ...doc,
-            score: cosineSimilarity(queryVector, doc.embedding)
-        })).sort((a, b) => b.score - a.score);
-        
-        // Return top results above a threshold
-        const top = results.slice(0, topK).filter(r => r.score > 0.65);
-        if (top.length > 0) {
-            console.log(`[KnowledgeBase] Found ${top.length} matches for query.`);
-            return top;
+        const queryTyped = new Float32Array(queryVector);
+        const results = [];
+
+        // Search most recent or top chunks first, but limit the total scan
+        const stmt = db.prepare("SELECT text, embedding, source_id FROM chunks LIMIT ?");
+        let processed = 0;
+
+        for (const c of stmt.iterate(MAX_CHUNKS_TO_SEARCH)) {
+            if (!c.embedding) continue;
+            processed++;
+            
+            const chunkVector = new Float32Array(c.embedding.buffer, c.embedding.byteOffset, c.embedding.byteLength / 4);
+            
+            let dotProduct = 0, normA = 0, normB = 0;
+            for (let i = 0; i < queryTyped.length; i++) {
+                const a = queryTyped[i];
+                const b = chunkVector[i];
+                dotProduct += a * b;
+                normA += a * a;
+                normB += b * b;
+            }
+            const score = dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
+
+            if (score > 0.65) {
+                results.push({ text: c.text, score, source_id: c.source_id });
+            }
         }
-    } catch(err) {
-        console.error("[KnowledgeBase] Search error:", err);
+
+        if (results.length > 0) {
+            results.sort((a, b) => b.score - a.score);
+            const top = results.slice(0, topK);
+            
+            const sourceIds = [...new Set(top.map(t => t.source_id))];
+            const sources = db.prepare(`SELECT id, name FROM sources WHERE id IN (${sourceIds.join(',')})`).all();
+            const sourceMap = Object.fromEntries(sources.map(s => [s.id, s.name]));
+            
+            console.log(`[RAG] Search took ${Date.now() - startTime}ms for ${processed} chunks.`);
+            return top.map(t => ({
+                text: t.text,
+                source: sourceMap[t.source_id],
+                score: t.score
+            }));
+        }
+    } catch (e) {
+        console.error("[RAG] Search Error:", e);
     }
     return null;
 }
 
-module.exports = { indexFile, searchKnowledge };
+
+module.exports = { indexFile, indexFolder, searchKnowledge };