feat(ts-sdk): add Upstash Vector vector store (#5811)

Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
AxelRay
2026-07-06 23:41:05 +05:30
committed by GitHub
co-authored by kartik-mem0
parent c944bed460
commit 9d36b2c94d
12 changed files with 501 additions and 1 deletions
+1
View File
@@ -16,6 +16,7 @@
"@anthropic-ai/sdk": "^0.18.0",
"@google/genai": "^0.7.0",
"@qdrant/js-client-rest": "^1.13.0",
"@upstash/vector": "^1.2.3",
"@types/node": "^20.11.19",
"@types/pg": "^8.11.0",
"@types/redis": "^4.0.10",
+3
View File
@@ -32,8 +32,11 @@ 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/upstash_vector";
export * from "./vector_stores/azure_mysql";
export * from "./vector_stores/cassandra";
export * from "./vector_stores/s3_vectors";
export * from "./vector_stores/vertex_ai_vector_search";
export * from "./vector_stores/pinecone";
export * from "./vector_stores/turbopuffer";
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 { UpstashVector } from "../vector_stores/upstash_vector";
import { AzureMySQLDB } from "../vector_stores/azure_mysql";
import { VertexAIVectorSearch } from "../vector_stores/vertex_ai_vector_search";
import { CassandraDB } from "../vector_stores/cassandra";
@@ -136,6 +137,8 @@ export class VectorStoreFactory {
return new VertexAIVectorSearch(config as any);
case "pgvector":
return new PGVector(config as any);
case "upstash_vector":
return new UpstashVector(config as any);
case "azure_mysql":
return new AzureMySQLDB(config as any);
case "cassandra":
@@ -0,0 +1,240 @@
import { Index, QueryResult, Vector } from "@upstash/vector";
import { VectorStore } from "./base";
import { SearchFilters, VectorStoreConfig, VectorStoreResult } from "../types";
interface UpstashVectorConfig extends VectorStoreConfig {
collectionName: string;
url?: string;
token?: string;
client?: Index<Record<string, unknown>>;
}
type UpstashMetadata = Record<string, unknown>;
export class UpstashVector implements VectorStore {
private readonly client: Index<UpstashMetadata>;
private readonly collectionName: string;
constructor(config: UpstashVectorConfig) {
if (!config.collectionName) {
throw new Error("collectionName is required for Upstash Vector.");
}
if (config.client) {
this.client = config.client;
} else if (config.url && config.token) {
this.client = new Index({
url: config.url,
token: config.token,
});
} else {
throw new Error("Either a client or url and token must be provided.");
}
this.collectionName = config.collectionName;
}
async initialize(): Promise<void> {
return;
}
async insert(
vectors: number[][],
ids: string[],
payloads: Record<string, any>[],
): Promise<void> {
const upsertData = vectors.map((vector, idx) => {
return {
id: ids[idx],
vector,
metadata: payloads[idx] ?? {},
};
});
await this.client.upsert(upsertData, { namespace: this.collectionName });
}
async search(
query: number[],
topK: number = 5,
filters?: SearchFilters,
): Promise<VectorStoreResult[]> {
const response = await this.client.query<UpstashMetadata>(
{
vector: query,
topK,
filter: this.convertFilters(filters),
includeMetadata: true,
},
{ namespace: this.collectionName },
);
return response.map((result) => this.parseResult(result));
}
async keywordSearch(
query: string,
topK: number = 5,
filters?: SearchFilters,
): Promise<VectorStoreResult[] | null> {
try {
const response = await this.client.query<UpstashMetadata>(
{
data: query,
topK,
filter: this.convertFilters(filters),
includeMetadata: true,
},
{ namespace: this.collectionName },
);
return response.map((result) => this.parseResult(result));
} catch (error) {
console.error(`Error during keyword search for query '${query}':`, error);
return null;
}
}
async get(vectorId: string): Promise<VectorStoreResult | null> {
const response = await this.client.fetch<UpstashMetadata>([vectorId], {
includeMetadata: true,
namespace: this.collectionName,
});
const vector = response[0];
if (!vector) {
return null;
}
return {
id: String(vector.id),
payload: (vector.metadata ?? {}) as Record<string, any>,
};
}
async update(
vectorId: string,
vector: number[],
payload: Record<string, any>,
): Promise<void> {
// Upstash's `update` can't set the vector and metadata in one call (its
// payload is a discriminated union of vector | data | metadata), so a
// single `upsert` replaces both atomically, the same way insert() writes.
await this.client.upsert(
{
id: vectorId,
vector,
metadata: payload,
},
{ namespace: this.collectionName },
);
}
async delete(vectorId: string): Promise<void> {
await this.client.delete(vectorId, { namespace: this.collectionName });
}
async deleteCol(): Promise<void> {
await this.client.reset({ namespace: this.collectionName });
}
async list(
filters?: SearchFilters,
topK: number = 100,
): Promise<[VectorStoreResult[], number]> {
const results: VectorStoreResult[] = [];
let cursor = "0";
do {
const response = await this.client.range<UpstashMetadata>(
{
cursor,
limit: Math.min(100, topK - results.length),
includeMetadata: true,
},
{ namespace: this.collectionName },
);
for (const vector of response.vectors) {
if (this.matchesFilters(vector, filters)) {
results.push({
id: String(vector.id),
payload: (vector.metadata ?? {}) as Record<string, any>,
});
}
if (results.length >= topK) {
break;
}
}
cursor = response.nextCursor;
// Upstash returns an empty-string cursor once the scan is exhausted (it
// never comes back as "0"), so "" is the termination sentinel. Checking
// for "0" here would re-scan from the start and return duplicates.
} while (cursor !== "" && results.length < topK);
return [results, results.length];
}
async getUserId(): Promise<string> {
return "anonymous-upstash-vector";
}
async setUserId(): Promise<void> {
return;
}
async reset(): Promise<void> {
await this.deleteCol();
}
private parseResult(result: QueryResult<UpstashMetadata>): VectorStoreResult {
return {
id: String(result.id),
payload: (result.metadata ?? {}) as Record<string, any>,
score: result.score,
};
}
private stringifyFilterValue(value: unknown): string {
if (typeof value === "string") {
return JSON.stringify(value);
}
if (typeof value === "boolean") {
return value ? "true" : "false";
}
return String(value);
}
private convertFilters(filters?: SearchFilters): string | undefined {
if (!filters) {
return undefined;
}
const expressions = Object.entries(filters)
.filter(([, value]) => value !== undefined && value !== null)
.map(([key, value]) => `${key} = ${this.stringifyFilterValue(value)}`);
return expressions.length > 0 ? expressions.join(" AND ") : undefined;
}
private matchesFilters(
vector: Vector<UpstashMetadata>,
filters?: SearchFilters,
): boolean {
if (!filters) {
return true;
}
return Object.entries(filters).every(([key, value]) => {
if (value === undefined || value === null) {
return true;
}
return vector.metadata?.[key] === value;
});
}
}
@@ -158,6 +158,11 @@ jest.mock("../src/vector_stores/pgvector", () => ({
.fn()
.mockImplementation((config) => ({ type: "pgvector", config })),
}));
jest.mock("../src/vector_stores/upstash_vector", () => ({
UpstashVector: jest
.fn()
.mockImplementation((config) => ({ type: "upstash-vector", config })),
}));
jest.mock("../src/vector_stores/azure_mysql", () => ({
AzureMySQLDB: jest
.fn()
@@ -299,6 +304,7 @@ describe("VectorStoreFactory", () => {
["vectorize"],
["azure-ai-search"],
["pgvector"],
["upstash_vector"],
["azure_mysql"],
["cassandra"],
["s3-vectors"],
@@ -0,0 +1,166 @@
import { UpstashVector } from "../src/vector_stores/upstash_vector";
describe("UpstashVector", () => {
const namespace = "memories";
function createClient(overrides: Record<string, jest.Mock> = {}) {
return {
upsert: jest.fn().mockResolvedValue("Success"),
query: jest.fn().mockResolvedValue([]),
fetch: jest.fn().mockResolvedValue([]),
update: jest.fn().mockResolvedValue({ updated: 1 }),
delete: jest.fn().mockResolvedValue({ deleted: 1 }),
reset: jest.fn().mockResolvedValue("Success"),
range: jest.fn().mockResolvedValue({ vectors: [], nextCursor: "" }),
...overrides,
};
}
it("upserts vectors into the collection namespace", async () => {
const client = createClient();
const store = new UpstashVector({
collectionName: namespace,
client: client as any,
});
await store.insert(
[[0.1, 0.2]],
["memory-1"],
[{ data: "hello", user_id: "user-1" }],
);
expect(client.upsert).toHaveBeenCalledWith(
[
{
id: "memory-1",
vector: [0.1, 0.2],
metadata: { data: "hello", user_id: "user-1" },
},
],
{ namespace },
);
});
it("queries vectors with converted filters", async () => {
const client = createClient({
query: jest.fn().mockResolvedValue([
{
id: "memory-1",
score: 0.9,
metadata: { data: "hello", user_id: "user-1" },
},
]),
});
const store = new UpstashVector({
collectionName: namespace,
client: client as any,
});
const results = await store.search([0.1, 0.2], 3, {
user_id: "user-1",
active: true,
});
expect(client.query).toHaveBeenCalledWith(
{
vector: [0.1, 0.2],
topK: 3,
filter: 'user_id = "user-1" AND active = true',
includeMetadata: true,
},
{ namespace },
);
expect(results).toEqual([
{
id: "memory-1",
payload: { data: "hello", user_id: "user-1" },
score: 0.9,
},
]);
});
it("fetches, updates, deletes, resets, and lists vectors in the namespace", async () => {
const client = createClient({
fetch: jest.fn().mockResolvedValue([
{
id: "memory-1",
metadata: { data: "hello" },
},
]),
range: jest
.fn()
.mockResolvedValueOnce({
vectors: [
{ id: "memory-1", metadata: { user_id: "user-1" } },
{ id: "memory-2", metadata: { user_id: "user-2" } },
],
nextCursor: "2",
})
.mockResolvedValueOnce({
vectors: [{ id: "memory-3", metadata: { user_id: "user-1" } }],
nextCursor: "",
}),
});
const store = new UpstashVector({
collectionName: namespace,
client: client as any,
});
await expect(store.get("memory-1")).resolves.toEqual({
id: "memory-1",
payload: { data: "hello" },
});
await store.update("memory-1", [0.3], { data: "updated" });
await store.delete("memory-1");
await store.deleteCol();
await store.reset();
await expect(store.list({ user_id: "user-1" }, 2)).resolves.toEqual([
[
{ id: "memory-1", payload: { user_id: "user-1" } },
{ id: "memory-3", payload: { user_id: "user-1" } },
],
2,
]);
expect(client.fetch).toHaveBeenCalledWith(["memory-1"], {
includeMetadata: true,
namespace,
});
expect(client.upsert).toHaveBeenCalledWith(
{ id: "memory-1", vector: [0.3], metadata: { data: "updated" } },
{ namespace },
);
expect(client.delete).toHaveBeenCalledWith("memory-1", { namespace });
expect(client.reset).toHaveBeenCalledTimes(2);
expect(client.reset).toHaveBeenCalledWith({ namespace });
});
it("stops paging when the cursor is exhausted instead of re-scanning", async () => {
// Upstash returns nextCursor "" at the end of a scan. If list() treated ""
// as "keep going" (e.g. by checking for "0"), it would re-fetch from the
// start and pile up duplicates until it hit topK. With fewer vectors than
// topK, a single page must end the scan: one range call, no duplicates.
const client = createClient({
range: jest.fn().mockResolvedValue({
vectors: [
{ id: "memory-1", metadata: { user_id: "user-1" } },
{ id: "memory-2", metadata: { user_id: "user-1" } },
],
nextCursor: "",
}),
});
const store = new UpstashVector({
collectionName: namespace,
client: client as any,
});
const [rows, count] = await store.list({ user_id: "user-1" }, 100);
expect(client.range).toHaveBeenCalledTimes(1);
expect(rows).toEqual([
{ id: "memory-1", payload: { user_id: "user-1" } },
{ id: "memory-2", payload: { user_id: "user-1" } },
]);
expect(count).toBe(2);
});
});