From 6a801bfe2f377ba34f12b3ef4260a5b84a4fc35b Mon Sep 17 00:00:00 2001 From: Kartik Date: Thu, 9 Jul 2026 22:18:38 +0530 Subject: [PATCH] feat(oss): add reranker + per-search rerank to the TypeScript OSS SDK (#6055) --- docs/components/rerankers/config.mdx | 29 +- docs/components/rerankers/models/cohere.mdx | 42 +- .../rerankers/models/huggingface.mdx | 34 ++ .../rerankers/models/llm_reranker.mdx | 37 ++ .../rerankers/models/sentence_transformer.mdx | 34 ++ .../rerankers/models/zero_entropy.mdx | 28 ++ docs/components/rerankers/optimization.mdx | 4 +- docs/components/rerankers/overview.mdx | 4 + docs/llms.txt | 3 +- docs/open-source/configuration.mdx | 2 +- docs/open-source/features/reranker-search.mdx | 132 +++++- mem0-ts/package.json | 3 + mem0-ts/pnpm-lock.yaml | 431 ++++++++++++++++++ mem0-ts/src/oss/src/config/manager.ts | 1 + mem0-ts/src/oss/src/index.ts | 5 + mem0-ts/src/oss/src/memory/index.ts | 33 +- mem0-ts/src/oss/src/memory/memory.types.ts | 5 + mem0-ts/src/oss/src/rerankers/base.ts | 21 + mem0-ts/src/oss/src/rerankers/cohere.test.ts | 140 ++++++ mem0-ts/src/oss/src/rerankers/cohere.ts | 61 +++ .../oss/src/rerankers/cross_encoder.test.ts | 189 ++++++++ .../src/oss/src/rerankers/cross_encoder.ts | 98 ++++ mem0-ts/src/oss/src/rerankers/llm.test.ts | 174 +++++++ mem0-ts/src/oss/src/rerankers/llm.ts | 87 ++++ .../src/oss/src/rerankers/zeroentropy.test.ts | 133 ++++++ mem0-ts/src/oss/src/rerankers/zeroentropy.ts | 58 +++ mem0-ts/src/oss/src/types/index.ts | 59 +++ mem0-ts/src/oss/src/utils/factory.test.ts | 86 ++++ mem0-ts/src/oss/src/utils/factory.ts | 69 +++ mem0-ts/src/oss/tests/memory.rerank.test.ts | 184 ++++++++ mem0-ts/tsup.config.ts | 3 + mem0/configs/rerankers/cohere.py | 3 +- 32 files changed, 2173 insertions(+), 19 deletions(-) create mode 100644 mem0-ts/src/oss/src/rerankers/base.ts create mode 100644 mem0-ts/src/oss/src/rerankers/cohere.test.ts create mode 100644 mem0-ts/src/oss/src/rerankers/cohere.ts create mode 100644 mem0-ts/src/oss/src/rerankers/cross_encoder.test.ts create mode 100644 mem0-ts/src/oss/src/rerankers/cross_encoder.ts create mode 100644 mem0-ts/src/oss/src/rerankers/llm.test.ts create mode 100644 mem0-ts/src/oss/src/rerankers/llm.ts create mode 100644 mem0-ts/src/oss/src/rerankers/zeroentropy.test.ts create mode 100644 mem0-ts/src/oss/src/rerankers/zeroentropy.ts create mode 100644 mem0-ts/src/oss/src/utils/factory.test.ts create mode 100644 mem0-ts/src/oss/tests/memory.rerank.test.ts diff --git a/docs/components/rerankers/config.mdx b/docs/components/rerankers/config.mdx index 475e8ca09..63d2fb3d3 100644 --- a/docs/components/rerankers/config.mdx +++ b/docs/components/rerankers/config.mdx @@ -26,7 +26,7 @@ All rerankers share these common configuration parameters: | Parameter | Description | Type | Default | | -------------------- | -------------------------------------------- | ------ | ----------------------- | -| `model` | Cohere rerank model | `str` | `"rerank-english-v3.0"` | +| `model` | Cohere rerank model | `str` | `"rerank-v3.5"` | | `api_key` | Cohere API key | `str` | `None` | | `return_documents` | Whether to return document texts in response | `bool` | `False` | | `max_chunks_per_doc` | Maximum chunks per document | `int` | `None` | @@ -103,3 +103,30 @@ config = { } } ``` + +## TypeScript SDK + +The self-hosted [TypeScript SDK](/open-source/features/reranker-search#typescript-sdk) (`mem0ai/oss`) supports the same five providers. Config keys are camelCase (`apiKey`, `topK`, `maxLength`) and each provider's SDK is a peer dependency you install per reranker. + +| Provider | Install | Default model | Key config fields | +| --- | --- | --- | --- | +| `cohere` | `pnpm add cohere-ai` | `rerank-v3.5` | `apiKey`, `model`, `topK` | +| `zero_entropy` | `pnpm add zeroentropy` | `zerank-1` | `apiKey`, `model`, `topK` | +| `sentence_transformer` | `pnpm add @huggingface/transformers` | `Xenova/ms-marco-MiniLM-L-6-v2` | `model`, `device`, `maxLength`, `normalize`, `topK` | +| `huggingface` | `pnpm add @huggingface/transformers` | `Xenova/bge-reranker-base` | `model`, `device`, `maxLength`, `normalize`, `topK` | +| `llm_reranker` | — (uses your LLM provider's own SDK) | `openai` / `gpt-4o-mini` | `provider`, `model`, `apiKey`, `llm` (nested override), `topK` | + +```typescript +import { Memory } from "mem0ai/oss"; + +const memory = new Memory({ + reranker: { + provider: "zero_entropy", + config: { apiKey: process.env.ZERO_ENTROPY_API_KEY, topK: 5 }, + }, +}); +``` + + + The local cross-encoder providers (`sentence_transformer`, `huggingface`) run on [Transformers.js](https://huggingface.co/docs/transformers.js) and default to ONNX (`Xenova/*`) model mirrors, so Python default model strings must be swapped for their ONNX equivalents. `batchSize` and `showProgressBar` are accepted for parity with Python but are no-ops in the TypeScript runtime. See the [reranker feature guide](/open-source/features/reranker-search#typescript-sdk) for full examples. + diff --git a/docs/components/rerankers/models/cohere.mdx b/docs/components/rerankers/models/cohere.mdx index 20e73a656..53ea0c28f 100644 --- a/docs/components/rerankers/models/cohere.mdx +++ b/docs/components/rerankers/models/cohere.mdx @@ -9,9 +9,9 @@ Cohere provides enterprise-grade reranking models with excellent multilingual su Cohere offers several reranking models: -- **`rerank-english-v3.0`**: Latest English reranker with best performance -- **`rerank-multilingual-v3.0`**: Multilingual support for global applications -- **`rerank-english-v2.0`**: Previous generation English reranker +- **`rerank-v3.5`** (default): Latest reranker, multilingual, best performance +- **`rerank-english-v3.0`**: Previous generation, English only +- **`rerank-multilingual-v3.0`**: Previous generation, multilingual ## Installation @@ -41,7 +41,7 @@ config = { "reranker": { "provider": "cohere", "config": { - "model": "rerank-english-v3.0", + "model": "rerank-v3.5", "api_key": "your-cohere-api-key", # or set COHERE_API_KEY "top_k": 5, "return_documents": False, @@ -53,6 +53,34 @@ config = { memory = Memory.from_config(config) ``` +## TypeScript (self-hosted) + +The [TypeScript OSS SDK](/open-source/features/reranker-search#typescript-sdk) (`mem0ai/oss`) ships the Cohere reranker. Config keys are camelCase, it defaults to the `rerank-v3.5` model, and you opt in per search with `rerank: true`. + +```bash +pnpm add cohere-ai +``` + +```typescript +import { Memory } from "mem0ai/oss"; + +const memory = new Memory({ + reranker: { + provider: "cohere", + config: { + apiKey: process.env.COHERE_API_KEY, // or set COHERE_API_KEY + // model: "rerank-v3.5", // default + topK: 5, + }, + }, +}); + +const results = await memory.search("What is the user's profession?", { + filters: { userId: "bob" }, + rerank: true, +}); +``` + ## Environment Variables Set your API key as an environment variable: @@ -77,7 +105,7 @@ config = { "rerank": { "provider": "cohere", "config": { - "model": "rerank-english-v3.0", + "model": "rerank-v3.5", "top_k": 3 } } @@ -124,7 +152,7 @@ config = { | Parameter | Description | Type | Default | | -------------------- | -------------------------------- | ------ | ----------------------- | -| `model` | Cohere rerank model to use | `str` | `"rerank-english-v3.0"` | +| `model` | Cohere rerank model to use | `str` | `"rerank-v3.5"` | | `api_key` | Cohere API key | `str` | `None` | | `top_k` | Maximum documents to return | `int` | `None` | | `return_documents` | Whether to return document texts | `bool` | `False` | @@ -139,7 +167,7 @@ config = { ## Best Practices -1. **Model Selection**: Use `rerank-english-v3.0` for English, `rerank-multilingual-v3.0` for other languages +1. **Model Selection**: `rerank-v3.5` handles English and multilingual workloads; pin an older `v3.0` model only if you need to reproduce prior results 2. **Batch Processing**: Process multiple queries efficiently 3. **Error Handling**: Implement retry logic for production systems 4. **Monitoring**: Track reranking performance and costs diff --git a/docs/components/rerankers/models/huggingface.mdx b/docs/components/rerankers/models/huggingface.mdx index 64828c095..95b7f6663 100644 --- a/docs/components/rerankers/models/huggingface.mdx +++ b/docs/components/rerankers/models/huggingface.mdx @@ -57,6 +57,40 @@ config = { } ``` +## TypeScript (self-hosted) + +The [TypeScript OSS SDK](/open-source/features/reranker-search#typescript-sdk) (`mem0ai/oss`) runs this reranker locally with [Transformers.js](https://huggingface.co/docs/transformers.js) — the same cross-encoder path as `sentence_transformer`, just a different default model. It executes ONNX weights, so the default is the ONNX mirror `Xenova/bge-reranker-base`. Point `model` at any ONNX-exported reranker on the Hub (a raw `BAAI/bge-reranker-*` PyTorch checkpoint will not load in this runtime). + +```bash +pnpm add @huggingface/transformers +``` + +```typescript +import { Memory } from "mem0ai/oss"; + +const memory = new Memory({ + reranker: { + provider: "huggingface", + config: { + // model: "Xenova/bge-reranker-base", // default (ONNX) + device: "cpu", // "cpu" | "wasm" | "webgpu" + maxLength: 512, // max tokens per query-document pair + normalize: true, // sigmoid-normalize logits to [0, 1] (default) + topK: 5, + }, + }, +}); + +const results = await memory.search("What are the user's interests?", { + filters: { userId: "alice" }, + rerank: true, +}); +``` + + + `batchSize` and `showProgressBar` are accepted for parity with the Python SDK but are no-ops in the TypeScript runtime. `trust_remote_code` and `model_kwargs` are Python-only. + + ## Popular Models ### BGE Rerankers (Recommended) diff --git a/docs/components/rerankers/models/llm_reranker.mdx b/docs/components/rerankers/models/llm_reranker.mdx index d23be45a4..2d3abf512 100644 --- a/docs/components/rerankers/models/llm_reranker.mdx +++ b/docs/components/rerankers/models/llm_reranker.mdx @@ -67,6 +67,43 @@ config = { } ``` +## TypeScript (self-hosted) + +The [TypeScript OSS SDK](/open-source/features/reranker-search#typescript-sdk) (`mem0ai/oss`) ships the LLM reranker under the provider name `llm_reranker`. It does **not** reuse the Memory's main `llm` instance — it builds its own LLM from the reranker's own config, defaulting to `openai` / `gpt-4o-mini`. Set `provider`/`model`/`apiKey` directly on `config`, or nest a fully separate `config.llm: { provider, config }` (its `provider`/`config` take priority over the top-level fields, which only backfill values missing from the nested config). + +```typescript +import { Memory } from "mem0ai/oss"; + +const memory = new Memory({ + reranker: { + provider: "llm_reranker", + config: { apiKey: process.env.OPENAI_API_KEY }, + }, +}); + +const results = await memory.search("What movies do I like?", { + filters: { userId: "alice" }, + rerank: true, +}); +``` + +To rerank with a different LLM provider than the Memory's main `llm`, nest it under `config.llm`: + +```typescript +const memory = new Memory({ + llm: { provider: "openai", config: { apiKey: process.env.OPENAI_API_KEY } }, + reranker: { + provider: "llm_reranker", + config: { + llm: { + provider: "anthropic", + config: { apiKey: process.env.ANTHROPIC_API_KEY }, + }, + }, + }, +}); +``` + ## Supported LLM Providers ### OpenAI diff --git a/docs/components/rerankers/models/sentence_transformer.mdx b/docs/components/rerankers/models/sentence_transformer.mdx index 8535e4627..26c02fe34 100644 --- a/docs/components/rerankers/models/sentence_transformer.mdx +++ b/docs/components/rerankers/models/sentence_transformer.mdx @@ -54,6 +54,40 @@ config = { memory = Memory.from_config(config) ``` +## TypeScript (self-hosted) + +The [TypeScript OSS SDK](/open-source/features/reranker-search#typescript-sdk) (`mem0ai/oss`) runs this reranker locally with [Transformers.js](https://huggingface.co/docs/transformers.js). Because it executes ONNX weights, the default model is the ONNX mirror of the Python default — `Xenova/ms-marco-MiniLM-L-6-v2`. Point `model` at any ONNX-exported cross-encoder on the Hub (a raw `cross-encoder/...` PyTorch checkpoint will not load in this runtime). + +```bash +pnpm add @huggingface/transformers +``` + +```typescript +import { Memory } from "mem0ai/oss"; + +const memory = new Memory({ + reranker: { + provider: "sentence_transformer", + config: { + // model: "Xenova/ms-marco-MiniLM-L-6-v2", // default (ONNX) + device: "cpu", // "cpu" | "wasm" | "webgpu" + maxLength: 512, // max tokens per query-document pair + normalize: true, // sigmoid-normalize logits to [0, 1] (default) + topK: 5, + }, + }, +}); + +const results = await memory.search("What books does the user like?", { + filters: { userId: "charlie" }, + rerank: true, +}); +``` + + + `batchSize` and `showProgressBar` are accepted for parity with the Python SDK but are no-ops in the TypeScript runtime — a search reranks a small candidate set in a single in-process forward pass. The model downloads once and is cached in-process. + + ## GPU Acceleration For better performance, use GPU acceleration: diff --git a/docs/components/rerankers/models/zero_entropy.mdx b/docs/components/rerankers/models/zero_entropy.mdx index 43973a323..e453ff748 100644 --- a/docs/components/rerankers/models/zero_entropy.mdx +++ b/docs/components/rerankers/models/zero_entropy.mdx @@ -50,6 +50,34 @@ config = { memory = Memory.from_config(config) ``` +## TypeScript (self-hosted) + +The [TypeScript OSS SDK](/open-source/features/reranker-search#typescript-sdk) (`mem0ai/oss`) ships the Zero Entropy reranker under the same provider name as Python, `zero_entropy`. It reads the key from config or `ZERO_ENTROPY_API_KEY` and defaults to the `zerank-1` model. + +```bash +pnpm add zeroentropy +``` + +```typescript +import { Memory } from "mem0ai/oss"; + +const memory = new Memory({ + reranker: { + provider: "zero_entropy", + config: { + apiKey: process.env.ZERO_ENTROPY_API_KEY, + // model: "zerank-1", // default (or "zerank-1-small") + topK: 5, + }, + }, +}); + +const results = await memory.search("What Italian food does the user like?", { + filters: { userId: "alice" }, + rerank: true, +}); +``` + ## Environment Variables Set your API key as an environment variable: diff --git a/docs/components/rerankers/optimization.mdx b/docs/components/rerankers/optimization.mdx index 5b396bf61..f85a28f4a 100644 --- a/docs/components/rerankers/optimization.mdx +++ b/docs/components/rerankers/optimization.mdx @@ -47,7 +47,7 @@ config = { "reranker": { "provider": "cohere", "config": { - "model": "rerank-english-v3.0", + "model": "rerank-v3.5", "top_n": 10, "max_chunks_per_doc": 10, # Limit chunk processing "return_documents": False # Reduce response size @@ -280,7 +280,7 @@ config = { ```python def benchmark_rerankers(): configs = [ - {"provider": "cohere", "model": "rerank-english-v3.0"}, + {"provider": "cohere", "model": "rerank-v3.5"}, {"provider": "sentence_transformer", "model": "cross-encoder/ms-marco-MiniLM-L-6-v2"}, {"provider": "huggingface", "model": "BAAI/bge-reranker-base"} ] diff --git a/docs/components/rerankers/overview.mdx b/docs/components/rerankers/overview.mdx index 58432159c..2fab7396f 100644 --- a/docs/components/rerankers/overview.mdx +++ b/docs/components/rerankers/overview.mdx @@ -19,6 +19,10 @@ Reranking trades extra latency for better precision. Start once you have baselin + +All five rerankers are available in both the Python and the [TypeScript](/open-source/features/reranker-search#typescript-sdk) self-hosted SDKs. Each provider page has a **TypeScript (self-hosted)** section with the camelCase config. + + ## Reranking Workflow diff --git a/docs/llms.txt b/docs/llms.txt index 304eb60fc..bca73bdbd 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -501,5 +501,6 @@ Everything below is OSS-only provider configuration. Skip this entire section wh - [Custom Reranker Prompts](https://docs.mem0.ai/components/rerankers/custom-prompts) [OSS]: Use when rewriting reranker prompts. - [Cohere Reranker](https://docs.mem0.ai/components/rerankers/models/cohere) [OSS]: Use for Cohere Rerank. - [Sentence Transformer Reranker](https://docs.mem0.ai/components/rerankers/models/sentence_transformer) [OSS]: Use for local cross-encoder rerankers. -- [Hugging Face Reranker](https://docs.mem0.ai/components/rerankers/models/huggingface) [OSS]: Use for HF-hosted reranker models.- [LLM Reranker](https://docs.mem0.ai/components/rerankers/models/llm_reranker) [OSS]: Use when the reranker is a prompted LLM (implementation reference). +- [Hugging Face Reranker](https://docs.mem0.ai/components/rerankers/models/huggingface) [OSS]: Use for HF-hosted reranker models. +- [LLM Reranker](https://docs.mem0.ai/components/rerankers/models/llm_reranker) [OSS]: Use when the reranker is a prompted LLM (implementation reference). - [Zero Entropy Reranker](https://docs.mem0.ai/components/rerankers/models/zero_entropy) [OSS]: Use for the Zero Entropy reranker. diff --git a/docs/open-source/configuration.mdx b/docs/open-source/configuration.mdx index 38bd074b8..0fbb72849 100644 --- a/docs/open-source/configuration.mdx +++ b/docs/open-source/configuration.mdx @@ -59,7 +59,7 @@ config = { }, "reranker": { "provider": "cohere", - "config": {"model": "rerank-english-v3.0"}, + "config": {"model": "rerank-v3.5"}, }, } diff --git a/docs/open-source/features/reranker-search.mdx b/docs/open-source/features/reranker-search.mdx index e1d747a0e..c2138b5c5 100644 --- a/docs/open-source/features/reranker-search.mdx +++ b/docs/open-source/features/reranker-search.mdx @@ -18,7 +18,129 @@ Reranker-enhanced search adds a second scoring pass after vector retrieval so Me - All configuration snippets translate directly to the TypeScript SDK: swap dictionaries for objects while keeping the same keys (`provider`, `config`, `rerank` flags). + The `Configure it` and `See it in action` snippets below use the Python SDK. The self-hosted **TypeScript SDK** supports the Cohere, Zero Entropy, Sentence Transformer, Hugging Face, and LLM rerankers — see [TypeScript SDK](#typescript-sdk). + + +--- + +## TypeScript SDK + +The self-hosted TypeScript SDK (`mem0ai/oss`) ships five rerankers: **Cohere**, **Zero Entropy**, **Sentence Transformer**, **Hugging Face**, and the **LLM reranker**. Configure one under `reranker`, then opt in per search with `rerank: true`. Keys are camelCase (`apiKey`, not `api_key`). + +Provider SDKs are peer dependencies — install the one your reranker needs: + +```bash +pnpm add cohere-ai # cohere +pnpm add zeroentropy # zero_entropy +pnpm add @huggingface/transformers # sentence_transformer, huggingface +# llm_reranker defaults to openai (already a core dependency) — install another +# provider's SDK only if you nest a different one under config.llm +``` + +### Hosted rerankers (Cohere, Zero Entropy) + +Both call a hosted API and read their key from config or the provider's environment variable (`COHERE_API_KEY`, `ZERO_ENTROPY_API_KEY`). + +```typescript +import { Memory } from "mem0ai/oss"; + +// Cohere reranker (defaults to the rerank-v3.5 model) +const memory = new Memory({ + reranker: { + provider: "cohere", + config: { apiKey: process.env.COHERE_API_KEY }, + }, +}); + +const results = await memory.search("What are my food preferences?", { + filters: { userId: "alice" }, + rerank: true, +}); +``` + +```typescript +// Zero Entropy reranker (defaults to the zerank-1 model) +const memory = new Memory({ + reranker: { + provider: "zero_entropy", + config: { apiKey: process.env.ZERO_ENTROPY_API_KEY }, + }, +}); +``` + +### Local cross-encoders (Sentence Transformer, Hugging Face) + +Both run a cross-encoder locally with [Transformers.js](https://huggingface.co/docs/transformers.js) — no API key, no network at inference time. Because Transformers.js runs ONNX weights, the default models are the ONNX mirrors of the Python SDK's defaults (`sentence_transformer` → `Xenova/ms-marco-MiniLM-L-6-v2`, `huggingface` → `Xenova/bge-reranker-base`). Point `model` at any ONNX-exported cross-encoder on the Hub to override. + +```typescript +const memory = new Memory({ + reranker: { + provider: "sentence_transformer", // or "huggingface" + config: { + // model: "Xenova/bge-reranker-base", // override the default + device: "cpu", // Transformers.js device: "cpu" | "wasm" | "webgpu" + maxLength: 512, // max tokens per query-document pair + normalize: true, // sigmoid-normalize logits to [0, 1] (default) + }, + }, +}); + +const results = await memory.search("What movies do I like?", { + filters: { userId: "alice" }, + rerank: true, +}); +``` + + + `batchSize` and `showProgressBar` are accepted for config parity with the Python SDK but are no-ops in this runtime — a memory search reranks a small candidate set in a single in-process forward pass. The model is downloaded once and cached in-process on first use. + + +### LLM reranker + +To score with an LLM instead of a dedicated reranker, use the `llm_reranker` provider. It builds its own LLM from the reranker's config — defaulting to `openai` / `gpt-4o-mini` — rather than reusing the Memory's main `llm`: + +```typescript +const memory = new Memory({ + reranker: { + provider: "llm_reranker", + config: { apiKey: process.env.OPENAI_API_KEY }, + }, +}); + +const results = await memory.search("What movies do I like?", { + filters: { userId: "alice" }, + rerank: true, +}); +``` + +Nest a different provider under `config.llm` to override the default: + +```typescript +const memory = new Memory({ + reranker: { + provider: "llm_reranker", + config: { + llm: { + provider: "anthropic", + config: { apiKey: process.env.ANTHROPIC_API_KEY }, + }, + }, + }, +}); +``` + +### Config reference + +| Provider | Default model | Key config fields | +| --- | --- | --- | +| `cohere` | `rerank-v3.5` | `apiKey`, `model`, `topK` | +| `zero_entropy` | `zerank-1` | `apiKey`, `model`, `topK` | +| `sentence_transformer` | `Xenova/ms-marco-MiniLM-L-6-v2` | `model`, `device`, `maxLength`, `normalize`, `topK` | +| `huggingface` | `Xenova/bge-reranker-base` | `model`, `device`, `maxLength`, `normalize`, `topK` | +| `llm_reranker` | `openai` / `gpt-4o-mini` | `provider`, `model`, `apiKey`, `llm` (nested override), `topK` | + + + `rerank` is opt-in per search and a no-op when no `reranker` is configured. If the reranker call fails, Mem0 logs a warning and returns the original vector-ranked results. --- @@ -61,7 +183,7 @@ config = { "reranker": { "provider": "cohere", "config": { - "model": "rerank-english-v3.0", + "model": "rerank-v3.5", "api_key": "your-cohere-api-key" } } @@ -86,7 +208,7 @@ config = { "reranker": { "provider": "cohere", "config": { - "model": "rerank-english-v3.0", + "model": "rerank-v3.5", "api_key": "your-cohere-api-key", "top_k": 10, "return_documents": True @@ -164,7 +286,7 @@ config = { "reranker": { "provider": "cohere", "config": { - "model": "rerank-english-v3.0", + "model": "rerank-v3.5", "api_key": "your-cohere-api-key", "top_k": 15, "return_documents": True @@ -338,7 +460,7 @@ config = { "reranker": { "provider": "cohere", "config": { - "model": "rerank-english-v3.0", + "model": "rerank-v3.5", "api_key": "your-cohere-api-key" } } diff --git a/mem0-ts/package.json b/mem0-ts/package.json index 48e019f73..206bd4f8d 100644 --- a/mem0-ts/package.json +++ b/mem0-ts/package.json @@ -114,6 +114,7 @@ "@cloudflare/workers-types": "^4.20250504.0", "@google-cloud/aiplatform": "^6.8.0", "@google/genai": "^1.40.0", + "@huggingface/transformers": "^3.0.0 || ^4.0.0", "@langchain/core": "^1.1.47", "@mistralai/mistralai": "^1.5.2", "@opensearch-project/opensearch": "^3.5.1", @@ -128,6 +129,7 @@ "cassandra-driver": "4.8.0", "chromadb": "^3.5.0", "cloudflare": "^4.2.0", + "cohere-ai": "^7.17.0 || ^8.0.0", "fastembed": "^2.1.0", "groq-sdk": "0.3.0", "mongodb": "^7.0.0", @@ -139,6 +141,7 @@ "iovalkey": "^0.3.3", "compromise": "^14.0.0", "natural": "^8.0.1", + "zeroentropy": "^0.1.0-alpha.10", "mysql2": "^3.0.0", "@zilliz/milvus2-sdk-node": "^2.4.0 || ^3.0.0" }, diff --git a/mem0-ts/pnpm-lock.yaml b/mem0-ts/pnpm-lock.yaml index 4a8b9311c..ae4b656ca 100644 --- a/mem0-ts/pnpm-lock.yaml +++ b/mem0-ts/pnpm-lock.yaml @@ -56,6 +56,9 @@ importers: '@google/genai': specifier: ^1.40.0 version: 1.52.0 + '@huggingface/transformers': + specifier: ^3.0.0 || ^4.0.0 + version: 4.2.0 '@langchain/core': specifier: ^1.1.47 version: 1.1.48(@opentelemetry/api@1.9.1)(openai@4.104.0(ws@5.2.5)(zod@3.25.76))(ws@5.2.5) @@ -104,6 +107,9 @@ importers: cloudflare: specifier: ^4.2.0 version: 4.5.0 + cohere-ai: + specifier: ^7.17.0 || ^8.0.0 + version: 8.0.0(@aws-crypto/sha256-js@5.2.0)(@smithy/protocol-http@5.5.2)(@smithy/signature-v4@5.5.2) compromise: specifier: ^14.0.0 version: 14.15.1 @@ -140,6 +146,9 @@ importers: weaviate-client: specifier: ^3.0.0 version: 3.13.1 + zeroentropy: + specifier: ^0.1.0-alpha.10 + version: 0.1.0-alpha.10 zod: specifier: ^3.24.1 version: 3.25.76 @@ -667,6 +676,9 @@ packages: resolution: {integrity: sha512-L38Ax21uF2OPUmCRWycZ/dZdMYf7gMrtClcxvVrqJVFmn8ET2M++GYmFGJpLqOHS1beATxOXLWe7y2ijSQz/ng==} engines: {node: '>=20'} + '@emnapi/runtime@1.11.2': + resolution: {integrity: sha512-kyOl3X0DuTiT1h2ft8r2fYO8JYtU9a9Xis/zBSiGArNaagCOWx90N1k2wxp18czFDH+OgcWGb5ZP/XMt3dcyPA==} + '@esbuild/aix-ppc64@0.28.1': resolution: {integrity: sha512-Svl7tq8k/08+p6CXPpRjQ1fKX+1odH/BQbb48fV6fj3CWHhsoIOoY87w1oHXm0qEpkIK3ZfVgp0hed3XBXzXMQ==} engines: {node: '>=18'} @@ -858,12 +870,159 @@ packages: engines: {node: '>=18'} hasBin: true + '@huggingface/jinja@0.5.9': + resolution: {integrity: sha512-uWTG+l3VJRsl7EXxYizuL3P+cCPoc3cRqbWWRcQN0FhejRfbdq0RNhCmbY/YDtnTcz9icdLYuLDjsnz4d8JMuw==} + engines: {node: '>=18'} + '@huggingface/tasks@0.21.21': resolution: {integrity: sha512-6cHtkeMvgbudI6p4+gPP/+xvcAG2uumF+ntojCSNnHy5xcSzuHXCMJHzh+rlX0TSZ/QdDZ4uWkPLQikbphNyEw==} + '@huggingface/tokenizers@0.1.3': + resolution: {integrity: sha512-8rF/RRT10u+kn7YuUbUg0OF30K8rjTc78aHpxT+qJ1uWSqxT1MHi8+9ltwYfkFYJzT/oS+qw3JVfHtNMGAdqyA==} + + '@huggingface/transformers@4.2.0': + resolution: {integrity: sha512-8BRCoBMH0XsWaEIamuR0LrJGAfftgHAfb2Vrffy0VKlSAE/MnUJ5/h/zTfEP3fDIft+nk7TqB8xXEyABGitBjQ==} + '@huggingface/xetchunk-wasm@0.1.0': resolution: {integrity: sha512-wWpp2qwPgf9kv1KLJjcDUk/OrpDOsFoQ3Qpz0U5LGn20csoymBf8eneOv6wm/GzPBzlWac1OYiR0aa1vT6aM2Q==} + '@img/colour@1.1.0': + resolution: {integrity: sha512-Td76q7j57o/tLVdgS746cYARfSyxk8iEfRxewL9h4OMzYhbW4TAcppl0mT4eyqXddh6L/jwoM75mo7ixa/pCeQ==} + engines: {node: '>=18'} + + '@img/sharp-darwin-arm64@0.34.5': + resolution: {integrity: sha512-imtQ3WMJXbMY4fxb/Ndp6HBTNVtWCUI0WdobyheGf5+ad6xX8VIDO8u2xE4qc/fr08CKG/7dDseFtn6M6g/r3w==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [arm64] + os: [darwin] + + '@img/sharp-darwin-x64@0.34.5': + resolution: {integrity: sha512-YNEFAF/4KQ/PeW0N+r+aVVsoIY0/qxxikF2SWdp+NRkmMB7y9LBZAVqQ4yhGCm/H3H270OSykqmQMKLBhBJDEw==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [x64] + os: [darwin] + + '@img/sharp-libvips-darwin-arm64@1.2.4': + resolution: {integrity: sha512-zqjjo7RatFfFoP0MkQ51jfuFZBnVE2pRiaydKJ1G/rHZvnsrHAOcQALIi9sA5co5xenQdTugCvtb1cuf78Vf4g==} + cpu: [arm64] + os: [darwin] + + '@img/sharp-libvips-darwin-x64@1.2.4': + resolution: {integrity: sha512-1IOd5xfVhlGwX+zXv2N93k0yMONvUlANylbJw1eTah8K/Jtpi15KC+WSiaX/nBmbm2HxRM1gZ0nSdjSsrZbGKg==} + cpu: [x64] + os: [darwin] + + '@img/sharp-libvips-linux-arm64@1.2.4': + resolution: {integrity: sha512-excjX8DfsIcJ10x1Kzr4RcWe1edC9PquDRRPx3YVCvQv+U5p7Yin2s32ftzikXojb1PIFc/9Mt28/y+iRklkrw==} + cpu: [arm64] + os: [linux] + + '@img/sharp-libvips-linux-arm@1.2.4': + resolution: {integrity: sha512-bFI7xcKFELdiNCVov8e44Ia4u2byA+l3XtsAj+Q8tfCwO6BQ8iDojYdvoPMqsKDkuoOo+X6HZA0s0q11ANMQ8A==} + cpu: [arm] + os: [linux] + + '@img/sharp-libvips-linux-ppc64@1.2.4': + resolution: {integrity: sha512-FMuvGijLDYG6lW+b/UvyilUWu5Ayu+3r2d1S8notiGCIyYU/76eig1UfMmkZ7vwgOrzKzlQbFSuQfgm7GYUPpA==} + cpu: [ppc64] + os: [linux] + + '@img/sharp-libvips-linux-riscv64@1.2.4': + resolution: {integrity: sha512-oVDbcR4zUC0ce82teubSm+x6ETixtKZBh/qbREIOcI3cULzDyb18Sr/Wcyx7NRQeQzOiHTNbZFF1UwPS2scyGA==} + cpu: [riscv64] + os: [linux] + + '@img/sharp-libvips-linux-s390x@1.2.4': + resolution: {integrity: sha512-qmp9VrzgPgMoGZyPvrQHqk02uyjA0/QrTO26Tqk6l4ZV0MPWIW6LTkqOIov+J1yEu7MbFQaDpwdwJKhbJvuRxQ==} + cpu: [s390x] + os: [linux] + + '@img/sharp-libvips-linux-x64@1.2.4': + resolution: {integrity: sha512-tJxiiLsmHc9Ax1bz3oaOYBURTXGIRDODBqhveVHonrHJ9/+k89qbLl0bcJns+e4t4rvaNBxaEZsFtSfAdquPrw==} + cpu: [x64] + os: [linux] + + '@img/sharp-libvips-linuxmusl-arm64@1.2.4': + resolution: {integrity: sha512-FVQHuwx1IIuNow9QAbYUzJ+En8KcVm9Lk5+uGUQJHaZmMECZmOlix9HnH7n1TRkXMS0pGxIJokIVB9SuqZGGXw==} + cpu: [arm64] + os: [linux] + + '@img/sharp-libvips-linuxmusl-x64@1.2.4': + resolution: {integrity: sha512-+LpyBk7L44ZIXwz/VYfglaX/okxezESc6UxDSoyo2Ks6Jxc4Y7sGjpgU9s4PMgqgjj1gZCylTieNamqA1MF7Dg==} + cpu: [x64] + os: [linux] + + '@img/sharp-linux-arm64@0.34.5': + resolution: {integrity: sha512-bKQzaJRY/bkPOXyKx5EVup7qkaojECG6NLYswgktOZjaXecSAeCWiZwwiFf3/Y+O1HrauiE3FVsGxFg8c24rZg==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [arm64] + os: [linux] + + '@img/sharp-linux-arm@0.34.5': + resolution: {integrity: sha512-9dLqsvwtg1uuXBGZKsxem9595+ujv0sJ6Vi8wcTANSFpwV/GONat5eCkzQo/1O6zRIkh0m/8+5BjrRr7jDUSZw==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [arm] + os: [linux] + + '@img/sharp-linux-ppc64@0.34.5': + resolution: {integrity: sha512-7zznwNaqW6YtsfrGGDA6BRkISKAAE1Jo0QdpNYXNMHu2+0dTrPflTLNkpc8l7MUP5M16ZJcUvysVWWrMefZquA==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [ppc64] + os: [linux] + + '@img/sharp-linux-riscv64@0.34.5': + resolution: {integrity: sha512-51gJuLPTKa7piYPaVs8GmByo7/U7/7TZOq+cnXJIHZKavIRHAP77e3N2HEl3dgiqdD/w0yUfiJnII77PuDDFdw==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [riscv64] + os: [linux] + + '@img/sharp-linux-s390x@0.34.5': + resolution: {integrity: sha512-nQtCk0PdKfho3eC5MrbQoigJ2gd1CgddUMkabUj+rBevs8tZ2cULOx46E7oyX+04WGfABgIwmMC0VqieTiR4jg==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [s390x] + os: [linux] + + '@img/sharp-linux-x64@0.34.5': + resolution: {integrity: sha512-MEzd8HPKxVxVenwAa+JRPwEC7QFjoPWuS5NZnBt6B3pu7EG2Ge0id1oLHZpPJdn3OQK+BQDiw9zStiHBTJQQQQ==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [x64] + os: [linux] + + '@img/sharp-linuxmusl-arm64@0.34.5': + resolution: {integrity: sha512-fprJR6GtRsMt6Kyfq44IsChVZeGN97gTD331weR1ex1c1rypDEABN6Tm2xa1wE6lYb5DdEnk03NZPqA7Id21yg==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [arm64] + os: [linux] + + '@img/sharp-linuxmusl-x64@0.34.5': + resolution: {integrity: sha512-Jg8wNT1MUzIvhBFxViqrEhWDGzqymo3sV7z7ZsaWbZNDLXRJZoRGrjulp60YYtV4wfY8VIKcWidjojlLcWrd8Q==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [x64] + os: [linux] + + '@img/sharp-wasm32@0.34.5': + resolution: {integrity: sha512-OdWTEiVkY2PHwqkbBI8frFxQQFekHaSSkUIJkwzclWZe64O1X4UlUjqqqLaPbUpMOQk6FBu/HtlGXNblIs0huw==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [wasm32] + + '@img/sharp-win32-arm64@0.34.5': + resolution: {integrity: sha512-WQ3AgWCWYSb2yt+IG8mnC6Jdk9Whs7O0gxphblsLvdhSpSTtmu69ZG1Gkb6NuvxsNACwiPV6cNSZNzt0KPsw7g==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [arm64] + os: [win32] + + '@img/sharp-win32-ia32@0.34.5': + resolution: {integrity: sha512-FV9m/7NmeCmSHDD5j4+4pNI8Cp3aW+JvLoXcTUo0IqyjSfAZJ8dIUmijx1qaJsIiU+Hosw6xM5KijAWRJCSgNg==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [ia32] + os: [win32] + + '@img/sharp-win32-x64@0.34.5': + resolution: {integrity: sha512-+29YMsqY2/9eFEiW93eqWnuLcWcufowXewwSNIT6UwZdUUCrM3oFjMWH/Z6/TMmb4hlFenmfAVbpWeup2jryCw==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + cpu: [x64] + os: [win32] + '@iovalkey/commands@0.1.0': resolution: {integrity: sha512-/B9W4qKSSITDii5nkBCHyPkIkAi+ealUtr1oqBJsLxjSRLka4pxun2VvMNSmcwgAMxgXtQfl0qRv7TE+udPJzg==} @@ -1923,6 +2082,24 @@ packages: resolution: {integrity: sha512-QVb0dM5HvG+uaxitm8wONl7jltx8dqhfU33DcqtOZcLSVIKSDDLDi7+0LbAKiyI8hD9u42m2YxXSkMGWThaecQ==} engines: {iojs: '>= 1.0.0', node: '>= 0.12.0'} + cohere-ai@8.0.0: + resolution: {integrity: sha512-j6gHcQfdK/r9i3e7FM6NtWx8OR270E/PyEyDEVQeeJT212bSJbWvWFfIg7WAPOUMHatDbxOFYP2aKKWnj8zSYA==} + engines: {node: '>=18.0.0'} + peerDependencies: + '@aws-crypto/sha256-js': ^5.2.0 + '@aws-sdk/credential-providers': ^3.583.0 + '@smithy/protocol-http': ^5.1.2 + '@smithy/signature-v4': ^5.1.2 + peerDependenciesMeta: + '@aws-crypto/sha256-js': + optional: true + '@aws-sdk/credential-providers': + optional: true + '@smithy/protocol-http': + optional: true + '@smithy/signature-v4': + optional: true + collect-v8-coverage@1.0.3: resolution: {integrity: sha512-1L5aqIkwPfiodaMgQunkF1zRhNqifHBmtbbbxcr6yVxxBnliw4TDOW6NxpO8DJLgJ16OT+Y4ztZqP6p/FtXnAw==} @@ -1974,6 +2151,10 @@ packages: convert-source-map@2.0.0: resolution: {integrity: sha512-Kvp459HrV2FEJ1CAsi1Ku+MY3kasH19TFykTz2xWmMeq6bk2NU3XXvfJ+Q61m0xktWwt+1HSYf3JZsTms3aRJg==} + convict@6.2.5: + resolution: {integrity: sha512-JtXpxqDqJ8P0UwEHwhxLzCIXQy97vlYBZR222Sbzb1q1Erex9ASrztJ29SyhWFQjod1AeFBaPzEEC8YvtZMIYg==} + engines: {node: '>=6'} + create-jest@29.7.0: resolution: {integrity: sha512-Adz2bdH0Vq3F53KEMJOoftQFutWCukm6J24wbPWRO4k1kMY7gS7ds/uoJkNuV8wDCtWWnuwGcJwpWcih+zEW1Q==} engines: {node: ^14.15.0 || ^16.10.0 || >=18.0.0} @@ -2251,6 +2432,9 @@ packages: resolution: {integrity: sha512-5z2nZxrnKxk+jLq5TyD0xbPXI2I18FF+knIZVG55e0CXWgXF/F4SpCBsiW7JTBPwghqXsC66T2yctnVT/sMO0g==} hasBin: true + flatbuffers@25.9.23: + resolution: {integrity: sha512-MI1qs7Lo4Syw0EOzUl0xjs2lsoeqFku44KpngfIduHBYvzm8h2+7K8YMQh1JtVVVrUvhLpNwqVi4DERegUJhPQ==} + fn.name@1.1.0: resolution: {integrity: sha512-GRnmB5gPyJpAhTQdSZTSp9uaPSvl09KoYcMQtsB9rQoOmzs9dH6ffeccH+Z+cv6P68Hu5bC6JjRh4Ah/mHSNRw==} @@ -2270,6 +2454,10 @@ packages: form-data-encoder@1.7.2: resolution: {integrity: sha512-qfqtYan3rxrnCk1VYaA4H+Ms9xdpPqvLZa6xmMgFvhO32x7/3J/ExcTd6qpxM0vH2GdMI+poehyBZvqfMTto8A==} + form-data-encoder@4.1.0: + resolution: {integrity: sha512-G6NsmEW15s0Uw9XnCg+33H3ViYRyiM0hMrMhhqQOR8NFc5GhYrI+6I3u7OTw7b91J2g8rtvMBZJDbcGb2YUniw==} + engines: {node: '>= 18'} + form-data@4.0.6: resolution: {integrity: sha512-vKatAh4SlVfgbv+YtmhiRjhEMJsYpsG1Y2rMQtR+SVSbytsSD1YGzDIcrAJmdFec88u/+VoGmxnl+80gL1tRCQ==} engines: {node: '>= 6'} @@ -2278,6 +2466,10 @@ packages: resolution: {integrity: sha512-0iirZp3uVDjVGt9p49aTaqjk84TrglENEDuqfdlZQ1roC9CWlPk6Avf8EEnZNcAqPonwkG35x4n3ww/1THYAeQ==} engines: {node: '>= 12.20'} + formdata-node@6.0.3: + resolution: {integrity: sha512-8e1++BCiTzUno9v5IZ2J6bv4RU+3UKDmqWUQD0MIMVCd9AdhWkO1gw57oo1mNEX1dMq2EGI+FbWz4B92pscSQg==} + engines: {node: '>= 18'} + formdata-polyfill@4.0.10: resolution: {integrity: sha512-buewHzMvYL29jdeQTVILecSaZKnt/RJWjoZCF5OW60Z67/GmSLBkOFM7qh1PI3zFNtJbaZL5eQu1vLfazOwj4g==} engines: {node: '>=12.20.0'} @@ -2405,6 +2597,9 @@ packages: resolution: {integrity: sha512-+CqsMbHPiSTdtSO14O51eMNlrp9N79gmeqmXeouJOhfucAedHw9noVe/n5uJk3tbKE6a+6ZCQg3RPhVhHByAIw==} engines: {node: '>=18'} + guid-typescript@1.0.9: + resolution: {integrity: sha512-Y8T4vYhEfwJOTbouREvG+3XDsjr8E3kIr7uf+JZ0BYloFsttiHU0WfvANVsR7TxNUJa/WpCnw/Ino/p+DeBhBQ==} + handlebars@4.7.9: resolution: {integrity: sha512-4E71E0rpOaQuJR2A3xDZ+GM1HyWYv1clR58tC8emQNeQe3RH7MAzSbat+V0wG78LQBo6m6bzSG/L4pBuCsgnUQ==} engines: {node: '>=0.4.7'} @@ -2839,6 +3034,9 @@ packages: lodash.camelcase@4.3.0: resolution: {integrity: sha512-TwuEnCnxbc3rAvhf/LbG7tJUDzhqXyFnv3dtzLOPgCG/hODL7WFnsbwktkD7yUV0RrreP/l1PALq/YSg6VvjlA==} + lodash.clonedeep@4.5.0: + resolution: {integrity: sha512-H5ZhCF25riFd9uB5UCkVKo61m3S/xZk1x4wA6yp/L3RFP6Z/eHH1ymQcGLo7J3GMPfm0V/7m1tryHuGVxpqEBQ==} + lodash.defaults@4.2.0: resolution: {integrity: sha512-qjxPLHd3r5DnsdGacqOMU6pb/avJzdh9tFX2ymgoZE27BmjXrNy/y4LoaiTeAb+O3gL8AfpJGtqfX/ae2leYYQ==} @@ -3144,10 +3342,23 @@ packages: onnxruntime-common@1.21.0: resolution: {integrity: sha512-Q632iLLrtCAVOTO65dh2+mNbQir/QNTVBG3h/QdZBpns7mZ0RYbLRBgGABPbpU9351AgYy7SJf1WaeVwMrBFPQ==} + onnxruntime-common@1.24.0-dev.20251116-b39e144322: + resolution: {integrity: sha512-BOoomdHYmNRL5r4iQ4bMvsl2t0/hzVQ3OM3PHD0gxeXu1PmggqBv3puZicEUVOA3AtHHYmqZtjMj9FOfGrATTw==} + + onnxruntime-common@1.24.3: + resolution: {integrity: sha512-GeuPZO6U/LBJXvwdaqHbuUmoXiEdeCjWi/EG7Y1HNnDwJYuk6WUbNXpF6luSUY8yASul3cmUlLGrCCL1ZgVXqA==} + onnxruntime-node@1.21.0: resolution: {integrity: sha512-NeaCX6WW2L8cRCSqy3bInlo5ojjQqu2fD3D+9W5qb5irwxhEyWKXeH2vZ8W9r6VxaMPUan+4/7NDwZMtouZxEw==} os: [win32, darwin, linux] + onnxruntime-node@1.24.3: + resolution: {integrity: sha512-JH7+czbc8ALA819vlTgcV+Q214/+VjGeBHDjX81+ZCD0PCVCIFGFNtT0V4sXG/1JXypKPgScQcB3ij/hk3YnTg==} + os: [win32, darwin, linux] + + onnxruntime-web@1.26.0-dev.20260416-b7804b056c: + resolution: {integrity: sha512-MD6Ss4GSpQBo6zqoJzyT9LRbKYs7x/JVN23FT24EcEvlqF4VuzPOeH6X38orZPKHQDbprn7K+SBpu0/mj2CQiw==} + open@10.2.0: resolution: {integrity: sha512-YgBpdJHPyQ2UE5x+hlSXcnejzAvD0b22U2OuAP+8OnlJT+PjWPxtgmGqKKc+RgTM63U9gN0YzrYc71R2WT/hTA==} engines: {node: '>=18'} @@ -3312,6 +3523,9 @@ packages: pkg-types@1.3.1: resolution: {integrity: sha512-/Jm5M4RvtBFVkKWRu2BLUTNP8/M2a+UwuAX+ae4770q1qVGtfjG+WTCupoZixokjmHiry8uI+dlY8KXYV5HVVQ==} + platform@1.3.6: + resolution: {integrity: sha512-fnWVljUchTro6RiCFvCXBbNhJc2NijN7oIQxbwsyL0buWJPG85v81ehlHI9fXrJsMNgTofEoWIQeClKpgxFLrg==} + postcss-load-config@6.0.1: resolution: {integrity: sha512-oPtTM4oerL+UXmx+93ytZVN82RrlY/wPUV8IeDxFrzIjXOLF1pN+EmKPLbubvKHT2HC20xXsCAH2Z+CKV6Oz/g==} engines: {node: '>= 18'} @@ -3555,6 +3769,10 @@ packages: resolution: {integrity: sha512-8I8TjW5KMOKsZQTvoxjuSIa7foAwPWGOts+6o7sgjz41/qMD9VQHEDxi6PBvK2l0MXUmqZyNpUK+T2tQaaElvw==} engines: {node: '>=10'} + sharp@0.34.5: + resolution: {integrity: sha512-Ou9I5Ft9WNcCbXrU9cMgPBcCK8LiwLqcbywW3t4oDV37n1pzpuNLsYiAV8eODnjbtQlSDwZ2cUEeQz4E54Hltg==} + engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} + shebang-command@2.0.0: resolution: {integrity: sha512-kHxr2zZpYtdmrN1qDjrrX/Z1rR1kG8Dx+gkpK1G4eXmvXswmcE1hTWBWYUzlraYw1/yZp6YuDY77YtvbN0dmDA==} engines: {node: '>=8'} @@ -4067,6 +4285,10 @@ packages: resolution: {integrity: sha512-YgvUTfwqyc7UXVMrB+SImsVYSmTS8X/tSrtdNZMImM+n7+QTriRXyXim0mBrTXNeqzVF0KWGgHPeiyViFFrNDw==} engines: {node: '>=18'} + yargs-parser@20.2.9: + resolution: {integrity: sha512-y11nGElTIV+CT3Zv9t7VKl+Q3hTQoT9a1Qzezhhl6Rp21gJ/IVTW7Z3y9EWXhuUBC2Shnf+DX0antecpAwSP8w==} + engines: {node: '>=10'} + yargs-parser@21.1.1: resolution: {integrity: sha512-tVpsJW7DdjecAiFpbIB1e3qxIQsE6NoPc5/eTdrbbIC4h0LVsWhnoa3g+m2HclBIujHzsxZ4VJVA+GUuc2/LBw==} engines: {node: '>=12'} @@ -4083,6 +4305,9 @@ packages: resolution: {integrity: sha512-rVksvsnNCdJ/ohGc6xgPwyN8eheCxsiLM8mxuE/t/mOVqJewPuO1miLpTHQiRgTKCLexL4MeAFVagts7HmNZ2Q==} engines: {node: '>=10'} + zeroentropy@0.1.0-alpha.10: + resolution: {integrity: sha512-4OwBUNuAQZzigzU75Nk+vbPGmL449U4XJtZTsZWu7cpztVHOx9gc9v0GPoyymmqtf0AnyUPPs8IWinNTUZLHhA==} + zod-to-json-schema@3.25.2: resolution: {integrity: sha512-O/PgfnpT1xKSDeQYSCfRI5Gy3hPf91mKVDuYLUHZJMiDFptvP41MSnWofm8dnCm0256ZNfZIM7DSzuSMAFnjHA==} peerDependencies: @@ -4997,6 +5222,11 @@ snapshots: transitivePeerDependencies: - supports-color + '@emnapi/runtime@1.11.2': + dependencies: + tslib: 2.8.1 + optional: true + '@esbuild/aix-ppc64@0.28.1': optional: true @@ -5117,13 +5347,121 @@ snapshots: optionalDependencies: cli-progress: 3.12.0 + '@huggingface/jinja@0.5.9': {} + '@huggingface/tasks@0.21.21': {} + '@huggingface/tokenizers@0.1.3': {} + + '@huggingface/transformers@4.2.0': + dependencies: + '@huggingface/jinja': 0.5.9 + '@huggingface/tokenizers': 0.1.3 + onnxruntime-node: 1.24.3 + onnxruntime-web: 1.26.0-dev.20260416-b7804b056c + sharp: 0.34.5 + '@huggingface/xetchunk-wasm@0.1.0': dependencies: '@huggingface/blake3-jit': 0.0.2 gearhash-jit: 1.0.2 + '@img/colour@1.1.0': {} + + '@img/sharp-darwin-arm64@0.34.5': + optionalDependencies: + '@img/sharp-libvips-darwin-arm64': 1.2.4 + optional: true + + '@img/sharp-darwin-x64@0.34.5': + optionalDependencies: + '@img/sharp-libvips-darwin-x64': 1.2.4 + optional: true + + '@img/sharp-libvips-darwin-arm64@1.2.4': + optional: true + + '@img/sharp-libvips-darwin-x64@1.2.4': + optional: true + + '@img/sharp-libvips-linux-arm64@1.2.4': + optional: true + + '@img/sharp-libvips-linux-arm@1.2.4': + optional: true + + '@img/sharp-libvips-linux-ppc64@1.2.4': + optional: true + + '@img/sharp-libvips-linux-riscv64@1.2.4': + optional: true + + '@img/sharp-libvips-linux-s390x@1.2.4': + optional: true + + '@img/sharp-libvips-linux-x64@1.2.4': + optional: true + + '@img/sharp-libvips-linuxmusl-arm64@1.2.4': + optional: true + + '@img/sharp-libvips-linuxmusl-x64@1.2.4': + optional: true + + '@img/sharp-linux-arm64@0.34.5': + optionalDependencies: + '@img/sharp-libvips-linux-arm64': 1.2.4 + optional: true + + '@img/sharp-linux-arm@0.34.5': + optionalDependencies: + '@img/sharp-libvips-linux-arm': 1.2.4 + optional: true + + '@img/sharp-linux-ppc64@0.34.5': + optionalDependencies: + '@img/sharp-libvips-linux-ppc64': 1.2.4 + optional: true + + '@img/sharp-linux-riscv64@0.34.5': + optionalDependencies: + '@img/sharp-libvips-linux-riscv64': 1.2.4 + optional: true + + '@img/sharp-linux-s390x@0.34.5': + optionalDependencies: + '@img/sharp-libvips-linux-s390x': 1.2.4 + optional: true + + '@img/sharp-linux-x64@0.34.5': + optionalDependencies: + '@img/sharp-libvips-linux-x64': 1.2.4 + optional: true + + '@img/sharp-linuxmusl-arm64@0.34.5': + optionalDependencies: + '@img/sharp-libvips-linuxmusl-arm64': 1.2.4 + optional: true + + '@img/sharp-linuxmusl-x64@0.34.5': + optionalDependencies: + '@img/sharp-libvips-linuxmusl-x64': 1.2.4 + optional: true + + '@img/sharp-wasm32@0.34.5': + dependencies: + '@emnapi/runtime': 1.11.2 + optional: true + + '@img/sharp-win32-arm64@0.34.5': + optional: true + + '@img/sharp-win32-ia32@0.34.5': + optional: true + + '@img/sharp-win32-x64@0.34.5': + optional: true + '@iovalkey/commands@0.1.0': {} '@isaacs/cliui@8.0.2': @@ -6328,6 +6666,18 @@ snapshots: co@4.6.0: {} + cohere-ai@8.0.0(@aws-crypto/sha256-js@5.2.0)(@smithy/protocol-http@5.5.2)(@smithy/signature-v4@5.5.2): + dependencies: + convict: 6.2.5 + form-data: 4.0.6 + form-data-encoder: 4.1.0 + formdata-node: 6.0.3 + readable-stream: 4.7.0 + optionalDependencies: + '@aws-crypto/sha256-js': 5.2.0 + '@smithy/protocol-http': 5.5.2 + '@smithy/signature-v4': 5.5.2 + collect-v8-coverage@1.0.3: {} color-convert@2.0.1: @@ -6371,6 +6721,11 @@ snapshots: convert-source-map@2.0.0: {} + convict@6.2.5: + dependencies: + lodash.clonedeep: 4.5.0 + yargs-parser: 20.2.9 + create-jest@29.7.0(@types/node@22.19.21)(ts-node@10.9.2(@types/node@22.19.21)(typescript@5.5.4)): dependencies: '@jest/types': 29.6.3 @@ -6666,6 +7021,8 @@ snapshots: kolorist: 1.8.0 read-pkg: 8.1.0 + flatbuffers@25.9.23: {} + fn.name@1.1.0: {} follow-redirects@1.16.0: {} @@ -6677,6 +7034,8 @@ snapshots: form-data-encoder@1.7.2: {} + form-data-encoder@4.1.0: {} + form-data@4.0.6: dependencies: asynckit: 0.4.0 @@ -6690,6 +7049,8 @@ snapshots: node-domexception: 1.0.0 web-streams-polyfill: 4.0.0-beta.3 + formdata-node@6.0.3: {} + formdata-polyfill@4.0.10: dependencies: fetch-blob: 3.2.0 @@ -6869,6 +7230,8 @@ snapshots: transitivePeerDependencies: - supports-color + guid-typescript@1.0.9: {} + handlebars@4.7.9: dependencies: minimist: 1.2.8 @@ -7463,6 +7826,8 @@ snapshots: lodash.camelcase@4.3.0: {} + lodash.clonedeep@4.5.0: {} + lodash.defaults@4.2.0: {} lodash.includes@4.3.0: {} @@ -7767,12 +8132,31 @@ snapshots: onnxruntime-common@1.21.0: {} + onnxruntime-common@1.24.0-dev.20251116-b39e144322: {} + + onnxruntime-common@1.24.3: {} + onnxruntime-node@1.21.0: dependencies: global-agent: 3.0.0 onnxruntime-common: 1.21.0 tar: 7.5.19 + onnxruntime-node@1.24.3: + dependencies: + adm-zip: 0.5.17 + global-agent: 3.0.0 + onnxruntime-common: 1.24.3 + + onnxruntime-web@1.26.0-dev.20260416-b7804b056c: + dependencies: + flatbuffers: 25.9.23 + guid-typescript: 1.0.9 + long: 5.3.2 + onnxruntime-common: 1.24.0-dev.20251116-b39e144322 + platform: 1.3.6 + protobufjs: 7.6.3 + open@10.2.0: dependencies: default-browser: 5.5.0 @@ -7944,6 +8328,8 @@ snapshots: mlly: 1.8.2 pathe: 2.0.3 + platform@1.3.6: {} + postcss-load-config@6.0.1: dependencies: lilconfig: 3.1.3 @@ -8198,6 +8584,37 @@ snapshots: dependencies: type-fest: 0.13.1 + sharp@0.34.5: + dependencies: + '@img/colour': 1.1.0 + detect-libc: 2.1.2 + semver: 7.8.4 + optionalDependencies: + '@img/sharp-darwin-arm64': 0.34.5 + '@img/sharp-darwin-x64': 0.34.5 + '@img/sharp-libvips-darwin-arm64': 1.2.4 + '@img/sharp-libvips-darwin-x64': 1.2.4 + '@img/sharp-libvips-linux-arm': 1.2.4 + '@img/sharp-libvips-linux-arm64': 1.2.4 + '@img/sharp-libvips-linux-ppc64': 1.2.4 + '@img/sharp-libvips-linux-riscv64': 1.2.4 + '@img/sharp-libvips-linux-s390x': 1.2.4 + '@img/sharp-libvips-linux-x64': 1.2.4 + '@img/sharp-libvips-linuxmusl-arm64': 1.2.4 + '@img/sharp-libvips-linuxmusl-x64': 1.2.4 + '@img/sharp-linux-arm': 0.34.5 + '@img/sharp-linux-arm64': 0.34.5 + '@img/sharp-linux-ppc64': 0.34.5 + '@img/sharp-linux-riscv64': 0.34.5 + '@img/sharp-linux-s390x': 0.34.5 + '@img/sharp-linux-x64': 0.34.5 + '@img/sharp-linuxmusl-arm64': 0.34.5 + '@img/sharp-linuxmusl-x64': 0.34.5 + '@img/sharp-wasm32': 0.34.5 + '@img/sharp-win32-arm64': 0.34.5 + '@img/sharp-win32-ia32': 0.34.5 + '@img/sharp-win32-x64': 0.34.5 + shebang-command@2.0.0: dependencies: shebang-regex: 3.0.0 @@ -8680,6 +9097,8 @@ snapshots: yallist@5.0.0: {} + yargs-parser@20.2.9: {} + yargs-parser@21.1.1: {} yargs@17.7.2: @@ -8696,6 +9115,18 @@ snapshots: yocto-queue@0.1.0: {} + zeroentropy@0.1.0-alpha.10: + dependencies: + '@types/node': 18.19.130 + '@types/node-fetch': 2.6.13 + abort-controller: 3.0.0 + agentkeepalive: 4.6.0 + form-data-encoder: 1.7.2 + formdata-node: 4.4.1 + node-fetch: 2.7.0 + transitivePeerDependencies: + - encoding + zod-to-json-schema@3.25.2(zod@3.25.76): dependencies: zod: 3.25.76 diff --git a/mem0-ts/src/oss/src/config/manager.ts b/mem0-ts/src/oss/src/config/manager.ts index c72692583..1fecc5154 100644 --- a/mem0-ts/src/oss/src/config/manager.ts +++ b/mem0-ts/src/oss/src/config/manager.ts @@ -177,6 +177,7 @@ export class ConfigManager { })(), disableHistory: userConfig.disableHistory || DEFAULT_MEMORY_CONFIG.disableHistory, + reranker: userConfig.reranker, }; // Validate the merged config diff --git a/mem0-ts/src/oss/src/index.ts b/mem0-ts/src/oss/src/index.ts index 8b772fbe1..ceaf8009c 100644 --- a/mem0-ts/src/oss/src/index.ts +++ b/mem0-ts/src/oss/src/index.ts @@ -45,4 +45,9 @@ export * from "./vector_stores/milvus"; export * from "./vector_stores/mongodb"; export * from "./vector_stores/opensearch"; export * from "./vector_stores/weaviate"; +export * from "./rerankers/base"; +export * from "./rerankers/cohere"; +export * from "./rerankers/llm"; +export * from "./rerankers/zeroentropy"; +export * from "./rerankers/cross_encoder"; export * from "./utils/factory"; diff --git a/mem0-ts/src/oss/src/memory/index.ts b/mem0-ts/src/oss/src/memory/index.ts index 5d8aea211..3173820fa 100644 --- a/mem0-ts/src/oss/src/memory/index.ts +++ b/mem0-ts/src/oss/src/memory/index.ts @@ -13,6 +13,7 @@ import { LLMFactory, VectorStoreFactory, HistoryManagerFactory, + RerankerFactory, } from "../utils/factory"; import { FactRetrievalSchema, @@ -28,6 +29,7 @@ import { import { DummyHistoryManager } from "../storage/DummyHistoryManager"; import { Embedder } from "../embeddings/base"; import { LLM } from "../llms/base"; +import { Reranker } from "../rerankers/base"; import { VectorStore } from "../vector_stores/base"; import { ConfigManager } from "../config/manager"; @@ -164,6 +166,7 @@ export class Memory { private embedder: Embedder; private vectorStore!: VectorStore; private llm: LLM; + private reranker: Reranker | null = null; private db: HistoryManager; private collectionName: string | undefined; private apiVersion: string; @@ -188,6 +191,12 @@ export class Memory { this.config.llm.provider, this.config.llm.config, ); + if (this.config.reranker) { + this.reranker = RerankerFactory.create( + this.config.reranker.provider, + this.config.reranker.config, + ); + } if (this.config.disableHistory) { this.db = new DummyHistoryManager(); } else { @@ -1541,8 +1550,30 @@ export class Memory { }; }); + // Step 10: Optionally re-rank with the configured reranker. Opt-in per + // search via `rerank: true`; a no-op when no reranker is configured. + const invokeReranker = Boolean( + config.rerank && this.reranker && results.length > 0, + ); + let finalResults = results; + if (invokeReranker) { + try { + const ranked = await this.reranker!.rerank( + query, + results.map((r) => r.memory), + topK, + ); + finalResults = ranked.map((r) => ({ + ...results[r.index], + rerankScore: r.rerankScore, + })); + } catch (e) { + console.warn(`Reranking failed, using original results: ${e}`); + } + } + const result = { - results, + results: finalResults, }; const searchElapsedMs = Date.now() - searchStartMs; if (temporalUsageNotice) { diff --git a/mem0-ts/src/oss/src/memory/memory.types.ts b/mem0-ts/src/oss/src/memory/memory.types.ts index 861db25f7..33ca2df14 100644 --- a/mem0-ts/src/oss/src/memory/memory.types.ts +++ b/mem0-ts/src/oss/src/memory/memory.types.ts @@ -36,6 +36,11 @@ export interface SearchMemoryOptions { threshold?: number; explain?: boolean; referenceDate?: number | string | Date | null; + /** + * Re-rank the results with the configured reranker before returning. No-op + * when no `reranker` is configured on the Memory. + */ + rerank?: boolean; /** Include expired memories in the results. Defaults to false. */ showExpired?: boolean; } diff --git a/mem0-ts/src/oss/src/rerankers/base.ts b/mem0-ts/src/oss/src/rerankers/base.ts new file mode 100644 index 000000000..a785558bb --- /dev/null +++ b/mem0-ts/src/oss/src/rerankers/base.ts @@ -0,0 +1,21 @@ +export interface RerankResult { + /** Index into the input `documents` array. */ + index: number; + /** Relevance of the document to the query, 0..1, higher = more relevant. */ + rerankScore: number; +} + +export interface Reranker { + /** + * Rank `documents` by relevance to `query`. + * + * Returns results sorted by descending relevance. When `topK` is given, at + * most that many results are returned. Each result's `index` points back into + * the input `documents` array so callers can recover the original item. + */ + rerank( + query: string, + documents: string[], + topK?: number, + ): Promise; +} diff --git a/mem0-ts/src/oss/src/rerankers/cohere.test.ts b/mem0-ts/src/oss/src/rerankers/cohere.test.ts new file mode 100644 index 000000000..33389a1b9 --- /dev/null +++ b/mem0-ts/src/oss/src/rerankers/cohere.test.ts @@ -0,0 +1,140 @@ +const mockRerank = jest.fn(); + +jest.mock("cohere-ai", () => ({ + CohereClient: jest.fn().mockImplementation(() => ({ + rerank: mockRerank, + })), +})); + +import { CohereClient } from "cohere-ai"; +import { CohereReranker } from "./cohere"; + +describe("CohereReranker", () => { + beforeEach(() => { + mockRerank.mockReset(); + (CohereClient as unknown as jest.Mock).mockClear(); + }); + + it("throws when no API key is provided or configured", () => { + const originalEnv = process.env.COHERE_API_KEY; + delete process.env.COHERE_API_KEY; + + expect(() => new CohereReranker({})).toThrow(/Cohere API key is required/); + + if (originalEnv !== undefined) process.env.COHERE_API_KEY = originalEnv; + }); + + it("sends the query, documents, topN, and default model to Cohere", async () => { + mockRerank.mockResolvedValue({ results: [] }); + const reranker = new CohereReranker({ apiKey: "key" }); + + await reranker.rerank("capital of US?", ["a", "b", "c"], 2); + + expect(mockRerank).toHaveBeenCalledWith({ + model: "rerank-v3.5", + query: "capital of US?", + documents: ["a", "b", "c"], + topN: 2, + returnDocuments: false, + maxChunksPerDoc: undefined, + }); + }); + + it("defaults topN to documents.length when neither the call nor config sets a top_k", async () => { + mockRerank.mockResolvedValue({ results: [] }); + const reranker = new CohereReranker({ apiKey: "key" }); + + await reranker.rerank("q", ["a", "b", "c"]); + + expect(mockRerank).toHaveBeenCalledWith( + expect.objectContaining({ topN: 3 }), + ); + }); + + it("forwards returnDocuments and maxChunksPerDoc from config", async () => { + mockRerank.mockResolvedValue({ results: [] }); + const reranker = new CohereReranker({ + apiKey: "key", + returnDocuments: true, + maxChunksPerDoc: 5, + }); + + await reranker.rerank("q", ["a"]); + + expect(mockRerank).toHaveBeenCalledWith( + expect.objectContaining({ returnDocuments: true, maxChunksPerDoc: 5 }), + ); + }); + + it("returns Cohere's ranked results as {index, rerankScore}", async () => { + mockRerank.mockResolvedValue({ + results: [ + { index: 2, relevanceScore: 0.9 }, + { index: 0, relevanceScore: 0.31 }, + ], + }); + const reranker = new CohereReranker({ apiKey: "key" }); + + const results = await reranker.rerank("q", ["x", "y", "z"]); + + expect(results).toEqual([ + { index: 2, rerankScore: 0.9 }, + { index: 0, rerankScore: 0.31 }, + ]); + }); + + it("uses a custom model when provided", async () => { + mockRerank.mockResolvedValue({ results: [] }); + const reranker = new CohereReranker({ + apiKey: "key", + model: "rerank-v4.0-pro", + }); + + await reranker.rerank("q", ["a"]); + + expect(mockRerank).toHaveBeenCalledWith( + expect.objectContaining({ model: "rerank-v4.0-pro" }), + ); + }); + + it("returns an empty array without calling Cohere when there are no documents", async () => { + const reranker = new CohereReranker({ apiKey: "key" }); + + const results = await reranker.rerank("q", []); + + expect(results).toEqual([]); + expect(mockRerank).not.toHaveBeenCalled(); + }); + + it("falls back to the original order with rerankScore 0.0 when the Cohere API call fails", async () => { + mockRerank.mockRejectedValue(new Error("cohere is down")); + const warnSpy = jest.spyOn(console, "warn").mockImplementation(() => {}); + const reranker = new CohereReranker({ apiKey: "key" }); + + const results = await reranker.rerank("q", ["a", "b", "c"]); + + expect(results).toEqual([ + { index: 0, rerankScore: 0.0 }, + { index: 1, rerankScore: 0.0 }, + { index: 2, rerankScore: 0.0 }, + ]); + expect(warnSpy).toHaveBeenCalled(); + + warnSpy.mockRestore(); + }); + + it("slices the fallback results by topK when the Cohere API call fails", async () => { + mockRerank.mockRejectedValue(new Error("cohere is down")); + jest.spyOn(console, "warn").mockImplementation(() => {}); + const reranker = new CohereReranker({ apiKey: "key", topK: 2 }); + + const results = await reranker.rerank("q", ["a", "b", "c"]); + + expect(results).toEqual([ + { index: 0, rerankScore: 0.0 }, + { index: 1, rerankScore: 0.0 }, + ]); + + (console.warn as jest.Mock).mockRestore(); + }); +}); diff --git a/mem0-ts/src/oss/src/rerankers/cohere.ts b/mem0-ts/src/oss/src/rerankers/cohere.ts new file mode 100644 index 000000000..886c2ec3c --- /dev/null +++ b/mem0-ts/src/oss/src/rerankers/cohere.ts @@ -0,0 +1,61 @@ +import { CohereClient } from "cohere-ai"; +import { RerankerConfig } from "../types"; +import { Reranker, RerankResult } from "./base"; + +const DEFAULT_MODEL = "rerank-v3.5"; + +export class CohereReranker implements Reranker { + private client: CohereClient; + private model: string; + private topK?: number; + private returnDocuments: boolean; + private maxChunksPerDoc?: number; + + constructor(config: RerankerConfig) { + const apiKey = config.apiKey || process.env.COHERE_API_KEY; + if (!apiKey) { + throw new Error( + "Cohere API key is required. Set COHERE_API_KEY environment variable or pass apiKey in config.", + ); + } + this.client = new CohereClient({ token: apiKey }); + this.model = config.model || DEFAULT_MODEL; + this.topK = config.topK; + this.returnDocuments = config.returnDocuments ?? false; + this.maxChunksPerDoc = config.maxChunksPerDoc; + } + + async rerank( + query: string, + documents: string[], + topK?: number, + ): Promise { + if (documents.length === 0) return []; + + try { + const response = await this.client.rerank({ + model: this.model, + query, + documents, + topN: topK || this.topK || documents.length, + returnDocuments: this.returnDocuments, + maxChunksPerDoc: this.maxChunksPerDoc, + }); + + return response.results.map((result) => ({ + index: result.index, + rerankScore: result.relevanceScore, + })); + } catch (e) { + console.warn( + `Cohere reranking failed, falling back to original order: ${e}`, + ); + const scored = documents.map((_, index) => ({ + index, + rerankScore: 0.0, + })); + const finalTopK = topK || this.topK; + return finalTopK ? scored.slice(0, finalTopK) : scored; + } + } +} diff --git a/mem0-ts/src/oss/src/rerankers/cross_encoder.test.ts b/mem0-ts/src/oss/src/rerankers/cross_encoder.test.ts new file mode 100644 index 000000000..552607c15 --- /dev/null +++ b/mem0-ts/src/oss/src/rerankers/cross_encoder.test.ts @@ -0,0 +1,189 @@ +const mockModelFromPretrained = jest.fn(); +const mockTokenizerFromPretrained = jest.fn(); + +jest.mock("@huggingface/transformers", () => ({ + AutoModelForSequenceClassification: { + from_pretrained: mockModelFromPretrained, + }, + AutoTokenizer: { from_pretrained: mockTokenizerFromPretrained }, +})); + +import { CrossEncoderReranker } from "./cross_encoder"; + +const sigmoid = (x: number) => 1 / (1 + Math.exp(-x)); + +/** Wire the mocked tokenizer + model so the model returns `logits` for a call. */ +function setupModel(logits: number[][]) { + const tokenizer = jest.fn().mockReturnValue({ input_ids: [] }); + mockTokenizerFromPretrained.mockResolvedValue(tokenizer); + const model = jest + .fn() + .mockResolvedValue({ logits: { tolist: () => logits } }); + mockModelFromPretrained.mockResolvedValue(model); + return { tokenizer, model }; +} + +describe("CrossEncoderReranker", () => { + beforeEach(() => { + mockModelFromPretrained.mockReset(); + mockTokenizerFromPretrained.mockReset(); + }); + + it("scores each document and returns them sorted by relevance, sigmoid-normalized to [0,1]", async () => { + setupModel([[0.0], [2.0], [-1.0]]); + const reranker = new CrossEncoderReranker({}, "default-model"); + + const results = await reranker.rerank("q", ["a", "b", "c"]); + + // sigmoid: b(2.0)=0.88 > a(0.0)=0.5 > c(-1.0)=0.27 + expect(results.map((r) => r.index)).toEqual([1, 0, 2]); + expect(results[0].rerankScore).toBeCloseTo(sigmoid(2.0), 5); + expect(results[1].rerankScore).toBeCloseTo(sigmoid(0.0), 5); + expect(results[2].rerankScore).toBeCloseTo(sigmoid(-1.0), 5); + }); + + it("pairs the query with each document via text_pair when tokenizing", async () => { + const { tokenizer } = setupModel([[0.1], [0.2]]); + const reranker = new CrossEncoderReranker( + { maxLength: 128 }, + "default-model", + ); + + await reranker.rerank("what is x", ["doc one", "doc two"]); + + expect(tokenizer).toHaveBeenCalledWith( + ["what is x", "what is x"], + expect.objectContaining({ + text_pair: ["doc one", "doc two"], + padding: true, + truncation: true, + max_length: 128, + }), + ); + }); + + it("applies the topK limit", async () => { + setupModel([[0.0], [2.0], [-1.0]]); + const reranker = new CrossEncoderReranker({}, "default-model"); + + const results = await reranker.rerank("q", ["a", "b", "c"], 2); + + expect(results).toHaveLength(2); + expect(results.map((r) => r.index)).toEqual([1, 0]); + }); + + it("falls back to config.topK when the rerank() call omits one", async () => { + setupModel([[0.0], [2.0], [-1.0]]); + const reranker = new CrossEncoderReranker({ topK: 1 }, "default-model"); + + const results = await reranker.rerank("q", ["a", "b", "c"]); + + expect(results).toHaveLength(1); + expect(results.map((r) => r.index)).toEqual([1]); + }); + + it("returns [] without loading the model when there are no documents", async () => { + const reranker = new CrossEncoderReranker({}, "default-model"); + + const results = await reranker.rerank("q", []); + + expect(results).toEqual([]); + expect(mockModelFromPretrained).not.toHaveBeenCalled(); + expect(mockTokenizerFromPretrained).not.toHaveBeenCalled(); + }); + + it("returns raw logits as scores when normalize is false", async () => { + setupModel([[2.0], [0.0]]); + const reranker = new CrossEncoderReranker( + { normalize: false }, + "default-model", + ); + + const results = await reranker.rerank("q", ["a", "b"]); + + expect(results.map((r) => r.index)).toEqual([0, 1]); + expect(results[0].rerankScore).toBe(2.0); + expect(results[1].rerankScore).toBe(0.0); + }); + + it("loads the model and tokenizer only once across multiple rerank calls", async () => { + setupModel([[0.5]]); + const reranker = new CrossEncoderReranker({}, "default-model"); + + await reranker.rerank("q", ["a"]); + await reranker.rerank("q2", ["b"]); + + expect(mockModelFromPretrained).toHaveBeenCalledTimes(1); + expect(mockTokenizerFromPretrained).toHaveBeenCalledTimes(1); + }); + + it("loads the default model, or the configured model when provided", async () => { + setupModel([[0.5]]); + + await new CrossEncoderReranker({}, "the-default").rerank("q", ["a"]); + expect(mockModelFromPretrained).toHaveBeenCalledWith( + "the-default", + expect.any(Object), + ); + + mockModelFromPretrained.mockClear(); + setupModel([[0.5]]); + await new CrossEncoderReranker( + { model: "custom/model" }, + "the-default", + ).rerank("q", ["a"]); + expect(mockModelFromPretrained).toHaveBeenCalledWith( + "custom/model", + expect.any(Object), + ); + }); + + it("applies a default maxLength (as the huggingface provider passes 512) when config omits one", async () => { + const { tokenizer } = setupModel([[0.5]]); + const reranker = new CrossEncoderReranker({}, "default-model", 512); + + await reranker.rerank("q", ["a"]); + + expect(tokenizer).toHaveBeenCalledWith( + ["q"], + expect.objectContaining({ max_length: 512 }), + ); + }); + + it("falls back to the original order with rerankScore 0.0 when the model fails to load", async () => { + mockModelFromPretrained.mockResolvedValue(jest.fn()); + mockTokenizerFromPretrained.mockRejectedValue( + new Error("model download failed"), + ); + const warnSpy = jest.spyOn(console, "warn").mockImplementation(() => {}); + const reranker = new CrossEncoderReranker({}, "default-model"); + + const results = await reranker.rerank("q", ["a", "b", "c"]); + + expect(results).toEqual([ + { index: 0, rerankScore: 0.0 }, + { index: 1, rerankScore: 0.0 }, + { index: 2, rerankScore: 0.0 }, + ]); + expect(warnSpy).toHaveBeenCalled(); + + warnSpy.mockRestore(); + }); + + it("falls back to the original order with rerankScore 0.0, sliced by topK, when scoring fails", async () => { + const tokenizer = jest.fn().mockReturnValue({ input_ids: [] }); + mockTokenizerFromPretrained.mockResolvedValue(tokenizer); + mockModelFromPretrained.mockResolvedValue( + jest.fn().mockRejectedValue(new Error("forward pass failed")), + ); + const warnSpy = jest.spyOn(console, "warn").mockImplementation(() => {}); + const reranker = new CrossEncoderReranker({}, "default-model"); + + const results = await reranker.rerank("q", ["a", "b"], 1); + + expect(results).toEqual([{ index: 0, rerankScore: 0.0 }]); + expect(warnSpy).toHaveBeenCalled(); + + warnSpy.mockRestore(); + }); +}); diff --git a/mem0-ts/src/oss/src/rerankers/cross_encoder.ts b/mem0-ts/src/oss/src/rerankers/cross_encoder.ts new file mode 100644 index 000000000..095455f11 --- /dev/null +++ b/mem0-ts/src/oss/src/rerankers/cross_encoder.ts @@ -0,0 +1,98 @@ +import { RerankerConfig } from "../types"; +import { Reranker, RerankResult } from "./base"; + +const sigmoid = (x: number) => 1 / (1 + Math.exp(-x)); + +export class CrossEncoderReranker implements Reranker { + private modelId: string; + private device?: string; + private maxLength?: number; + private normalize: boolean; + private topK?: number; + // ponytail: batchSize/showProgressBar are accepted for config parity with the + // Python SDK but are no-ops here — a memory search reranks a small candidate + // set in a single forward pass. Chunk by batchSize if that ever grows. + private loaded?: Promise<{ model: any; tokenizer: any }>; + + constructor( + config: RerankerConfig, + defaultModel: string, + defaultMaxLength?: number, + ) { + this.modelId = config.model || defaultModel; + this.device = config.device; + this.maxLength = config.maxLength ?? defaultMaxLength; + this.normalize = config.normalize ?? true; + this.topK = config.topK; + } + + private load() { + if (!this.loaded) { + this.loaded = (async () => { + // Lazy-load Transformers.js (and its onnxruntime native binding) only + // when a rerank actually runs. A static import would pull onnxruntime + // into every `new Memory()`, colliding on Linux with fastembed's + // separate onnxruntime version — see the merge with the FastEmbed + // embedder. Deferring it keeps memory construction free of ONNX. + const { AutoModelForSequenceClassification, AutoTokenizer } = + await import("@huggingface/transformers"); + const options: any = {}; + if (this.device) options.device = this.device; + const model = await AutoModelForSequenceClassification.from_pretrained( + this.modelId, + options, + ); + const tokenizer = await AutoTokenizer.from_pretrained(this.modelId); + return { model, tokenizer }; + })(); + } + return this.loaded; + } + + async rerank( + query: string, + documents: string[], + topK?: number, + ): Promise { + if (documents.length === 0) return []; + + try { + const { model, tokenizer } = await this.load(); + + const inputs = tokenizer( + documents.map(() => query), + { + text_pair: documents, + padding: true, + truncation: true, + ...(this.maxLength ? { max_length: this.maxLength } : {}), + }, + ); + + const { logits } = await model(inputs); + const rows: unknown[] = logits.tolist(); + + const scored = rows.map((row, index) => { + const logit = Array.isArray(row) ? (row[0] as number) : (row as number); + return { + index, + rerankScore: this.normalize ? sigmoid(logit) : logit, + }; + }); + + scored.sort((a, b) => b.rerankScore - a.rerankScore); + const finalTopK = topK || this.topK; + return finalTopK ? scored.slice(0, finalTopK) : scored; + } catch (e) { + console.warn( + `Cross-encoder reranking failed, falling back to original order: ${e}`, + ); + const scored = documents.map((_, index) => ({ + index, + rerankScore: 0.0, + })); + const finalTopK = topK || this.topK; + return finalTopK ? scored.slice(0, finalTopK) : scored; + } + } +} diff --git a/mem0-ts/src/oss/src/rerankers/llm.test.ts b/mem0-ts/src/oss/src/rerankers/llm.test.ts new file mode 100644 index 000000000..0b953d260 --- /dev/null +++ b/mem0-ts/src/oss/src/rerankers/llm.test.ts @@ -0,0 +1,174 @@ +import { LLM } from "../llms/base"; +import { LLMReranker } from "./llm"; + +// Duplicated rather than imported from ./llm.ts so this test catches drift. +const EXPECTED_SYSTEM_PROMPT = `You are a relevance scoring assistant. Given a query and a document, score how relevant the document is to the query. + +Score the relevance on a scale from 0.0 to 1.0, where: +- 1.0 = Perfectly relevant and directly answers the query +- 0.8-0.9 = Highly relevant with good information +- 0.6-0.7 = Moderately relevant with some useful information +- 0.4-0.5 = Slightly relevant with limited useful information +- 0.0-0.3 = Not relevant or no useful information + +Respond with only a single numerical score between 0.0 and 1.0. Do not include any explanation or additional text.`; + +/** + * Fake LLM that scores a document by looking up the document text inside the + * prompt. Test document tokens must be distinct and must not be substrings of + * the prompt boilerplate (e.g. avoid "a"/"b"), or the lookup resolves the + * wrong doc. + */ +function makeLLM(scoreByDoc: Record): LLM { + return { + generateResponse: async ( + messages: Array<{ role: string; content: string }>, + ) => { + const prompt = messages.map((m) => m.content).join("\n"); + const doc = Object.keys(scoreByDoc).find((d) => prompt.includes(d)); + return doc ? scoreByDoc[doc] : "no number here"; + }, + generateChat: async () => ({ content: "", role: "assistant" }), + }; +} + +describe("LLMReranker", () => { + it("sorts documents by descending relevance score", async () => { + const llm = makeLLM({ cats: "0.2", dogs: "0.9", fish: "0.5" }); + const reranker = new LLMReranker({}, llm); + + const results = await reranker.rerank("pets", ["cats", "dogs", "fish"]); + + expect(results.map((r) => r.index)).toEqual([1, 2, 0]); + expect(results.map((r) => r.rerankScore)).toEqual([0.9, 0.5, 0.2]); + }); + + it("clamps scores to the [0, 1] range", async () => { + const llm = makeLLM({ zebra: "1.5", walrus: "-0.3" }); + const reranker = new LLMReranker({}, llm); + + const results = await reranker.rerank("q", ["zebra", "walrus"]); + + const byIndex = new Map(results.map((r) => [r.index, r.rerankScore])); + expect(byIndex.get(0)).toBe(1); // "zebra" 1.5 -> clamped to 1 + expect(byIndex.get(1)).toBe(0); // "walrus" -0.3 -> clamped to 0 + }); + + it("truncates results to topK", async () => { + const llm = makeLLM({ alpha: "0.1", bravo: "0.8", charlie: "0.5" }); + const reranker = new LLMReranker({}, llm); + + const results = await reranker.rerank( + "q", + ["alpha", "bravo", "charlie"], + 2, + ); + + expect(results).toHaveLength(2); + expect(results.map((r) => r.index)).toEqual([1, 2]); // bravo(0.8), charlie(0.5) + }); + + it("falls back to config.topK when the rerank() call omits one", async () => { + const llm = makeLLM({ alpha: "0.1", bravo: "0.8", charlie: "0.5" }); + const reranker = new LLMReranker({ topK: 1 }, llm); + + const results = await reranker.rerank("q", ["alpha", "bravo", "charlie"]); + + expect(results).toHaveLength(1); + expect(results[0].index).toBe(1); // bravo(0.8) + }); + + it("falls back to a neutral score of 0.5 (not 0) when the LLM output has no number", async () => { + const llm = makeLLM({ junk: "I cannot rate this" }); + const reranker = new LLMReranker({}, llm); + + const results = await reranker.rerank("q", ["junk"]); + + expect(results[0].rerankScore).toBe(0.5); + }); + + it("prefers a decimal match over an integer match when extracting the score", async () => { + const llm = makeLLM({ item: "The score is 0.73 out of 1" }); + const reranker = new LLMReranker({}, llm); + + const results = await reranker.rerank("q", ["item"]); + + expect(results[0].rerankScore).toBe(0.73); + }); + + it("falls back to an integer match when no decimal is present", async () => { + const llm = makeLLM({ item: "I'd say this is a solid 1" }); + const reranker = new LLMReranker({}, llm); + + const results = await reranker.rerank("q", ["item"]); + + expect(results[0].rerankScore).toBe(1); + }); + + it("assigns a neutral 0.5 score (not 0.0) when a per-document LLM call fails, and still returns that document", async () => { + const llm: LLM = { + generateResponse: jest + .fn() + .mockResolvedValueOnce("0.9") // scores "good" + .mockRejectedValueOnce(new Error("rate limited")), // scores "bad" + generateChat: async () => ({ content: "", role: "assistant" }), + }; + const warnSpy = jest.spyOn(console, "warn").mockImplementation(() => {}); + const reranker = new LLMReranker({}, llm); + + const results = await reranker.rerank("q", ["good", "bad"]); + + expect(results).toHaveLength(2); + const byIndex = new Map(results.map((r) => [r.index, r.rerankScore])); + expect(byIndex.get(0)).toBe(0.9); + expect(byIndex.get(1)).toBe(0.5); + expect(warnSpy).toHaveBeenCalled(); + + warnSpy.mockRestore(); + }); + + it("sends the exact system prompt and a separate user message with the query and document", async () => { + const generateResponse = jest.fn().mockResolvedValue("0.5"); + const llm: LLM = { + generateResponse, + generateChat: async () => ({ content: "", role: "assistant" }), + }; + const reranker = new LLMReranker({}, llm); + + await reranker.rerank("what is the capital?", [ + "Paris is the capital of France.", + ]); + + expect(generateResponse).toHaveBeenCalledWith([ + { role: "system", content: EXPECTED_SYSTEM_PROMPT }, + { + role: "user", + content: + "Query: what is the capital?\n\nDocument: Paris is the capital of France.", + }, + ]); + }); + + it("truncates the query and document to 4000 characters before sending", async () => { + const generateResponse = jest.fn().mockResolvedValue("0.5"); + const llm: LLM = { + generateResponse, + generateChat: async () => ({ content: "", role: "assistant" }), + }; + const reranker = new LLMReranker({}, llm); + const longQuery = "q".repeat(5000); + const longDoc = "d".repeat(5000); + + await reranker.rerank(longQuery, [longDoc]); + + const userMessage = generateResponse.mock.calls[0][0][1]; + const sentQuery = userMessage.content.match(/^Query: (q+)/)[1]; + const sentDoc = userMessage.content.match(/Document: (d+)/)[1]; + expect(sentQuery).toHaveLength(4000); + expect(sentDoc).toHaveLength(4000); + }); + + it("throws when no LLM is provided", () => { + expect(() => new LLMReranker({}, undefined as unknown as LLM)).toThrow(); + }); +}); diff --git a/mem0-ts/src/oss/src/rerankers/llm.ts b/mem0-ts/src/oss/src/rerankers/llm.ts new file mode 100644 index 000000000..e2285705b --- /dev/null +++ b/mem0-ts/src/oss/src/rerankers/llm.ts @@ -0,0 +1,87 @@ +import { RerankerConfig } from "../types"; +import { LLM, LLMResponse } from "../llms/base"; +import { Reranker, RerankResult } from "./base"; + +const SYSTEM_PROMPT = `You are a relevance scoring assistant. Given a query and a document, score how relevant the document is to the query. + +Score the relevance on a scale from 0.0 to 1.0, where: +- 1.0 = Perfectly relevant and directly answers the query +- 0.8-0.9 = Highly relevant with good information +- 0.6-0.7 = Moderately relevant with some useful information +- 0.4-0.5 = Slightly relevant with limited useful information +- 0.0-0.3 = Not relevant or no useful information + +Respond with only a single numerical score between 0.0 and 1.0. Do not include any explanation or additional text.`; + +const MAX_INPUT_LEN = 4000; + +export class LLMReranker implements Reranker { + private llm: LLM; + private topK?: number; + + constructor(config: RerankerConfig, llm: LLM) { + if (!llm) { + throw new Error( + "LLMReranker requires an LLM instance; RerankerFactory should always provide one for the llm_reranker provider.", + ); + } + this.llm = llm; + this.topK = config.topK; + } + + async rerank( + query: string, + documents: string[], + topK?: number, + ): Promise { + if (documents.length === 0) return []; + + const scored = await Promise.all( + documents.map(async (document, index) => { + try { + const rerankScore = await this.score(query, document); + return { index, rerankScore }; + } catch (e) { + console.warn( + `LLM reranking failed for a document, assigning neutral score: ${e}`, + ); + return { index, rerankScore: 0.5 }; + } + }), + ); + + scored.sort((a, b) => b.rerankScore - a.rerankScore); + const finalTopK = topK || this.topK; + return finalTopK ? scored.slice(0, finalTopK) : scored; + } + + private async score(query: string, document: string): Promise { + const safeQuery = query.slice(0, MAX_INPUT_LEN); + const safeDoc = document.slice(0, MAX_INPUT_LEN); + const userMessage = `Query: ${safeQuery}\n\nDocument: ${safeDoc}`; + + const response = await this.llm.generateResponse([ + { role: "system", content: SYSTEM_PROMPT }, + { role: "user", content: userMessage }, + ]); + + const text = + typeof response === "string" + ? response + : ((response as LLMResponse)?.content ?? ""); + + return this.extractScore(text); + } + + private extractScore(responseText: string): number { + const matches = + responseText.match(/-?\d+\.\d+/g) || responseText.match(/-?\d+/g); + + if (matches && matches.length > 0) { + const score = parseFloat(matches[0]); + return Math.min(Math.max(score, 0.0), 1.0); + } + + return 0.5; + } +} diff --git a/mem0-ts/src/oss/src/rerankers/zeroentropy.test.ts b/mem0-ts/src/oss/src/rerankers/zeroentropy.test.ts new file mode 100644 index 000000000..a18202373 --- /dev/null +++ b/mem0-ts/src/oss/src/rerankers/zeroentropy.test.ts @@ -0,0 +1,133 @@ +const mockRerank = jest.fn(); + +jest.mock("zeroentropy", () => ({ + ZeroEntropy: jest.fn().mockImplementation(() => ({ + models: { rerank: mockRerank }, + })), +})); + +import { ZeroEntropy } from "zeroentropy"; +import { ZeroEntropyReranker } from "./zeroentropy"; + +describe("ZeroEntropyReranker", () => { + beforeEach(() => { + mockRerank.mockReset(); + (ZeroEntropy as unknown as jest.Mock).mockClear(); + }); + + it("throws when no API key is provided or configured", () => { + const originalEnv = process.env.ZERO_ENTROPY_API_KEY; + delete process.env.ZERO_ENTROPY_API_KEY; + + expect(() => new ZeroEntropyReranker({})).toThrow( + /Zero Entropy API key is required/, + ); + + if (originalEnv !== undefined) + process.env.ZERO_ENTROPY_API_KEY = originalEnv; + }); + + it("sends the query, documents, and default model to ZeroEntropy without a top_n parameter", async () => { + mockRerank.mockResolvedValue({ results: [] }); + const reranker = new ZeroEntropyReranker({ apiKey: "key" }); + + await reranker.rerank("capital of US?", ["a", "b", "c"], 2); + + expect(mockRerank).toHaveBeenCalledWith({ + model: "zerank-1", + query: "capital of US?", + documents: ["a", "b", "c"], + }); + }); + + it("maps ZeroEntropy's results (relevance_score) to {index, rerankScore}", async () => { + mockRerank.mockResolvedValue({ + results: [ + { index: 2, relevance_score: 0.9 }, + { index: 0, relevance_score: 0.31 }, + ], + }); + const reranker = new ZeroEntropyReranker({ apiKey: "key" }); + + const results = await reranker.rerank("q", ["x", "y", "z"]); + + expect(results).toEqual([ + { index: 2, rerankScore: 0.9 }, + { index: 0, rerankScore: 0.31 }, + ]); + }); + + it("sorts unsorted API results by descending relevance score client-side", async () => { + mockRerank.mockResolvedValue({ + results: [ + { index: 0, relevance_score: 0.2 }, + { index: 1, relevance_score: 0.9 }, + { index: 2, relevance_score: 0.5 }, + ], + }); + const reranker = new ZeroEntropyReranker({ apiKey: "key" }); + + const results = await reranker.rerank("q", ["a", "b", "c"]); + + expect(results.map((r) => r.index)).toEqual([1, 2, 0]); + expect(results.map((r) => r.rerankScore)).toEqual([0.9, 0.5, 0.2]); + }); + + it("slices to topK client-side after sorting", async () => { + mockRerank.mockResolvedValue({ + results: [ + { index: 0, relevance_score: 0.2 }, + { index: 1, relevance_score: 0.9 }, + { index: 2, relevance_score: 0.5 }, + ], + }); + const reranker = new ZeroEntropyReranker({ apiKey: "key" }); + + const results = await reranker.rerank("q", ["a", "b", "c"], 2); + + expect(results).toEqual([ + { index: 1, rerankScore: 0.9 }, + { index: 2, rerankScore: 0.5 }, + ]); + }); + + it("uses a custom model when provided", async () => { + mockRerank.mockResolvedValue({ results: [] }); + const reranker = new ZeroEntropyReranker({ + apiKey: "key", + model: "zerank-1-small", + }); + + await reranker.rerank("q", ["a"]); + + expect(mockRerank).toHaveBeenCalledWith( + expect.objectContaining({ model: "zerank-1-small" }), + ); + }); + + it("returns an empty array without calling ZeroEntropy when there are no documents", async () => { + const reranker = new ZeroEntropyReranker({ apiKey: "key" }); + + const results = await reranker.rerank("q", []); + + expect(results).toEqual([]); + expect(mockRerank).not.toHaveBeenCalled(); + }); + + it("falls back to the original order with rerankScore 0.0 when the API call fails", async () => { + mockRerank.mockRejectedValue(new Error("zero entropy is down")); + const warnSpy = jest.spyOn(console, "warn").mockImplementation(() => {}); + const reranker = new ZeroEntropyReranker({ apiKey: "key" }); + + const results = await reranker.rerank("q", ["a", "b", "c"]); + + expect(results).toEqual([ + { index: 0, rerankScore: 0.0 }, + { index: 1, rerankScore: 0.0 }, + { index: 2, rerankScore: 0.0 }, + ]); + expect(warnSpy).toHaveBeenCalled(); + + warnSpy.mockRestore(); + }); +}); diff --git a/mem0-ts/src/oss/src/rerankers/zeroentropy.ts b/mem0-ts/src/oss/src/rerankers/zeroentropy.ts new file mode 100644 index 000000000..86a4e7de3 --- /dev/null +++ b/mem0-ts/src/oss/src/rerankers/zeroentropy.ts @@ -0,0 +1,58 @@ +import { ZeroEntropy } from "zeroentropy"; +import { RerankerConfig } from "../types"; +import { Reranker, RerankResult } from "./base"; + +const DEFAULT_MODEL = "zerank-1"; + +export class ZeroEntropyReranker implements Reranker { + private client: ZeroEntropy; + private model: string; + private topK?: number; + + constructor(config: RerankerConfig) { + const apiKey = config.apiKey || process.env.ZERO_ENTROPY_API_KEY; + if (!apiKey) { + throw new Error( + "Zero Entropy API key is required. Set ZERO_ENTROPY_API_KEY environment variable or pass apiKey in config.", + ); + } + this.client = new ZeroEntropy({ apiKey }); + this.model = config.model || DEFAULT_MODEL; + this.topK = config.topK; + } + + async rerank( + query: string, + documents: string[], + topK?: number, + ): Promise { + if (documents.length === 0) return []; + + try { + const response = await this.client.models.rerank({ + model: this.model, + query, + documents, + }); + + const scored = response.results.map((result) => ({ + index: result.index, + rerankScore: result.relevance_score, + })); + scored.sort((a, b) => b.rerankScore - a.rerankScore); + + const finalTopK = topK || this.topK; + return finalTopK ? scored.slice(0, finalTopK) : scored; + } catch (e) { + console.warn( + `Zero Entropy reranking failed, falling back to original order: ${e}`, + ); + const scored = documents.map((_, index) => ({ + index, + rerankScore: 0.0, + })); + const finalTopK = topK || this.topK; + return finalTopK ? scored.slice(0, finalTopK) : scored; + } + } +} diff --git a/mem0-ts/src/oss/src/types/index.ts b/mem0-ts/src/oss/src/types/index.ts index 82f7bcc17..484aad462 100644 --- a/mem0-ts/src/oss/src/types/index.ts +++ b/mem0-ts/src/oss/src/types/index.ts @@ -60,6 +60,53 @@ export interface LLMConfig { maxTokens?: number; } +export interface RerankerConfig { + apiKey?: string; + /** The reranker model to use. Default varies by provider. */ + model?: string; + /** Maximum number of documents to return after reranking. Default: unset (return all). */ + topK?: number; + /** `cohere` only. Return document texts in the response. Default: `false`. */ + returnDocuments?: boolean; + /** `cohere` only. Maximum number of chunks per document. Default: unset. */ + maxChunksPerDoc?: number; + /** + * `sentence_transformer` / `huggingface` only. Transformers.js device, e.g. + * `"cpu"`, `"wasm"`, `"webgpu"`. Default: unset (auto-detect). + */ + device?: string; + /** `huggingface` only. Max token length per query-document pair. Default: `512`. */ + maxLength?: number; + /** + * `sentence_transformer` / `huggingface` only. Sigmoid-normalize raw logits + * to `[0, 1]`. Default: `true`; set `false` to surface raw logits. + */ + normalize?: boolean; + /** No-op: a search reranks a small candidate set in one forward pass. */ + batchSize?: number; + /** No-op in this runtime. */ + showProgressBar?: boolean; + /** + * `llm_reranker` only. LLM provider used to build the scoring LLM when + * `llm` is not set. Default: `"openai"`. + */ + provider?: string; + /** `llm_reranker` only. Temperature for LLM generation. Default: `0.0`. */ + temperature?: number; + /** `llm_reranker` only. Maximum tokens for the LLM response. Default: `100`. */ + maxTokens?: number; + /** + * `llm_reranker` only. Nested LLM configuration. When set, it overrides the + * top-level `provider`/`model`/`temperature`/`maxTokens`/`apiKey`, which + * then only act as defaults for fields missing from `llm.config`. + */ + llm?: { + provider: string; + config: LLMConfig; + }; + [key: string]: any; +} + export interface MemoryConfig { version?: string; embedder: { @@ -74,6 +121,10 @@ export interface MemoryConfig { provider: string; config: LLMConfig; }; + reranker?: { + provider: string; + config: RerankerConfig; + }; historyStore?: HistoryStoreConfig; disableHistory?: boolean; historyDbPath?: string; @@ -87,6 +138,8 @@ export interface MemoryItem { createdAt?: string; updatedAt?: string; score?: number; + /** Relevance score added by the reranker, alongside (not replacing) `score`. */ + rerankScore?: number; metadata?: Record; attributedTo?: string; } @@ -156,5 +209,11 @@ export const MemoryConfigSchema = z.object({ config: z.record(z.string(), z.any()), }) .optional(), + reranker: z + .object({ + provider: z.string(), + config: z.record(z.string(), z.any()), + }) + .optional(), disableHistory: z.boolean().optional(), }); diff --git a/mem0-ts/src/oss/src/utils/factory.test.ts b/mem0-ts/src/oss/src/utils/factory.test.ts new file mode 100644 index 000000000..0a358ca44 --- /dev/null +++ b/mem0-ts/src/oss/src/utils/factory.test.ts @@ -0,0 +1,86 @@ +jest.mock("zeroentropy", () => ({ ZeroEntropy: jest.fn() })); +jest.mock("@huggingface/transformers", () => ({ + AutoModelForSequenceClassification: { from_pretrained: jest.fn() }, + AutoTokenizer: { from_pretrained: jest.fn() }, +})); + +import { RerankerFactory } from "./factory"; +import { CohereReranker } from "../rerankers/cohere"; +import { LLMReranker } from "../rerankers/llm"; +import { ZeroEntropyReranker } from "../rerankers/zeroentropy"; +import { CrossEncoderReranker } from "../rerankers/cross_encoder"; + +describe("RerankerFactory", () => { + const originalOpenAiKey = process.env.OPENAI_API_KEY; + + afterEach(() => { + if (originalOpenAiKey === undefined) { + delete process.env.OPENAI_API_KEY; + } else { + process.env.OPENAI_API_KEY = originalOpenAiKey; + } + }); + + it("creates a CohereReranker for provider 'cohere'", () => { + const reranker = RerankerFactory.create("cohere", { apiKey: "key" }); + expect(reranker).toBeInstanceOf(CohereReranker); + }); + + it("matches the provider name case-insensitively", () => { + const reranker = RerankerFactory.create("Cohere", { apiKey: "key" }); + expect(reranker).toBeInstanceOf(CohereReranker); + }); + + it("creates a ZeroEntropyReranker for provider 'zero_entropy'", () => { + const reranker = RerankerFactory.create("zero_entropy", { + apiKey: "key", + }); + expect(reranker).toBeInstanceOf(ZeroEntropyReranker); + }); + + it("creates a CrossEncoderReranker for provider 'sentence_transformer'", () => { + const reranker = RerankerFactory.create("sentence_transformer", {}); + expect(reranker).toBeInstanceOf(CrossEncoderReranker); + }); + + it("creates a CrossEncoderReranker for provider 'huggingface'", () => { + const reranker = RerankerFactory.create("huggingface", {}); + expect(reranker).toBeInstanceOf(CrossEncoderReranker); + }); + + it("creates an LLMReranker for provider 'llm_reranker', building a default openai LLM from top-level config", () => { + const reranker = RerankerFactory.create("llm_reranker", { + apiKey: "key", + }); + expect(reranker).toBeInstanceOf(LLMReranker); + }); + + it("creates an LLMReranker that builds its own LLM from a nested config.llm", () => { + const reranker = RerankerFactory.create("llm_reranker", { + llm: { provider: "openai", config: { apiKey: "x" } }, + }); + expect(reranker).toBeInstanceOf(LLMReranker); + }); + + it("prefers the nested llm.provider over the top-level provider when building the llm_reranker's LLM", () => { + // If the top-level `provider` were used instead of the nested one, this + // would throw ("Unsupported LLM provider: not-a-real-provider"). + const reranker = RerankerFactory.create("llm_reranker", { + provider: "not-a-real-provider", + llm: { provider: "openai", config: { apiKey: "key" } }, + }); + expect(reranker).toBeInstanceOf(LLMReranker); + }); + + it("throws for the 'llm_reranker' provider when the default LLM has no API key available", () => { + delete process.env.OPENAI_API_KEY; + + expect(() => RerankerFactory.create("llm_reranker", {})).toThrow(); + }); + + it("throws for an unsupported provider", () => { + expect(() => RerankerFactory.create("banana", {})).toThrow( + /unsupported reranker provider/i, + ); + }); +}); diff --git a/mem0-ts/src/oss/src/utils/factory.ts b/mem0-ts/src/oss/src/utils/factory.ts index 50f487e23..9039ba668 100644 --- a/mem0-ts/src/oss/src/utils/factory.ts +++ b/mem0-ts/src/oss/src/utils/factory.ts @@ -12,8 +12,14 @@ import { EmbeddingConfig, HistoryStoreConfig, LLMConfig, + RerankerConfig, VectorStoreConfig, } from "../types"; +import { Reranker } from "../rerankers/base"; +import { CohereReranker } from "../rerankers/cohere"; +import { LLMReranker } from "../rerankers/llm"; +import { ZeroEntropyReranker } from "../rerankers/zeroentropy"; +import { CrossEncoderReranker } from "../rerankers/cross_encoder"; import { Embedder } from "../embeddings/base"; import { LLM } from "../llms/base"; import { VectorStore } from "../vector_stores/base"; @@ -186,6 +192,69 @@ export class VectorStoreFactory { } } +export class RerankerFactory { + static create(provider: string, config: RerankerConfig): Reranker { + switch (provider.toLowerCase()) { + case "cohere": + return new CohereReranker(config); + case "zero_entropy": + return new ZeroEntropyReranker(config); + case "sentence_transformer": + return new CrossEncoderReranker( + config, + "Xenova/ms-marco-MiniLM-L-6-v2", + ); + case "huggingface": + return new CrossEncoderReranker( + config, + "Xenova/bge-reranker-base", + 512, + ); + case "llm_reranker": { + const llm = RerankerFactory.buildLLMRerankerLLM(config); + return new LLMReranker(config, llm); + } + default: + throw new Error(`Unsupported reranker provider: ${provider}`); + } + } + + private static buildLLMRerankerLLM(config: RerankerConfig): LLM { + const nested = config.llm; + let llmProvider: string; + let llmConfig: LLMConfig; + + if (nested) { + llmProvider = nested.provider || config.provider || "openai"; + llmConfig = { ...(nested.config || {}) }; + if (llmConfig.model === undefined) { + llmConfig.model = config.model ?? "gpt-4o-mini"; + } + if (llmConfig.temperature === undefined) { + llmConfig.temperature = config.temperature ?? 0.0; + } + if (llmConfig.maxTokens === undefined) { + llmConfig.maxTokens = config.maxTokens ?? 100; + } + if (config.apiKey && llmConfig.apiKey === undefined) { + llmConfig.apiKey = config.apiKey; + } + } else { + llmProvider = config.provider || "openai"; + llmConfig = { + model: config.model ?? "gpt-4o-mini", + temperature: config.temperature ?? 0.0, + maxTokens: config.maxTokens ?? 100, + }; + if (config.apiKey) { + llmConfig.apiKey = config.apiKey; + } + } + + return LLMFactory.create(llmProvider, llmConfig); + } +} + export class HistoryManagerFactory { static create(provider: string, config: HistoryStoreConfig): HistoryManager { switch (provider.toLowerCase()) { diff --git a/mem0-ts/src/oss/tests/memory.rerank.test.ts b/mem0-ts/src/oss/tests/memory.rerank.test.ts new file mode 100644 index 000000000..2ad477580 --- /dev/null +++ b/mem0-ts/src/oss/tests/memory.rerank.test.ts @@ -0,0 +1,184 @@ +/** + * Reranker integration tests for Memory.search(). + * + * Verifies the per-search `rerank` flag: when a reranker is configured and + * `rerank: true` is passed, search reorders results by the reranker's output; + * otherwise results pass through unchanged. Failures degrade gracefully. + */ +/// +import { Memory } from "../src/memory"; +import { CohereReranker } from "../src/rerankers/cohere"; +import type { RerankResult } from "../src/rerankers/base"; + +jest.setTimeout(15000); + +jest.mock("../src/embeddings/google", () => ({ + GoogleEmbedder: jest.fn(), +})); +jest.mock("../src/llms/google", () => ({ + GoogleLLM: jest.fn(), +})); +jest.mock("../src/llms/openai", () => ({ + OpenAILLM: jest.fn().mockImplementation(() => ({ + generateResponse: jest + .fn() + .mockResolvedValue(JSON.stringify({ memory: [] })), + })), +})); + +const mockEmbedding = new Array(1536).fill(0.1); +jest.mock("../src/embeddings/openai", () => ({ + OpenAIEmbedder: jest.fn().mockImplementation(() => ({ + embed: jest.fn().mockResolvedValue(mockEmbedding), + embedBatch: jest + .fn() + .mockImplementation((texts: string[]) => + Promise.resolve(texts.map(() => mockEmbedding)), + ), + embeddingDims: 1536, + })), +})); + +function createMemory(config: Record = {}): Memory { + return new Memory({ + version: "v1.1", + embedder: { + provider: "openai", + config: { apiKey: "test-key", model: "text-embedding-3-small" }, + }, + vectorStore: { + provider: "memory", + config: { + collectionName: `test-rerank-${Date.now()}-${Math.random()}`, + dimension: 1536, + dbPath: ":memory:", + }, + }, + llm: { + provider: "openai", + config: { apiKey: "test-key", model: "gpt-5-mini" }, + }, + historyDbPath: ":memory:", + ...config, + }); +} + +// Two semantic results whose natural (score-sorted) order is [alpha, bravo]. +async function primeSearch(m: any) { + await m._ensureInitialized(); + m.embedder = { embed: jest.fn().mockResolvedValue(mockEmbedding) }; + m.vectorStore.search = jest.fn().mockResolvedValue([ + { id: "a", score: 0.9, payload: { data: "alpha" } }, + { id: "b", score: 0.8, payload: { data: "bravo" } }, + ]); + m.vectorStore.keywordSearch = jest.fn().mockResolvedValue(null); +} + +describe("Memory.search reranking", () => { + it("reorders results by the reranker when rerank:true, adding rerankScore while preserving the original vector score", async () => { + const memory = createMemory(); + const m = memory as any; + await primeSearch(m); + + const rerank = jest + .fn, [string, string[], number?]>() + .mockResolvedValue([ + { index: 1, rerankScore: 0.99 }, // bravo + { index: 0, rerankScore: 0.4 }, // alpha + ]); + m.reranker = { rerank }; + + const result = await m.search("what did i eat", { + filters: { user_id: "u1" }, + rerank: true, + }); + + expect(rerank).toHaveBeenCalledWith( + "what did i eat", + ["alpha", "bravo"], + expect.any(Number), + ); + expect(result.results.map((r: any) => r.memory)).toEqual([ + "bravo", + "alpha", + ]); + expect(result.results[0].rerankScore).toBe(0.99); + expect(result.results[1].rerankScore).toBe(0.4); + // The original vector similarity `score` must survive reranking. + expect(result.results[0].score).toBe(0.8); // bravo's original vector score + expect(result.results[1].score).toBe(0.9); // alpha's original vector score + + await memory.reset(); + }); + + it("leaves results untouched and does not call the reranker when rerank is omitted", async () => { + const memory = createMemory(); + const m = memory as any; + await primeSearch(m); + const rerank = jest.fn(); + m.reranker = { rerank }; + + const result = await m.search("what did i eat", { + filters: { user_id: "u1" }, + }); + + expect(rerank).not.toHaveBeenCalled(); + expect(result.results.map((r: any) => r.memory)).toEqual([ + "alpha", + "bravo", + ]); + expect(result.results[0].rerankScore).toBeUndefined(); + + await memory.reset(); + }); + + it("is a no-op (no throw) when rerank:true but no reranker is configured", async () => { + const memory = createMemory(); + const m = memory as any; + await primeSearch(m); + + const result = await m.search("what did i eat", { + filters: { user_id: "u1" }, + rerank: true, + }); + + expect(result.results.map((r: any) => r.memory)).toEqual([ + "alpha", + "bravo", + ]); + + await memory.reset(); + }); + + it("falls back to the original results when the reranker throws", async () => { + const memory = createMemory(); + const m = memory as any; + await primeSearch(m); + m.reranker = { + rerank: jest.fn().mockRejectedValue(new Error("provider down")), + }; + const warnSpy = jest.spyOn(console, "warn").mockImplementation(() => {}); + + const result = await m.search("what did i eat", { + filters: { user_id: "u1" }, + rerank: true, + }); + + expect(result.results.map((r: any) => r.memory)).toEqual([ + "alpha", + "bravo", + ]); + expect(warnSpy).toHaveBeenCalled(); + + warnSpy.mockRestore(); + await memory.reset(); + }); + + it("wires a reranker from config in the constructor", () => { + const memory = createMemory({ + reranker: { provider: "cohere", config: { apiKey: "test-key" } }, + }); + + expect((memory as any).reranker).toBeInstanceOf(CohereReranker); + }); +}); diff --git a/mem0-ts/tsup.config.ts b/mem0-ts/tsup.config.ts index d79558a09..863339632 100644 --- a/mem0-ts/tsup.config.ts +++ b/mem0-ts/tsup.config.ts @@ -6,6 +6,9 @@ const external = [ "@anthropic-ai/sdk", "@aws-sdk/client-s3vectors", "groq-sdk", + "cohere-ai", + "@huggingface/transformers", + "zeroentropy", "uuid", "pg", "zod", diff --git a/mem0/configs/rerankers/cohere.py b/mem0/configs/rerankers/cohere.py index 3a2f27fec..a7d408829 100644 --- a/mem0/configs/rerankers/cohere.py +++ b/mem0/configs/rerankers/cohere.py @@ -1,4 +1,5 @@ from typing import Optional + from pydantic import Field from mem0.configs.rerankers.base import BaseRerankerConfig @@ -10,6 +11,6 @@ class CohereRerankerConfig(BaseRerankerConfig): Inherits from BaseRerankerConfig and adds Cohere-specific settings. """ - model: Optional[str] = Field(default="rerank-english-v3.0", description="The Cohere rerank model to use") + model: Optional[str] = Field(default="rerank-v3.5", description="The Cohere rerank model to use") return_documents: bool = Field(default=False, description="Whether to return the document texts in the response") max_chunks_per_doc: Optional[int] = Field(default=None, description="Maximum number of chunks per document")