diff --git a/docs/components/vectordbs/dbs/vertex_ai.mdx b/docs/components/vectordbs/dbs/vertex_ai.mdx index 6df1b5b1d..5679136e3 100644 --- a/docs/components/vectordbs/dbs/vertex_ai.mdx +++ b/docs/components/vectordbs/dbs/vertex_ai.mdx @@ -8,8 +8,8 @@ description: "Use Google Cloud Vertex AI Vector Search as a managed vector store To use Google Cloud Vertex AI Vector Search with `mem0`, you need to configure the `vector_store` in your `mem0` config: - -```python + +```python Python import os from mem0 import Memory @@ -20,7 +20,7 @@ config = { "provider": "vertex_ai_vector_search", "config": { "endpoint_id": "YOUR_ENDPOINT_ID", # Required: Vector Search endpoint ID - "index_id": "YOUR_INDEX_ID", # Required: Vector Search index ID + "index_id": "YOUR_INDEX_ID", # Required: Vector Search index ID "deployment_index_id": "YOUR_DEPLOYMENT_INDEX_ID", # Required: Deployment-specific ID "project_id": "YOUR_PROJECT_ID", # Required: Google Cloud project ID "project_number": "YOUR_PROJECT_NUMBER", # Required: Google Cloud project number @@ -34,9 +34,40 @@ m = Memory.from_config(config) m.add("Your text here", user_id="user", metadata={"category": "example"}) ``` +```typescript TypeScript +import { Memory } from "mem0ai/oss"; + +// Authenticate with GOOGLE_APPLICATION_CREDENTIALS in your environment, +// or pass credentialsPath / serviceAccountJson in the config below. +const config = { + vectorStore: { + provider: "vertex_ai_vector_search", + config: { + endpointId: "YOUR_ENDPOINT_ID", // Required: Vector Search endpoint ID + indexId: "YOUR_INDEX_ID", // Required: Vector Search index ID + deploymentIndexId: "YOUR_DEPLOYMENT_INDEX_ID", // Required: Deployment-specific ID + projectId: "YOUR_PROJECT_ID", // Required: Google Cloud project ID + projectNumber: "YOUR_PROJECT_NUMBER", // Required: Google Cloud project number + region: "YOUR_REGION", // Required: Google Cloud region + credentialsPath: "path/to/credentials.json", // Optional: defaults to GOOGLE_APPLICATION_CREDENTIALS + vectorSearchApiEndpoint: "YOUR_API_ENDPOINT", // Required for search/get operations + }, + }, +}; + +const memory = new Memory(config); +await memory.add("Your text here", { + userId: "user", + metadata: { category: "example" }, +}); +``` + + ### Required Parameters + + | Parameter | Description | Required | |-----------|-------------|----------| | `endpoint_id` | Vector Search endpoint ID | Yes | @@ -48,3 +79,18 @@ m.add("Your text here", user_id="user", metadata={"category": "example"}) | `region` | Google Cloud region | Yes | | `credentials_path` | Path to service account credentials | No (defaults to GOOGLE_APPLICATION_CREDENTIALS) | | `service_account_json` | Service account credentials as a dictionary (alternative to `credentials_path`) | `None` | + + +| Parameter | Description | Required | +|-----------|-------------|----------| +| `endpointId` | Vector Search endpoint ID | Yes | +| `indexId` | Vector Search index ID | Yes | +| `deploymentIndexId` | Deployment-specific index ID | Yes | +| `projectId` | Google Cloud project ID | Yes | +| `projectNumber` | Google Cloud project number | Yes | +| `vectorSearchApiEndpoint` | Vector search API endpoint | Yes (for get operations) | +| `region` | Google Cloud region | Yes | +| `credentialsPath` | Path to service account credentials | No (defaults to GOOGLE_APPLICATION_CREDENTIALS) | +| `serviceAccountJson` | Service account credentials as an object (alternative to `credentialsPath`) | No | + + diff --git a/mem0-ts/package.json b/mem0-ts/package.json index fd5d653f8..e8824e67e 100644 --- a/mem0-ts/package.json +++ b/mem0-ts/package.json @@ -112,6 +112,7 @@ "@azure/identity": "^4.0.0", "@azure/search-documents": "^12.0.0", "@cloudflare/workers-types": "^4.20250504.0", + "@google-cloud/aiplatform": "^6.8.0", "@google/genai": "^1.40.0", "@langchain/core": "^1.1.47", "@mistralai/mistralai": "^1.5.2", diff --git a/mem0-ts/pnpm-lock.yaml b/mem0-ts/pnpm-lock.yaml index 05b260a13..27a8d89bf 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 + '@google-cloud/aiplatform': + specifier: ^6.8.0 + version: 6.8.1 '@google/genai': specifier: ^1.40.0 version: 1.52.0 @@ -695,6 +698,10 @@ packages: cpu: [x64] os: [win32] + '@google-cloud/aiplatform@6.8.1': + resolution: {integrity: sha512-8O64485bzBzQQ1VlgEhaLJcOvVlfC/V8MB5nQ5y2Dcl8UAfGXW5LkXK1CEfZ7c6IZYxpkJ9U818ldc8pg7eteg==} + engines: {node: '>=18'} + '@google/genai@1.52.0': resolution: {integrity: sha512-gwSvbpiN/17O9TbsqSsE/OzZcpv5Fo4RQjdngGgogtuB9RsyJ8ZHhX5KjHj1bp5N9snN2eK8LDGXSaWW2hof8Q==} engines: {node: '>=20.0.0'} @@ -704,6 +711,15 @@ packages: '@modelcontextprotocol/sdk': optional: true + '@grpc/grpc-js@1.14.4': + resolution: {integrity: sha512-k9Dj3DV/itK9D06Y8f190Qgop7/Ui+D0njFV3LHMPwPT75DpXLQohE9Wmz0QElrJnzsjB7KPWiKJbOl7IPDArQ==} + engines: {node: '>=12.10.0'} + + '@grpc/proto-loader@0.8.1': + resolution: {integrity: sha512-wtF6h+DY6M3YaDBPAmvuuA6jV8Sif9MjtOI5euKFWRgCDl5PeDpPsHR9u2l6St5ceY8AZgoNDww5+HvEsXFsGg==} + engines: {node: '>=6'} + hasBin: true + '@huggingface/blake3-jit@0.0.2': resolution: {integrity: sha512-Bq7B5qabyjrJfhBsl85Jd2QBtf+HzRD7h7A9GfN2lzrrsABhOa5evVPgzoCTxR7Ub0QFj7YDK1YkYRWBU25+2w==} @@ -822,6 +838,9 @@ packages: '@jridgewell/trace-mapping@0.3.9': resolution: {integrity: sha512-3Belt6tdc8bPgAtbcmdtNJlirVoTmEb5e2gC94PnkwEW9jI6CAHUeoG85tjWP5WquqfavoMtMwiG4P926ZKKuQ==} + '@js-sdsl/ordered-map@4.4.2': + resolution: {integrity: sha512-iUKgm52T8HOE/makSxjqoWhe95ZJA1/G1sYsGev2JDKUSS14KAgg1LHb+Ba+IPow0xflbnSkOsZcO08C7w1gYw==} + '@langchain/core@1.1.48': resolution: {integrity: sha512-fQU6Guyb1pwc2fEplmA8FPbKfOMAofjnyJzExevro0FxEiuGHE18Ov/ZHmT9trWCDTZRI9eW1VIc6aChxV8pAQ==} engines: {node: '>=20'} @@ -1771,6 +1790,9 @@ packages: resolution: {integrity: sha512-KIN/nDJBQRcXw0MLVhZE9iQHmG68qAVIBg9CqmUYjmQIhgij9U5MFvrqkUL5FbtyyzZuOeOt0zdeRe4UY7ct+A==} engines: {node: '>= 0.4'} + duplexify@4.1.3: + resolution: {integrity: sha512-M3BmBhwJRZsSx38lZyhE53Csddgzl5R7xGJNk7CVddZD6CcmwMCH8J+7AprIrQKH7TonKxaCjcv27Qmf+sQ+oA==} + eastasianwidth@0.2.0: resolution: {integrity: sha512-I88TYZWc9XiYHRQ4/3c5rjjfgkjhLyW2luGIheGERbNQ6OY7yTybanSpDXZa8y7VUP9YmDcYa+eyq4ca7iLqWA==} @@ -2026,10 +2048,18 @@ packages: resolution: {integrity: sha512-DpLKbNU4WylpxJykQujfCcwYWiV/Jhm50Goo0wrVILAv5jOr9d+H+UR3PhSCD2rCCEIg0uc+G+muBTwD54JhDQ==} engines: {node: '>= 0.4'} + google-auth-library@10.5.0: + resolution: {integrity: sha512-7ABviyMOlX5hIVD60YOfHw4/CxOfBhyduaYB+wbFWCWoni4N7SLcV46hrVRktuBbZjFC9ONyqamZITN7q3n32w==} + engines: {node: '>=18'} + google-auth-library@10.7.0: resolution: {integrity: sha512-QpTAbNJ36TliZLx3TTtahR8HG0hN9RllL1e3FymOvQSIKK8JmgV58H924ub2wa2DsS3ANjjP1Aw1N+Ramc8hqQ==} engines: {node: '>=18'} + google-gax@5.0.7: + resolution: {integrity: sha512-EhiqaWWJ+9h7sCcKJTsoo6tMcjokVHhWsbSuWCnZJT4vIBP3y4mAoFLnt9SzgkVZeq24ZsFaArr06nnYYku2yA==} + engines: {node: '>=18'} + google-logging-utils@1.1.3: resolution: {integrity: sha512-eAmLkjDjAFCVXg7A1unxHsLf961m6y17QFqXqAXGj/gVkKFrEICfStRfwUlGNfeCEjNRa32JEWOUTlYXPyyKvA==} engines: {node: '>=14'} @@ -2048,6 +2078,10 @@ packages: groq-sdk@0.3.0: resolution: {integrity: sha512-Cdgjh4YoSBE2X4S9sxPGXaAy1dlN4bRtAaDZ3cnq+XsxhhN9WSBeHF64l7LWwuD5ntmw7YC5Vf4Ff1oHCg1LOg==} + gtoken@8.0.0: + resolution: {integrity: sha512-+CqsMbHPiSTdtSO14O51eMNlrp9N79gmeqmXeouJOhfucAedHw9noVe/n5uJk3tbKE6a+6ZCQg3RPhVhHByAIw==} + engines: {node: '>=18'} + handlebars@4.7.9: resolution: {integrity: sha512-4E71E0rpOaQuJR2A3xDZ+GM1HyWYv1clR58tC8emQNeQe3RH7MAzSbat+V0wG78LQBo6m6bzSG/L4pBuCsgnUQ==} engines: {node: '>=0.4.7'} @@ -2447,6 +2481,9 @@ packages: resolution: {integrity: sha512-t7hw9pI+WvuwNJXwk5zVHpyhIqzg2qTlklJOf0mVxGSbe3Fp2VieZcduNYjaLDoy6p9uGpQEGWG87WpMKlNq8g==} engines: {node: '>=8'} + lodash.camelcase@4.3.0: + resolution: {integrity: sha512-TwuEnCnxbc3rAvhf/LbG7tJUDzhqXyFnv3dtzLOPgCG/hODL7WFnsbwktkD7yUV0RrreP/l1PALq/YSg6VvjlA==} + lodash.defaults@4.2.0: resolution: {integrity: sha512-qjxPLHd3r5DnsdGacqOMU6pb/avJzdh9tFX2ymgoZE27BmjXrNy/y4LoaiTeAb+O3gL8AfpJGtqfX/ae2leYYQ==} @@ -2708,6 +2745,10 @@ packages: resolution: {integrity: sha512-rJgTQnkUnH1sFw8yT6VSU3zD3sWmu6sZhIseY8VX+GRu3P6F7Fu+JNDoXfklElbLJSnc3FUQHVe4cU5hj+BcUg==} engines: {node: '>=0.10.0'} + object-hash@3.0.0: + resolution: {integrity: sha512-RSn9F68PjH9HqtltsSnqYC1XXoWe9Bju5+213R98cNGttag9q9yAOTzdbsqvIa7aNm5WffBZFpWYr2aWrklWAw==} + engines: {node: '>= 6'} + object-keys@1.1.1: resolution: {integrity: sha512-NuAESUOUMrlIXOfHKzD6bpPu3tYt3xvjNdRIQ+FeT0lNb4K8WR70CaDxhuNguS2XG+GjkyMwOzsN5ZktImfhLA==} engines: {node: '>= 0.4'} @@ -2965,6 +3006,10 @@ packages: resolution: {integrity: sha512-NxNv/kLguCA7p3jE8oL2aEBsrJWgAakBpgmgK6lpPWV+WuOmY6r2/zbAVnP+T8bQlA0nzHXSJSJW0Hq7ylaD2Q==} engines: {node: '>= 6'} + proto3-json-serializer@3.0.4: + resolution: {integrity: sha512-E1sbAYg3aEbXrq0n1ojJkRHQJGE1kaE/O6GLA94y8rnJBfgvOPTOd1b9hOceQK1FFZI9qMh1vBERCyO2ifubcw==} + engines: {node: '>=18'} + protobufjs@7.6.3: resolution: {integrity: sha512-+k0vdJKNdW+Vu+dYe8tZA/VvQb6XKNWexC6URwBFXxNnjLJz9nQJCemGyNgRAWD+B7+nGNc9qMPGwcD7s4nzUw==} engines: {node: '>=12.0.0'} @@ -3048,6 +3093,10 @@ packages: engines: {node: '>= 0.4'} hasBin: true + retry-request@8.0.3: + resolution: {integrity: sha512-qqoc4kkGgP9cmQDWELlOpAmfgJOg0Yi7MT82ZjiPWu451ayju4itwomjM4/dBEliify8C1b3tSaeCOldugtwPQ==} + engines: {node: '>=18'} + retry@0.13.1: resolution: {integrity: sha512-XQBQ3I8W1Cge0Seh+6gjj03LbmRFWuoszgK9ooCpwYIrhhoO80pfq4cUkU5DkknwfOfFteRwlZ56PYOGYyFWdg==} engines: {node: '>= 4'} @@ -3178,6 +3227,12 @@ packages: resolution: {integrity: sha512-I6GPS/E0zyieHehMRPQcqkiBMJKGgLta+1hREixhoLPqEA0AlVFiC43dl8uPpmkkeRdDMzYRWFWk5/l9x7nmNg==} engines: {node: '>=0.10.0'} + stream-events@1.0.5: + resolution: {integrity: sha512-E1GUzBSgvct8Jsb3v2X15pjzN1tYebtbLaMg+eBOUOAxgbLoSbT2NS91ckc5lJD1KfLjId+jXJRgo0qnV5Nerg==} + + stream-shift@1.0.3: + resolution: {integrity: sha512-76ORR0DO1o1hlKwTbi/DM3EXWGf3ZJYO8cXX5RJwnul2DEg2oyoZyjLNoQM8WsvZiFKCRfC1O0J7iCvie3RZmQ==} + string-length@4.0.2: resolution: {integrity: sha512-+l6rNN5fYHNhZZy41RXsYptCjA2Igmq4EG7kZAYFQI1E1VTXarr6ZPXBg6eq7Y6eK4FEhY6AJlyuFIb/v/S0VQ==} engines: {node: '>=10'} @@ -3220,6 +3275,9 @@ packages: strnum@2.4.1: resolution: {integrity: sha512-M9eUSMT2dCB2cTNPG7UYj6KuK7RJR2SN2+yCV/fTW3xzTCS6EaGZ5pSMgDIjB7r8zSfTGk+dvvn9rTjpVS9Mwg==} + stubs@3.0.0: + resolution: {integrity: sha512-PdHt7hHUJKxvTCgbKX9C1V/ftOcjJQgz8BZwNfV5c4B6dcGqlpelTbJ999jBGZ2jYiPAwcX5dP6oBwVlBlUbxw==} + sucrase@3.35.1: resolution: {integrity: sha512-DhuTmvZWux4H1UOnWMB3sk0sbaCVOoQZjv8u1rDoTV0HTdGem9hkAZtl4JZy8P2z4Bg0nT+YMeOFyVr4zcG5Tw==} engines: {node: '>=16 || 14 >=14.17'} @@ -3264,6 +3322,10 @@ packages: resolution: {integrity: sha512-4LeEWl96twnS2Q7Bz4MGqgazLqO+hJN63GZxXoIqh1T3VweYD997gbU1ItNsQafqqXTXd5WFyFdReLtwvRBNiw==} engines: {node: '>=18'} + teeny-request@10.1.3: + resolution: {integrity: sha512-5yDliI1uWkYPo7W+Zvrxg6YmoWuj5iC5EydewqrRTvc68nyMTZhlPPlLg6cptUGfbQAb+N9XDPDPzF6N081lug==} + engines: {node: '>=18'} + test-exclude@6.0.0: resolution: {integrity: sha512-cAGWPIyOHU6zlmg88jwm7VRyXnMN7iV68OGAbYDk/Mh/xC/pzVPlQtY6ngoIH/5/tciuhGfvESU8GrHrcxD56w==} engines: {node: '>=8'} @@ -4351,6 +4413,12 @@ snapshots: '@esbuild/win32-x64@0.28.1': optional: true + '@google-cloud/aiplatform@6.8.1': + dependencies: + google-gax: 5.0.7 + transitivePeerDependencies: + - supports-color + '@google/genai@1.52.0': dependencies: google-auth-library: 10.7.0 @@ -4362,6 +4430,18 @@ snapshots: - supports-color - utf-8-validate + '@grpc/grpc-js@1.14.4': + dependencies: + '@grpc/proto-loader': 0.8.1 + '@js-sdsl/ordered-map': 4.4.2 + + '@grpc/proto-loader@0.8.1': + dependencies: + lodash.camelcase: 4.3.0 + long: 5.3.2 + protobufjs: 7.6.3 + yargs: 17.7.2 + '@huggingface/blake3-jit@0.0.2': {} '@huggingface/hub@2.13.2': @@ -4589,6 +4669,8 @@ snapshots: '@jridgewell/resolve-uri': 3.1.2 '@jridgewell/sourcemap-codec': 1.5.5 + '@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)': dependencies: '@cfworker/json-schema': 4.1.1 @@ -5520,6 +5602,13 @@ snapshots: es-errors: 1.3.0 gopd: 1.2.0 + duplexify@4.1.3: + dependencies: + end-of-stream: 1.4.5 + inherits: 2.0.4 + readable-stream: 3.6.2 + stream-shift: 1.0.3 + eastasianwidth@0.2.0: {} ecdsa-sig-formatter@1.0.11: @@ -5814,6 +5903,18 @@ snapshots: define-properties: 1.2.1 gopd: 1.2.0 + google-auth-library@10.5.0: + dependencies: + base64-js: 1.5.1 + ecdsa-sig-formatter: 1.0.11 + gaxios: 7.1.5 + gcp-metadata: 8.1.2 + google-logging-utils: 1.1.3 + gtoken: 8.0.0 + jws: 4.0.1 + transitivePeerDependencies: + - supports-color + google-auth-library@10.7.0: dependencies: base64-js: 1.5.1 @@ -5825,6 +5926,22 @@ snapshots: transitivePeerDependencies: - supports-color + google-gax@5.0.7: + dependencies: + '@grpc/grpc-js': 1.14.4 + '@grpc/proto-loader': 0.8.1 + duplexify: 4.1.3 + google-auth-library: 10.5.0 + google-logging-utils: 1.1.3 + node-fetch: 3.3.2 + object-hash: 3.0.0 + proto3-json-serializer: 3.0.4 + protobufjs: 7.6.3 + retry-request: 8.0.3 + rimraf: 5.0.10 + transitivePeerDependencies: + - supports-color + google-logging-utils@1.1.3: {} gopd@1.2.0: {} @@ -5847,6 +5964,13 @@ snapshots: transitivePeerDependencies: - encoding + gtoken@8.0.0: + dependencies: + gaxios: 7.1.5 + jws: 4.0.1 + transitivePeerDependencies: + - supports-color + handlebars@4.7.9: dependencies: minimist: 1.2.8 @@ -6414,6 +6538,8 @@ snapshots: dependencies: p-locate: 4.1.0 + lodash.camelcase@4.3.0: {} + lodash.defaults@4.2.0: {} lodash.includes@4.3.0: {} @@ -6658,6 +6784,8 @@ snapshots: object-assign@4.1.1: {} + object-hash@3.0.0: {} + object-keys@1.1.1: {} obuf@1.1.2: {} @@ -6905,6 +7033,10 @@ snapshots: kleur: 3.0.3 sisteransi: 1.0.5 + proto3-json-serializer@3.0.4: + dependencies: + protobufjs: 7.6.3 + protobufjs@7.6.3: dependencies: '@protobufjs/aspromise': 1.1.2 @@ -7006,6 +7138,13 @@ snapshots: path-parse: 1.0.7 supports-preserve-symlinks-flag: 1.0.0 + retry-request@8.0.3: + dependencies: + extend: 3.0.2 + teeny-request: 10.1.3 + transitivePeerDependencies: + - supports-color + retry@0.13.1: {} reusify@1.1.0: {} @@ -7141,6 +7280,12 @@ snapshots: stopwords-iso@1.1.0: {} + stream-events@1.0.5: + dependencies: + stubs: 3.0.0 + + stream-shift@1.0.3: {} + string-length@4.0.2: dependencies: char-regex: 1.0.2 @@ -7182,6 +7327,8 @@ snapshots: dependencies: anynum: 1.0.1 + stubs@3.0.0: {} + sucrase@3.35.1: dependencies: '@jridgewell/gen-mapping': 0.3.13 @@ -7242,6 +7389,15 @@ snapshots: minizlib: 3.1.0 yallist: 5.0.0 + teeny-request@10.1.3: + dependencies: + http-proxy-agent: 7.0.2 + https-proxy-agent: 7.0.6 + node-fetch: 3.3.2 + stream-events: 1.0.5 + transitivePeerDependencies: + - supports-color + test-exclude@6.0.0: dependencies: '@istanbuljs/schema': 0.1.6 diff --git a/mem0-ts/src/oss/src/index.ts b/mem0-ts/src/oss/src/index.ts index 41c15cbe8..842a0bc2c 100644 --- a/mem0-ts/src/oss/src/index.ts +++ b/mem0-ts/src/oss/src/index.ts @@ -34,4 +34,5 @@ export * from "./vector_stores/azure_ai_search"; export * from "./vector_stores/pgvector"; export * from "./vector_stores/cassandra"; export * from "./vector_stores/s3_vectors"; +export * from "./vector_stores/vertex_ai_vector_search"; export * from "./utils/factory"; diff --git a/mem0-ts/src/oss/src/utils/factory.ts b/mem0-ts/src/oss/src/utils/factory.ts index 956cd665d..5a3e6b8e3 100644 --- a/mem0-ts/src/oss/src/utils/factory.ts +++ b/mem0-ts/src/oss/src/utils/factory.ts @@ -42,6 +42,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 { VertexAIVectorSearch } from "../vector_stores/vertex_ai_vector_search"; import { CassandraDB } from "../vector_stores/cassandra"; import { PineconeDB } from "../vector_stores/pinecone"; import { S3Vectors } from "../vector_stores/s3_vectors"; @@ -129,6 +130,8 @@ export class VectorStoreFactory { return new VectorizeDB(config as any); case "azure-ai-search": return new AzureAISearch(config as any); + case "vertex_ai_vector_search": + return new VertexAIVectorSearch(config as any); case "pgvector": return new PGVector(config as any); case "cassandra": diff --git a/mem0-ts/src/oss/src/vector_stores/vertex_ai_vector_search.ts b/mem0-ts/src/oss/src/vector_stores/vertex_ai_vector_search.ts new file mode 100644 index 000000000..23089fa4c --- /dev/null +++ b/mem0-ts/src/oss/src/vector_stores/vertex_ai_vector_search.ts @@ -0,0 +1,360 @@ +import { VectorStore } from "./base"; +import { SearchFilters, VectorStoreConfig, VectorStoreResult } from "../types"; + +export interface GoogleMatchingEngineConfig extends VectorStoreConfig { + projectId: string; + projectNumber: string; + region: string; + endpointId: string; + indexId: string; + deploymentIndexId: string; + collectionName?: string; + credentialsPath?: string; + serviceAccountJson?: Record; + vectorSearchApiEndpoint?: string; +} + +export class VertexAIVectorSearch implements VectorStore { + private config: GoogleMatchingEngineConfig; + private matchClient: any; + private indexClient: any; + private _initPromise?: Promise; + + constructor(config: GoogleMatchingEngineConfig) { + this.config = { ...config }; + if (!this.config.collectionName) { + this.config.collectionName = + this.config.indexId || this.config.deploymentIndexId; + } + this.initialize().catch(console.error); + } + + async initialize(): Promise { + if (!this._initPromise) { + this._initPromise = this._doInitialize(); + } + return this._initPromise; + } + + private async _doInitialize(): Promise { + try { + const aiplatform = await import("@google-cloud/aiplatform"); + const { MatchServiceClient, IndexServiceClient } = aiplatform.v1; + + const clientOptions: any = { + projectId: this.config.projectId, + apiEndpoint: + this.config.vectorSearchApiEndpoint || + `${this.config.region}-aiplatform.googleapis.com`, + }; + + if (this.config.credentialsPath) { + clientOptions.keyFilename = this.config.credentialsPath; + } else if (this.config.serviceAccountJson) { + clientOptions.credentials = { + client_email: this.config.serviceAccountJson.client_email, + private_key: this.config.serviceAccountJson.private_key, + }; + } + + this.matchClient = new MatchServiceClient(clientOptions); + this.indexClient = new IndexServiceClient(clientOptions); + } catch (error) { + console.error( + "Failed to initialize Vertex AI client. Make sure @google-cloud/aiplatform is installed.", + error, + ); + throw error; + } + } + + private _createRestriction(key: string, value: any): any { + return { + namespace: key, + allowList: [String(value)], + }; + } + + private _createDatapoint( + vectorId: string, + vector: number[], + payload: Record = {}, + ): any { + const restricts = Object.entries(payload).map(([key, value]) => + this._createRestriction(key, value), + ); + return { + datapointId: vectorId, + featureVector: vector, + restricts: restricts.length > 0 ? restricts : undefined, + }; + } + + private get indexPath(): string { + return `projects/${this.config.projectNumber}/locations/${this.config.region}/indexes/${this.config.indexId}`; + } + + private get indexEndpointPath(): string { + return `projects/${this.config.projectNumber}/locations/${this.config.region}/indexEndpoints/${this.config.endpointId}`; + } + + async insert( + vectors: number[][], + ids: string[], + payloads: Record[], + ): Promise { + await this.initialize(); + const datapoints = vectors.map((vector, i) => + this._createDatapoint(ids[i], vector, payloads[i] || {}), + ); + + await this.indexClient.upsertDatapoints({ + index: this.indexPath, + datapoints, + }); + } + + async search( + query: number[], + topK: number = 5, + filters?: SearchFilters, + ): Promise { + if (!this.config.vectorSearchApiEndpoint) { + throw new Error( + "vectorSearchApiEndpoint is required for search operation", + ); + } + await this.initialize(); + + const restricts: any[] = []; + if (filters) { + for (const [key, value] of Object.entries(filters)) { + if (typeof value === "object" && value !== null) { + const includes = value.include || []; + const excludes = value.exclude || []; + restricts.push({ + namespace: key, + allowList: includes.map(String), + denyList: excludes.map(String), + }); + } else { + restricts.push({ + namespace: key, + allowList: [String(value)], + }); + } + } + } + + const request = { + indexEndpoint: this.indexEndpointPath, + deployedIndexId: this.config.deploymentIndexId, + queries: [ + { + datapoint: { + featureVector: query, + restricts: restricts.length > 0 ? restricts : undefined, + }, + neighborCount: topK, + }, + ], + returnFullDatapoint: true, + }; + + const [response] = await this.matchClient.findNeighbors(request); + + if ( + !response || + !response.nearestNeighbors || + response.nearestNeighbors.length === 0 + ) { + return []; + } + + const neighbors = response.nearestNeighbors[0].neighbors || []; + return neighbors + .filter( + (neighbor: any) => + neighbor.datapoint?.datapointId !== "mem0-user-id-record", + ) + .map((neighbor: any) => { + const payload: Record = {}; + if (neighbor.datapoint?.restricts) { + for (const restrict of neighbor.datapoint.restricts) { + if (restrict.allowList && restrict.allowList.length > 0) { + payload[restrict.namespace] = restrict.allowList[0]; + } + } + } + + const score = + neighbor.distance !== undefined + ? Math.max(0.0, 1.0 - neighbor.distance) + : undefined; + return { + id: neighbor.datapoint.datapointId, + payload, + score, + }; + }); + } + + async get(vectorId: string): Promise { + if (!this.config.vectorSearchApiEndpoint) { + throw new Error("vectorSearchApiEndpoint is required for get operation"); + } + await this.initialize(); + + const request = { + indexEndpoint: this.indexEndpointPath, + deployedIndexId: this.config.deploymentIndexId, + queries: [ + { + datapoint: { + datapointId: vectorId, + }, + neighborCount: 1, + }, + ], + returnFullDatapoint: true, + }; + + const [response] = await this.matchClient.findNeighbors(request); + + if ( + !response || + !response.nearestNeighbors || + response.nearestNeighbors.length === 0 + ) { + return null; + } + + const neighbors = response.nearestNeighbors[0].neighbors || []; + if (neighbors.length === 0) { + return null; + } + + const neighbor = neighbors[0]; + if (neighbor.datapoint?.datapointId !== vectorId) { + return null; + } + + const payload: Record = {}; + if (neighbor.datapoint?.restricts) { + for (const restrict of neighbor.datapoint.restricts) { + if (restrict.allowList && restrict.allowList.length > 0) { + payload[restrict.namespace] = restrict.allowList[0]; + } + } + } + + const score = + neighbor.distance !== undefined + ? Math.max(0.0, 1.0 - neighbor.distance) + : undefined; + return { + id: neighbor.datapoint.datapointId, + payload, + score, + }; + } + + async keywordSearch( + query: string, + topK?: number, + filters?: SearchFilters, + ): Promise { + return null; + } + + async update( + vectorId: string, + vector: number[], + payload: Record, + ): Promise { + await this.initialize(); + + // Verify existence first + const existing = await this.get(vectorId); + if (!existing) { + console.warn(`Vector not found for id: ${vectorId}`); + return; + } + + const datapoint = this._createDatapoint(vectorId, vector, payload); + await this.indexClient.upsertDatapoints({ + index: this.indexPath, + datapoints: [datapoint], + }); + } + + async delete(vectorId: string): Promise { + await this.initialize(); + try { + await this.indexClient.removeDatapoints({ + index: this.indexPath, + datapointIds: [vectorId], + }); + } catch (error: any) { + // Ignore if not found + if (error.code !== 5) { + // 5 is NOT_FOUND in gRPC + throw error; + } + } + } + + async deleteCol(): Promise { + console.warn( + "Delete collection operation is not supported for Google Matching Engine", + ); + } + + async list( + filters?: SearchFilters, + topK: number = 10000, + ): Promise<[VectorStoreResult[], number]> { + await this.initialize(); + + // We do not have dimension natively, usually 768 or 1536. + // Vertex AI returns error if vector size is wrong, but some setups might allow empty vector. + // In Python SDK, it uses a zero vector of size 768. + const dimension = this.config.dimension || 768; + const zeroVector = Array(dimension).fill(0.0); + + const results = await this.search(zeroVector, topK, filters); + return [results, results.length]; + } + + async getUserId(): Promise { + // Vertex AI doesn't easily let us create arbitrary new indexes dynamically. + // So we use a special datapoint in the same index, with a specific ID. + const userIdDatapointId = "mem0-user-id-record"; + const existing = await this.get(userIdDatapointId); + if (existing && existing.payload?.user_id) { + return existing.payload.user_id; + } + + const randomUserId = + Math.random().toString(36).substring(2, 15) + + Math.random().toString(36).substring(2, 15); + + await this.setUserId(randomUserId); + return randomUserId; + } + + async setUserId(userId: string): Promise { + await this.initialize(); + const userIdDatapointId = "mem0-user-id-record"; + const dimension = this.config.dimension || 768; + const zeroVector = Array(dimension).fill(0.0); + + const datapoint = this._createDatapoint(userIdDatapointId, zeroVector, { + user_id: userId, + }); + await this.indexClient.upsertDatapoints({ + index: this.indexPath, + datapoints: [datapoint], + }); + } +} diff --git a/mem0-ts/src/oss/tests/vertex-ai-vector-search.test.ts b/mem0-ts/src/oss/tests/vertex-ai-vector-search.test.ts new file mode 100644 index 000000000..b1c24fdce --- /dev/null +++ b/mem0-ts/src/oss/tests/vertex-ai-vector-search.test.ts @@ -0,0 +1,242 @@ +/// +import { VertexAIVectorSearch } from "../src/vector_stores/vertex_ai_vector_search"; + +jest.mock("@google-cloud/aiplatform", () => { + const MatchServiceClient = jest.fn().mockImplementation(() => ({ + findNeighbors: jest.fn().mockResolvedValue([{ nearestNeighbors: [] }]), + })); + const IndexServiceClient = jest.fn().mockImplementation(() => ({ + upsertDatapoints: jest.fn().mockResolvedValue([{}]), + removeDatapoints: jest.fn().mockResolvedValue([{}]), + })); + return { + v1: { + MatchServiceClient, + IndexServiceClient, + }, + }; +}); + +describe("VertexAIVectorSearch", () => { + let store: VertexAIVectorSearch; + + beforeEach(() => { + store = new VertexAIVectorSearch({ + projectId: "test-project", + projectNumber: "123456789", + region: "us-central1", + endpointId: "test-endpoint", + indexId: "test-index", + deploymentIndexId: "test-deployment", + vectorSearchApiEndpoint: "test-api-endpoint", + }); + }); + + it("should initialize with correct collection name", () => { + expect((store as any).config.collectionName).toBe("test-index"); + }); + + it("should insert vectors", async () => { + await store.insert([[1, 2, 3]], ["id1"], [{ key: "value" }]); + expect((store as any).indexClient.upsertDatapoints).toHaveBeenCalled(); + }); + + it("should search vectors", async () => { + const results = await store.search([1, 2, 3], 5); + expect((store as any).matchClient.findNeighbors).toHaveBeenCalled(); + expect(results).toEqual([]); + }); + + it("should search vectors and return populated results", async () => { + const mockFindNeighbors = jest.fn().mockResolvedValue([ + { + nearestNeighbors: [ + { + neighbors: [ + { + datapoint: { + datapointId: "id1", + restricts: [{ namespace: "key", allowList: ["value"] }], + }, + distance: 0.1, + }, + ], + }, + ], + }, + ]); + (store as any).matchClient.findNeighbors = mockFindNeighbors; + + const results = await store.search([1, 2, 3], 5, { key: "value" }); + + expect(mockFindNeighbors).toHaveBeenCalled(); + // It should map payload correctly and score should be 1.0 - distance + expect(results).toEqual([ + { + id: "id1", + payload: { key: "value" }, + score: 0.9, + }, + ]); + }); + + it("should get vector by id", async () => { + const result = await store.get("id1"); + expect((store as any).matchClient.findNeighbors).toHaveBeenCalled(); + expect(result).toBeNull(); + }); + + it("should build allowList/denyList restricts from include/exclude filters", async () => { + const spy = jest.fn().mockResolvedValue([{ nearestNeighbors: [] }]); + (store as any).matchClient.findNeighbors = spy; + + await store.search([1, 2, 3], 5, { + key: { include: ["a"], exclude: ["b"] }, + }); + + const request = spy.mock.calls[0][0]; + expect(request.queries[0].datapoint.restricts).toEqual([ + { namespace: "key", allowList: ["a"], denyList: ["b"] }, + ]); + }); + + it("should exclude the mem0-user-id-record sentinel from search results", async () => { + (store as any).matchClient.findNeighbors = jest.fn().mockResolvedValue([ + { + nearestNeighbors: [ + { + neighbors: [ + { + datapoint: { datapointId: "mem0-user-id-record" }, + distance: 0.0, + }, + { datapoint: { datapointId: "id1" }, distance: 0.1 }, + ], + }, + ], + }, + ]); + + const results = await store.search([1, 2, 3], 5); + expect(results.map((r) => r.id)).toEqual(["id1"]); + }); + + it("should update an existing vector", async () => { + (store as any).matchClient.findNeighbors = jest.fn().mockResolvedValue([ + { + nearestNeighbors: [ + { + neighbors: [{ datapoint: { datapointId: "id1" }, distance: 0.1 }], + }, + ], + }, + ]); + + await store.update("id1", [4, 5, 6], { key: "new" }); + expect((store as any).indexClient.upsertDatapoints).toHaveBeenCalled(); + }); + + it("should skip update when the vector does not exist", async () => { + const warnSpy = jest.spyOn(console, "warn").mockImplementation(() => {}); + + // default findNeighbors returns empty → get() resolves null + await store.update("missing", [4, 5, 6], { key: "new" }); + + expect((store as any).indexClient.upsertDatapoints).not.toHaveBeenCalled(); + expect(warnSpy).toHaveBeenCalled(); + warnSpy.mockRestore(); + }); + + it("should delete a vector by id", async () => { + await store.delete("id1"); + expect((store as any).indexClient.removeDatapoints).toHaveBeenCalledWith({ + index: expect.any(String), + datapointIds: ["id1"], + }); + }); + + it("should ignore NOT_FOUND (gRPC code 5) on delete", async () => { + (store as any).indexClient.removeDatapoints = jest + .fn() + .mockRejectedValue({ code: 5 }); + await expect(store.delete("missing")).resolves.toBeUndefined(); + }); + + it("should rethrow non-NOT_FOUND errors on delete", async () => { + (store as any).indexClient.removeDatapoints = jest + .fn() + .mockRejectedValue({ code: 13 }); + await expect(store.delete("id1")).rejects.toEqual({ code: 13 }); + }); + + it("should list vectors via a zero-vector search", async () => { + (store as any).matchClient.findNeighbors = jest.fn().mockResolvedValue([ + { + nearestNeighbors: [ + { + neighbors: [{ datapoint: { datapointId: "id1" }, distance: 0.2 }], + }, + ], + }, + ]); + + const [results, count] = await store.list(); + expect((store as any).matchClient.findNeighbors).toHaveBeenCalled(); + expect(count).toBe(1); + expect(results[0].id).toBe("id1"); + }); + + it("should return null for keywordSearch (unsupported)", async () => { + await expect(store.keywordSearch("hello")).resolves.toBeNull(); + }); + + it("should warn and no-op on deleteCol (unsupported)", async () => { + const warnSpy = jest.spyOn(console, "warn").mockImplementation(() => {}); + await store.deleteCol(); + expect(warnSpy).toHaveBeenCalled(); + warnSpy.mockRestore(); + }); + + it("should set a user id via the sentinel datapoint", async () => { + await store.setUserId("user-123"); + + expect((store as any).indexClient.upsertDatapoints).toHaveBeenCalled(); + const request = (store as any).indexClient.upsertDatapoints.mock + .calls[0][0]; + expect(request.datapoints[0].datapointId).toBe("mem0-user-id-record"); + }); + + it("should return an existing user id from the sentinel record", async () => { + (store as any).matchClient.findNeighbors = jest.fn().mockResolvedValue([ + { + nearestNeighbors: [ + { + neighbors: [ + { + datapoint: { + datapointId: "mem0-user-id-record", + restricts: [ + { namespace: "user_id", allowList: ["existing-user"] }, + ], + }, + distance: 0.0, + }, + ], + }, + ], + }, + ]); + + const userId = await store.getUserId(); + expect(userId).toBe("existing-user"); + }); + + it("should generate and persist a new user id when none exists", async () => { + // default findNeighbors empty → get() null → generate + setUserId + const userId = await store.getUserId(); + + expect(typeof userId).toBe("string"); + expect(userId.length).toBeGreaterThan(0); + expect((store as any).indexClient.upsertDatapoints).toHaveBeenCalled(); + }); +}); diff --git a/mem0-ts/tsup.config.ts b/mem0-ts/tsup.config.ts index 3d70b13d9..f4f7216ee 100644 --- a/mem0-ts/tsup.config.ts +++ b/mem0-ts/tsup.config.ts @@ -17,6 +17,7 @@ const external = [ "iovalkey", "ollama", "@google/genai", + "@google-cloud/aiplatform", "@mistralai/mistralai", "@supabase/supabase-js", "@azure/search-documents",