feat: add Google Vertex AI Vector Search support to vector store factory (#5791)

Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
Co-authored-by: divyansh-1009 <divyansh-1009@users.noreply.github.com>
This commit is contained in:
Div
2026-07-06 22:54:24 +05:30
committed by GitHub
parent 7fb3feb5cd
commit 2bc2f763d9
8 changed files with 813 additions and 3 deletions
+49 -3
View File
@@ -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
<CodeGroup>
```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" },
});
```
</CodeGroup>
### Required Parameters
<Tabs>
<Tab title="Python">
| 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` |
</Tab>
<Tab title="TypeScript">
| 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 |
</Tab>
</Tabs>
+1
View File
@@ -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",
+156
View File
@@ -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
+1
View File
@@ -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";
+3
View File
@@ -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":
@@ -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<string, any>;
vectorSearchApiEndpoint?: string;
}
export class VertexAIVectorSearch implements VectorStore {
private config: GoogleMatchingEngineConfig;
private matchClient: any;
private indexClient: any;
private _initPromise?: Promise<void>;
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<void> {
if (!this._initPromise) {
this._initPromise = this._doInitialize();
}
return this._initPromise;
}
private async _doInitialize(): Promise<void> {
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<string, any> = {},
): 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<string, any>[],
): Promise<void> {
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<VectorStoreResult[]> {
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<string, any> = {};
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<VectorStoreResult | null> {
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<string, any> = {};
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<VectorStoreResult[] | null> {
return null;
}
async update(
vectorId: string,
vector: number[],
payload: Record<string, any>,
): Promise<void> {
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<void> {
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<void> {
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<string> {
// 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<void> {
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],
});
}
}
@@ -0,0 +1,242 @@
/// <reference types="jest" />
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();
});
});
+1
View File
@@ -17,6 +17,7 @@ const external = [
"iovalkey",
"ollama",
"@google/genai",
"@google-cloud/aiplatform",
"@mistralai/mistralai",
"@supabase/supabase-js",
"@azure/search-documents",