feat(ts-sdk): add OpenSearch vector store (#5810)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
AxelRay
@@ -15,6 +15,7 @@
|
||||
"dependencies": {
|
||||
"@anthropic-ai/sdk": "^0.18.0",
|
||||
"@google/genai": "^0.7.0",
|
||||
"@opensearch-project/opensearch": "^3.5.1",
|
||||
"@qdrant/js-client-rest": "^1.13.0",
|
||||
"@upstash/vector": "^1.2.3",
|
||||
"@types/node": "^20.11.19",
|
||||
|
||||
@@ -39,4 +39,5 @@ 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 "./vector_stores/opensearch";
|
||||
export * from "./utils/factory";
|
||||
|
||||
@@ -43,6 +43,7 @@ import { LangchainEmbedder } from "../embeddings/langchain";
|
||||
import { LangchainVectorStore } from "../vector_stores/langchain";
|
||||
import { AzureAISearch } from "../vector_stores/azure_ai_search";
|
||||
import { PGVector } from "../vector_stores/pgvector";
|
||||
import { OpenSearchDB } from "../vector_stores/opensearch";
|
||||
import { UpstashVector } from "../vector_stores/upstash_vector";
|
||||
import { AzureMySQLDB } from "../vector_stores/azure_mysql";
|
||||
import { VertexAIVectorSearch } from "../vector_stores/vertex_ai_vector_search";
|
||||
@@ -140,6 +141,8 @@ export class VectorStoreFactory {
|
||||
return new VertexAIVectorSearch(config as any);
|
||||
case "pgvector":
|
||||
return new PGVector(config as any);
|
||||
case "opensearch":
|
||||
return new OpenSearchDB(config as any);
|
||||
case "upstash_vector":
|
||||
return new UpstashVector(config as any);
|
||||
case "azure_mysql":
|
||||
|
||||
@@ -0,0 +1,621 @@
|
||||
import { Client } from "@opensearch-project/opensearch";
|
||||
import { VectorStore } from "./base";
|
||||
import { SearchFilters, VectorStoreConfig, VectorStoreResult } from "../types";
|
||||
|
||||
type OpenSearchAuth =
|
||||
| {
|
||||
username?: string;
|
||||
password?: string;
|
||||
}
|
||||
| {
|
||||
user?: string;
|
||||
password?: string;
|
||||
}
|
||||
| Record<string, any>;
|
||||
|
||||
interface OpenSearchConfig extends VectorStoreConfig {
|
||||
client?: Client;
|
||||
host?: string;
|
||||
port?: number;
|
||||
httpAuth?: OpenSearchAuth | [string, string];
|
||||
user?: string;
|
||||
password?: string;
|
||||
useSSL?: boolean;
|
||||
verifyCerts?: boolean;
|
||||
collectionName: string;
|
||||
embeddingModelDims: number;
|
||||
autoRefresh?: boolean;
|
||||
}
|
||||
|
||||
interface OpenSearchHit {
|
||||
_id?: string;
|
||||
_score?: number;
|
||||
_source?: {
|
||||
id?: string;
|
||||
vector_field?: number[];
|
||||
payload?: Record<string, any>;
|
||||
};
|
||||
}
|
||||
|
||||
const KEY_MAP: Record<string, string> = {
|
||||
$and: "AND",
|
||||
$or: "OR",
|
||||
$not: "NOT",
|
||||
};
|
||||
|
||||
function responseBody<T>(response: T | { body: T }): T {
|
||||
if (
|
||||
response &&
|
||||
typeof response === "object" &&
|
||||
"body" in response &&
|
||||
(response as { body: T }).body !== undefined
|
||||
) {
|
||||
return (response as { body: T }).body;
|
||||
}
|
||||
|
||||
return response as T;
|
||||
}
|
||||
|
||||
function escapeWildcard(value: string): string {
|
||||
return value.replace(/([\\*?])/g, "\\$1");
|
||||
}
|
||||
|
||||
export class OpenSearchDB implements VectorStore {
|
||||
private client: Client;
|
||||
private readonly collectionName: string;
|
||||
private readonly embeddingModelDims: number;
|
||||
private readonly autoRefresh: boolean;
|
||||
private _initPromise?: Promise<void>;
|
||||
|
||||
constructor(config: OpenSearchConfig) {
|
||||
this.collectionName = config.collectionName;
|
||||
this.embeddingModelDims = config.embeddingModelDims;
|
||||
this.autoRefresh = config.autoRefresh ?? false;
|
||||
|
||||
if (config.client) {
|
||||
this.client = config.client;
|
||||
} else {
|
||||
const useSSL = config.useSSL ?? false;
|
||||
const host = config.host || "localhost";
|
||||
const port = config.port || 9200;
|
||||
const auth =
|
||||
config.httpAuth ||
|
||||
(config.user && config.password
|
||||
? { username: config.user, password: config.password }
|
||||
: undefined);
|
||||
|
||||
this.client = new Client({
|
||||
node: `${useSSL ? "https" : "http"}://${host}:${port}`,
|
||||
auth: this.normalizeAuth(auth),
|
||||
ssl: {
|
||||
// Default false to match the Python SDK: self-hosted OpenSearch
|
||||
// commonly uses self-signed certs, so verification is opt-in.
|
||||
rejectUnauthorized: config.verifyCerts ?? false,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
this.initialize().catch(console.error);
|
||||
}
|
||||
|
||||
async initialize(): Promise<void> {
|
||||
if (!this._initPromise) {
|
||||
this._initPromise = this._doInitialize();
|
||||
}
|
||||
|
||||
return this._initPromise;
|
||||
}
|
||||
|
||||
private async _doInitialize(): Promise<void> {
|
||||
await this.createCol(this.collectionName, this.embeddingModelDims);
|
||||
await this.ensureMigrationIndex();
|
||||
}
|
||||
|
||||
private normalizeAuth(auth?: OpenSearchAuth | [string, string]): any {
|
||||
if (!auth) return undefined;
|
||||
|
||||
if (Array.isArray(auth)) {
|
||||
return {
|
||||
username: auth[0],
|
||||
password: auth[1],
|
||||
};
|
||||
}
|
||||
|
||||
if ("user" in auth && auth.user) {
|
||||
const { user, ...rest } = auth;
|
||||
return {
|
||||
...rest,
|
||||
username: user,
|
||||
};
|
||||
}
|
||||
|
||||
return auth;
|
||||
}
|
||||
|
||||
private async indexExists(index: string): Promise<boolean> {
|
||||
const response = await this.client.indices.exists({ index });
|
||||
return Boolean(responseBody<boolean>(response));
|
||||
}
|
||||
|
||||
private async createCol(name: string, vectorSize: number): Promise<void> {
|
||||
if (await this.indexExists(name)) return;
|
||||
|
||||
await this.client.indices.create({
|
||||
index: name,
|
||||
body: {
|
||||
settings: {
|
||||
index: {
|
||||
knn: true,
|
||||
},
|
||||
},
|
||||
mappings: {
|
||||
properties: {
|
||||
vector_field: {
|
||||
type: "knn_vector",
|
||||
dimension: vectorSize,
|
||||
method: {
|
||||
engine: "nmslib",
|
||||
name: "hnsw",
|
||||
space_type: "cosinesimil",
|
||||
},
|
||||
},
|
||||
// payload is a dynamic object so that string sub-fields (user_id,
|
||||
// agent_id, run_id, and any user metadata) each get an automatic
|
||||
// `.keyword` sub-field, which the filter builder queries for exact
|
||||
// match term/terms clauses. Mapping these as explicit `keyword`
|
||||
// fields would remove the `.keyword` sub-field and make every
|
||||
// scoped filter silently match nothing.
|
||||
payload: { type: "object" },
|
||||
id: { type: "keyword" },
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
private async ensureMigrationIndex(): Promise<void> {
|
||||
if (await this.indexExists("memory_migrations")) return;
|
||||
|
||||
await this.client.indices.create({
|
||||
index: "memory_migrations",
|
||||
body: {
|
||||
mappings: {
|
||||
properties: {
|
||||
user_id: { type: "keyword" },
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
private validateVector(vector: number[], index: number): void {
|
||||
if (!vector) {
|
||||
throw new Error(`Vector at index ${index} is null or undefined.`);
|
||||
}
|
||||
if (vector.length === 0) {
|
||||
throw new Error(
|
||||
`Vector at index ${index} is empty. Expected dimension ${this.embeddingModelDims}.`,
|
||||
);
|
||||
}
|
||||
if (vector.length !== this.embeddingModelDims) {
|
||||
throw new Error(
|
||||
`Vector at index ${index} has dimension ${vector.length}, but index ` +
|
||||
`'${this.collectionName}' expects dimension ${this.embeddingModelDims}.`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async insert(
|
||||
vectors: number[][],
|
||||
ids: string[],
|
||||
payloads: Record<string, any>[],
|
||||
): Promise<void> {
|
||||
vectors.forEach((vector, index) => this.validateVector(vector, index));
|
||||
|
||||
const operations = vectors.flatMap((vector, index) => {
|
||||
const id = ids[index] || String(index);
|
||||
return [
|
||||
{ index: { _index: this.collectionName, _id: id } },
|
||||
{
|
||||
vector_field: vector,
|
||||
payload: payloads[index] || {},
|
||||
id,
|
||||
},
|
||||
];
|
||||
});
|
||||
|
||||
if (operations.length === 0) return;
|
||||
|
||||
const response = responseBody<{ errors?: boolean; items?: any[] }>(
|
||||
await this.client.bulk({
|
||||
refresh: this.autoRefresh,
|
||||
body: operations,
|
||||
}),
|
||||
);
|
||||
|
||||
if (response.errors) {
|
||||
const failedItem = response.items?.find((item) => {
|
||||
const action = item.index || item.create || item.update || item.delete;
|
||||
return action?.error;
|
||||
});
|
||||
|
||||
throw new Error(
|
||||
`OpenSearch bulk insert failed: ${JSON.stringify(failedItem)}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async keywordSearch(
|
||||
query: string,
|
||||
topK: number = 5,
|
||||
filters?: SearchFilters,
|
||||
): Promise<VectorStoreResult[] | null> {
|
||||
const boolQuery: Record<string, any> = {
|
||||
should: [
|
||||
{ match: { "payload.data": query } },
|
||||
{ match: { "payload.text_lemmatized": query } },
|
||||
],
|
||||
minimum_should_match: 1,
|
||||
};
|
||||
|
||||
const filter = this.buildFilterClauses(filters);
|
||||
if (filter.length) {
|
||||
boolQuery.filter = filter;
|
||||
}
|
||||
|
||||
const response = responseBody<{ hits: { hits: OpenSearchHit[] } }>(
|
||||
(await this.client.search({
|
||||
index: this.collectionName,
|
||||
body: {
|
||||
size: topK,
|
||||
query: { bool: boolQuery },
|
||||
},
|
||||
})) as any,
|
||||
);
|
||||
|
||||
return response.hits.hits.map((hit) => this.hitToResult(hit));
|
||||
}
|
||||
|
||||
async search(
|
||||
query: number[],
|
||||
topK: number = 5,
|
||||
filters?: SearchFilters,
|
||||
): Promise<VectorStoreResult[]> {
|
||||
const knnQuery = {
|
||||
knn: {
|
||||
vector_field: {
|
||||
vector: query,
|
||||
k: topK * 2,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const filter = this.buildFilterClauses(filters);
|
||||
const searchQuery = filter.length
|
||||
? {
|
||||
bool: {
|
||||
must: knnQuery,
|
||||
filter,
|
||||
},
|
||||
}
|
||||
: knnQuery;
|
||||
|
||||
const response = responseBody<{ hits: { hits: OpenSearchHit[] } }>(
|
||||
(await this.client.search({
|
||||
index: this.collectionName,
|
||||
body: {
|
||||
size: topK * 2,
|
||||
query: searchQuery,
|
||||
},
|
||||
})) as any,
|
||||
);
|
||||
|
||||
return response.hits.hits
|
||||
.slice(0, topK)
|
||||
.map((hit) => this.hitToResult(hit));
|
||||
}
|
||||
|
||||
async get(vectorId: string): Promise<VectorStoreResult | null> {
|
||||
try {
|
||||
const response = responseBody<{ _source?: OpenSearchHit["_source"] }>(
|
||||
await this.client.get({
|
||||
index: this.collectionName,
|
||||
id: vectorId,
|
||||
}),
|
||||
);
|
||||
|
||||
if (!response._source) return null;
|
||||
|
||||
return {
|
||||
id: response._source.id || vectorId,
|
||||
payload: response._source.payload || {},
|
||||
};
|
||||
} catch (error: any) {
|
||||
if (error?.statusCode === 404 || error?.meta?.statusCode === 404) {
|
||||
return null;
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
async update(
|
||||
vectorId: string,
|
||||
vector: number[],
|
||||
payload: Record<string, any>,
|
||||
): Promise<void> {
|
||||
if (vector) {
|
||||
this.validateVector(vector, 0);
|
||||
}
|
||||
|
||||
await this.client.update({
|
||||
index: this.collectionName,
|
||||
id: vectorId,
|
||||
body: {
|
||||
doc: {
|
||||
...(vector && { vector_field: vector }),
|
||||
...(payload && { payload }),
|
||||
id: vectorId,
|
||||
},
|
||||
},
|
||||
refresh: this.autoRefresh,
|
||||
});
|
||||
}
|
||||
|
||||
async delete(vectorId: string): Promise<void> {
|
||||
try {
|
||||
await this.client.delete({
|
||||
index: this.collectionName,
|
||||
id: vectorId,
|
||||
refresh: this.autoRefresh,
|
||||
});
|
||||
} catch (error: any) {
|
||||
if (error?.statusCode === 404 || error?.meta?.statusCode === 404) {
|
||||
return;
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
async deleteCol(): Promise<void> {
|
||||
if (!(await this.indexExists(this.collectionName))) return;
|
||||
await this.client.indices.delete({ index: this.collectionName });
|
||||
}
|
||||
|
||||
async list(
|
||||
filters?: SearchFilters,
|
||||
topK: number = 100,
|
||||
): Promise<[VectorStoreResult[], number]> {
|
||||
const filter = this.buildFilterClauses(filters);
|
||||
const query = filter.length ? { bool: { filter } } : { match_all: {} };
|
||||
|
||||
const response = responseBody<{
|
||||
hits: {
|
||||
total?: number | { value: number };
|
||||
hits: OpenSearchHit[];
|
||||
};
|
||||
}>(
|
||||
(await this.client.search({
|
||||
index: this.collectionName,
|
||||
body: {
|
||||
size: topK,
|
||||
query,
|
||||
},
|
||||
})) as any,
|
||||
);
|
||||
|
||||
const results = response.hits.hits.map((hit) => this.hitToResult(hit));
|
||||
const total =
|
||||
typeof response.hits.total === "number"
|
||||
? response.hits.total
|
||||
: response.hits.total?.value || results.length;
|
||||
|
||||
return [results, total];
|
||||
}
|
||||
|
||||
async reset(): Promise<void> {
|
||||
await this.deleteCol();
|
||||
await this.createCol(this.collectionName, this.embeddingModelDims);
|
||||
}
|
||||
|
||||
async getUserId(): Promise<string> {
|
||||
await this.ensureMigrationIndex();
|
||||
|
||||
const response = responseBody<{ hits: { hits: OpenSearchHit[] } }>(
|
||||
(await this.client.search({
|
||||
index: "memory_migrations",
|
||||
body: {
|
||||
size: 1,
|
||||
query: { match_all: {} },
|
||||
},
|
||||
})) as any,
|
||||
);
|
||||
|
||||
const existing = response.hits.hits[0]?._source?.payload?.user_id;
|
||||
if (existing) return String(existing);
|
||||
|
||||
const userId =
|
||||
Math.random().toString(36).substring(2, 15) +
|
||||
Math.random().toString(36).substring(2, 15);
|
||||
|
||||
await this.setUserId(userId);
|
||||
return userId;
|
||||
}
|
||||
|
||||
async setUserId(userId: string): Promise<void> {
|
||||
await this.ensureMigrationIndex();
|
||||
await this.client.index({
|
||||
index: "memory_migrations",
|
||||
id: "user_id",
|
||||
body: {
|
||||
user_id: userId,
|
||||
payload: { user_id: userId },
|
||||
},
|
||||
refresh: this.autoRefresh,
|
||||
});
|
||||
}
|
||||
|
||||
private hitToResult(hit: OpenSearchHit): VectorStoreResult {
|
||||
return {
|
||||
id: hit._source?.id || hit._id || "",
|
||||
payload: hit._source?.payload || {},
|
||||
score: hit._score,
|
||||
};
|
||||
}
|
||||
|
||||
private buildFilterClauses(filters?: SearchFilters): any[] {
|
||||
if (!filters || Object.keys(filters).length === 0) return [];
|
||||
|
||||
const normalized: Record<string, any> = {};
|
||||
for (const [key, value] of Object.entries(filters)) {
|
||||
normalized[KEY_MAP[key] || key] = value;
|
||||
}
|
||||
|
||||
return Object.entries(normalized)
|
||||
.map(([key, value]) => this.buildFilterClause(key, value))
|
||||
.filter((clause): clause is Record<string, any> => Boolean(clause));
|
||||
}
|
||||
|
||||
private buildFilterClause(
|
||||
key: string,
|
||||
value: any,
|
||||
): Record<string, any> | null {
|
||||
if (value === null || value === undefined) return null;
|
||||
|
||||
if (key === "AND" || key === "OR" || key === "NOT") {
|
||||
if (!Array.isArray(value)) {
|
||||
throw new Error(`${key} filter value must be an array.`);
|
||||
}
|
||||
|
||||
const clauses = value
|
||||
.flatMap((filter) => this.buildFilterClauses(filter))
|
||||
.filter(Boolean);
|
||||
|
||||
if (clauses.length === 0) return null;
|
||||
|
||||
if (key === "AND") {
|
||||
return { bool: { filter: clauses } };
|
||||
}
|
||||
if (key === "OR") {
|
||||
return { bool: { should: clauses, minimum_should_match: 1 } };
|
||||
}
|
||||
return { bool: { must_not: clauses } };
|
||||
}
|
||||
|
||||
if (value === "*") {
|
||||
return { exists: { field: this.payloadField(key, false) } };
|
||||
}
|
||||
|
||||
if (Array.isArray(value)) {
|
||||
value.forEach((item) => this.assertScalarValue(key, item));
|
||||
return { terms: { [this.payloadField(key, true)]: value } };
|
||||
}
|
||||
|
||||
if (typeof value === "object") {
|
||||
return this.buildOperatorClause(key, value);
|
||||
}
|
||||
|
||||
return {
|
||||
term: { [this.payloadField(key, typeof value === "string")]: value },
|
||||
};
|
||||
}
|
||||
|
||||
private buildOperatorClause(
|
||||
key: string,
|
||||
value: Record<string, any>,
|
||||
): Record<string, any> {
|
||||
const clauses = Object.entries(value).map(([operator, operatorValue]) => {
|
||||
switch (operator) {
|
||||
case "eq":
|
||||
this.assertScalarValue(key, operatorValue);
|
||||
return {
|
||||
term: {
|
||||
[this.payloadField(key, typeof operatorValue === "string")]:
|
||||
operatorValue,
|
||||
},
|
||||
};
|
||||
case "ne":
|
||||
this.assertScalarValue(key, operatorValue);
|
||||
return {
|
||||
bool: {
|
||||
must_not: [
|
||||
{
|
||||
term: {
|
||||
[this.payloadField(key, typeof operatorValue === "string")]:
|
||||
operatorValue,
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
case "in":
|
||||
this.assertScalarArray(key, operatorValue);
|
||||
return { terms: { [this.payloadField(key, true)]: operatorValue } };
|
||||
case "nin":
|
||||
this.assertScalarArray(key, operatorValue);
|
||||
return {
|
||||
bool: {
|
||||
must_not: [
|
||||
{ terms: { [this.payloadField(key, true)]: operatorValue } },
|
||||
],
|
||||
},
|
||||
};
|
||||
case "gt":
|
||||
case "gte":
|
||||
case "lt":
|
||||
case "lte":
|
||||
this.assertScalarValue(key, operatorValue);
|
||||
return {
|
||||
range: {
|
||||
[this.payloadField(key, false)]: {
|
||||
[operator]: operatorValue,
|
||||
},
|
||||
},
|
||||
};
|
||||
case "contains":
|
||||
case "icontains":
|
||||
this.assertScalarValue(key, operatorValue);
|
||||
return {
|
||||
wildcard: {
|
||||
[this.payloadField(key, true)]: {
|
||||
value: `*${escapeWildcard(String(operatorValue))}*`,
|
||||
case_insensitive: operator === "icontains",
|
||||
},
|
||||
},
|
||||
};
|
||||
default:
|
||||
throw new Error(`Unsupported filter operator: ${operator}`);
|
||||
}
|
||||
});
|
||||
|
||||
return clauses.length === 1 ? clauses[0] : { bool: { filter: clauses } };
|
||||
}
|
||||
|
||||
private payloadField(key: string, keyword: boolean): string {
|
||||
if (key.startsWith("payload.")) {
|
||||
return keyword && !key.endsWith(".keyword") ? `${key}.keyword` : key;
|
||||
}
|
||||
|
||||
const field = `payload.${key}`;
|
||||
return keyword ? `${field}.keyword` : field;
|
||||
}
|
||||
|
||||
// Filter values become OpenSearch term/terms/range/wildcard leaves. Allowing
|
||||
// an object here lets a caller inject raw query parameters (e.g. a `term`
|
||||
// object form with `boost`/`case_insensitive`), so reject non-scalar leaves.
|
||||
// Mirrors the Python SDK's `_validate_filter` scalar allow-list (PR #5986).
|
||||
private assertScalarValue(key: string, value: any): void {
|
||||
if (value !== null && typeof value === "object") {
|
||||
throw new Error(
|
||||
`Filter value for '${key}' must be a string, number, or boolean, got ` +
|
||||
`${Array.isArray(value) ? "an array" : "an object"}.`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
private assertScalarArray(key: string, value: any): void {
|
||||
if (!Array.isArray(value)) {
|
||||
throw new Error(`Filter value for '${key}' must be an array.`);
|
||||
}
|
||||
value.forEach((item) => this.assertScalarValue(key, item));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,156 @@
|
||||
import { OpenSearchDB } from "../src/vector_stores/opensearch";
|
||||
|
||||
// These tests run entirely against a mocked OpenSearch client (no live cluster).
|
||||
// They lock the index mapping and the filter query paths together: `payload` must
|
||||
// stay a dynamic object so that a `.keyword` sub-field exists for the exact-match
|
||||
// term clauses the filter builder emits. Mapping payload sub-keys as explicit
|
||||
// `keyword` fields removes that sub-field and makes every scoped filter match
|
||||
// nothing, so the createCol + search assertions here guard that regression.
|
||||
|
||||
describe("OpenSearchDB", () => {
|
||||
const collectionName = "memories";
|
||||
const dims = 3;
|
||||
|
||||
function createClient(overrides: Record<string, any> = {}) {
|
||||
const { indices: indicesOverride, ...rest } = overrides;
|
||||
return {
|
||||
indices: {
|
||||
exists: jest.fn().mockResolvedValue(false),
|
||||
create: jest.fn().mockResolvedValue({ body: { acknowledged: true } }),
|
||||
delete: jest.fn().mockResolvedValue({ body: { acknowledged: true } }),
|
||||
...(indicesOverride || {}),
|
||||
},
|
||||
bulk: jest.fn().mockResolvedValue({ body: { errors: false, items: [] } }),
|
||||
search: jest.fn().mockResolvedValue({
|
||||
body: { hits: { hits: [], total: { value: 0 } } },
|
||||
}),
|
||||
get: jest.fn().mockResolvedValue({ body: { _source: null } }),
|
||||
update: jest.fn().mockResolvedValue({ body: {} }),
|
||||
delete: jest.fn().mockResolvedValue({ body: {} }),
|
||||
index: jest.fn().mockResolvedValue({ body: {} }),
|
||||
...rest,
|
||||
};
|
||||
}
|
||||
|
||||
async function createStore(client: any) {
|
||||
const store = new OpenSearchDB({
|
||||
collectionName,
|
||||
embeddingModelDims: dims,
|
||||
client: client as any,
|
||||
});
|
||||
await store.initialize();
|
||||
return store;
|
||||
}
|
||||
|
||||
it("maps payload as a dynamic object and enables knn without a slow refresh override", async () => {
|
||||
const client = createClient();
|
||||
await createStore(client);
|
||||
|
||||
const createCall = client.indices.create.mock.calls.find(
|
||||
([arg]: any[]) => arg.index === collectionName,
|
||||
);
|
||||
expect(createCall).toBeDefined();
|
||||
|
||||
const body = createCall[0].body;
|
||||
// Dynamic payload object is what gives user_id/agent_id/run_id their
|
||||
// `.keyword` sub-field. Enumerated keyword sub-properties would break filters.
|
||||
expect(body.mappings.properties.payload).toEqual({ type: "object" });
|
||||
expect(body.mappings.properties.payload.properties).toBeUndefined();
|
||||
// No dead metadata mirror of payload.
|
||||
expect(body.mappings.properties.metadata).toBeUndefined();
|
||||
// knn on, and no 10s refresh_interval that would hide freshly added memories.
|
||||
expect(body.settings.index.knn).toBe(true);
|
||||
expect(body.settings.index.refresh_interval).toBeUndefined();
|
||||
});
|
||||
|
||||
it("inserts payload and id without a duplicate metadata field", async () => {
|
||||
const client = createClient();
|
||||
const store = await createStore(client);
|
||||
|
||||
await store.insert(
|
||||
[[0.1, 0.2, 0.3]],
|
||||
["mem-1"],
|
||||
[{ user_id: "alice", data: "Alice likes pizza" }],
|
||||
);
|
||||
|
||||
expect(client.bulk).toHaveBeenCalledTimes(1);
|
||||
const operations = client.bulk.mock.calls[0][0].body;
|
||||
const [action, doc] = operations;
|
||||
expect(action).toEqual({ index: { _index: collectionName, _id: "mem-1" } });
|
||||
expect(doc.payload).toEqual({
|
||||
user_id: "alice",
|
||||
data: "Alice likes pizza",
|
||||
});
|
||||
expect(doc.id).toBe("mem-1");
|
||||
// No dead metadata/text mirror fields (Python live insert writes neither).
|
||||
expect(doc).not.toHaveProperty("metadata");
|
||||
expect(doc).not.toHaveProperty("text");
|
||||
});
|
||||
|
||||
it("scopes search by user_id via the payload.<key>.keyword term path", async () => {
|
||||
const client = createClient();
|
||||
const store = await createStore(client);
|
||||
|
||||
await store.search([0.1, 0.2, 0.3], 5, { user_id: "alice" });
|
||||
|
||||
const searchCall = client.search.mock.calls.find(
|
||||
([arg]: any[]) => arg.index === collectionName,
|
||||
);
|
||||
expect(searchCall).toBeDefined();
|
||||
const filter = searchCall[0].body.query.bool.filter;
|
||||
expect(filter).toContainEqual({
|
||||
term: { "payload.user_id.keyword": "alice" },
|
||||
});
|
||||
});
|
||||
|
||||
it("lists filtered results with a total count", async () => {
|
||||
const client = createClient({
|
||||
search: jest.fn().mockResolvedValue({
|
||||
body: {
|
||||
hits: {
|
||||
total: { value: 1 },
|
||||
hits: [
|
||||
{
|
||||
_id: "mem-1",
|
||||
_score: 1,
|
||||
_source: { id: "mem-1", payload: { user_id: "alice" } },
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
}),
|
||||
});
|
||||
const store = await createStore(client);
|
||||
|
||||
const [results, total] = await store.list({ user_id: "alice" });
|
||||
|
||||
expect(total).toBe(1);
|
||||
expect(results).toEqual([
|
||||
{ id: "mem-1", payload: { user_id: "alice" }, score: 1 },
|
||||
]);
|
||||
});
|
||||
|
||||
it("rejects object filter values that could inject query parameters", async () => {
|
||||
const client = createClient();
|
||||
const store = await createStore(client);
|
||||
|
||||
// A term query accepts an object form ({value, boost, case_insensitive}),
|
||||
// so an object leaf would let a caller inject raw query params. Guarded.
|
||||
await expect(
|
||||
store.search([0.1, 0.2, 0.3], 5, {
|
||||
user_id: { eq: { boost: 999, value: "x" } },
|
||||
} as any),
|
||||
).rejects.toThrow(/must be a string, number, or boolean/);
|
||||
|
||||
expect(client.search).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("returns null from get when the document is missing (404)", async () => {
|
||||
const client = createClient({
|
||||
get: jest.fn().mockRejectedValue({ statusCode: 404 }),
|
||||
});
|
||||
const store = await createStore(client);
|
||||
|
||||
await expect(store.get("missing")).resolves.toBeNull();
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user