diff --git a/docs/components/vectordbs/dbs/elasticsearch.mdx b/docs/components/vectordbs/dbs/elasticsearch.mdx index 1c6961c3f..2701b9597 100644 --- a/docs/components/vectordbs/dbs/elasticsearch.mdx +++ b/docs/components/vectordbs/dbs/elasticsearch.mdx @@ -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 + +```bash Python pip install elasticsearch>=8.0.0 ``` +```bash TypeScript +npm install mem0ai @elastic/elasticsearch +``` + + ### Usage -```python + +```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" } }); +``` + + + + 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. + + ### Config Here are the parameters available for configuring Elasticsearch: @@ -74,6 +121,10 @@ Here are the parameters available for configuring Elasticsearch: ### Custom Search Query + + `custom_search_query` is available in the Python SDK only. The TypeScript SDK runs a fixed k-NN query with optional metadata filters. + + The `custom_search_query` parameter allows you to customize the search query when `Memory.search` is called. __Example__ diff --git a/mem0-ts/package.json b/mem0-ts/package.json index 2a6d8be01..2398063fc 100644 --- a/mem0-ts/package.json +++ b/mem0-ts/package.json @@ -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", diff --git a/mem0-ts/pnpm-lock.yaml b/mem0-ts/pnpm-lock.yaml index 5b5c00438..145bb3988 100644 --- a/mem0-ts/pnpm-lock.yaml +++ b/mem0-ts/pnpm-lock.yaml @@ -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: {} diff --git a/mem0-ts/src/oss/src/index.ts b/mem0-ts/src/oss/src/index.ts index b40cf01ce..d902cbdcc 100644 --- a/mem0-ts/src/oss/src/index.ts +++ b/mem0-ts/src/oss/src/index.ts @@ -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"; diff --git a/mem0-ts/src/oss/src/utils/factory.ts b/mem0-ts/src/oss/src/utils/factory.ts index 7aa372da2..f41a05f5b 100644 --- a/mem0-ts/src/oss/src/utils/factory.ts +++ b/mem0-ts/src/oss/src/utils/factory.ts @@ -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": diff --git a/mem0-ts/src/oss/src/vector_stores/elasticsearch.ts b/mem0-ts/src/oss/src/vector_stores/elasticsearch.ts new file mode 100644 index 000000000..4b036ae52 --- /dev/null +++ b/mem0-ts/src/oss/src/vector_stores/elasticsearch.ts @@ -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; +} + +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; + + 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 = {}; + + 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 { + if (!this._initPromise) { + this._initPromise = this._doInitialize(); + } + return this._initPromise; + } + + private async _doInitialize(): Promise { + 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 { + 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[], + ): Promise { + 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 { + const searchBody: Record = { + 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 { + 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, + ): Promise { + const doc: Record = {}; + 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 { + await this.client.delete({ + index: this.collectionName, + id: vectorId, + }); + } + + async deleteCol(): Promise { + await this.client.indices.delete({ index: this.collectionName }); + } + + async list( + filters?: SearchFilters, + topK: number = 100, + ): Promise<[VectorStoreResult[], number]> { + const query: Record = { 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 { + 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 { + 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; + } + } +} diff --git a/mem0-ts/src/oss/tests/elasticsearch.unit.test.ts b/mem0-ts/src/oss/tests/elasticsearch.unit.test.ts new file mode 100644 index 000000000..69a78a27c --- /dev/null +++ b/mem0-ts/src/oss/tests/elasticsearch.unit.test.ts @@ -0,0 +1,524 @@ +/// + +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); + }); + }); +}); diff --git a/mem0-ts/tsup.config.ts b/mem0-ts/tsup.config.ts index dec748219..96ed43523 100644 --- a/mem0-ts/tsup.config.ts +++ b/mem0-ts/tsup.config.ts @@ -32,6 +32,7 @@ const external = [ "mysql2", "@turbopuffer/turbopuffer", "@opensearch-project/opensearch", + "@elastic/elasticsearch", ]; const define = {