cozo-memory 1.0.9 → 1.1.1

package/dist/embedding-service.js

@@ -94,9 +94,16 @@ class EmbeddingService {
     tokenizer = null;
     modelId;
     dimensions;
+    useOllama;
+    ollamaModel;
+    ollamaBaseUrl;
     queue = Promise.resolve();
     constructor() {
         this.cache = new LRUCache(1000, 3600000); // 1000 entries, 1h TTL
+        // Check if Ollama should be used
+        this.useOllama = process.env.USE_OLLAMA === 'true';
+        this.ollamaModel = process.env.OLLAMA_EMBEDDING_MODEL || 'argus-ai/pplx-embed-v1-0.6b:q8_0';
+        this.ollamaBaseUrl = process.env.OLLAMA_BASE_URL || 'http://localhost:11434';
         // Support multiple embedding models via environment variable
         this.modelId = process.env.EMBEDDING_MODEL || "Xenova/bge-m3";
         // Set dimensions based on model
@@ -106,9 +113,20 @@ class EmbeddingService {
             "Xenova/bge-small-en-v1.5": 384,
             "Xenova/nomic-embed-text-v1": 768,
             "onnx-community/Qwen3-Embedding-0.6B-ONNX": 1024,
+            // Note: perplexity-ai models require manual ONNX file placement
+            // See PPLX_EMBED_INTEGRATION.md for instructions
+            "perplexity-ai/pplx-embed-v1-0.6b": 1024,
+            "perplexity-ai/pplx-embed-v1-4b": 2560,
+            // Ollama models
+            "argus-ai/pplx-embed-v1-0.6b:q8_0": 1024,
         };
-        this.dimensions = dimensionMap[this.modelId] || 1024;
-        console.error(`[EmbeddingService] Using model: ${this.modelId} (${this.dimensions} dimensions)`);
+        this.dimensions = dimensionMap[this.useOllama ? this.ollamaModel : this.modelId] || 1024;
+        if (this.useOllama) {
+            console.error(`[EmbeddingService] Using Ollama: ${this.ollamaModel} @ ${this.ollamaBaseUrl} (${this.dimensions} dimensions)`);
+        }
+        else {
+            console.error(`[EmbeddingService] Using ONNX model: ${this.modelId} (${this.dimensions} dimensions)`);
+        }
     }
     // Public getter for dimensions
     getDimensions() {
@@ -125,6 +143,11 @@ class EmbeddingService {
     async init() {
         if (this.session && this.tokenizer)
             return;
+        // Skip ONNX initialization if using Ollama
+        if (this.useOllama) {
+            console.error('[EmbeddingService] Using Ollama backend, skipping ONNX initialization');
+            return;
+        }
         try {
             // 1. Check if model needs to be downloaded
             // Extract namespace and model name from modelId (e.g., "Xenova/bge-m3" or "onnx-community/Qwen3-Embedding-0.6B-ONNX")
@@ -139,10 +162,23 @@ class EmbeddingService {
             if (!fs.existsSync(fp32Path) && !fs.existsSync(quantizedPath)) {
                 console.log(`[EmbeddingService] Model not found, downloading ${this.modelId}...`);
                 console.log(`[EmbeddingService] This may take a few minutes on first run.`);
-                // Import AutoModel dynamically to trigger download
-                const { AutoModel } = await import("@xenova/transformers");
-                await AutoModel.from_pretrained(this.modelId, { quantized: false });
-                console.log(`[EmbeddingService] Model download completed.`);
+                // Check if this is a Xenova-compatible model
+                if (namespace === 'Xenova' || namespace === 'onnx-community') {
+                    // Import AutoModel dynamically to trigger download
+                    const { AutoModel } = await import("@xenova/transformers");
+                    await AutoModel.from_pretrained(this.modelId, { quantized: false });
+                    console.log(`[EmbeddingService] Model download completed.`);
+                }
+                else {
+                    // For non-Xenova models (like perplexity-ai), provide manual download instructions
+                    console.error(`[EmbeddingService] ERROR: Model ${this.modelId} is not available via @xenova/transformers`);
+                    console.error(`[EmbeddingService] Please download the model manually:`);
+                    console.error(`[EmbeddingService] 1. Visit: https://huggingface.co/${this.modelId}`);
+                    console.error(`[EmbeddingService] 2. Download the 'onnx' folder contents`);
+                    console.error(`[EmbeddingService] 3. Place files in: ${baseDir}`);
+                    console.error(`[EmbeddingService] See PPLX_EMBED_INTEGRATION.md for detailed instructions`);
+                    throw new Error(`Model ${this.modelId} requires manual download. See error messages above.`);
+                }
             }
             // 2. Load Tokenizer
             if (!this.tokenizer) {
@@ -188,6 +224,10 @@ class EmbeddingService {
                 return cached;
             }
             try {
+                // Use Ollama if enabled
+                if (this.useOllama) {
+                    return await this.embedWithOllama(textStr);
+                }
                 await this.init();
                 if (!this.session || !this.tokenizer)
                     throw new Error("Session/Tokenizer not initialized");
@@ -240,6 +280,37 @@ class EmbeddingService {
             }
         });
     }
+    async embedWithOllama(text) {
+        try {
+            const response = await fetch(`${this.ollamaBaseUrl}/api/embeddings`, {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                },
+                body: JSON.stringify({
+                    model: this.ollamaModel,
+                    prompt: text,
+                }),
+            });
+            if (!response.ok) {
+                throw new Error(`Ollama API error: ${response.status} ${response.statusText}`);
+            }
+            const data = await response.json();
+            if (!data.embedding || !Array.isArray(data.embedding)) {
+                throw new Error('Invalid response from Ollama API');
+            }
+            const embedding = data.embedding;
+            // Normalize the embedding
+            const normalized = this.normalize(embedding);
+            // Cache it
+            this.cache.set(text, normalized);
+            return normalized;
+        }
+        catch (error) {
+            console.error(`[EmbeddingService] Ollama error for "${text.substring(0, 20)}...":`, error?.message || error);
+            return new Array(this.dimensions).fill(0);
+        }
+    }
     // Batch-Embeddings
     async embedBatch(texts) {
         // For now, process sequentially via serialized queue to avoid overloading
@@ -306,7 +377,8 @@ class EmbeddingService {
         return {
             size: this.cache.size(),
             maxSize: 1000,
-            model: this.modelId,
+            model: this.useOllama ? this.ollamaModel : this.modelId,
+            backend: this.useOllama ? 'ollama' : 'onnx',
             dimensions: this.dimensions
         };
     }

package/dist/eval-suite.js

@@ -88,10 +88,14 @@ async function runEvaluation() {
     }
     const server = new index_1.MemoryServer(EVAL_DB_PATH);
     await server.embeddingService.embed("warmup");
+    // Warmup reranker
+    await server.hybridSearch.advancedSearch({ query: "warmup", limit: 1, rerank: true });
     await setupEvalData(server);
     const methods = [
         { name: "Hybrid Search", func: (q) => server.hybridSearch.search({ query: q, limit: 10 }) },
+        { name: "Reranked Search", func: (q) => server.hybridSearch.search({ query: q, limit: 10, rerank: true }) },
         { name: "Graph-RAG", func: (q) => server.hybridSearch.graphRag({ query: q, limit: 10, graphConstraints: { maxDepth: 2 } }) },
+        { name: "Graph-RAG (Reranked)", func: (q) => server.hybridSearch.graphRag({ query: q, limit: 10, graphConstraints: { maxDepth: 2 }, rerank: true }) },
         { name: "Graph-Walking", func: (q) => server.graph_walking({ query: q, limit: 10, max_depth: 3 }) }
     ];
     const summary = [];
@@ -101,6 +105,9 @@ async function runEvaluation() {
         let totalRecall10 = 0;
         let totalMRR = 0;
         let totalLatency = 0;
+        const n = EVAL_DATASET.length;
+        // Reset cache between methods to get accurate latency
+        await server.hybridSearch.clearCache();
         for (const task of EVAL_DATASET) {
             const t0 = perf_hooks_1.performance.now();
             const results = await method.func(task.query);
@@ -113,7 +120,6 @@ async function runEvaluation() {
             totalMRR += mrr;
             totalLatency += (t1 - t0);
         }
-        const n = EVAL_DATASET.length;
         summary.push({
             Method: method.name,
             "Recall@3": (totalRecall3 / n).toFixed(3),

package/dist/hybrid-search.js

@@ -5,15 +5,18 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.HybridSearch = void 0;
 const crypto_1 = __importDefault(require("crypto"));
+const reranker_service_1 = require("./reranker-service");
 const SEMANTIC_CACHE_THRESHOLD = 0.95;
 class HybridSearch {
     db;
     embeddingService;
+    rerankerService;
     searchCache = new Map();
     CACHE_TTL = 300000; // 5 minutes cache
     constructor(db, embeddingService) {
         this.db = db;
         this.embeddingService = embeddingService;
+        this.rerankerService = new reranker_service_1.RerankerService();
     }
     getCacheKey(options) {
         const str = JSON.stringify({
@@ -75,6 +78,65 @@ class HybridSearch {
             return { ...r, score };
         });
     }
+    applyContextBoost(results, options) {
+        const { session_id, task_id } = options;
+        if (!session_id && !task_id)
+            return results;
+        return results.map(result => {
+            let boost = 1.0;
+            let reasons = [];
+            const metadata = result.metadata || {};
+            if (task_id && metadata.task_id === task_id) {
+                boost += 0.5;
+                reasons.push("Task Match");
+            }
+            if (session_id && metadata.session_id === session_id) {
+                boost += 0.3;
+                reasons.push("Session Match");
+            }
+            if (boost > 1.0) {
+                // Cap the score at 1.0 to stay within standard search score range
+                const newScore = Math.min(1.0, result.score * boost);
+                return {
+                    ...result,
+                    score: newScore,
+                    explanation: (typeof result.explanation === 'string' ? result.explanation : JSON.stringify(result.explanation)) +
+                        ` | Context Boost (x${boost.toFixed(1)}): ${reasons.join(', ')}`
+                };
+            }
+            return result;
+        });
+    }
+    async applyReranking(query, results) {
+        if (results.length <= 1)
+            return results;
+        console.error(`[HybridSearch] Reranking ${results.length} candidates...`);
+        const documents = results.map(r => {
+            const parts = [
+                r.name ? `Name: ${r.name}` : '',
+                r.type ? `Type: ${r.type}` : '',
+                r.text ? `Description: ${r.text}` : '',
+                r.metadata ? `Details: ${JSON.stringify(r.metadata)}` : ''
+            ].filter(p => p !== '');
+            return parts.join(' | ');
+        });
+        try {
+            const rerankedOrder = await this.rerankerService.rerank(query, documents);
+            return rerankedOrder.map((item, i) => {
+                const original = results[item.index];
+                return {
+                    ...original,
+                    score: (item.score + 1.0) / 2.0, // Normalize to 0-1 range if it's logits, or just use as is
+                    explanation: (typeof original.explanation === 'string' ? original.explanation : JSON.stringify(original.explanation)) +
+                        ` | Reranked (Rank ${i + 1}, Cross-Encoder Score: ${item.score.toFixed(4)})`
+                };
+            });
+        }
+        catch (e) {
+            console.error(`[HybridSearch] Reranking failed, returning original results:`, e);
+            return results;
+        }
+    }
     async advancedSearch(options) {
         console.error("[HybridSearch] Starting advancedSearch with options:", JSON.stringify(options, null, 2));
         const { query, limit = 10, filters, graphConstraints, vectorParams } = options;
@@ -181,10 +243,6 @@ class HybridSearch {
             `:sort -score`,
             `:limit $limit`
         ].join('\n').trim();
-        console.error('--- DEBUG: Cozo Datalog Query ---');
-        console.error(datalogQuery);
-        console.error('--- DEBUG: Params ---');
-        console.error(JSON.stringify(params, null, 2));
         try {
             const results = await this.db.run(datalogQuery, params);
             let searchResults = results.rows.map((r) => ({
@@ -211,7 +269,14 @@ class HybridSearch {
                     return Object.entries(filters.metadata).every(([key, val]) => r.metadata[key] === val);
                 });
             }
-            const finalResults = this.applyTimeDecay(searchResults);
+            const timeDecayedResults = this.applyTimeDecay(searchResults);
+            const finalResults = this.applyContextBoost(timeDecayedResults, options);
+            // Phase 3: Reranking
+            if (options.rerank) {
+                const rerankedResults = await this.applyReranking(options.query, finalResults);
+                await this.updateCache(options, queryEmbedding, rerankedResults);
+                return rerankedResults;
+            }
             await this.updateCache(options, queryEmbedding, finalResults);
             return finalResults;
         }
@@ -300,7 +365,6 @@ class HybridSearch {
       :sort -score
       :limit $limit
     `.trim();
-        console.error("[HybridSearch] Graph-RAG Datalog Query:\n", datalogQuery);
         try {
             const results = await this.db.run(datalogQuery, params);
             let searchResults = results.rows.map((r) => ({
@@ -330,7 +394,11 @@ class HybridSearch {
                     return Object.entries(filters.metadata).every(([key, val]) => r.metadata[key] === val);
                 });
             }
-            return this.applyTimeDecay(searchResults);
+            const decayedResults = this.applyTimeDecay(searchResults);
+            if (options.rerank) {
+                return await this.applyReranking(options.query, decayedResults);
+            }
+            return decayedResults;
         }
         catch (e) {
             console.error("[HybridSearch] Error in graphRag:", e.message);
@@ -414,7 +482,8 @@ No markdown, no explanation. Just the JSON.`;
                     filters: {
                         ...options.filters,
                         entityTypes: ["CommunitySummary"]
-                    }
+                    },
+                    rerank: options.rerank
                 });
                 // If no community summaries found, fallback to standard search
                 if (results.length === 0) {
@@ -437,5 +506,15 @@ No markdown, no explanation. Just the JSON.`;
             }
         }));
     }
+    async clearCache() {
+        this.searchCache.clear();
+        try {
+            await this.db.run(`{ ?[query_hash] := *search_cache{query_hash} :rm search_cache {query_hash} }`);
+            console.error("[HybridSearch] Cache cleared successfully.");
+        }
+        catch (e) {
+            console.error("[HybridSearch] Error clearing cache:", e);
+        }
+    }
 }
 exports.HybridSearch = HybridSearch;