feat(oss): add reranker + per-search rerank to the TypeScript OSS SDK (#6055)
This commit is contained in:
Kartik
@@ -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 },
|
||||
},
|
||||
});
|
||||
```
|
||||
|
||||
<Note>
|
||||
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.
|
||||
</Note>
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
});
|
||||
```
|
||||
|
||||
<Note>
|
||||
`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.
|
||||
</Note>
|
||||
|
||||
## Popular Models
|
||||
|
||||
### BGE Rerankers (Recommended)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
});
|
||||
```
|
||||
|
||||
<Note>
|
||||
`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.
|
||||
</Note>
|
||||
|
||||
## GPU Acceleration
|
||||
|
||||
For better performance, use GPU acceleration:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"}
|
||||
]
|
||||
|
||||
@@ -19,6 +19,10 @@ Reranking trades extra latency for better precision. Start once you have baselin
|
||||
<Card title="Zero Entropy" icon="/images/provider-icons/zeroentropy.svg" href="/components/rerankers/models/zero_entropy" />
|
||||
</CardGroup>
|
||||
|
||||
<Note>
|
||||
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.
|
||||
</Note>
|
||||
|
||||
## Reranking Workflow
|
||||
|
||||
<CardGroup cols={3}>
|
||||
|
||||
+2
-1
@@ -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.
|
||||
|
||||
@@ -59,7 +59,7 @@ config = {
|
||||
},
|
||||
"reranker": {
|
||||
"provider": "cohere",
|
||||
"config": {"model": "rerank-english-v3.0"},
|
||||
"config": {"model": "rerank-v3.5"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -18,7 +18,129 @@ Reranker-enhanced search adds a second scoring pass after vector retrieval so Me
|
||||
</Warning>
|
||||
|
||||
<Note>
|
||||
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).
|
||||
</Note>
|
||||
|
||||
---
|
||||
|
||||
## 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,
|
||||
});
|
||||
```
|
||||
|
||||
<Note>
|
||||
`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.
|
||||
</Note>
|
||||
|
||||
### 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` |
|
||||
|
||||
<Note>
|
||||
`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.
|
||||
</Note>
|
||||
|
||||
---
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
},
|
||||
|
||||
Generated
+431
@@ -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
|
||||
|
||||
@@ -177,6 +177,7 @@ export class ConfigManager {
|
||||
})(),
|
||||
disableHistory:
|
||||
userConfig.disableHistory || DEFAULT_MEMORY_CONFIG.disableHistory,
|
||||
reranker: userConfig.reranker,
|
||||
};
|
||||
|
||||
// Validate the merged config
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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<RerankResult[]>;
|
||||
}
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
@@ -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<RerankResult[]> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
@@ -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<RerankResult[]> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<string, string>): 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();
|
||||
});
|
||||
});
|
||||
@@ -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<RerankResult[]> {
|
||||
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<number> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
@@ -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<RerankResult[]> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<string, any>;
|
||||
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(),
|
||||
});
|
||||
|
||||
@@ -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,
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -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()) {
|
||||
|
||||
@@ -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.
|
||||
*/
|
||||
/// <reference types="jest" />
|
||||
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<string, any> = {}): 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<Promise<RerankResult[]>, [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);
|
||||
});
|
||||
});
|
||||
@@ -6,6 +6,9 @@ const external = [
|
||||
"@anthropic-ai/sdk",
|
||||
"@aws-sdk/client-s3vectors",
|
||||
"groq-sdk",
|
||||
"cohere-ai",
|
||||
"@huggingface/transformers",
|
||||
"zeroentropy",
|
||||
"uuid",
|
||||
"pg",
|
||||
"zod",
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user