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
co-authored by
kartik-mem0
divyansh-1009
kartik-mem0
divyansh-1009
parent
7fb3feb5cd
commit
2bc2f763d9
@@ -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";
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user