Feat/valkey vector store (#5826)

Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
Yash
2026-07-06 20:20:59 +05:30
committed by GitHub
parent 6dc4606dcf
commit a7ecf781cd
11 changed files with 1020 additions and 2 deletions
+46 -1
View File
@@ -14,7 +14,8 @@ pip install mem0ai[vector-stores]
## Usage
```python
<CodeGroup>
```python Python
config = {
"vector_store": {
"provider": "valkey",
@@ -37,8 +38,36 @@ messages = [
m.add(messages, user_id="alice", metadata={"category": "movies"})
```
```typescript TypeScript
import { Memory } from 'mem0ai/oss';
const config = {
vectorStore: {
provider: 'valkey',
config: {
collectionName: 'test',
valkeyUrl: 'valkey://localhost:6379',
embeddingModelDims: 1536,
indexType: 'flat',
},
},
};
const memory = new Memory(config);
const messages = [
{ role: 'user', content: "I'm planning to watch a movie tonight. Any recommendations?" },
{ role: 'assistant', content: 'How about thriller movies? They can be quite engaging.' },
{ role: 'user', content: "I'm not a big fan of thriller movies but I love sci-fi movies." },
{ role: 'assistant', content: "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future." },
];
await memory.add(messages, { userId: 'alice', metadata: { category: 'movies' } });
```
</CodeGroup>
## Parameters
<Tabs>
<Tab title="Python">
Here are the parameters available for configuring Valkey:
| Parameter | Description | Default Value |
@@ -52,6 +81,22 @@ Here are the parameters available for configuring Valkey:
| `hnsw_ef_runtime` | Size of dynamic candidate list for search | `10` |
| `cluster_mode` | Enable cluster mode for Valkey cluster (CME) deployments | `false` |
| `timezone` | Timezone for timestamp handling | `UTC` |
</Tab>
<Tab title="TypeScript">
| Parameter | Description | Default Value |
| --- | --- | --- |
| `collectionName` | The name of the collection to store the vectors | `mem0` |
| `valkeyUrl` | Connection URL for the Valkey server | `valkey://localhost:6379` |
| `embeddingModelDims` | Dimensions of the embedding model | `1536` |
| `indexType` | Vector index algorithm (`hnsw` or `flat`) | `hnsw` |
| `hnswM` | Number of bi-directional links for HNSW | `16` |
| `hnswEfConstruction` | Size of dynamic candidate list for HNSW | `200` |
| `hnswEfRuntime` | Size of dynamic candidate list for search | `10` |
| `clusterMode` | Enable cluster mode for Valkey cluster (CME) deployments | `false` |
| `timezone` | Timezone for timestamp handling | `UTC` |
</Tab>
</Tabs>
## Cluster Mode
+3 -1
View File
@@ -97,7 +97,8 @@
"ts-jest": "^29.2.6",
"ts-node": "^10.9.2",
"tsup": "^8.3.0",
"typescript": "5.5.4"
"typescript": "5.5.4",
"iovalkey": "^0.3.3"
},
"dependencies": {
"axios": "^1.16.0",
@@ -124,6 +125,7 @@
"ollama": "^0.5.14",
"pg": "8.11.3",
"redis": "^4.6.13",
"iovalkey": "^0.3.3",
"compromise": "^14.0.0",
"natural": "^8.0.1"
},
+61
View File
@@ -117,6 +117,9 @@ importers:
fix-tsup-cjs:
specifier: ^1.2.0
version: 1.2.0
iovalkey:
specifier: ^0.3.3
version: 0.3.3
jest:
specifier: ^29.7.0
version: 29.7.0(@types/node@22.19.21)(ts-node@10.9.2(@types/node@22.19.21)(typescript@5.5.4))
@@ -671,6 +674,9 @@ packages:
'@modelcontextprotocol/sdk':
optional: true
'@iovalkey/commands@0.1.0':
resolution: {integrity: sha512-/B9W4qKSSITDii5nkBCHyPkIkAi+ealUtr1oqBJsLxjSRLka4pxun2VvMNSmcwgAMxgXtQfl0qRv7TE+udPJzg==}
'@isaacs/cliui@8.0.2':
resolution: {integrity: sha512-O8jcjabXaleOG9DQ0+ARXWZBTfnP4WNAqzuiJK7ll44AmxGKv/J2M4TPjxjY3znBCfvBXFzucm1twdyFybFqEA==}
engines: {node: '>=12'}
@@ -1643,6 +1649,10 @@ packages:
resolution: {integrity: sha512-ZySD7Nf91aLB0RxL4KGrKHBXl7Eds1DAmEdcoVawXnLD7SDhpNgtuII2aAkg7a7QS41jxPSZ17p4VdGnMHk3MQ==}
engines: {node: '>=0.4.0'}
denque@2.1.0:
resolution: {integrity: sha512-HVQE3AAb/pxF8fQAoiqpvg9i3evqug3hoiwakOyZAwJm+6vZehbkYXZ0l4JxS+I3QxM97v5aaRNhj8v5oBhekw==}
engines: {node: '>=0.10'}
detect-libc@2.1.2:
resolution: {integrity: sha512-Btj2BOOO83o3WyH59e8MgXsxEQVcarkUOpEYrubB0urwnN10yQ364rsiByU11nZlqWYZm05i/of7io4mzihBtQ==}
engines: {node: '>=8'}
@@ -2006,6 +2016,10 @@ packages:
ini@1.3.8:
resolution: {integrity: sha512-JV/yugV2uzW5iMRSiZAyDtQd+nxtUnjeLt0acNdw98kKLrvuRVyB80tsREOE7yvGVgalhZ6RNXCmEHkUKBKxew==}
iovalkey@0.3.3:
resolution: {integrity: sha512-4rTJX6Q5wTYEvxboXi8DsEiUo+OvqJGtLYOSGm37KpdRXsG5XJjbVtYKGJpPSWP+QT7rWscA4vsrdmzbEbenpw==}
engines: {node: '>=18.12.0'}
is-arrayish@0.2.1:
resolution: {integrity: sha512-zz06S8t0ozoDXMG+ube26zeCTNXcKIPJZJi8hBrF4idCLms4CG9QtK7qBl1boi5ODzFpjswb5JPmHCbMpjaYzg==}
@@ -2315,9 +2329,15 @@ packages:
resolution: {integrity: sha512-t7hw9pI+WvuwNJXwk5zVHpyhIqzg2qTlklJOf0mVxGSbe3Fp2VieZcduNYjaLDoy6p9uGpQEGWG87WpMKlNq8g==}
engines: {node: '>=8'}
lodash.defaults@4.2.0:
resolution: {integrity: sha512-qjxPLHd3r5DnsdGacqOMU6pb/avJzdh9tFX2ymgoZE27BmjXrNy/y4LoaiTeAb+O3gL8AfpJGtqfX/ae2leYYQ==}
lodash.includes@4.3.0:
resolution: {integrity: sha512-W3Bx6mdkRTGtlJISOvVD/lbqjTlPPUDTMnlXZFnVwi9NKJ6tiAk6LVdlhZMm17VZisqhKcgzpO5Wz91PCt5b0w==}
lodash.isarguments@3.1.0:
resolution: {integrity: sha512-chi4NHZlZqZD18a0imDHnZPrDeBbTtVN7GXMwuGdRH9qotxAjYs3aVLKc7zNOG9eddR5Ksd8rvFEBc9SsggPpg==}
lodash.isboolean@3.0.3:
resolution: {integrity: sha512-Bz5mupy2SVbPHURB98VAcw+aHh4vRV5IPNhILUCsOzRmsTmSQ17jIuqopAentWoehktxGd9e/hbIXq980/1QJg==}
@@ -2831,6 +2851,14 @@ packages:
resolution: {integrity: sha512-GDhwkLfywWL2s6vEjyhri+eXmfH6j1L7JE27WhqLeYzoh/A3DBaYGEj2H/HFZCn/kMfim73FXxEJTw06WtxQwg==}
engines: {node: '>= 14.18.0'}
redis-errors@1.2.0:
resolution: {integrity: sha512-1qny3OExCf0UvUV/5wpYKf2YwPcOqXzkwKKSmKHiE6ZMQs5heeE/c8eXK+PNllPvmjgAbfnsbpkGZWy8cBpn9w==}
engines: {node: '>=4'}
redis-parser@3.0.0:
resolution: {integrity: sha512-DJnGAeenTdpMEH6uAJRK/uiyEIH9WVsUmoLwzudwGJUwZPp80PDBWPHXSAGNPwNvIXAbe7MSUB1zQFugFml66A==}
engines: {node: '>=4'}
redis@4.7.1:
resolution: {integrity: sha512-S1bJDnqLftzHXHP8JsT5II/CtHWQrASX5K96REjWjlmWKrviSOLWmM7QnRLstAWsu1VBBV1ffV6DzCvxNP0UJQ==}
@@ -2968,6 +2996,9 @@ packages:
resolution: {integrity: sha512-XlkWvfIm6RmsWtNJx+uqtKLS8eqFbxUg0ZzLXqY0caEy9l7hruX8IpiDnjsLavoBgqCCR71TqWO8MaXYheJ3RQ==}
engines: {node: '>=10'}
standard-as-callback@2.1.0:
resolution: {integrity: sha512-qoRRSyROncaz1z0mvYqIE4lCd9p2R90i6GxW3uZv5ucSu8tU7B5HXUP1gG8pVZsYNVaXjk8ClXHPttLyxAL48A==}
stopwords-iso@1.1.0:
resolution: {integrity: sha512-I6GPS/E0zyieHehMRPQcqkiBMJKGgLta+1hREixhoLPqEA0AlVFiC43dl8uPpmkkeRdDMzYRWFWk5/l9x7nmNg==}
engines: {node: '>=0.10.0'}
@@ -4124,6 +4155,8 @@ snapshots:
- supports-color
- utf-8-validate
'@iovalkey/commands@0.1.0': {}
'@isaacs/cliui@8.0.2':
dependencies:
string-width: 5.1.2
@@ -5200,6 +5233,8 @@ snapshots:
delayed-stream@1.0.0: {}
denque@2.1.0: {}
detect-libc@2.1.2: {}
detect-newline@3.1.0: {}
@@ -5596,6 +5631,20 @@ snapshots:
ini@1.3.8: {}
iovalkey@0.3.3:
dependencies:
'@iovalkey/commands': 0.1.0
cluster-key-slot: 1.1.2
debug: 4.4.3(supports-color@5.5.0)
denque: 2.1.0
lodash.defaults: 4.2.0
lodash.isarguments: 3.1.0
redis-errors: 1.2.0
redis-parser: 3.0.0
standard-as-callback: 2.1.0
transitivePeerDependencies:
- supports-color
is-arrayish@0.2.1: {}
is-binary-path@2.1.0:
@@ -6065,8 +6114,12 @@ snapshots:
dependencies:
p-locate: 4.1.0
lodash.defaults@4.2.0: {}
lodash.includes@4.3.0: {}
lodash.isarguments@3.1.0: {}
lodash.isboolean@3.0.3: {}
lodash.isinteger@4.0.4: {}
@@ -6575,6 +6628,12 @@ snapshots:
readdirp@4.1.2: {}
redis-errors@1.2.0: {}
redis-parser@3.0.0:
dependencies:
redis-errors: 1.2.0
redis@4.7.1:
dependencies:
'@redis/bloom': 1.2.0(@redis/client@1.6.1)
@@ -6726,6 +6785,8 @@ snapshots:
dependencies:
escape-string-regexp: 2.0.0
standard-as-callback@2.1.0: {}
stopwords-iso@1.1.0: {}
string-length@4.0.2:
+1
View File
@@ -23,6 +23,7 @@ export * from "./vector_stores/base";
export * from "./vector_stores/memory";
export * from "./vector_stores/qdrant";
export * from "./vector_stores/redis";
export * from "./vector_stores/valkey";
export * from "./vector_stores/supabase";
export * from "./vector_stores/langchain";
export * from "./vector_stores/vectorize";
+2
View File
@@ -21,6 +21,8 @@ export interface EmbeddingConfig {
modelProperties?: Record<string, any>;
}
export type { ValkeyConfig } from "./valkey";
export interface VectorStoreConfig {
collectionName?: string;
dimension?: number;
+13
View File
@@ -0,0 +1,13 @@
import { VectorStoreConfig } from "./index";
export interface ValkeyConfig extends VectorStoreConfig {
valkeyUrl: string;
collectionName: string;
embeddingModelDims: number;
timezone?: string;
indexType?: "hnsw" | "flat";
hnswM?: number;
hnswEfConstruction?: number;
hnswEfRuntime?: number;
clusterMode?: boolean;
}
+3
View File
@@ -19,6 +19,7 @@ import { VectorStore } from "../vector_stores/base";
import { Qdrant } from "../vector_stores/qdrant";
import { VectorizeDB } from "../vector_stores/vectorize";
import { RedisDB } from "../vector_stores/redis";
import { ValkeyDB } from "../vector_stores/valkey";
import { OllamaLLM } from "../llms/ollama";
import { LMStudioLLM } from "../llms/lmstudio";
import { DeepSeekLLM } from "../llms/deepseek";
@@ -107,6 +108,8 @@ export class VectorStoreFactory {
return new Qdrant(config as any);
case "redis":
return new RedisDB(config as any);
case "valkey":
return new ValkeyDB(config as any);
case "supabase":
return new SupabaseDB(config as any);
case "langchain":
+608
View File
@@ -0,0 +1,608 @@
import { VectorStore } from "./base";
import { SearchFilters, VectorStoreResult } from "../types";
import { ValkeyConfig } from "../types/valkey";
interface ValkeyClient {
call: (...args: (string | number | Buffer)[]) => Promise<unknown>;
hset: (key: string, data: Record<string, string | Buffer>) => Promise<number>;
hgetall: (key: string) => Promise<Record<string, string>>;
exists: (key: string) => Promise<number>;
del: (key: string) => Promise<number>;
get: (key: string) => Promise<string | null>;
set: (key: string, value: string) => Promise<"OK">;
quit: () => Promise<"OK">;
on: (event: string, listener: (...args: any[]) => void) => void;
}
/**
* Escape Valkey Search TAG filter special characters.
*/
function escapeTagValue(value: unknown): string {
return String(value).replace(
/([,.<>{}\[\]"':;!@#$%^&*()\-+=~|/\\\s])/g,
"\\$1",
);
}
const EXCLUDED_KEYS = new Set([
"user_id",
"agent_id",
"run_id",
"hash",
"data",
"created_at",
"updated_at",
]);
function toSnakeCase(obj: Record<string, any>): Record<string, any> {
if (typeof obj !== "object" || obj === null) return obj;
return Object.fromEntries(
Object.entries(obj).map(([key, value]) => [
key.replace(/[A-Z]/g, (letter) => `_${letter.toLowerCase()}`),
value,
]),
);
}
function toCamelCase(obj: Record<string, any>): Record<string, any> {
if (typeof obj !== "object" || obj === null) return obj;
return Object.fromEntries(
Object.entries(obj).map(([key, value]) => [
key.replace(/_([a-z])/g, (_, letter) => letter.toUpperCase()),
value,
]),
);
}
interface ValkeySearchDoc {
memory_id?: string;
hash?: string;
memory?: string;
created_at?: string;
updated_at?: string;
agent_id?: string;
run_id?: string;
user_id?: string;
metadata?: string;
vector_score?: string;
}
function parseFtSearchResults(result: unknown[]): {
total: number;
docs: ValkeySearchDoc[];
} {
const total = Number(result[0]) || 0;
const docs: ValkeySearchDoc[] = [];
for (let i = 1; i < result.length; i += 2) {
const fields = result[i + 1] as string[];
const doc: ValkeySearchDoc = {};
for (let j = 0; j < fields.length; j += 2) {
const key = fields[j] as keyof ValkeySearchDoc;
(doc as Record<string, string>)[key] = fields[j + 1] as string;
}
docs.push(doc);
}
return { total, docs };
}
function parseValkeyUrl(url: string): {
host: string;
port: number;
username?: string;
password?: string;
} {
const normalized = url.replace(/^valkey:\/\//, "redis://");
const parsed = new URL(normalized);
return {
host: parsed.hostname,
port: parsed.port ? parseInt(parsed.port, 10) : 6379,
username: parsed.username ? decodeURIComponent(parsed.username) : undefined,
password: parsed.password ? decodeURIComponent(parsed.password) : undefined,
};
}
function formatTimestamp(timestamp: number, timezone: string = "UTC"): string {
const date = new Date(timestamp * 1000);
if (timezone === "UTC") {
return date.toISOString();
}
// Mirror Python's datetime.fromtimestamp(ts, tz).isoformat(): render the
// instant as an ISO-8601 string carrying the target IANA timezone's offset.
const parts = new Intl.DateTimeFormat("en-US", {
timeZone: timezone,
year: "numeric",
month: "2-digit",
day: "2-digit",
hour: "2-digit",
minute: "2-digit",
second: "2-digit",
hourCycle: "h23",
}).formatToParts(date);
const get = (type: string) =>
parts.find((part) => part.type === type)?.value ?? "00";
const yyyy = get("year");
const MM = get("month");
const dd = get("day");
const HH = get("hour");
const mm = get("minute");
const ss = get("second");
const asIfUtc = Date.UTC(
Number(yyyy),
Number(MM) - 1,
Number(dd),
Number(HH),
Number(mm),
Number(ss),
);
const offsetMinutes = Math.round((asIfUtc - date.getTime()) / 60000);
const sign = offsetMinutes >= 0 ? "+" : "-";
const absOffset = Math.abs(offsetMinutes);
const offHH = String(Math.floor(absOffset / 60)).padStart(2, "0");
const offMM = String(absOffset % 60).padStart(2, "0");
return `${yyyy}-${MM}-${dd}T${HH}:${mm}:${ss}${sign}${offHH}:${offMM}`;
}
async function loadIovalkey(): Promise<typeof import("iovalkey")> {
try {
return await import("iovalkey");
} catch {
throw new Error(
"iovalkey is required for the Valkey vector store. Install it with: npm install iovalkey",
);
}
}
export class ValkeyDB implements VectorStore {
private client!: ValkeyClient;
private readonly collectionName: string;
private readonly indexPrefix: string;
private readonly embeddingModelDims: number;
private readonly timezone: string;
private readonly indexType: "hnsw" | "flat";
private readonly hnswM: number;
private readonly hnswEfConstruction: number;
private readonly hnswEfRuntime: number;
private readonly clusterMode: boolean;
private readonly valkeyUrl: string;
private _initPromise?: Promise<void>;
constructor(config: ValkeyConfig) {
this.collectionName = config.collectionName;
this.indexPrefix = `mem0:${config.collectionName}`;
this.embeddingModelDims = config.embeddingModelDims;
this.timezone = config.timezone ?? "UTC";
this.indexType = (config.indexType ?? "hnsw").toLowerCase() as
| "hnsw"
| "flat";
this.hnswM = config.hnswM ?? 16;
this.hnswEfConstruction = config.hnswEfConstruction ?? 200;
this.hnswEfRuntime = config.hnswEfRuntime ?? 10;
this.clusterMode = config.clusterMode ?? false;
this.valkeyUrl = config.valkeyUrl;
if (this.indexType !== "hnsw" && this.indexType !== "flat") {
throw new Error(
`Invalid indexType: ${config.indexType}. Must be 'hnsw' or 'flat'`,
);
}
this.initialize().catch((err) => {
console.error("Failed to initialize Valkey:", err);
});
}
private buildIndexCreateCommand(
collectionName: string,
embeddingDims: number,
distanceMetric: string,
prefix: string,
): (string | number)[] {
const vectorConfig =
this.indexType === "hnsw"
? [
"embedding",
"VECTOR",
"HNSW",
"12",
"TYPE",
"FLOAT32",
"DIM",
String(embeddingDims),
"DISTANCE_METRIC",
distanceMetric,
"M",
String(this.hnswM),
"EF_CONSTRUCTION",
String(this.hnswEfConstruction),
"EF_RUNTIME",
String(this.hnswEfRuntime),
]
: [
"embedding",
"VECTOR",
"FLAT",
"6",
"TYPE",
"FLOAT32",
"DIM",
String(embeddingDims),
"DISTANCE_METRIC",
distanceMetric,
];
return [
"FT.CREATE",
collectionName,
"ON",
"HASH",
"PREFIX",
"1",
prefix,
"SCHEMA",
"memory_id",
"TAG",
"hash",
"TAG",
"agent_id",
"TAG",
"run_id",
"TAG",
"user_id",
"TAG",
"memory",
"TEXT",
"metadata",
"TAG",
"created_at",
"NUMERIC",
"updated_at",
"NUMERIC",
...vectorConfig,
];
}
private async ensureSearchModule(): Promise<void> {
try {
await this.client.call("FT._LIST");
} catch (error: any) {
const message = String(error?.message ?? error).toLowerCase();
if (message.includes("unknown command")) {
throw new Error(
"Valkey search module is not available. Please ensure Valkey is running with the search module enabled.",
);
}
throw error;
}
}
private async createIndex(): Promise<void> {
await this.ensureSearchModule();
try {
await this.client.call("FT.INFO", this.collectionName);
return;
} catch (error: any) {
const message = String(error?.message ?? error).toLowerCase();
if (
!message.includes("not found") &&
!message.includes("unknown index")
) {
throw error;
}
}
const cmd = this.buildIndexCreateCommand(
this.collectionName,
this.embeddingModelDims,
"COSINE",
this.indexPrefix,
);
await this.client.call(...cmd);
}
private async connectClient(): Promise<void> {
const iovalkey = await loadIovalkey();
const Valkey = iovalkey.default;
if (this.clusterMode) {
const { Cluster } = iovalkey;
const { host, port, username, password } = parseValkeyUrl(this.valkeyUrl);
// Standalone `new Valkey(url)` keeps credentials embedded in the URL, but
// Cluster takes discrete nodes, so pass the parsed auth through redisOptions
// (mirrors Python's ValkeyCluster.from_url, which preserves URL creds).
const redisOptions: { username?: string; password?: string } = {};
if (username) redisOptions.username = username;
if (password) redisOptions.password = password;
this.client = new Cluster([{ host, port }], {
redisOptions,
}) as unknown as ValkeyClient;
} else {
this.client = new Valkey(this.valkeyUrl) as unknown as ValkeyClient;
}
this.client.on("error", (err) =>
console.error("Valkey Client Error:", err),
);
this.client.on("connect", () => console.log("Valkey Client Connected"));
}
async initialize(): Promise<void> {
if (!this._initPromise) {
this._initPromise = this._doInitialize();
}
return this._initPromise;
}
private async _doInitialize(): Promise<void> {
await this.connectClient();
await this.createIndex();
}
private buildSearchQuery(knnPart: string, filters?: SearchFilters): string {
const snakeFilters = filters ? toSnakeCase(filters) : undefined;
if (
!snakeFilters ||
!Object.entries(snakeFilters).some(
([, value]) => value !== null && value !== undefined,
)
) {
return `*=>${knnPart}`;
}
const filterParts = Object.entries(snakeFilters)
.filter(([, value]) => value !== null && value !== undefined)
.map(([key, value]) => `@${key}:{${escapeTagValue(value)}}`);
if (!filterParts.length) {
return `*=>${knnPart}`;
}
return `${filterParts.join(" ")} =>${knnPart}`;
}
private docToResult(doc: ValkeySearchDoc): VectorStoreResult {
const rawDistance = doc.vector_score ? Number(doc.vector_score) : undefined;
const score =
rawDistance !== undefined ? Math.max(0, 1 - rawDistance) : undefined;
const resultPayload: Record<string, any> = {
hash: doc.hash ?? "",
data: doc.memory ?? "",
created_at: doc.created_at
? formatTimestamp(Number(doc.created_at), this.timezone)
: undefined,
};
if (doc.updated_at) {
resultPayload.updated_at = formatTimestamp(
Number(doc.updated_at),
this.timezone,
);
}
if (doc.agent_id) resultPayload.agent_id = doc.agent_id;
if (doc.run_id) resultPayload.run_id = doc.run_id;
if (doc.user_id) resultPayload.user_id = doc.user_id;
if (doc.metadata) {
try {
Object.assign(resultPayload, JSON.parse(doc.metadata));
} catch {
console.warn("Failed to parse Valkey metadata:", doc.metadata);
}
}
return {
id: doc.memory_id ?? "",
payload: toCamelCase(resultPayload),
score,
};
}
async insert(
vectors: number[][],
ids: string[],
payloads: Record<string, any>[],
): Promise<void> {
await Promise.all(
vectors.map(async (vector, idx) => {
const payload = toSnakeCase(payloads[idx] ?? {});
const id = ids[idx];
const key = `${this.indexPrefix}:${id}`;
if (!payload.created_at) {
payload.created_at = new Date().toISOString();
}
const hashData: Record<string, string | Buffer> = {
memory_id: id,
hash: payload.hash ?? `hash_${id}`,
memory: payload.data ?? `data_${id}`,
created_at: String(
Math.floor(new Date(payload.created_at).getTime() / 1000),
),
embedding: Buffer.from(new Float32Array(vector).buffer),
metadata: JSON.stringify(
Object.fromEntries(
Object.entries(payload).filter(([k]) => !EXCLUDED_KEYS.has(k)),
),
),
};
for (const field of ["agent_id", "run_id", "user_id"]) {
if (field in payload) {
hashData[field] = String(payload[field]);
}
}
await this.client.hset(key, hashData);
}),
);
}
async keywordSearch(): Promise<null> {
return null;
}
async search(
query: number[],
topK: number = 5,
filters?: SearchFilters,
): Promise<VectorStoreResult[]> {
const vectorBytes = Buffer.from(new Float32Array(query).buffer);
const knnPart =
this.indexType === "hnsw"
? `[KNN ${topK} @embedding $vec_param EF_RUNTIME ${this.hnswEfRuntime} AS vector_score]`
: `[KNN ${topK} @embedding $vec_param AS vector_score]`;
const searchQuery = this.buildSearchQuery(knnPart, filters);
const result = (await this.client.call(
"FT.SEARCH",
this.collectionName,
searchQuery,
"PARAMS",
"2",
"vec_param",
vectorBytes,
"RETURN",
"10",
"memory_id",
"hash",
"agent_id",
"run_id",
"user_id",
"memory",
"metadata",
"created_at",
"updated_at",
"vector_score",
"DIALECT",
"2",
"LIMIT",
"0",
String(topK),
)) as unknown[];
const { docs } = parseFtSearchResults(result);
return docs.map((doc) => this.docToResult(doc));
}
async get(vectorId: string): Promise<VectorStoreResult | null> {
const key = `${this.indexPrefix}:${vectorId}`;
const exists = await this.client.exists(key);
if (!exists) {
return null;
}
const result = await this.client.hgetall(key);
if (!Object.keys(result).length) {
return null;
}
const doc: ValkeySearchDoc = {
memory_id: result.memory_id,
hash: result.hash,
memory: result.memory,
created_at: result.created_at,
updated_at: result.updated_at,
agent_id: result.agent_id,
run_id: result.run_id,
user_id: result.user_id,
metadata: result.metadata,
};
return this.docToResult(doc);
}
async update(
vectorId: string,
vector: number[],
payload: Record<string, any>,
): Promise<void> {
const snakePayload = toSnakeCase(payload);
const key = `${this.indexPrefix}:${vectorId}`;
if (!snakePayload.created_at) {
snakePayload.created_at = new Date().toISOString();
}
const hashData: Record<string, string | Buffer> = {
memory_id: vectorId,
hash: snakePayload.hash ?? `hash_${vectorId}`,
memory: snakePayload.data ?? `data_${vectorId}`,
created_at: String(
Math.floor(new Date(snakePayload.created_at).getTime() / 1000),
),
embedding: Buffer.from(new Float32Array(vector).buffer),
metadata: JSON.stringify(
Object.fromEntries(
Object.entries(snakePayload).filter(([k]) => !EXCLUDED_KEYS.has(k)),
),
),
};
if (snakePayload.updated_at) {
hashData.updated_at = String(
Math.floor(new Date(snakePayload.updated_at).getTime() / 1000),
);
}
for (const field of ["agent_id", "run_id", "user_id"]) {
if (field in snakePayload) {
hashData[field] = String(snakePayload[field]);
}
}
await this.client.hset(key, hashData);
}
async delete(vectorId: string): Promise<void> {
const key = `${this.indexPrefix}:${vectorId}`;
const exists = await this.client.exists(key);
if (!exists) {
console.warn(`Memory with ID ${vectorId} does not exist`);
return;
}
await this.client.del(key);
}
async deleteCol(): Promise<void> {
try {
await this.client.call("FT.DROPINDEX", this.collectionName);
} catch (error: any) {
const message = String(error?.message ?? error);
if (!message.includes("Unknown index name")) {
throw error;
}
}
}
async list(
filters?: SearchFilters,
topK: number = 100,
): Promise<[VectorStoreResult[], number]> {
const dummyVector = new Array(this.embeddingModelDims).fill(0);
const results = await this.search(dummyVector, topK, filters);
return [results, results.length];
}
async close(): Promise<void> {
await this.client.quit();
}
async getUserId(): Promise<string> {
const userId = await this.client.get("memory_migrations:1");
if (userId) {
return userId;
}
const randomUserId =
Math.random().toString(36).substring(2, 15) +
Math.random().toString(36).substring(2, 15);
await this.client.set("memory_migrations:1", randomUserId);
return randomUserId;
}
async setUserId(userId: string): Promise<void> {
await this.client.set("memory_migrations:1", userId);
}
}
@@ -113,6 +113,11 @@ jest.mock("../src/vector_stores/redis", () => ({
.fn()
.mockImplementation((config) => ({ type: "redis", config })),
}));
jest.mock("../src/vector_stores/valkey", () => ({
ValkeyDB: jest
.fn()
.mockImplementation((config) => ({ type: "valkey", config })),
}));
jest.mock("../src/vector_stores/supabase", () => ({
SupabaseDB: jest
.fn()
@@ -260,6 +265,7 @@ describe("VectorStoreFactory", () => {
test.each([
["qdrant"],
["redis"],
["valkey"],
["supabase"],
["langchain"],
["vectorize"],
+276
View File
@@ -0,0 +1,276 @@
/**
* Valkey vector store unit tests with mocked iovalkey client.
*/
/// <reference types="jest" />
describe("Valkey – mocked iovalkey client", () => {
let ValkeyDB: any;
beforeEach(() => {
jest.resetModules();
jest.doMock("iovalkey", () => {
const store = new Map<string, Record<string, string>>();
const mockClient = {
on: jest.fn(),
call: jest.fn().mockImplementation(async (...args: any[]) => {
const command = args[0];
if (command === "FT._LIST") {
return [];
}
if (command === "FT.INFO") {
throw new Error("Unknown index name");
}
if (command === "FT.CREATE") {
return "OK";
}
if (command === "FT.SEARCH") {
return [0];
}
if (command === "FT.DROPINDEX") {
return "OK";
}
return "OK";
}),
hset: jest.fn().mockImplementation(async (key: string, obj: any) => {
const existing = store.get(key) ?? {};
const normalized: Record<string, string> = { ...existing };
for (const [field, value] of Object.entries(obj)) {
normalized[field] =
value instanceof Buffer ? value.toString("hex") : String(value);
}
store.set(key, normalized);
return 1;
}),
hgetall: jest.fn().mockImplementation(async (key: string) => {
return store.get(key) ?? {};
}),
exists: jest.fn().mockImplementation(async (key: string) => {
return store.has(key) ? 1 : 0;
}),
del: jest.fn().mockImplementation(async (key: string) => {
store.delete(key);
return 1;
}),
get: jest.fn().mockResolvedValue(null),
set: jest.fn().mockResolvedValue("OK"),
quit: jest.fn().mockResolvedValue("OK"),
};
const Valkey = jest.fn().mockImplementation(() => mockClient);
const Cluster = jest.fn().mockImplementation(() => mockClient);
return {
__esModule: true,
default: Valkey,
Cluster,
__mockClient: mockClient,
};
});
ValkeyDB = require("../src/vector_stores/valkey").ValkeyDB;
});
afterEach(() => {
jest.restoreAllMocks();
jest.resetModules();
});
it("implements full VectorStore interface", () => {
const store = new ValkeyDB({
collectionName: "test",
embeddingModelDims: 4,
valkeyUrl: "valkey://localhost:6379",
});
expect(typeof store.insert).toBe("function");
expect(typeof store.search).toBe("function");
expect(typeof store.get).toBe("function");
expect(typeof store.update).toBe("function");
expect(typeof store.delete).toBe("function");
expect(typeof store.deleteCol).toBe("function");
expect(typeof store.list).toBe("function");
expect(typeof store.getUserId).toBe("function");
expect(typeof store.setUserId).toBe("function");
expect(typeof store.initialize).toBe("function");
});
it("initialize() is idempotent", async () => {
const store = new ValkeyDB({
collectionName: "test",
embeddingModelDims: 4,
valkeyUrl: "valkey://localhost:6379",
});
const p1 = store.initialize();
const p2 = store.initialize();
await Promise.all([p1, p2]);
const iovalkey = require("iovalkey");
expect(iovalkey.default).toHaveBeenCalledTimes(1);
});
it("creates HNSW index when indexType is hnsw", async () => {
const store = new ValkeyDB({
collectionName: "test",
embeddingModelDims: 4,
valkeyUrl: "valkey://localhost:6379",
indexType: "hnsw",
});
await store.initialize();
const iovalkey = require("iovalkey");
const mockClient = iovalkey.__mockClient;
const createCall = mockClient.call.mock.calls.find(
(call: any[]) => call[0] === "FT.CREATE",
);
expect(createCall).toBeDefined();
expect(createCall).toContain("HNSW");
});
it("inserts and retrieves a vector", async () => {
const store = new ValkeyDB({
collectionName: "test",
embeddingModelDims: 4,
valkeyUrl: "valkey://localhost:6379",
});
await store.initialize();
await store.insert(
[[0.1, 0.2, 0.3, 0.4]],
["mem-1"],
[
{
data: "hello valkey",
hash: "hash-1",
created_at: "2024-01-01T00:00:00.000Z",
user_id: "alice",
},
],
);
const iovalkey = require("iovalkey");
const mockClient = iovalkey.__mockClient;
expect(mockClient.hset).toHaveBeenCalledWith(
"mem0:test:mem-1",
expect.objectContaining({
memory_id: "mem-1",
memory: "hello valkey",
hash: "hash-1",
user_id: "alice",
}),
);
// Read back through the real stateful mock store (populated by the hset
// above) rather than a hand-rolled hgetall override, so the insert→get
// round-trip and timestamp rendering are genuinely exercised.
const result = await store.get("mem-1");
expect(result?.id).toBe("mem-1");
expect(result?.payload.data).toBe("hello valkey");
expect(result?.payload.userId).toBe("alice");
// created_at is persisted as unix seconds and rendered back to its ISO instant.
expect(result?.payload.createdAt).toBe("2024-01-01T00:00:00.000Z");
});
it("uses Cluster client when clusterMode is enabled", async () => {
const store = new ValkeyDB({
collectionName: "test",
embeddingModelDims: 4,
valkeyUrl: "valkey://cluster.example:6379",
clusterMode: true,
});
await store.initialize();
const iovalkey = require("iovalkey");
expect(iovalkey.Cluster).toHaveBeenCalledTimes(1);
expect(iovalkey.default).not.toHaveBeenCalled();
});
it("passes URL credentials to Cluster via redisOptions in cluster mode", async () => {
const store = new ValkeyDB({
collectionName: "test",
embeddingModelDims: 4,
valkeyUrl: "valkey://user:s3cret@cluster.example:6379",
clusterMode: true,
});
await store.initialize();
const iovalkey = require("iovalkey");
// Cluster ignores URL-embedded auth, so credentials must be forwarded
// explicitly via redisOptions — otherwise every cluster connection is
// silently unauthenticated.
expect(iovalkey.Cluster).toHaveBeenCalledWith(
[{ host: "cluster.example", port: 6379 }],
{ redisOptions: { username: "user", password: "s3cret" } },
);
});
it("renders timestamps in the configured timezone", async () => {
const store = new ValkeyDB({
collectionName: "test",
embeddingModelDims: 4,
valkeyUrl: "valkey://localhost:6379",
timezone: "America/New_York",
});
await store.initialize();
await store.insert(
[[0.1, 0.2, 0.3, 0.4]],
["mem-tz"],
[{ data: "tz", created_at: "2024-01-01T00:00:00.000Z" }],
);
const result = await store.get("mem-tz");
// 2024-01-01T00:00:00Z is 2023-12-31T19:00:00 in America/New_York (UTC-5).
expect(result?.payload.createdAt).toBe("2023-12-31T19:00:00-05:00");
});
it("escapes special characters in filter values (query-injection safety)", async () => {
const store = new ValkeyDB({
collectionName: "test",
embeddingModelDims: 4,
valkeyUrl: "valkey://localhost:6379",
});
await store.initialize();
const iovalkey = require("iovalkey");
const mockClient = iovalkey.__mockClient;
await store.search([0.1, 0.2, 0.3, 0.4], 5, { user_id: "a|b c" });
const searchCall = mockClient.call.mock.calls.find(
(call: any[]) => call[0] === "FT.SEARCH",
);
expect(searchCall).toBeDefined();
// `|` and whitespace must be escaped so a filter value can't rewrite the query.
expect(searchCall[2]).toContain("@user_id:{a\\|b\\ c}");
});
it("does not raise an unhandled rejection when initialization fails", async () => {
const iovalkey = require("iovalkey");
iovalkey.__mockClient.call.mockImplementationOnce(async () => {
throw new Error("ERR unknown command 'FT._LIST'");
});
const errorSpy = jest.spyOn(console, "error").mockImplementation(() => {});
const unhandled: unknown[] = [];
const onUnhandled = (reason: unknown) => unhandled.push(reason);
process.on("unhandledRejection", onUnhandled);
// The constructor kicks off initialize() in a detached .catch; it must log
// and swallow, never re-throw — a re-throw surfaces as an unhandled promise
// rejection that can crash the Node process.
const store = new ValkeyDB({
collectionName: "test",
embeddingModelDims: 4,
valkeyUrl: "valkey://localhost:6379",
});
await expect(store.initialize()).rejects.toThrow(/search module/i);
// Give Node a macrotask to surface any unhandled rejection from the catch.
await new Promise((resolve) => setTimeout(resolve, 10));
process.off("unhandledRejection", onUnhandled);
expect(unhandled).toHaveLength(0);
expect(errorSpy).toHaveBeenCalled();
errorSpy.mockRestore();
});
});
+1
View File
@@ -12,6 +12,7 @@ const external = [
"better-sqlite3",
"@qdrant/js-client-rest",
"redis",
"iovalkey",
"ollama",
"@google/genai",
"@mistralai/mistralai",