@smyslenny/agent-memory 2.1.0 → 2.2.0

package/dist/mcp/server.js

@@ -14580,6 +14580,26 @@ function getStrategy(intent) {
 }
 
 // src/search/rerank.ts
+async function rerankWithProvider(results, query, provider) {
+  if (results.length === 0) return results;
+  const documents = results.map((r) => r.memory.content);
+  try {
+    const apiResults = await provider.rerank(query, documents);
+    const scoreMap = new Map(apiResults.map((r) => [r.index, r.relevance_score]));
+    return results.map((r, i) => {
+      const score = scoreMap.get(i);
+      if (score === void 0) return r;
+      return {
+        ...r,
+        score,
+        matchReason: `${r.matchReason}+rerank`
+      };
+    });
+  } catch (err) {
+    console.warn("[agent-memory] External rerank failed, falling back:", err);
+    return results;
+  }
+}
 function rerank(results, opts) {
   const now2 = Date.now();
   const scored = results.map((r) => {
@@ -14741,7 +14761,8 @@ async function searchHybrid(db, query, opts) {
   if (!provider || !model) {
     return bm25.slice(0, limit);
   }
-  const qVec = Float32Array.from(await provider.embed(query));
+  const embedFn = provider.embedQuery ?? provider.embed;
+  const qVec = Float32Array.from(await embedFn.call(provider, query));
   const embeddings = listEmbeddings(db, agentId, model);
   const scored = [];
   for (const e of embeddings) {
@@ -14774,6 +14795,29 @@ async function searchHybrid(db, query, opts) {
 }
 
 // src/search/providers.ts
+var QWEN_DEFAULT_INSTRUCTION = "Given a query, retrieve the most semantically relevant document";
+function getDefaultInstruction(model) {
+  const m = model.toLowerCase();
+  if (m.includes("qwen")) return QWEN_DEFAULT_INSTRUCTION;
+  if (m.includes("gemini")) return null;
+  return null;
+}
+function resolveInstruction(model) {
+  const override = process.env.AGENT_MEMORY_EMBEDDINGS_INSTRUCTION;
+  if (override !== void 0) {
+    const normalized = override.trim();
+    if (!normalized) return null;
+    const lowered = normalized.toLowerCase();
+    if (lowered === "none" || lowered === "off" || lowered === "false" || lowered === "null") return null;
+    return normalized;
+  }
+  return getDefaultInstruction(model);
+}
+function buildQueryInput(query, instructionPrefix) {
+  if (!instructionPrefix) return query;
+  return `Instruct: ${instructionPrefix}
+Query: ${query}`;
+}
 function getEmbeddingProviderFromEnv() {
   const provider = (process.env.AGENT_MEMORY_EMBEDDINGS_PROVIDER ?? "none").toLowerCase();
   if (provider === "none" || provider === "off" || provider === "false") return null;
@@ -14782,14 +14826,24 @@ function getEmbeddingProviderFromEnv() {
     const model = process.env.AGENT_MEMORY_EMBEDDINGS_MODEL ?? "text-embedding-3-small";
     const baseUrl = process.env.OPENAI_BASE_URL ?? "https://api.openai.com/v1";
     if (!apiKey) return null;
-    return createOpenAIProvider({ apiKey, model, baseUrl });
+    const instruction = resolveInstruction(model);
+    return createOpenAIProvider({ apiKey, model, baseUrl, instruction });
+  }
+  if (provider === "gemini" || provider === "google") {
+    const apiKey = process.env.GEMINI_API_KEY ?? process.env.OPENAI_API_KEY;
+    const model = process.env.AGENT_MEMORY_EMBEDDINGS_MODEL ?? "gemini-embedding-001";
+    const baseUrl = process.env.GEMINI_BASE_URL ?? process.env.OPENAI_BASE_URL ?? "https://generativelanguage.googleapis.com/v1beta";
+    if (!apiKey) return null;
+    const instruction = resolveInstruction(model);
+    return createOpenAIProvider({ id: "gemini", apiKey, model, baseUrl, instruction });
   }
   if (provider === "qwen" || provider === "dashscope" || provider === "tongyi") {
     const apiKey = process.env.DASHSCOPE_API_KEY;
     const model = process.env.AGENT_MEMORY_EMBEDDINGS_MODEL ?? "text-embedding-v3";
     const baseUrl = process.env.DASHSCOPE_BASE_URL ?? "https://dashscope.aliyuncs.com";
     if (!apiKey) return null;
-    return createDashScopeProvider({ apiKey, model, baseUrl });
+    const instruction = resolveInstruction(model);
+    return createDashScopeProvider({ apiKey, model, baseUrl, instruction });
   }
   return null;
 }
@@ -14806,52 +14860,112 @@ function normalizeEmbedding(e) {
 }
 function createOpenAIProvider(opts) {
   const baseUrl = opts.baseUrl ?? "https://api.openai.com/v1";
+  const instructionPrefix = opts.instruction ?? null;
+  async function requestEmbedding(input) {
+    const resp = await fetch(`${baseUrl.replace(/\/$/, "")}/embeddings`, {
+      method: "POST",
+      headers: {
+        "content-type": "application/json",
+        authorization: authHeader(opts.apiKey)
+      },
+      body: JSON.stringify({ model: opts.model, input })
+    });
+    if (!resp.ok) {
+      const body = await resp.text().catch(() => "");
+      throw new Error(`OpenAI embeddings failed: ${resp.status} ${resp.statusText} ${body}`.trim());
+    }
+    const data = await resp.json();
+    return normalizeEmbedding(data.data?.[0]?.embedding);
+  }
   return {
-    id: "openai",
+    id: opts.id ?? "openai",
     model: opts.model,
+    instructionPrefix,
     async embed(text) {
-      const resp = await fetch(`${baseUrl.replace(/\/$/, "")}/embeddings`, {
-        method: "POST",
-        headers: {
-          "content-type": "application/json",
-          authorization: authHeader(opts.apiKey)
-        },
-        body: JSON.stringify({ model: opts.model, input: text })
-      });
-      if (!resp.ok) {
-        const body = await resp.text().catch(() => "");
-        throw new Error(`OpenAI embeddings failed: ${resp.status} ${resp.statusText} ${body}`.trim());
-      }
-      const data = await resp.json();
-      const emb = data.data?.[0]?.embedding;
-      return normalizeEmbedding(emb);
+      return requestEmbedding(text);
+    },
+    async embedQuery(query) {
+      return requestEmbedding(buildQueryInput(query, instructionPrefix));
     }
   };
 }
 function createDashScopeProvider(opts) {
   const baseUrl = opts.baseUrl ?? "https://dashscope.aliyuncs.com";
+  const instructionPrefix = opts.instruction ?? null;
+  async function requestEmbedding(text) {
+    const resp = await fetch(`${baseUrl.replace(/\/$/, "")}/api/v1/services/embeddings/text-embedding/text-embedding`, {
+      method: "POST",
+      headers: {
+        "content-type": "application/json",
+        authorization: authHeader(opts.apiKey)
+      },
+      body: JSON.stringify({
+        model: opts.model,
+        input: { texts: [text] }
+      })
+    });
+    if (!resp.ok) {
+      const body = await resp.text().catch(() => "");
+      throw new Error(`DashScope embeddings failed: ${resp.status} ${resp.statusText} ${body}`.trim());
+    }
+    const data = await resp.json();
+    const emb = data.output?.embeddings?.[0]?.embedding ?? data.output?.embeddings?.[0]?.vector ?? data.output?.embedding ?? data.data?.[0]?.embedding;
+    return normalizeEmbedding(emb);
+  }
   return {
     id: "dashscope",
     model: opts.model,
+    instructionPrefix,
     async embed(text) {
-      const resp = await fetch(`${baseUrl.replace(/\/$/, "")}/api/v1/services/embeddings/text-embedding/text-embedding`, {
+      return requestEmbedding(text);
+    },
+    async embedQuery(query) {
+      return requestEmbedding(buildQueryInput(query, instructionPrefix));
+    }
+  };
+}
+
+// src/search/rerank-provider.ts
+function authHeader2(apiKey) {
+  return apiKey.startsWith("Bearer ") ? apiKey : `Bearer ${apiKey}`;
+}
+function getRerankerProviderFromEnv() {
+  const provider = (process.env.AGENT_MEMORY_RERANK_PROVIDER ?? "none").toLowerCase();
+  if (provider === "none" || provider === "off") return null;
+  if (provider === "openai" || provider === "jina" || provider === "cohere") {
+    const apiKey = process.env.AGENT_MEMORY_RERANK_API_KEY ?? process.env.OPENAI_API_KEY;
+    const model = process.env.AGENT_MEMORY_RERANK_MODEL ?? "Qwen/Qwen3-Reranker-8B";
+    const baseUrl = process.env.AGENT_MEMORY_RERANK_BASE_URL ?? process.env.OPENAI_BASE_URL ?? "https://api.openai.com/v1";
+    if (!apiKey) return null;
+    return createOpenAIRerankProvider({ apiKey, model, baseUrl });
+  }
+  return null;
+}
+function createOpenAIRerankProvider(opts) {
+  const baseUrl = opts.baseUrl ?? "https://api.openai.com/v1";
+  return {
+    id: "openai-rerank",
+    model: opts.model,
+    async rerank(query, documents) {
+      const resp = await fetch(`${baseUrl.replace(/\/$/, "")}/rerank`, {
         method: "POST",
         headers: {
           "content-type": "application/json",
-          authorization: authHeader(opts.apiKey)
+          authorization: authHeader2(opts.apiKey)
         },
-        body: JSON.stringify({
-          model: opts.model,
-          input: { texts: [text] }
-        })
+        body: JSON.stringify({ model: opts.model, query, documents })
       });
       if (!resp.ok) {
         const body = await resp.text().catch(() => "");
-        throw new Error(`DashScope embeddings failed: ${resp.status} ${resp.statusText} ${body}`.trim());
+        throw new Error(`Rerank API failed: ${resp.status} ${resp.statusText} ${body}`.trim());
       }
       const data = await resp.json();
-      const emb = data?.output?.embeddings?.[0]?.embedding ?? data?.output?.embeddings?.[0]?.vector ?? data?.output?.embedding ?? data?.data?.[0]?.embedding;
-      return normalizeEmbedding(emb);
+      const results = data.results ?? [];
+      return results.map((r) => {
+        const index = typeof r.index === "number" ? r.index : Number.NaN;
+        const relevance = typeof r.relevance_score === "number" ? r.relevance_score : Number.NaN;
+        return { index, relevance_score: relevance };
+      }).filter((r) => Number.isInteger(r.index) && Number.isFinite(r.relevance_score));
     }
   };
 }
@@ -15226,6 +15340,7 @@ function createMcpServer(dbPath, agentId) {
   const db = openDatabase({ path: dbPath ?? DB_PATH });
   const aid = agentId ?? AGENT_ID;
   const embeddingProvider = getEmbeddingProviderFromEnv();
+  const rerankerProvider = getRerankerProviderFromEnv();
   const server = new McpServer({
     name: "agent-memory",
     version: "2.1.0"
@@ -15263,7 +15378,10 @@ function createMcpServer(dbPath, agentId) {
     async ({ query, limit }) => {
       const { intent, confidence } = classifyIntent(query);
       const strategy = getStrategy(intent);
-      const raw = await searchHybrid(db, query, { agent_id: aid, embeddingProvider, limit: limit * 2 });
+      let raw = await searchHybrid(db, query, { agent_id: aid, embeddingProvider, limit: limit * 2 });
+      if (rerankerProvider) {
+        raw = await rerankWithProvider(raw, query, rerankerProvider);
+      }
       const results = rerank(raw, { ...strategy, limit });
       const output = {
         intent,