feat(ts-sdk): add Elasticsearch vector store provider (#5866)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
404 Ameyy
@@ -6,15 +6,22 @@ description: "Use Elasticsearch as a vector database in Mem0 for distributed vec
|
||||
|
||||
### Installation
|
||||
|
||||
Elasticsearch support requires additional dependencies. Install them with:
|
||||
Elasticsearch support requires the Elasticsearch client as an extra dependency.
|
||||
|
||||
```bash
|
||||
<CodeGroup>
|
||||
```bash Python
|
||||
pip install elasticsearch>=8.0.0
|
||||
```
|
||||
|
||||
```bash TypeScript
|
||||
npm install mem0ai @elastic/elasticsearch
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
### Usage
|
||||
|
||||
```python
|
||||
<CodeGroup>
|
||||
```python Python
|
||||
import os
|
||||
from mem0 import Memory
|
||||
|
||||
@@ -36,12 +43,52 @@ m = Memory.from_config(config)
|
||||
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": "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."}
|
||||
]
|
||||
m.add(messages, user_id="alice", metadata={"category": "movies"})
|
||||
```
|
||||
|
||||
```typescript TypeScript
|
||||
import { Memory } from "mem0ai/oss";
|
||||
|
||||
// Set OPENAI_API_KEY in your environment.
|
||||
const config = {
|
||||
embedder: {
|
||||
provider: "openai",
|
||||
config: {
|
||||
apiKey: process.env.OPENAI_API_KEY,
|
||||
model: "text-embedding-3-small",
|
||||
},
|
||||
},
|
||||
vectorStore: {
|
||||
provider: "elasticsearch",
|
||||
config: {
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 1536,
|
||||
host: "localhost",
|
||||
port: 9200,
|
||||
// For Elastic Cloud, pass cloudId and apiKey instead of host/port.
|
||||
// For basic auth, pass username and password.
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
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>
|
||||
|
||||
<Note>
|
||||
The TypeScript SDK uses camelCase config keys: `collectionName`, `embeddingModelDims`, `cloudId`, `apiKey`, `useSsl`, `verifyCerts`, `caCerts`, `autoCreateIndex`, and `username` (in place of the Python `user`). `collectionName` and `embeddingModelDims` are required. Because the vector store embeds text with your configured embedder before writing, set an `embedder` in the config as shown above.
|
||||
</Note>
|
||||
|
||||
### Config
|
||||
|
||||
Here are the parameters available for configuring Elasticsearch:
|
||||
@@ -74,6 +121,10 @@ Here are the parameters available for configuring Elasticsearch:
|
||||
|
||||
### Custom Search Query
|
||||
|
||||
<Note>
|
||||
`custom_search_query` is available in the Python SDK only. The TypeScript SDK runs a fixed k-NN query with optional metadata filters.
|
||||
</Note>
|
||||
|
||||
The `custom_search_query` parameter allows you to customize the search query when `Memory.search` is called.
|
||||
|
||||
__Example__
|
||||
|
||||
@@ -132,6 +132,7 @@
|
||||
"ollama": "^0.5.14",
|
||||
"pg": "8.11.3",
|
||||
"redis": "^4.6.13",
|
||||
"@elastic/elasticsearch": "^9.0.0",
|
||||
"iovalkey": "^0.3.3",
|
||||
"compromise": "^14.0.0",
|
||||
"natural": "^8.0.1",
|
||||
|
||||
Generated
+89
-22
@@ -44,6 +44,9 @@ importers:
|
||||
'@cloudflare/workers-types':
|
||||
specifier: ^4.20250504.0
|
||||
version: 4.20260611.1
|
||||
'@elastic/elasticsearch':
|
||||
specifier: ^9.0.0
|
||||
version: 9.4.2
|
||||
'@google-cloud/aiplatform':
|
||||
specifier: ^6.8.0
|
||||
version: 6.8.1
|
||||
@@ -52,7 +55,7 @@ importers:
|
||||
version: 1.52.0
|
||||
'@langchain/core':
|
||||
specifier: ^1.1.47
|
||||
version: 1.1.48(openai@4.104.0(ws@8.21.0)(zod@3.25.76))(ws@8.21.0)
|
||||
version: 1.1.48(@opentelemetry/api@1.9.1)(openai@4.104.0(ws@8.21.0)(zod@3.25.76))(ws@8.21.0)
|
||||
'@mistralai/mistralai':
|
||||
specifier: ^1.5.2
|
||||
version: 1.15.1
|
||||
@@ -106,7 +109,7 @@ importers:
|
||||
version: 3.22.5(@types/node@22.19.21)
|
||||
natural:
|
||||
specifier: ^8.0.1
|
||||
version: 8.1.1
|
||||
version: 8.1.1(@opentelemetry/api@1.9.1)
|
||||
ollama:
|
||||
specifier: ^0.5.14
|
||||
version: 0.5.18
|
||||
@@ -554,6 +557,19 @@ packages:
|
||||
resolution: {integrity: sha512-IchNf6dN4tHoMFIn/7OE8LWZ19Y6q/67Bmf6vnGREv8RSbBVb9LPJxEcnwrcwX6ixSvaiGoomAUvu4YSxXrVgw==}
|
||||
engines: {node: '>=12'}
|
||||
|
||||
'@elastic/elasticsearch@9.4.2':
|
||||
resolution: {integrity: sha512-H9myMlLUeotkZhZ4ppinoMGDFxmW3lY8/s+4TIk1vFHyCvWU1Ej4T7azX5buCzemyFApgN0ywnEuvOtpel2VZg==}
|
||||
engines: {node: '>=20'}
|
||||
peerDependencies:
|
||||
apache-arrow: 18.x - 21.x
|
||||
peerDependenciesMeta:
|
||||
apache-arrow:
|
||||
optional: true
|
||||
|
||||
'@elastic/transport@9.3.7':
|
||||
resolution: {integrity: sha512-L38Ax21uF2OPUmCRWycZ/dZdMYf7gMrtClcxvVrqJVFmn8ET2M++GYmFGJpLqOHS1beATxOXLWe7y2ijSQz/ng==}
|
||||
engines: {node: '>=20'}
|
||||
|
||||
'@esbuild/aix-ppc64@0.28.1':
|
||||
resolution: {integrity: sha512-Svl7tq8k/08+p6CXPpRjQ1fKX+1odH/BQbb48fV6fj3CWHhsoIOoY87w1oHXm0qEpkIK3ZfVgp0hed3XBXzXMQ==}
|
||||
engines: {node: '>=18'}
|
||||
@@ -879,6 +895,20 @@ packages:
|
||||
resolution: {integrity: sha512-ow1A/Z7MBlNL/JZzTY4+M+8psiVh66g0Z7EQSl/qbT0U6zLBsVDYh6nvm8D4e4tIzQtn6ODbA2OsqYGvqZj76g==}
|
||||
engines: {node: '>=14', yarn: ^1.22.10}
|
||||
|
||||
'@opentelemetry/api@1.9.1':
|
||||
resolution: {integrity: sha512-gLyJlPHPZYdAk1JENA9LeHejZe1Ti77/pTeFm/nMXmQH/HFZlcS/O2XJB+L8fkbrNSqhdtlvjBVjxwUYanNH5Q==}
|
||||
engines: {node: '>=8.0.0'}
|
||||
|
||||
'@opentelemetry/core@2.9.0':
|
||||
resolution: {integrity: sha512-m2nckMT80NnmjTYSPjJQObBJ+8dgkoajEOUbznL8AHZ3T3yHRk2P7gI1PhEBc1+lOnrYE9UWrWHqJDsmqjmNbw==}
|
||||
engines: {node: ^18.19.0 || >=20.6.0}
|
||||
peerDependencies:
|
||||
'@opentelemetry/api': '>=1.0.0 <1.10.0'
|
||||
|
||||
'@opentelemetry/semantic-conventions@1.42.0':
|
||||
resolution: {integrity: sha512-icc5xCzndZfhuJMy5oqk5AvloWquR7jtae74qzpkKkhGp8BivK+oCcEXgGnjCdTfp8hA44l+w8gE8yYJbocJJw==}
|
||||
engines: {node: '>=14'}
|
||||
|
||||
'@pinecone-database/pinecone@8.0.0':
|
||||
resolution: {integrity: sha512-dItFqLdis2Pd5lC67aKn8HhvXajzlOz4+0RyK1CRcZMdSwG8YxUCBtH4yXPC8/6uLxfr2dvMWnRFuKgtPwwajQ==}
|
||||
engines: {node: '>=20.0.0'}
|
||||
@@ -3202,6 +3232,9 @@ packages:
|
||||
secure-json-parse@2.7.0:
|
||||
resolution: {integrity: sha512-6aU+Rwsezw7VR8/nyvKTx8QpWH9FrcYiXXlqC4z5d5XQBDRqtbfsRjnwGyqbi3gddNtWHuEk9OANUotL26qKUw==}
|
||||
|
||||
secure-json-parse@4.1.0:
|
||||
resolution: {integrity: sha512-l4KnYfEyqYJxDwlNVyRfO2E4NTHfMKAWdUuA8J0yve2Dz/E/PdBepY03RvyJpssIpRFwJoCD55wA+mEDs6ByWA==}
|
||||
|
||||
semver-compare@1.0.0:
|
||||
resolution: {integrity: sha512-YM3/ITh2MJ5MtzaM429anh+x2jiLVjqILF4m4oyQB18W7Ggea7BfqdH/wGMK7dDiMghv/6WG7znWMwUDzJiXow==}
|
||||
|
||||
@@ -4409,6 +4442,26 @@ snapshots:
|
||||
dependencies:
|
||||
'@jridgewell/trace-mapping': 0.3.9
|
||||
|
||||
'@elastic/elasticsearch@9.4.2':
|
||||
dependencies:
|
||||
'@elastic/transport': 9.3.7
|
||||
tslib: 2.8.1
|
||||
transitivePeerDependencies:
|
||||
- supports-color
|
||||
|
||||
'@elastic/transport@9.3.7':
|
||||
dependencies:
|
||||
'@opentelemetry/api': 1.9.1
|
||||
'@opentelemetry/core': 2.9.0(@opentelemetry/api@1.9.1)
|
||||
debug: 4.4.3(supports-color@5.5.0)
|
||||
hpagent: 1.2.0
|
||||
ms: 2.1.3
|
||||
secure-json-parse: 4.1.0
|
||||
tslib: 2.8.1
|
||||
undici: 7.28.0
|
||||
transitivePeerDependencies:
|
||||
- supports-color
|
||||
|
||||
'@esbuild/aix-ppc64@0.28.1':
|
||||
optional: true
|
||||
|
||||
@@ -4745,12 +4798,12 @@ snapshots:
|
||||
|
||||
'@js-sdsl/ordered-map@4.4.2': {}
|
||||
|
||||
'@langchain/core@1.1.48(openai@4.104.0(ws@8.21.0)(zod@3.25.76))(ws@8.21.0)':
|
||||
'@langchain/core@1.1.48(@opentelemetry/api@1.9.1)(openai@4.104.0(ws@8.21.0)(zod@3.25.76))(ws@8.21.0)':
|
||||
dependencies:
|
||||
'@cfworker/json-schema': 4.1.1
|
||||
'@standard-schema/spec': 1.1.0
|
||||
js-tiktoken: 1.0.21
|
||||
langsmith: 0.6.3(openai@4.104.0(ws@8.21.0)(zod@3.25.76))(ws@8.21.0)
|
||||
langsmith: 0.6.3(@opentelemetry/api@1.9.1)(openai@4.104.0(ws@8.21.0)(zod@3.25.76))(ws@8.21.0)
|
||||
mustache: 4.2.0
|
||||
p-queue: 6.6.2
|
||||
zod: 3.25.76
|
||||
@@ -4797,6 +4850,15 @@ snapshots:
|
||||
transitivePeerDependencies:
|
||||
- supports-color
|
||||
|
||||
'@opentelemetry/api@1.9.1': {}
|
||||
|
||||
'@opentelemetry/core@2.9.0(@opentelemetry/api@1.9.1)':
|
||||
dependencies:
|
||||
'@opentelemetry/api': 1.9.1
|
||||
'@opentelemetry/semantic-conventions': 1.42.0
|
||||
|
||||
'@opentelemetry/semantic-conventions@1.42.0': {}
|
||||
|
||||
'@pinecone-database/pinecone@8.0.0': {}
|
||||
|
||||
'@pkgjs/parseargs@0.11.0':
|
||||
@@ -4836,9 +4898,9 @@ snapshots:
|
||||
dependencies:
|
||||
'@redis/client': 1.6.1
|
||||
|
||||
'@redis/bloom@5.12.1(@redis/client@5.12.1)':
|
||||
'@redis/bloom@5.12.1(@redis/client@5.12.1(@opentelemetry/api@1.9.1))':
|
||||
dependencies:
|
||||
'@redis/client': 5.12.1
|
||||
'@redis/client': 5.12.1(@opentelemetry/api@1.9.1)
|
||||
|
||||
'@redis/client@1.6.1':
|
||||
dependencies:
|
||||
@@ -4846,9 +4908,11 @@ snapshots:
|
||||
generic-pool: 3.9.0
|
||||
yallist: 4.0.0
|
||||
|
||||
'@redis/client@5.12.1':
|
||||
'@redis/client@5.12.1(@opentelemetry/api@1.9.1)':
|
||||
dependencies:
|
||||
cluster-key-slot: 1.1.2
|
||||
optionalDependencies:
|
||||
'@opentelemetry/api': 1.9.1
|
||||
|
||||
'@redis/graph@1.1.1(@redis/client@1.6.1)':
|
||||
dependencies:
|
||||
@@ -4858,25 +4922,25 @@ snapshots:
|
||||
dependencies:
|
||||
'@redis/client': 1.6.1
|
||||
|
||||
'@redis/json@5.12.1(@redis/client@5.12.1)':
|
||||
'@redis/json@5.12.1(@redis/client@5.12.1(@opentelemetry/api@1.9.1))':
|
||||
dependencies:
|
||||
'@redis/client': 5.12.1
|
||||
'@redis/client': 5.12.1(@opentelemetry/api@1.9.1)
|
||||
|
||||
'@redis/search@1.2.0(@redis/client@1.6.1)':
|
||||
dependencies:
|
||||
'@redis/client': 1.6.1
|
||||
|
||||
'@redis/search@5.12.1(@redis/client@5.12.1)':
|
||||
'@redis/search@5.12.1(@redis/client@5.12.1(@opentelemetry/api@1.9.1))':
|
||||
dependencies:
|
||||
'@redis/client': 5.12.1
|
||||
'@redis/client': 5.12.1(@opentelemetry/api@1.9.1)
|
||||
|
||||
'@redis/time-series@1.1.0(@redis/client@1.6.1)':
|
||||
dependencies:
|
||||
'@redis/client': 1.6.1
|
||||
|
||||
'@redis/time-series@5.12.1(@redis/client@5.12.1)':
|
||||
'@redis/time-series@5.12.1(@redis/client@5.12.1(@opentelemetry/api@1.9.1))':
|
||||
dependencies:
|
||||
'@redis/client': 5.12.1
|
||||
'@redis/client': 5.12.1(@opentelemetry/api@1.9.1)
|
||||
|
||||
'@rollup/rollup-android-arm-eabi@4.61.1':
|
||||
optional: true
|
||||
@@ -6627,10 +6691,11 @@ snapshots:
|
||||
|
||||
kolorist@1.8.0: {}
|
||||
|
||||
langsmith@0.6.3(openai@4.104.0(ws@8.21.0)(zod@3.25.76))(ws@8.21.0):
|
||||
langsmith@0.6.3(@opentelemetry/api@1.9.1)(openai@4.104.0(ws@8.21.0)(zod@3.25.76))(ws@8.21.0):
|
||||
dependencies:
|
||||
p-queue: 6.6.2
|
||||
optionalDependencies:
|
||||
'@opentelemetry/api': 1.9.1
|
||||
openai: 4.104.0(ws@8.21.0)(zod@3.25.76)
|
||||
ws: 8.21.0
|
||||
|
||||
@@ -6835,7 +6900,7 @@ snapshots:
|
||||
|
||||
natural-compare@1.4.0: {}
|
||||
|
||||
natural@8.1.1:
|
||||
natural@8.1.1(@opentelemetry/api@1.9.1):
|
||||
dependencies:
|
||||
afinn-165: 2.0.2
|
||||
afinn-165-financialmarketnews: 3.0.0
|
||||
@@ -6844,7 +6909,7 @@ snapshots:
|
||||
memjs: 1.3.2
|
||||
mongoose: 9.7.0
|
||||
pg: 8.21.0
|
||||
redis: 5.12.1
|
||||
redis: 5.12.1(@opentelemetry/api@1.9.1)
|
||||
safe-stable-stringify: 2.5.0
|
||||
stopwords-iso: 1.1.0
|
||||
sylvester: 0.0.21
|
||||
@@ -7240,13 +7305,13 @@ snapshots:
|
||||
'@redis/search': 1.2.0(@redis/client@1.6.1)
|
||||
'@redis/time-series': 1.1.0(@redis/client@1.6.1)
|
||||
|
||||
redis@5.12.1:
|
||||
redis@5.12.1(@opentelemetry/api@1.9.1):
|
||||
dependencies:
|
||||
'@redis/bloom': 5.12.1(@redis/client@5.12.1)
|
||||
'@redis/client': 5.12.1
|
||||
'@redis/json': 5.12.1(@redis/client@5.12.1)
|
||||
'@redis/search': 5.12.1(@redis/client@5.12.1)
|
||||
'@redis/time-series': 5.12.1(@redis/client@5.12.1)
|
||||
'@redis/bloom': 5.12.1(@redis/client@5.12.1(@opentelemetry/api@1.9.1))
|
||||
'@redis/client': 5.12.1(@opentelemetry/api@1.9.1)
|
||||
'@redis/json': 5.12.1(@redis/client@5.12.1(@opentelemetry/api@1.9.1))
|
||||
'@redis/search': 5.12.1(@redis/client@5.12.1(@opentelemetry/api@1.9.1))
|
||||
'@redis/time-series': 5.12.1(@redis/client@5.12.1(@opentelemetry/api@1.9.1))
|
||||
transitivePeerDependencies:
|
||||
- '@node-rs/xxhash'
|
||||
- '@opentelemetry/api'
|
||||
@@ -7337,6 +7402,8 @@ snapshots:
|
||||
|
||||
secure-json-parse@2.7.0: {}
|
||||
|
||||
secure-json-parse@4.1.0: {}
|
||||
|
||||
semver-compare@1.0.0: {}
|
||||
|
||||
semver@6.3.1: {}
|
||||
|
||||
@@ -32,6 +32,7 @@ export * from "./vector_stores/langchain";
|
||||
export * from "./vector_stores/vectorize";
|
||||
export * from "./vector_stores/azure_ai_search";
|
||||
export * from "./vector_stores/pgvector";
|
||||
export * from "./vector_stores/elasticsearch";
|
||||
export * from "./vector_stores/upstash_vector";
|
||||
export * from "./vector_stores/azure_mysql";
|
||||
export * from "./vector_stores/cassandra";
|
||||
|
||||
@@ -43,6 +43,7 @@ import { LangchainEmbedder } from "../embeddings/langchain";
|
||||
import { LangchainVectorStore } from "../vector_stores/langchain";
|
||||
import { AzureAISearch } from "../vector_stores/azure_ai_search";
|
||||
import { PGVector } from "../vector_stores/pgvector";
|
||||
import { ElasticsearchDB } from "../vector_stores/elasticsearch";
|
||||
import { OpenSearchDB } from "../vector_stores/opensearch";
|
||||
import { UpstashVector } from "../vector_stores/upstash_vector";
|
||||
import { AzureMySQLDB } from "../vector_stores/azure_mysql";
|
||||
@@ -141,6 +142,8 @@ export class VectorStoreFactory {
|
||||
return new VertexAIVectorSearch(config as any);
|
||||
case "pgvector":
|
||||
return new PGVector(config as any);
|
||||
case "elasticsearch":
|
||||
return new ElasticsearchDB(config as any);
|
||||
case "opensearch":
|
||||
return new OpenSearchDB(config as any);
|
||||
case "upstash_vector":
|
||||
|
||||
@@ -0,0 +1,333 @@
|
||||
import { Client } from "@elastic/elasticsearch";
|
||||
import { VectorStore } from "./base";
|
||||
import { SearchFilters, VectorStoreConfig, VectorStoreResult } from "../types";
|
||||
|
||||
interface ElasticsearchConfig extends VectorStoreConfig {
|
||||
client?: Client;
|
||||
host?: string;
|
||||
port?: number;
|
||||
cloudId?: string;
|
||||
apiKey?: string;
|
||||
username?: string;
|
||||
password?: string;
|
||||
collectionName: string;
|
||||
embeddingModelDims: number;
|
||||
dimension?: number;
|
||||
useSsl?: boolean;
|
||||
caCerts?: string;
|
||||
verifyCerts?: boolean;
|
||||
autoCreateIndex?: boolean;
|
||||
headers?: Record<string, string>;
|
||||
}
|
||||
|
||||
const SAFE_FILTER_KEY = /^[a-zA-Z_][a-zA-Z0-9_]*$/;
|
||||
|
||||
// Mirrors the Python provider's _validate_filter. Filter keys are interpolated
|
||||
// into `metadata.${key}` field paths, so reject anything that is not a plain
|
||||
// identifier and reject non-scalar values before they reach the query.
|
||||
function validateFilter(key: string, value: unknown): void {
|
||||
if (typeof key !== "string" || !SAFE_FILTER_KEY.test(key)) {
|
||||
throw new Error(`Invalid filter key: ${JSON.stringify(key)}`);
|
||||
}
|
||||
if (
|
||||
typeof value !== "string" &&
|
||||
typeof value !== "number" &&
|
||||
typeof value !== "boolean"
|
||||
) {
|
||||
throw new Error(
|
||||
`Filter value for ${JSON.stringify(key)} must be string, number, or boolean`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
export class ElasticsearchDB implements VectorStore {
|
||||
private client: Client;
|
||||
private readonly collectionName: string;
|
||||
private readonly dimension: number;
|
||||
private readonly autoCreateIndex: boolean;
|
||||
private _initPromise?: Promise<void>;
|
||||
|
||||
constructor(config: ElasticsearchConfig) {
|
||||
this.collectionName = config.collectionName;
|
||||
this.dimension = config.dimension || config.embeddingModelDims || 1536;
|
||||
this.autoCreateIndex = config.autoCreateIndex !== false;
|
||||
|
||||
if (config.client) {
|
||||
this.client = config.client;
|
||||
} else {
|
||||
const params: Record<string, any> = {};
|
||||
|
||||
if (config.cloudId) {
|
||||
params.cloud = { id: config.cloudId };
|
||||
if (config.apiKey) {
|
||||
params.auth = { apiKey: config.apiKey };
|
||||
}
|
||||
} else {
|
||||
const host = config.host || "localhost";
|
||||
const port = config.port || 9200;
|
||||
params.node = `${config.useSsl !== false ? "https" : "http"}://${host}:${port}`;
|
||||
|
||||
if (config.apiKey) {
|
||||
params.auth = { apiKey: config.apiKey };
|
||||
} else if (config.username && config.password) {
|
||||
params.auth = {
|
||||
username: config.username,
|
||||
password: config.password,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
if (config.verifyCerts !== undefined) {
|
||||
params.tls = { rejectUnauthorized: config.verifyCerts };
|
||||
}
|
||||
if (config.caCerts) {
|
||||
params.tls = { ...params.tls, ca: config.caCerts };
|
||||
}
|
||||
if (config.headers) {
|
||||
params.headers = config.headers;
|
||||
}
|
||||
|
||||
this.client = new Client(params);
|
||||
}
|
||||
|
||||
this.initialize().catch(console.error);
|
||||
}
|
||||
|
||||
async initialize(): Promise<void> {
|
||||
if (!this._initPromise) {
|
||||
this._initPromise = this._doInitialize();
|
||||
}
|
||||
return this._initPromise;
|
||||
}
|
||||
|
||||
private async _doInitialize(): Promise<void> {
|
||||
try {
|
||||
if (this.autoCreateIndex) {
|
||||
await this.ensureIndex(this.collectionName, this.dimension);
|
||||
}
|
||||
await this.ensureIndex("memory_migrations", 1);
|
||||
} catch (error) {
|
||||
console.error("Error initializing Elasticsearch:", error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
private async ensureIndex(name: string, size: number): Promise<void> {
|
||||
const exists = await this.client.indices.exists({ index: name });
|
||||
if (!exists) {
|
||||
await this.client.indices.create({
|
||||
index: name,
|
||||
settings: {
|
||||
number_of_replicas: 1,
|
||||
number_of_shards: 5,
|
||||
refresh_interval: "1s",
|
||||
},
|
||||
mappings: {
|
||||
properties: {
|
||||
vector: {
|
||||
type: "dense_vector",
|
||||
dims: size,
|
||||
index: true,
|
||||
similarity: "cosine",
|
||||
},
|
||||
metadata: {
|
||||
type: "object",
|
||||
properties: {
|
||||
user_id: { type: "keyword" },
|
||||
agent_id: { type: "keyword" },
|
||||
run_id: { type: "keyword" },
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async insert(
|
||||
vectors: number[][],
|
||||
ids: string[],
|
||||
payloads: Record<string, any>[],
|
||||
): Promise<void> {
|
||||
const operations: any[] = [];
|
||||
for (let i = 0; i < vectors.length; i++) {
|
||||
operations.push(
|
||||
{ index: { _index: this.collectionName, _id: ids[i] } },
|
||||
{ vector: vectors[i], metadata: payloads[i] || {} },
|
||||
);
|
||||
}
|
||||
|
||||
await this.client.bulk({ operations });
|
||||
}
|
||||
|
||||
async search(
|
||||
query: number[],
|
||||
topK: number = 5,
|
||||
filters?: SearchFilters,
|
||||
): Promise<VectorStoreResult[]> {
|
||||
const searchBody: Record<string, any> = {
|
||||
knn: {
|
||||
field: "vector",
|
||||
query_vector: query,
|
||||
k: topK,
|
||||
num_candidates: topK * 2,
|
||||
},
|
||||
};
|
||||
|
||||
if (filters && Object.keys(filters).length > 0) {
|
||||
const filterConditions = Object.entries(filters).map(([key, value]) => {
|
||||
validateFilter(key, value);
|
||||
return { term: { [`metadata.${key}`]: value } };
|
||||
});
|
||||
searchBody.knn.filter = { bool: { must: filterConditions } };
|
||||
}
|
||||
|
||||
const response = await this.client.search({
|
||||
index: this.collectionName,
|
||||
...searchBody,
|
||||
});
|
||||
|
||||
return response.hits.hits.map((hit: any) => ({
|
||||
id: hit._id,
|
||||
score: hit._score,
|
||||
payload: hit._source?.metadata || {},
|
||||
}));
|
||||
}
|
||||
|
||||
async get(vectorId: string): Promise<VectorStoreResult | null> {
|
||||
try {
|
||||
const response = await this.client.get({
|
||||
index: this.collectionName,
|
||||
id: vectorId,
|
||||
});
|
||||
return {
|
||||
id: response._id,
|
||||
payload: (response._source as any)?.metadata || {},
|
||||
};
|
||||
} catch (error: any) {
|
||||
if (error?.meta?.statusCode === 404) return null;
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
async update(
|
||||
vectorId: string,
|
||||
vector: number[],
|
||||
payload: Record<string, any>,
|
||||
): Promise<void> {
|
||||
const doc: Record<string, any> = {};
|
||||
if (vector) doc.vector = vector;
|
||||
if (payload) doc.metadata = payload;
|
||||
|
||||
await this.client.update({
|
||||
index: this.collectionName,
|
||||
id: vectorId,
|
||||
doc,
|
||||
});
|
||||
}
|
||||
|
||||
async delete(vectorId: string): Promise<void> {
|
||||
await this.client.delete({
|
||||
index: this.collectionName,
|
||||
id: vectorId,
|
||||
});
|
||||
}
|
||||
|
||||
async deleteCol(): Promise<void> {
|
||||
await this.client.indices.delete({ index: this.collectionName });
|
||||
}
|
||||
|
||||
async list(
|
||||
filters?: SearchFilters,
|
||||
topK: number = 100,
|
||||
): Promise<[VectorStoreResult[], number]> {
|
||||
const query: Record<string, any> = { query: { match_all: {} } };
|
||||
|
||||
if (filters && Object.keys(filters).length > 0) {
|
||||
const filterConditions = Object.entries(filters).map(([key, value]) => {
|
||||
validateFilter(key, value);
|
||||
return { term: { [`metadata.${key}`]: value } };
|
||||
});
|
||||
query.query = { bool: { must: filterConditions } };
|
||||
}
|
||||
|
||||
query.size = topK;
|
||||
|
||||
const response = await this.client.search({
|
||||
index: this.collectionName,
|
||||
...query,
|
||||
});
|
||||
|
||||
const results = response.hits.hits.map((hit: any) => ({
|
||||
id: hit._id,
|
||||
payload: hit._source?.metadata || {},
|
||||
}));
|
||||
|
||||
return [results, results.length];
|
||||
}
|
||||
|
||||
private generateUUID(): string {
|
||||
return "xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx".replace(/[xy]/g, (c) => {
|
||||
const r = (Math.random() * 16) | 0;
|
||||
const v = c === "x" ? r : (r & 0x3) | 0x8;
|
||||
return v.toString(16);
|
||||
});
|
||||
}
|
||||
|
||||
async getUserId(): Promise<string> {
|
||||
try {
|
||||
const response = await this.client.search({
|
||||
index: "memory_migrations",
|
||||
query: { match_all: {} },
|
||||
size: 1,
|
||||
});
|
||||
|
||||
if (response.hits.hits.length > 0) {
|
||||
return (response.hits.hits[0]._source as any)?.metadata
|
||||
?.user_id as string;
|
||||
}
|
||||
|
||||
const randomUserId =
|
||||
Math.random().toString(36).substring(2, 15) +
|
||||
Math.random().toString(36).substring(2, 15);
|
||||
|
||||
// memory_migrations is only ever read via match_all, never by vector,
|
||||
// so store no vector. A zero vector would be rejected by the cosine
|
||||
// dense_vector mapping ("does not support vectors with zero magnitude").
|
||||
await this.client.index({
|
||||
index: "memory_migrations",
|
||||
id: this.generateUUID(),
|
||||
document: { metadata: { user_id: randomUserId } },
|
||||
});
|
||||
|
||||
return randomUserId;
|
||||
} catch (error) {
|
||||
console.error("Error getting user ID:", error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
async setUserId(userId: string): Promise<void> {
|
||||
try {
|
||||
const response = await this.client.search({
|
||||
index: "memory_migrations",
|
||||
query: { match_all: {} },
|
||||
size: 1,
|
||||
});
|
||||
|
||||
const docId =
|
||||
response.hits.hits.length > 0
|
||||
? response.hits.hits[0]._id
|
||||
: this.generateUUID();
|
||||
|
||||
await this.client.index({
|
||||
index: "memory_migrations",
|
||||
id: docId,
|
||||
document: { metadata: { user_id: userId } },
|
||||
});
|
||||
} catch (error) {
|
||||
console.error("Error setting user ID:", error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,524 @@
|
||||
/// <reference types="jest" />
|
||||
|
||||
type MockFn = jest.Mock;
|
||||
|
||||
const mockIndicesExists = jest.fn();
|
||||
const mockIndicesCreate = jest.fn();
|
||||
const mockBulk = jest.fn();
|
||||
const mockSearch = jest.fn();
|
||||
const mockGet = jest.fn();
|
||||
const mockUpdate = jest.fn();
|
||||
const mockDelete = jest.fn();
|
||||
const mockIndicesDelete = jest.fn();
|
||||
const mockIndex_ = jest.fn();
|
||||
|
||||
const mockClient = jest.fn().mockImplementation(() => ({
|
||||
indices: {
|
||||
exists: mockIndicesExists,
|
||||
create: mockIndicesCreate,
|
||||
delete: mockIndicesDelete,
|
||||
},
|
||||
bulk: mockBulk,
|
||||
search: mockSearch,
|
||||
get: mockGet,
|
||||
update: mockUpdate,
|
||||
delete: mockDelete,
|
||||
index: mockIndex_,
|
||||
}));
|
||||
|
||||
jest.mock("@elastic/elasticsearch", () => ({
|
||||
Client: mockClient,
|
||||
}));
|
||||
|
||||
const { ElasticsearchDB } = require("../src/vector_stores/elasticsearch");
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
mockIndicesExists.mockResolvedValue(false);
|
||||
mockIndicesCreate.mockResolvedValue({});
|
||||
mockBulk.mockResolvedValue({});
|
||||
mockSearch.mockResolvedValue({ hits: { hits: [] } });
|
||||
mockGet.mockResolvedValue({
|
||||
_id: "test-id",
|
||||
_source: { metadata: { key: "val" } },
|
||||
});
|
||||
mockUpdate.mockResolvedValue({});
|
||||
mockDelete.mockResolvedValue({});
|
||||
mockIndicesDelete.mockResolvedValue({});
|
||||
mockIndex_.mockResolvedValue({ _id: "new-id" });
|
||||
});
|
||||
|
||||
describe("ElasticsearchDB", () => {
|
||||
describe("constructor", () => {
|
||||
it("creates client with self-hosted host:port", () => {
|
||||
new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 768,
|
||||
host: "localhost",
|
||||
port: 9200,
|
||||
username: "user",
|
||||
password: "pass",
|
||||
});
|
||||
|
||||
expect(mockClient).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
node: "https://localhost:9200",
|
||||
auth: { username: "user", password: "pass" },
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("creates client with cloud config", () => {
|
||||
new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 1536,
|
||||
cloudId:
|
||||
"my-cloud:dXMtZWFzdDQuZ2NwLmVsYXN0aWMtY2xvdWQuY29tOjQ0MyQxMjM0NTY3ODkw",
|
||||
apiKey: "base64-key",
|
||||
});
|
||||
|
||||
expect(mockClient).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
cloud: {
|
||||
id: "my-cloud:dXMtZWFzdDQuZ2NwLmVsYXN0aWMtY2xvdWQuY29tOjQ0MyQxMjM0NTY3ODkw",
|
||||
},
|
||||
auth: { apiKey: "base64-key" },
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("accepts pre-configured client", () => {
|
||||
const preExisting = new mockClient();
|
||||
new ElasticsearchDB({
|
||||
client: preExisting,
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 1536,
|
||||
});
|
||||
|
||||
expect(mockClient).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("defaults to port 9200 and https", () => {
|
||||
new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 384,
|
||||
host: "es.example.com",
|
||||
});
|
||||
|
||||
expect(mockClient).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
node: "https://es.example.com:9200",
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("initialize", () => {
|
||||
it("creates index with correct dense_vector mapping", async () => {
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "my_memories",
|
||||
embeddingModelDims: 768,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
await store.initialize();
|
||||
|
||||
expect(mockIndicesExists).toHaveBeenCalledWith({ index: "my_memories" });
|
||||
expect(mockIndicesCreate).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
index: "my_memories",
|
||||
mappings: {
|
||||
properties: {
|
||||
vector: {
|
||||
type: "dense_vector",
|
||||
dims: 768,
|
||||
index: true,
|
||||
similarity: "cosine",
|
||||
},
|
||||
metadata: {
|
||||
type: "object",
|
||||
properties: {
|
||||
user_id: { type: "keyword" },
|
||||
agent_id: { type: "keyword" },
|
||||
run_id: { type: "keyword" },
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("creates memory_migrations index with dims 1", async () => {
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "test",
|
||||
embeddingModelDims: 1536,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
await store.initialize();
|
||||
|
||||
expect(mockIndicesCreate).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
index: "memory_migrations",
|
||||
mappings: {
|
||||
properties: {
|
||||
vector: {
|
||||
type: "dense_vector",
|
||||
dims: 1,
|
||||
index: true,
|
||||
similarity: "cosine",
|
||||
},
|
||||
metadata: {
|
||||
type: "object",
|
||||
properties: {
|
||||
user_id: { type: "keyword" },
|
||||
agent_id: { type: "keyword" },
|
||||
run_id: { type: "keyword" },
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("is idempotent (initialize does not re-run)", async () => {
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "test",
|
||||
embeddingModelDims: 1536,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
await store.initialize();
|
||||
|
||||
mockIndicesExists.mockClear();
|
||||
|
||||
await store.initialize();
|
||||
expect(mockIndicesExists).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("skips index creation when autoCreateIndex is false", async () => {
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "test",
|
||||
embeddingModelDims: 1536,
|
||||
host: "localhost",
|
||||
autoCreateIndex: false,
|
||||
});
|
||||
|
||||
await store.initialize();
|
||||
|
||||
expect(mockIndicesExists).not.toHaveBeenCalledWith({ index: "test" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("insert", () => {
|
||||
it("bulk indexes vectors with ids and payloads", async () => {
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
await store.insert(
|
||||
[
|
||||
[0.1, 0.2, 0.3, 0.4],
|
||||
[0.5, 0.6, 0.7, 0.8],
|
||||
],
|
||||
["id-1", "id-2"],
|
||||
[{ user_id: "u1" }, { user_id: "u2" }],
|
||||
);
|
||||
|
||||
expect(mockBulk).toHaveBeenCalledWith({
|
||||
operations: [
|
||||
{ index: { _index: "mem0", _id: "id-1" } },
|
||||
{ vector: [0.1, 0.2, 0.3, 0.4], metadata: { user_id: "u1" } },
|
||||
{ index: { _index: "mem0", _id: "id-2" } },
|
||||
{ vector: [0.5, 0.6, 0.7, 0.8], metadata: { user_id: "u2" } },
|
||||
],
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("search", () => {
|
||||
it("builds correct knn query", async () => {
|
||||
mockSearch.mockResolvedValue({
|
||||
hits: {
|
||||
hits: [
|
||||
{
|
||||
_id: "id-1",
|
||||
_score: 0.95,
|
||||
_source: { metadata: { user_id: "u1" } },
|
||||
},
|
||||
],
|
||||
},
|
||||
});
|
||||
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
const results = await store.search([0.1, 0.2, 0.3, 0.4], 3);
|
||||
|
||||
expect(mockSearch).toHaveBeenCalledWith({
|
||||
index: "mem0",
|
||||
knn: {
|
||||
field: "vector",
|
||||
query_vector: [0.1, 0.2, 0.3, 0.4],
|
||||
k: 3,
|
||||
num_candidates: 6,
|
||||
},
|
||||
});
|
||||
expect(results).toEqual([
|
||||
{ id: "id-1", score: 0.95, payload: { user_id: "u1" } },
|
||||
]);
|
||||
});
|
||||
|
||||
it("applies metadata filters", async () => {
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
await store.search([0.1, 0.2, 0.3, 0.4], 5, {
|
||||
user_id: "u1",
|
||||
agent_id: "a1",
|
||||
});
|
||||
|
||||
expect(mockSearch).toHaveBeenCalledWith({
|
||||
index: "mem0",
|
||||
knn: {
|
||||
field: "vector",
|
||||
query_vector: [0.1, 0.2, 0.3, 0.4],
|
||||
k: 5,
|
||||
num_candidates: 10,
|
||||
filter: {
|
||||
bool: {
|
||||
must: [
|
||||
{ term: { "metadata.user_id": "u1" } },
|
||||
{ term: { "metadata.agent_id": "a1" } },
|
||||
],
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
it("rejects filter keys and values that are not safe scalars", async () => {
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
await expect(
|
||||
store.search([0.1, 0.2, 0.3, 0.4], 5, { "bad.key": "x" }),
|
||||
).rejects.toThrow(/Invalid filter key/);
|
||||
await expect(
|
||||
store.search([0.1, 0.2, 0.3, 0.4], 5, { user_id: { $ne: null } }),
|
||||
).rejects.toThrow(/must be string, number, or boolean/);
|
||||
expect(mockSearch).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe("get", () => {
|
||||
it("returns document by id", async () => {
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
const result = await store.get("id-1");
|
||||
|
||||
expect(mockGet).toHaveBeenCalledWith({ index: "mem0", id: "id-1" });
|
||||
expect(result).toEqual({ id: "test-id", payload: { key: "val" } });
|
||||
});
|
||||
|
||||
it("returns null on 404", async () => {
|
||||
mockGet.mockRejectedValue({ meta: { statusCode: 404 } });
|
||||
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
const result = await store.get("missing");
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("update", () => {
|
||||
it("updates vector and payload", async () => {
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
await store.update("id-1", [1, 2, 3, 4], { key: "val" });
|
||||
|
||||
expect(mockUpdate).toHaveBeenCalledWith({
|
||||
index: "mem0",
|
||||
id: "id-1",
|
||||
doc: { vector: [1, 2, 3, 4], metadata: { key: "val" } },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("delete", () => {
|
||||
it("deletes document by id", async () => {
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
await store.delete("id-1");
|
||||
|
||||
expect(mockDelete).toHaveBeenCalledWith({ index: "mem0", id: "id-1" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("deleteCol", () => {
|
||||
it("deletes the index", async () => {
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
await store.deleteCol();
|
||||
|
||||
expect(mockIndicesDelete).toHaveBeenCalledWith({ index: "mem0" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("list", () => {
|
||||
it("returns all documents with count", async () => {
|
||||
mockSearch.mockResolvedValue({
|
||||
hits: {
|
||||
hits: [
|
||||
{ _id: "a", _source: { metadata: { user_id: "u1" } } },
|
||||
{ _id: "b", _source: { metadata: { user_id: "u2" } } },
|
||||
],
|
||||
},
|
||||
});
|
||||
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
const [results, count] = await store.list();
|
||||
|
||||
expect(mockSearch).toHaveBeenCalledWith({
|
||||
index: "mem0",
|
||||
query: { match_all: {} },
|
||||
size: 100,
|
||||
});
|
||||
expect(results).toHaveLength(2);
|
||||
expect(count).toBe(2);
|
||||
});
|
||||
|
||||
it("applies filters to list query", async () => {
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
await store.list({ user_id: "u1" }, 10);
|
||||
|
||||
expect(mockSearch).toHaveBeenCalledWith({
|
||||
index: "mem0",
|
||||
query: { bool: { must: [{ term: { "metadata.user_id": "u1" } }] } },
|
||||
size: 10,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("getUserId / setUserId", () => {
|
||||
it("generates new userId when memory_migrations is empty", async () => {
|
||||
mockSearch.mockResolvedValue({ hits: { hits: [] } });
|
||||
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
const userId = await store.getUserId();
|
||||
|
||||
expect(mockSearch).toHaveBeenCalledWith({
|
||||
index: "memory_migrations",
|
||||
query: { match_all: {} },
|
||||
size: 1,
|
||||
});
|
||||
expect(mockIndex_).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
index: "memory_migrations",
|
||||
document: { metadata: { user_id: expect.any(String) } },
|
||||
}),
|
||||
);
|
||||
expect(userId).toBeTruthy();
|
||||
});
|
||||
|
||||
it("returns existing userId", async () => {
|
||||
mockSearch.mockResolvedValue({
|
||||
hits: {
|
||||
hits: [
|
||||
{
|
||||
_id: "doc-1",
|
||||
_source: { metadata: { user_id: "existing-user" } },
|
||||
},
|
||||
],
|
||||
},
|
||||
});
|
||||
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
const userId = await store.getUserId();
|
||||
expect(userId).toBe("existing-user");
|
||||
});
|
||||
|
||||
it("setUserId upserts into memory_migrations", async () => {
|
||||
const store = new ElasticsearchDB({
|
||||
collectionName: "mem0",
|
||||
embeddingModelDims: 4,
|
||||
host: "localhost",
|
||||
});
|
||||
|
||||
await store.setUserId("custom-user");
|
||||
|
||||
expect(mockIndex_).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
index: "memory_migrations",
|
||||
document: { metadata: { user_id: "custom-user" } },
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("factory registration", () => {
|
||||
it("is registered in VectorStoreFactory", () => {
|
||||
const { VectorStoreFactory } = require("../src/utils/factory");
|
||||
const store = VectorStoreFactory.create("elasticsearch", {
|
||||
collectionName: "test",
|
||||
embeddingModelDims: 128,
|
||||
host: "localhost",
|
||||
client: new mockClient(),
|
||||
});
|
||||
expect(store).toBeInstanceOf(ElasticsearchDB);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -32,6 +32,7 @@ const external = [
|
||||
"mysql2",
|
||||
"@turbopuffer/turbopuffer",
|
||||
"@opensearch-project/opensearch",
|
||||
"@elastic/elasticsearch",
|
||||
];
|
||||
|
||||
const define = {
|
||||
|
||||
Reference in New Issue
Block a user