Support external rerank endpoints

.env.example

@@ -95,6 +95,11 @@
 
 # BM25_WEIGHT=0.4                                # Hybrid search weight for BM25 leg
 # VECTOR_WEIGHT=0.6                              # Hybrid search weight for vector leg
+# RERANK_ENABLED=false                           # Rerank the top hybrid candidates
+# RERANK_BASE_URL=http://localhost:4000          # Optional Cohere-compatible /rerank endpoint
+# RERANK_API_KEY=                                # Optional; falls back to OPENAI_API_KEY
+# RERANK_MODEL=bge-reranker-v2-m3                # External reranker model name
+# RERANK_TIMEOUT_MS=30000                        # External rerank request timeout
 # AGENTMEMORY_GRAPH_WEIGHT=0.2                   # Graph traversal bonus on smart-search ranking
 # TOKEN_BUDGET=2000                              # Max tokens injected via mem::context per session
 # MAX_OBS_PER_SESSION=500                        # Per-session observation cap before consolidation kicks in

README.md

@@ -1425,6 +1425,11 @@ Create `~/.agentmemory/.env`:
 # Search tuning
 # BM25_WEIGHT=0.4
 # VECTOR_WEIGHT=0.6
+# RERANK_ENABLED=false
+# RERANK_BASE_URL=http://localhost:4000
+# RERANK_API_KEY=your-reranker-key  # optional for unauthenticated local endpoints
+# RERANK_MODEL=bge-reranker-v2-m3
+# RERANK_TIMEOUT_MS=30000
 # TOKEN_BUDGET=2000
 
 # Auth

src/state/reranker.ts

@@ -1,9 +1,114 @@
 import type { HybridSearchResult } from "../types.js";
+import { logger } from "../logger.js";
 
 let pipeline: any = null;
 let pipelineLoading: Promise<any> | null = null;
 let pipelineUnavailable = false;
 
+async function rerankExternal(
+  query: string,
+  results: HybridSearchResult[],
+  topK: number,
+): Promise<HybridSearchResult[] | null> {
+  const baseUrl = process.env.RERANK_BASE_URL?.replace(/\/+$/, "");
+  if (!baseUrl) return null;
+  const openaiBaseUrl = (
+    process.env.OPENAI_BASE_URL || "https://api.openai.com"
+  ).replace(/\/+$/, "");
+  const apiKey =
+    process.env.RERANK_API_KEY ||
+    (baseUrl === openaiBaseUrl ? process.env.OPENAI_API_KEY : undefined);
+  const configuredTimeout = Number(process.env.RERANK_TIMEOUT_MS);
+  const timeout =
+    Number.isFinite(configuredTimeout) && configuredTimeout > 0
+      ? configuredTimeout
+      : 30_000;
+
+  const candidates = results.slice(0, Math.min(results.length, topK));
+  try {
+    const response = await fetch(`${baseUrl}/rerank`, {
+      method: "POST",
+      headers: {
+        ...(apiKey ? { authorization: `Bearer ${apiKey}` } : {}),
+        "content-type": "application/json",
+      },
+      body: JSON.stringify({
+        model: process.env.RERANK_MODEL || "bge-reranker-v2-m3",
+        query,
+        documents: candidates.map(
+          (result) =>
+            `${result.observation.title || ""} ${result.observation.narrative || ""}`,
+        ),
+        top_n: candidates.length,
+      }),
+      signal: AbortSignal.timeout(timeout),
+    });
+    if (!response.ok) {
+      logger.warn("External reranker request failed", {
+        status: response.status,
+      });
+      return results;
+    }
+    const payload = (await response.json()) as {
+      results?: Array<{
+        index: number;
+        relevance_score?: number;
+        score?: number;
+      }>;
+    };
+    const providerResults = payload.results;
+    const seen = new Set<number>();
+    const valid =
+      Array.isArray(providerResults) &&
+      providerResults.length > 0 &&
+      providerResults.every((item) => {
+        const score = item.relevance_score ?? item.score;
+        if (
+          !Number.isInteger(item.index) ||
+          item.index < 0 ||
+          item.index >= candidates.length ||
+          seen.has(item.index) ||
+          typeof score !== "number" ||
+          !Number.isFinite(score)
+        ) {
+          return false;
+        }
+        seen.add(item.index);
+        return true;
+      });
+    if (!valid) {
+      logger.warn("External reranker returned invalid results", {
+        resultCount: Array.isArray(providerResults)
+          ? providerResults.length
+          : 0,
+      });
+      return results;
+    }
+    const scores = new Map(
+      providerResults.map((item) => [
+        item.index,
+        item.relevance_score ?? item.score ?? 0,
+      ]),
+    );
+    return candidates
+      .map((result, index) => ({
+        result,
+        score: scores.get(index) ?? result.combinedScore,
+      }))
+      .sort((a, b) => b.score - a.score)
+      .map(({ result, score }, index) => ({
+        ...result,
+        combinedScore: score,
+        rerankPosition: index + 1,
+      }));
+  } catch (error) {
+    logger.warn("External reranker request failed", {
+      error: error instanceof Error ? error.message : String(error),
+    });
+    return results;
+  }
+}
+
 async function loadPipeline(): Promise<any> {
   if (pipelineUnavailable) return null;
   if (pipeline) return pipeline;
@@ -38,6 +143,9 @@ export async function rerank(
 ): Promise<HybridSearchResult[]> {
   if (results.length <= 1) return results;
 
+  const external = await rerankExternal(query, results, topK);
+  if (external) return external;
+
   const reranker = await loadPipeline();
   if (!reranker) return results;
 
@@ -70,5 +178,5 @@ export async function rerank(
 }
 
 export function isRerankerAvailable(): boolean {
-  return pipeline !== null;
+  return Boolean(process.env.RERANK_BASE_URL || pipeline);
 }