diff --git a/.env.example b/.env.example index 77ca0f3a3..2ee7207ce 100644 --- a/.env.example +++ b/.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 diff --git a/README.md b/README.md index 84d747f9b..e727c9fe2 100644 --- a/README.md +++ b/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 diff --git a/src/state/reranker.ts b/src/state/reranker.ts index d0aae68cc..f27e5f4d4 100644 --- a/src/state/reranker.ts +++ b/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 | null = null; let pipelineUnavailable = false; +async function rerankExternal( + query: string, + results: HybridSearchResult[], + topK: number, +): Promise { + 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(); + 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 { if (pipelineUnavailable) return null; if (pipeline) return pipeline; @@ -38,6 +143,9 @@ export async function rerank( ): Promise { 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); } diff --git a/test/reranker.test.ts b/test/reranker.test.ts index 0694e6d50..ba5116b02 100644 --- a/test/reranker.test.ts +++ b/test/reranker.test.ts @@ -1,39 +1,63 @@ -import { describe, it, expect, vi } from "vitest"; +import { afterEach, beforeEach, describe, it, expect, vi } from "vitest"; vi.mock("@xenova/transformers", () => { throw new Error("not installed"); }); import { rerank, isRerankerAvailable } from "../src/state/reranker.js"; +import { logger } from "../src/logger.js"; +import type { HybridSearchResult } from "../src/types.js"; + +function makeResult( + id: string, + title: string, + narrative: string, + combinedScore: number, +): HybridSearchResult { + return { + observation: { + id, + sessionId: "s1", + timestamp: "2026-07-19T00:00:00Z", + type: "decision", + title, + facts: [], + narrative, + concepts: [], + files: [], + importance: 5, + }, + bm25Score: combinedScore, + vectorScore: 0, + graphScore: 0, + combinedScore, + sessionId: "s1", + }; +} describe("reranker", () => { + const resetEnvironment = () => { + delete process.env.RERANK_BASE_URL; + delete process.env.RERANK_API_KEY; + delete process.env.RERANK_MODEL; + delete process.env.RERANK_TIMEOUT_MS; + delete process.env.OPENAI_API_KEY; + delete process.env.OPENAI_BASE_URL; + }; + + beforeEach(resetEnvironment); + + afterEach(() => { + resetEnvironment(); + vi.restoreAllMocks(); + vi.unstubAllGlobals(); + }); + it("returns results unchanged when @xenova/transformers is unavailable", async () => { const results = [ - { - observation: { - id: "o1", - title: "First", - narrative: "First result", - }, - bm25Score: 0.5, - vectorScore: 0.6, - graphScore: 0, - combinedScore: 0.8, - sessionId: "s1", - }, - { - observation: { - id: "o2", - title: "Second", - narrative: "Second result", - }, - bm25Score: 0.3, - vectorScore: 0.4, - graphScore: 0, - combinedScore: 0.5, - sessionId: "s1", - }, - ] as any; + makeResult("o1", "First", "First result", 0.8), + makeResult("o2", "Second", "Second result", 0.5), + ]; const reranked = await rerank("test query", results); expect(reranked).toEqual(results); @@ -44,12 +68,7 @@ describe("reranker", () => { }); it("handles single result gracefully", async () => { - const results = [ - { - observation: { id: "o1", title: "Only" }, - combinedScore: 1.0, - }, - ] as any; + const results = [makeResult("o1", "Only", "", 1)]; const reranked = await rerank("query", results); expect(reranked).toHaveLength(1); @@ -59,4 +78,237 @@ describe("reranker", () => { const reranked = await rerank("query", []); expect(reranked).toHaveLength(0); }); + + it("uses an OpenAI-compatible external rerank endpoint when configured", async () => { + process.env.RERANK_BASE_URL = "http://reranker.test"; + process.env.RERANK_API_KEY = "test-key"; + process.env.RERANK_MODEL = "bge-reranker-v2-m3"; + const fetchMock = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + results: [ + { index: 1, relevance_score: 0.9 }, + { index: 0, relevance_score: 0.1 }, + ], + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ); + vi.stubGlobal("fetch", fetchMock); + const results = [ + makeResult("o1", "First", "First result", 0.8), + makeResult("o2", "Second", "Second result", 0.5), + ]; + + const reranked = await rerank("test query", results); + + expect(reranked.map((result) => result.observation.id)).toEqual(["o2", "o1"]); + expect(fetchMock).toHaveBeenCalledWith( + "http://reranker.test/rerank", + expect.objectContaining({ method: "POST" }), + ); + const headers = fetchMock.mock.calls[0][1]?.headers as Record; + expect(headers.authorization).toBe("Bearer test-key"); + }); + + it("supports local external rerank endpoints without an API key", async () => { + process.env.RERANK_BASE_URL = "http://reranker.test"; + const fetchMock = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + results: [ + { index: 1, relevance_score: 0.9 }, + { index: 0, relevance_score: 0.1 }, + ], + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ); + vi.stubGlobal("fetch", fetchMock); + const results = [ + makeResult("o1", "First", "First result", 0.8), + makeResult("o2", "Second", "Second result", 0.5), + ]; + + const reranked = await rerank("test query", results); + + expect(reranked.map((result) => result.observation.id)).toEqual(["o2", "o1"]); + const headers = fetchMock.mock.calls[0][1]?.headers as Record; + expect(headers).not.toHaveProperty("authorization"); + }); + + it("does not send OPENAI_API_KEY to a different rerank endpoint", async () => { + process.env.RERANK_BASE_URL = "https://reranker.test"; + process.env.OPENAI_BASE_URL = "https://api.openai.com"; + process.env.OPENAI_API_KEY = "openai-secret"; + const fetchMock = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + results: [ + { index: 1, relevance_score: 0.9 }, + { index: 0, relevance_score: 0.1 }, + ], + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ); + vi.stubGlobal("fetch", fetchMock); + const results = [ + makeResult("o1", "First", "First result", 0.8), + makeResult("o2", "Second", "Second result", 0.5), + ]; + + await rerank("test query", results); + + const headers = fetchMock.mock.calls[0][1]?.headers as Record; + expect(headers).not.toHaveProperty("authorization"); + }); + + it("uses OPENAI_API_KEY when the rerank endpoint matches OPENAI_BASE_URL", async () => { + process.env.RERANK_BASE_URL = "https://api.openai.com/"; + process.env.OPENAI_BASE_URL = "https://api.openai.com"; + process.env.OPENAI_API_KEY = "openai-secret"; + const fetchMock = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + results: [ + { index: 1, relevance_score: 0.9 }, + { index: 0, relevance_score: 0.1 }, + ], + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ); + vi.stubGlobal("fetch", fetchMock); + const results = [ + makeResult("o1", "First", "First result", 0.8), + makeResult("o2", "Second", "Second result", 0.5), + ]; + + await rerank("test query", results); + + const headers = fetchMock.mock.calls[0][1]?.headers as Record; + expect(headers.authorization).toBe("Bearer openai-secret"); + }); + + it.each([undefined, "invalid", "0", "-1", "Infinity"])( + "defaults invalid timeout %s to 30 seconds", + async (timeout) => { + process.env.RERANK_BASE_URL = "http://reranker.test"; + if (timeout !== undefined) process.env.RERANK_TIMEOUT_MS = timeout; + const timeoutSpy = vi + .spyOn(AbortSignal, "timeout") + .mockReturnValue(new AbortController().signal); + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue(new Response("unavailable", { status: 503 })), + ); + const results = [ + makeResult("o1", "First", "First result", 0.8), + makeResult("o2", "Second", "Second result", 0.5), + ]; + + await rerank("test query", results); + + expect(timeoutSpy).toHaveBeenCalledWith(30_000); + }, + ); + + it("fails open and logs non-successful external responses", async () => { + process.env.RERANK_BASE_URL = "http://reranker.test"; + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue(new Response("unavailable", { status: 503 })), + ); + const warn = vi.spyOn(logger, "warn").mockImplementation(() => {}); + const results = [ + makeResult("o1", "First", "First result", 0.8), + makeResult("o2", "Second", "Second result", 0.5), + ]; + + expect(await rerank("test query", results)).toEqual(results); + expect(warn).toHaveBeenCalledWith( + "External reranker request failed", + expect.objectContaining({ status: 503 }), + ); + }); + + it("fails open and logs malformed external responses", async () => { + process.env.RERANK_BASE_URL = "http://reranker.test"; + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + new Response("{", { + status: 200, + headers: { "content-type": "application/json" }, + }), + ), + ); + const warn = vi.spyOn(logger, "warn").mockImplementation(() => {}); + const results = [ + makeResult("o1", "First", "First result", 0.8), + makeResult("o2", "Second", "Second result", 0.5), + ]; + + expect(await rerank("test query", results)).toEqual(results); + expect(warn).toHaveBeenCalledWith( + "External reranker request failed", + expect.objectContaining({ error: expect.any(String) }), + ); + }); + + it("fails open and logs timeout or network errors", async () => { + process.env.RERANK_BASE_URL = "http://reranker.test"; + vi.stubGlobal( + "fetch", + vi.fn().mockRejectedValue(new DOMException("timed out", "TimeoutError")), + ); + const warn = vi.spyOn(logger, "warn").mockImplementation(() => {}); + const results = [ + makeResult("o1", "First", "First result", 0.8), + makeResult("o2", "Second", "Second result", 0.5), + ]; + + expect(await rerank("test query", results)).toEqual(results); + expect(warn).toHaveBeenCalledWith( + "External reranker request failed", + expect.objectContaining({ error: "timed out" }), + ); + }); + + it.each([ + { + name: "out-of-range", + providerResults: [{ index: 7, relevance_score: 0.9 }], + }, + { + name: "duplicate", + providerResults: [ + { index: 0, relevance_score: 0.9 }, + { index: 0, relevance_score: 0.1 }, + ], + }, + ])("fails open for $name result indices", async ({ providerResults }) => { + process.env.RERANK_BASE_URL = "http://reranker.test"; + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + new Response(JSON.stringify({ results: providerResults }), { + status: 200, + headers: { "content-type": "application/json" }, + }), + ), + ); + const warn = vi.spyOn(logger, "warn").mockImplementation(() => {}); + const results = [ + makeResult("o1", "First", "First result", 0.8), + makeResult("o2", "Second", "Second result", 0.5), + ]; + + expect(await rerank("test query", results)).toEqual(results); + expect(warn).toHaveBeenCalledWith( + "External reranker returned invalid results", + expect.any(Object), + ); + }); });