feat(mem0-ts): add Baidu vector store provider (#5790)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
Aditya
@@ -5,10 +5,22 @@ description: "Use Baidu Mochow as an enterprise vector database in Mem0 for high
|
||||
|
||||
[Baidu VectorDB](https://cloud.baidu.com/doc/VDB/index.html) is an enterprise-level distributed vector database service developed by Baidu Intelligent Cloud. It is powered by Baidu's proprietary "Mochow" vector database kernel, providing high performance, availability, and security for vector search.
|
||||
|
||||
### Installation
|
||||
|
||||
<CodeGroup>
|
||||
```bash Python
|
||||
pip install pymochow
|
||||
```
|
||||
|
||||
```bash TypeScript
|
||||
npm install @mochow/mochow-sdk-node
|
||||
```
|
||||
|
||||
</CodeGroup>
|
||||
|
||||
### Usage
|
||||
|
||||
```python
|
||||
import os
|
||||
from mem0 import Memory
|
||||
|
||||
config = {
|
||||
@@ -36,19 +48,63 @@ messages = [
|
||||
m.add(messages, user_id="alice", metadata={"category": "movies"})
|
||||
```
|
||||
|
||||
```typescript
|
||||
import { Memory } from "mem0ai/oss";
|
||||
|
||||
const memory = new Memory({
|
||||
embedder: {
|
||||
provider: "openai",
|
||||
config: {
|
||||
apiKey: process.env.OPENAI_API_KEY || "",
|
||||
model: "text-embedding-3-small",
|
||||
embeddingDims: 1536,
|
||||
},
|
||||
},
|
||||
vectorStore: {
|
||||
provider: "baidu",
|
||||
config: {
|
||||
endpoint: process.env.BAIDU_ENDPOINT || "",
|
||||
account: process.env.BAIDU_ACCOUNT || "root",
|
||||
apiKey: process.env.BAIDU_API_KEY || "",
|
||||
databaseName: "mem0",
|
||||
tableName: "mem0_table",
|
||||
embeddingModelDims: 1536,
|
||||
metricType: "COSINE",
|
||||
},
|
||||
},
|
||||
llm: {
|
||||
provider: "openai",
|
||||
config: {
|
||||
apiKey: process.env.OPENAI_API_KEY || "",
|
||||
model: "gpt-5-mini",
|
||||
},
|
||||
},
|
||||
});
|
||||
```
|
||||
|
||||
### Config
|
||||
|
||||
Here are the parameters available for configuring Baidu VectorDB:
|
||||
|
||||
| Parameter | Description | Default Value |
|
||||
| --- | --- | --- |
|
||||
| `endpoint` | Endpoint URL for your Baidu VectorDB instance | Required |
|
||||
| `account` | Baidu VectorDB account name | `root` |
|
||||
| `api_key` | API key for accessing Baidu VectorDB | Required |
|
||||
| `database_name` | Name of the database | `mem0` |
|
||||
| `table_name` | Name of the table | `mem0` |
|
||||
| `embedding_model_dims` | Dimensions of the embedding model | `1536` |
|
||||
| `metric_type` | Distance metric for similarity search | `L2` |
|
||||
| Parameter | Description | Default Value |
|
||||
| ---------------------- | --------------------------------------------- | ------------- |
|
||||
| `endpoint` | Endpoint URL for your Baidu VectorDB instance | Required |
|
||||
| `account` | Baidu VectorDB account name | `root` |
|
||||
| `api_key` | API key for accessing Baidu VectorDB | Required |
|
||||
| `database_name` | Name of the database | `mem0` |
|
||||
| `table_name` | Name of the table | `mem0` |
|
||||
| `embedding_model_dims` | Dimensions of the embedding model | `1536` |
|
||||
| `metric_type` | Distance metric for similarity search | `L2` |
|
||||
| `client` | Prebuilt Mochow client (TypeScript SDK only) | `None` |
|
||||
|
||||
For the TypeScript OSS SDK, use the camelCase equivalents:
|
||||
|
||||
- `databaseName`
|
||||
- `tableName`
|
||||
- `embeddingModelDims`
|
||||
- `metricType`
|
||||
|
||||
For OSS TS usage, `endpoint`, `account`, `apiKey`, `databaseName`, `tableName`, and `embeddingModelDims` are required unless you inject a prebuilt client. `metricType` defaults to `L2`, matching the Python SDK.
|
||||
|
||||
### Distance Metrics
|
||||
|
||||
@@ -66,3 +122,5 @@ The vector index is automatically configured with the following HNSW parameters:
|
||||
- `efconstruction`: 200 (size of the dynamic candidate list)
|
||||
- `auto_build`: true (automatically build index)
|
||||
- `auto_build_index_policy`: Incremental build with 10000 rows increment
|
||||
|
||||
The TypeScript provider also creates a BM25 inverted index over a `textLemmatized` column so `keywordSearch()` runs against a real full-text index. Mem0 lemmatizes the query before it reaches the vector store, so only the lemmatized form of each memory is indexed. If you point `tableName` at a table created before this index existed, `keywordSearch()` returns `null` and search falls back to vector similarity alone; recreate the table to enable it.
|
||||
|
||||
@@ -98,7 +98,8 @@
|
||||
"ts-node": "^10.9.2",
|
||||
"tsup": "^8.3.0",
|
||||
"typescript": "5.5.4",
|
||||
"iovalkey": "^0.3.3"
|
||||
"iovalkey": "^0.3.3",
|
||||
"@mochow/mochow-sdk-node": "^2.1.5"
|
||||
},
|
||||
"dependencies": {
|
||||
"axios": "^1.16.0",
|
||||
@@ -109,6 +110,7 @@
|
||||
"peerDependencies": {
|
||||
"@anthropic-ai/sdk": "^0.40.1",
|
||||
"@aws-sdk/client-s3vectors": "3.967.0",
|
||||
"@mochow/mochow-sdk-node": "^2.1.5",
|
||||
"@azure/identity": "^4.0.0",
|
||||
"@azure/search-documents": "^12.0.0",
|
||||
"@cloudflare/workers-types": "^4.20250504.0",
|
||||
@@ -151,6 +153,9 @@
|
||||
},
|
||||
"@zilliz/milvus2-sdk-node": {
|
||||
"optional": true
|
||||
},
|
||||
"@mochow/mochow-sdk-node": {
|
||||
"optional": true
|
||||
}
|
||||
},
|
||||
"engines": {
|
||||
|
||||
Generated
+31
-7
@@ -109,7 +109,7 @@ importers:
|
||||
version: 4.5.0
|
||||
cohere-ai:
|
||||
specifier: ^7.17.0 || ^8.0.0
|
||||
version: 8.0.0(@aws-crypto/sha256-js@5.2.0)(@smithy/protocol-http@5.5.2)(@smithy/signature-v4@5.5.2)
|
||||
version: 8.0.0(@aws-crypto/sha256-js@5.2.0)(@smithy/protocol-http@5.5.2)(@smithy/signature-v4@5.6.2)
|
||||
compromise:
|
||||
specifier: ^14.0.0
|
||||
version: 14.15.1
|
||||
@@ -153,6 +153,9 @@ importers:
|
||||
specifier: ^3.24.1
|
||||
version: 3.25.76
|
||||
devDependencies:
|
||||
'@mochow/mochow-sdk-node':
|
||||
specifier: ^2.1.5
|
||||
version: 2.1.5
|
||||
'@types/better-sqlite3':
|
||||
specifier: ^7.6.13
|
||||
version: 7.6.13
|
||||
@@ -191,7 +194,7 @@ importers:
|
||||
version: 10.9.2(@types/node@22.19.21)(typescript@5.5.4)
|
||||
tsup:
|
||||
specifier: ^8.3.0
|
||||
version: 8.5.1(typescript@5.5.4)
|
||||
version: 8.5.1(tsx@4.23.0)(typescript@5.5.4)
|
||||
typescript:
|
||||
specifier: 5.5.4
|
||||
version: 5.5.4
|
||||
@@ -1137,6 +1140,9 @@ packages:
|
||||
'@mistralai/mistralai@1.15.1':
|
||||
resolution: {integrity: sha512-fb995eiz3r0KsBGtRjFV+/iLbX+UpfalxpF+YitT3R6ukrPD4PN+FGwwmYcRFhNAzVzDUtTVxQYnjQWEnwV5nw==}
|
||||
|
||||
'@mochow/mochow-sdk-node@2.1.5':
|
||||
resolution: {integrity: sha512-IYluqAf50wH51uMGDqRni1Mmuc8GCNDHibJf0x3YlJFQkTJkU+l/PvLsmrer9qV1gMj6wHA+LE0OkNevW7KsCw==}
|
||||
|
||||
'@mongodb-js/saslprep@1.4.11':
|
||||
resolution: {integrity: sha512-o9rAHc0IpIjuPSxRutWpE1F62x7n+4mVS4rCNHkzhIUMQcc18bb6xEq5wd2NdN0WjepIyXIppRshYI2kQDOZVA==}
|
||||
|
||||
@@ -4076,6 +4082,11 @@ packages:
|
||||
typescript:
|
||||
optional: true
|
||||
|
||||
tsx@4.23.0:
|
||||
resolution: {integrity: sha512-eUdUIaCr963q2h5u3+QwvYp0+eqPvn+egeqZUm0hwERCqqx1E3kK5ehbGCvqSE5MQAULr67ww0cA3jKc3YkM1w==}
|
||||
engines: {node: '>=18.0.0'}
|
||||
hasBin: true
|
||||
|
||||
tunnel-agent@0.6.0:
|
||||
resolution: {integrity: sha512-McnNiV1l8RYeY8tBgEpuodCC1mLUdbSN+CYBL7kJsJNInOP8UjDDEwdk6Mw60vdLLrr5NHKZhMAOSrR2NZuQ+w==}
|
||||
|
||||
@@ -5700,6 +5711,11 @@ snapshots:
|
||||
- bufferutil
|
||||
- utf-8-validate
|
||||
|
||||
'@mochow/mochow-sdk-node@2.1.5':
|
||||
dependencies:
|
||||
tsx: 4.23.0
|
||||
winston: 3.19.0
|
||||
|
||||
'@mongodb-js/saslprep@1.4.11':
|
||||
dependencies:
|
||||
sparse-bitfield: 3.0.3
|
||||
@@ -6666,7 +6682,7 @@ snapshots:
|
||||
|
||||
co@4.6.0: {}
|
||||
|
||||
cohere-ai@8.0.0(@aws-crypto/sha256-js@5.2.0)(@smithy/protocol-http@5.5.2)(@smithy/signature-v4@5.5.2):
|
||||
cohere-ai@8.0.0(@aws-crypto/sha256-js@5.2.0)(@smithy/protocol-http@5.5.2)(@smithy/signature-v4@5.6.2):
|
||||
dependencies:
|
||||
convict: 6.2.5
|
||||
form-data: 4.0.6
|
||||
@@ -6676,7 +6692,7 @@ snapshots:
|
||||
optionalDependencies:
|
||||
'@aws-crypto/sha256-js': 5.2.0
|
||||
'@smithy/protocol-http': 5.5.2
|
||||
'@smithy/signature-v4': 5.5.2
|
||||
'@smithy/signature-v4': 5.6.2
|
||||
|
||||
collect-v8-coverage@1.0.3: {}
|
||||
|
||||
@@ -8330,9 +8346,11 @@ snapshots:
|
||||
|
||||
platform@1.3.6: {}
|
||||
|
||||
postcss-load-config@6.0.1:
|
||||
postcss-load-config@6.0.1(tsx@4.23.0):
|
||||
dependencies:
|
||||
lilconfig: 3.1.3
|
||||
optionalDependencies:
|
||||
tsx: 4.23.0
|
||||
|
||||
postgres-array@2.0.0: {}
|
||||
|
||||
@@ -8895,7 +8913,7 @@ snapshots:
|
||||
|
||||
tslib@2.8.1: {}
|
||||
|
||||
tsup@8.5.1(typescript@5.5.4):
|
||||
tsup@8.5.1(tsx@4.23.0)(typescript@5.5.4):
|
||||
dependencies:
|
||||
bundle-require: 5.1.0(esbuild@0.28.1)
|
||||
cac: 6.7.14
|
||||
@@ -8906,7 +8924,7 @@ snapshots:
|
||||
fix-dts-default-cjs-exports: 1.0.1
|
||||
joycon: 3.1.1
|
||||
picocolors: 1.1.1
|
||||
postcss-load-config: 6.0.1
|
||||
postcss-load-config: 6.0.1(tsx@4.23.0)
|
||||
resolve-from: 5.0.0
|
||||
rollup: 4.61.1
|
||||
source-map: 0.7.6
|
||||
@@ -8922,6 +8940,12 @@ snapshots:
|
||||
- tsx
|
||||
- yaml
|
||||
|
||||
tsx@4.23.0:
|
||||
dependencies:
|
||||
esbuild: 0.28.1
|
||||
optionalDependencies:
|
||||
fsevents: 2.3.3
|
||||
|
||||
tunnel-agent@0.6.0:
|
||||
dependencies:
|
||||
safe-buffer: 5.2.1
|
||||
|
||||
@@ -25,6 +25,7 @@ export * from "./llms/litellm";
|
||||
export * from "./llms/vllm";
|
||||
export * from "./vector_stores/base";
|
||||
export * from "./vector_stores/memory";
|
||||
export * from "./vector_stores/baidu";
|
||||
export * from "./vector_stores/qdrant";
|
||||
export * from "./vector_stores/redis";
|
||||
export * from "./vector_stores/valkey";
|
||||
|
||||
@@ -23,6 +23,7 @@ import { CrossEncoderReranker } from "../rerankers/cross_encoder";
|
||||
import { Embedder } from "../embeddings/base";
|
||||
import { LLM } from "../llms/base";
|
||||
import { VectorStore } from "../vector_stores/base";
|
||||
import { BaiduDB } from "../vector_stores/baidu";
|
||||
import { Qdrant } from "../vector_stores/qdrant";
|
||||
import { ChromaDB } from "../vector_stores/chroma";
|
||||
import { VectorizeDB } from "../vector_stores/vectorize";
|
||||
@@ -143,6 +144,8 @@ export class VectorStoreFactory {
|
||||
switch (provider.toLowerCase()) {
|
||||
case "memory":
|
||||
return new MemoryVectorStore(config);
|
||||
case "baidu":
|
||||
return new BaiduDB(config as any);
|
||||
case "qdrant":
|
||||
return new Qdrant(config as any);
|
||||
case "chroma":
|
||||
|
||||
@@ -0,0 +1,613 @@
|
||||
import type {
|
||||
AutoBuildIncrementPolicy,
|
||||
CommonResponse,
|
||||
DescTableResponse,
|
||||
FieldType,
|
||||
IndexSchema,
|
||||
MochowClient,
|
||||
QueryResponse,
|
||||
SearchResponse,
|
||||
SelectResponse,
|
||||
TableSchema,
|
||||
} from "@mochow/mochow-sdk-node";
|
||||
import { VectorStore } from "./base";
|
||||
import { SearchFilters, VectorStoreConfig, VectorStoreResult } from "../types";
|
||||
|
||||
type MochowSdk = typeof import("@mochow/mochow-sdk-node");
|
||||
|
||||
export interface BaiduConfig extends VectorStoreConfig {
|
||||
endpoint: string;
|
||||
account: string;
|
||||
apiKey: string;
|
||||
databaseName: string;
|
||||
tableName: string;
|
||||
embeddingModelDims: number;
|
||||
metricType?: "L2" | "IP" | "COSINE";
|
||||
client?: MochowClient;
|
||||
}
|
||||
|
||||
const VECTOR_INDEX = "vector_idx";
|
||||
const FILTERING_INDEX = "metadata_filtering_idx";
|
||||
// Named after the column it actually indexes, and deliberately not Python's "data_bm25_idx".
|
||||
// This index holds Porter-stemmed text, but mem0/vector_stores/baidu.py's keyword_search()
|
||||
// sends a raw, unstemmed query to that name. Sharing it would let Python find an index whose
|
||||
// contents it cannot match properly, silently returning degraded hits instead of None.
|
||||
const BM25_INDEX = "text_lemmatized_bm25_idx";
|
||||
const PROJECTIONS = ["id", "data", "metadata"];
|
||||
const TABLE_POLL_INTERVAL_MS = 2000;
|
||||
const TABLE_POLL_ATTEMPTS = 60;
|
||||
|
||||
// Mochow's server accepts JSON columns, but the Node SDK's FieldType enum predates them
|
||||
// (pymochow 2.4.1 ships FieldType.JSON == "JSON"). The wire value is the bare string.
|
||||
const JSON_FIELD_TYPE = "JSON" as unknown as FieldType;
|
||||
|
||||
// Querying a primary key that isn't there answers with this code, not an empty row. The Node
|
||||
// SDK's ServerErrCode stops at 100, but its siblings against the same server name it:
|
||||
// pymochow ROW_KEY_NOT_FOUND = 101, mochow-sdk-go RowKeyNotFound = 101.
|
||||
const ROW_KEY_NOT_FOUND = 101;
|
||||
|
||||
const SAFE_FILTER_KEY = /^[a-zA-Z_][a-zA-Z0-9_]*$/;
|
||||
|
||||
function escapeFilterString(value: string): string {
|
||||
return value.replace(/\\/g, "\\\\").replace(/"/g, '\\"');
|
||||
}
|
||||
|
||||
function sleep(ms: number): Promise<void> {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
|
||||
// Mochow resolves with a {code, msg} envelope instead of rejecting, so every call site
|
||||
// has to inspect the code. `tolerated` lets callers accept the idempotent outcomes
|
||||
// (database/table already exists, table already dropped).
|
||||
function check(
|
||||
response: CommonResponse,
|
||||
action: string,
|
||||
...tolerated: number[]
|
||||
): number {
|
||||
if (response.code !== 0 && !tolerated.includes(response.code)) {
|
||||
throw new Error(
|
||||
`Baidu Mochow ${action} failed (code ${response.code}): ${response.msg}`,
|
||||
);
|
||||
}
|
||||
return response.code;
|
||||
}
|
||||
|
||||
function lemmatizedText(payload: Record<string, any>): string {
|
||||
const data = typeof payload.data === "string" ? payload.data : "";
|
||||
return typeof payload.textLemmatized === "string" &&
|
||||
payload.textLemmatized.length > 0
|
||||
? payload.textLemmatized
|
||||
: data;
|
||||
}
|
||||
|
||||
function memoryData(payload: Record<string, any>): string {
|
||||
return typeof payload.data === "string" ? payload.data : "";
|
||||
}
|
||||
|
||||
function metadataPayload(payload: Record<string, any>): Record<string, any> {
|
||||
const { data: _data, textLemmatized: _textLemmatized, ...metadata } = payload;
|
||||
return metadata;
|
||||
}
|
||||
|
||||
function resultPayload(row: Record<string, any>): Record<string, any> {
|
||||
return {
|
||||
...(row.metadata || {}),
|
||||
...(typeof row.data === "string" ? { data: row.data } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
export class BaiduDB implements VectorStore {
|
||||
private client: MochowClient | null = null;
|
||||
private sdk: MochowSdk | null = null;
|
||||
private readonly endpoint: string;
|
||||
private readonly account: string;
|
||||
private readonly apiKey: string;
|
||||
private readonly databaseName: string;
|
||||
private readonly tableName: string;
|
||||
private readonly embeddingModelDims: number;
|
||||
private readonly metricType: "L2" | "IP" | "COSINE";
|
||||
// Fails closed: keyword search stays off until an inverted index is observed.
|
||||
private supportsKeywordSearch = false;
|
||||
private storeUserId = "anonymous-baidu-user";
|
||||
private _initPromise?: Promise<void>;
|
||||
|
||||
constructor(config: BaiduConfig) {
|
||||
this.endpoint = config.endpoint;
|
||||
this.account = config.account;
|
||||
this.apiKey = config.apiKey;
|
||||
this.databaseName = config.databaseName;
|
||||
this.tableName = config.tableName;
|
||||
this.embeddingModelDims = config.embeddingModelDims;
|
||||
this.metricType = config.metricType || "L2";
|
||||
this.client = config.client || null;
|
||||
|
||||
const requiredFields: Array<
|
||||
readonly [string, string | number | undefined]
|
||||
> = [
|
||||
["databaseName", this.databaseName],
|
||||
["tableName", this.tableName],
|
||||
["embeddingModelDims", this.embeddingModelDims],
|
||||
];
|
||||
|
||||
if (!this.client) {
|
||||
requiredFields.unshift(
|
||||
["endpoint", this.endpoint],
|
||||
["account", this.account],
|
||||
["apiKey", this.apiKey],
|
||||
);
|
||||
}
|
||||
|
||||
for (const [name, value] of requiredFields) {
|
||||
if (value === undefined || value === null || value === "") {
|
||||
throw new Error(
|
||||
`Baidu vector store requires a non-empty '${name}' config value.`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
this.initialize().catch(console.error);
|
||||
}
|
||||
|
||||
private get ns(): { database: string; table: string } {
|
||||
return { database: this.databaseName, table: this.tableName };
|
||||
}
|
||||
|
||||
// Loaded dynamically: @mochow/mochow-sdk-node is an optional peer dependency, so a static
|
||||
// value import would break `import { Memory } from "mem0ai/oss"` for everyone else.
|
||||
private async loadSdk(): Promise<MochowSdk> {
|
||||
if (!this.sdk) {
|
||||
let module: MochowSdk & { default?: MochowSdk };
|
||||
try {
|
||||
module = await import("@mochow/mochow-sdk-node");
|
||||
} catch {
|
||||
throw new Error(
|
||||
"The Baidu vector store requires the '@mochow/mochow-sdk-node' package. Install it with: npm install @mochow/mochow-sdk-node",
|
||||
);
|
||||
}
|
||||
this.sdk = module.default ?? module;
|
||||
}
|
||||
return this.sdk;
|
||||
}
|
||||
|
||||
private async ensureClient(): Promise<MochowClient> {
|
||||
if (!this.client) {
|
||||
const sdk = await this.loadSdk();
|
||||
this.client = new sdk.MochowClient({
|
||||
endpoint: this.endpoint,
|
||||
credential: { account: this.account, apiKey: this.apiKey },
|
||||
});
|
||||
}
|
||||
return this.client;
|
||||
}
|
||||
|
||||
private async ready(): Promise<{ client: MochowClient; sdk: MochowSdk }> {
|
||||
await this.initialize();
|
||||
return { client: await this.ensureClient(), sdk: await this.loadSdk() };
|
||||
}
|
||||
|
||||
private buildSchema(sdk: MochowSdk): TableSchema {
|
||||
const {
|
||||
AutoBuildPolicyType,
|
||||
FieldType,
|
||||
IndexType,
|
||||
InvertedIndexAnalyzer,
|
||||
InvertedIndexFieldAttribute,
|
||||
InvertedIndexParseMode,
|
||||
MetricType,
|
||||
} = sdk;
|
||||
|
||||
// sdk.AutoBuildIncrement() stamps policyType "TIMING" (bug in 2.1.5), so build the
|
||||
// increment policy by hand.
|
||||
const autoBuildPolicy: AutoBuildIncrementPolicy = {
|
||||
policyType: AutoBuildPolicyType.Increment,
|
||||
rowCountIncrement: 10000,
|
||||
};
|
||||
|
||||
const vectorIndex: IndexSchema = {
|
||||
indexName: VECTOR_INDEX,
|
||||
indexType: IndexType.HNSW,
|
||||
field: "vector",
|
||||
metricType: MetricType[this.metricType],
|
||||
params: { M: 16, efConstruction: 200 },
|
||||
autoBuild: true,
|
||||
autoBuildPolicy,
|
||||
};
|
||||
|
||||
return {
|
||||
fields: [
|
||||
{
|
||||
fieldName: "id",
|
||||
fieldType: FieldType.String,
|
||||
primaryKey: true,
|
||||
partitionKey: true,
|
||||
autoIncrement: false,
|
||||
notNull: true,
|
||||
},
|
||||
{
|
||||
fieldName: "data",
|
||||
fieldType: FieldType.Text,
|
||||
},
|
||||
{
|
||||
fieldName: "vector",
|
||||
fieldType: FieldType.FloatVector,
|
||||
notNull: true,
|
||||
dimension: this.embeddingModelDims,
|
||||
},
|
||||
// Stored outside `metadata` because Mochow cannot build an inverted index on a
|
||||
// field inside a JSON column. Memory.search() passes an already-lemmatized query,
|
||||
// so only the lemmatized form is worth indexing.
|
||||
{ fieldName: "textLemmatized", fieldType: FieldType.Text },
|
||||
{ fieldName: "metadata", fieldType: JSON_FIELD_TYPE },
|
||||
],
|
||||
indexes: [
|
||||
vectorIndex,
|
||||
{
|
||||
indexName: FILTERING_INDEX,
|
||||
indexType: IndexType.FilteringIndex,
|
||||
fields: ["metadata"],
|
||||
},
|
||||
{
|
||||
indexName: BM25_INDEX,
|
||||
indexType: IndexType.InvertedIndex,
|
||||
fields: ["textLemmatized"],
|
||||
fieldAttributes: [InvertedIndexFieldAttribute.Analyzed],
|
||||
params: {
|
||||
analyzer: InvertedIndexAnalyzer.EnglishAnalyzer,
|
||||
parseMode: InvertedIndexParseMode.FineMode,
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
private buildFilter(filters: SearchFilters): string {
|
||||
const conditions: string[] = [];
|
||||
|
||||
for (const [key, value] of Object.entries(filters)) {
|
||||
if (!SAFE_FILTER_KEY.test(key)) {
|
||||
throw new Error(`Invalid filter key: ${key}`);
|
||||
}
|
||||
|
||||
if (typeof value === "string") {
|
||||
conditions.push(`metadata["${key}"] = "${escapeFilterString(value)}"`);
|
||||
continue;
|
||||
}
|
||||
|
||||
if (typeof value === "number" || typeof value === "boolean") {
|
||||
conditions.push(`metadata["${key}"] = ${value}`);
|
||||
continue;
|
||||
}
|
||||
|
||||
throw new Error(
|
||||
`Filter value for ${key} must be str, int, float, or bool, got ${Array.isArray(value) ? "array" : typeof value}`,
|
||||
);
|
||||
}
|
||||
|
||||
return conditions.join(" AND ");
|
||||
}
|
||||
|
||||
private filterOf(filters?: SearchFilters): string | undefined {
|
||||
return filters && Object.keys(filters).length > 0
|
||||
? this.buildFilter(filters)
|
||||
: undefined;
|
||||
}
|
||||
|
||||
private async pollTable(
|
||||
client: MochowClient,
|
||||
settled: (response: DescTableResponse) => boolean,
|
||||
what: string,
|
||||
): Promise<void> {
|
||||
for (let attempt = 0; attempt < TABLE_POLL_ATTEMPTS; attempt++) {
|
||||
if (settled(await client.descTable(this.databaseName, this.tableName))) {
|
||||
return;
|
||||
}
|
||||
await sleep(TABLE_POLL_INTERVAL_MS);
|
||||
}
|
||||
|
||||
throw new Error(
|
||||
`Baidu Mochow table '${this.tableName}' was not ${what} after ${(TABLE_POLL_ATTEMPTS * TABLE_POLL_INTERVAL_MS) / 1000}s.`,
|
||||
);
|
||||
}
|
||||
|
||||
private async ensureTable(): Promise<void> {
|
||||
const sdk = await this.loadSdk();
|
||||
const client = await this.ensureClient();
|
||||
const { ServerErrCode, TableState } = sdk;
|
||||
|
||||
check(
|
||||
await client.createDatabase(this.databaseName),
|
||||
`createDatabase '${this.databaseName}'`,
|
||||
ServerErrCode.DBAlreadyExist,
|
||||
);
|
||||
|
||||
const created = check(
|
||||
await client.createTable({
|
||||
...this.ns,
|
||||
description: "mem0 memories",
|
||||
replication: 3,
|
||||
partition: { partitionType: sdk.PartitionType.HASH, partitionNum: 1 },
|
||||
enableDynamicField: false,
|
||||
schema: this.buildSchema(sdk),
|
||||
}),
|
||||
`createTable '${this.tableName}'`,
|
||||
ServerErrCode.TableAlreadyExist,
|
||||
);
|
||||
|
||||
// A table is CREATING until its indexes are built; writing to it before then fails.
|
||||
let description: DescTableResponse | undefined;
|
||||
await this.pollTable(
|
||||
client,
|
||||
(response) => {
|
||||
check(response, `descTable '${this.tableName}'`);
|
||||
description = response;
|
||||
return response.table.state === TableState.Normal;
|
||||
},
|
||||
"ready",
|
||||
);
|
||||
|
||||
this.applySchema(
|
||||
created === ServerErrCode.TableAlreadyExist,
|
||||
description!.table.schema,
|
||||
);
|
||||
}
|
||||
|
||||
private applySchema(preexisting: boolean, schema: TableSchema): void {
|
||||
if (!preexisting) {
|
||||
this.supportsKeywordSearch = true;
|
||||
return;
|
||||
}
|
||||
|
||||
const fields = schema?.fields ?? [];
|
||||
const indexes = schema?.indexes ?? [];
|
||||
const field = (name: string) => fields.find((f) => f.fieldName === name);
|
||||
const typeOf = (name: string) => String(field(name)?.fieldType ?? "");
|
||||
const label = `${this.databaseName}.${this.tableName}`;
|
||||
|
||||
if (
|
||||
typeOf("id") !== "STRING" ||
|
||||
!typeOf("data").startsWith("TEXT") ||
|
||||
typeOf("vector") !== "FLOAT_VECTOR" ||
|
||||
typeOf("metadata") !== "JSON"
|
||||
) {
|
||||
throw new Error(
|
||||
`Baidu Mochow table '${label}' exists but is missing the id/data/vector/metadata schema mem0 requires. Drop it, or point 'tableName' at an unused table.`,
|
||||
);
|
||||
}
|
||||
|
||||
const dimension = field("vector")?.dimension;
|
||||
if (dimension !== undefined && dimension !== this.embeddingModelDims) {
|
||||
throw new Error(
|
||||
`Baidu Mochow table '${label}' stores ${dimension}-dimensional vectors, but 'embeddingModelDims' is ${this.embeddingModelDims}.`,
|
||||
);
|
||||
}
|
||||
|
||||
this.supportsKeywordSearch =
|
||||
typeOf("textLemmatized").startsWith("TEXT") &&
|
||||
indexes.some((index) => index.indexName === BM25_INDEX);
|
||||
|
||||
if (!this.supportsKeywordSearch) {
|
||||
console.warn(
|
||||
`Baidu Mochow table '${label}' has no '${BM25_INDEX}' inverted index. keywordSearch() will return null until the table is recreated.`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async initialize(): Promise<void> {
|
||||
if (!this._initPromise) {
|
||||
this._initPromise = this.ensureTable().catch((error) => {
|
||||
this._initPromise = undefined;
|
||||
throw error;
|
||||
});
|
||||
}
|
||||
|
||||
return this._initPromise;
|
||||
}
|
||||
|
||||
async insert(
|
||||
vectors: number[][],
|
||||
ids: string[],
|
||||
payloads: Record<string, any>[],
|
||||
): Promise<void> {
|
||||
const { client } = await this.ready();
|
||||
|
||||
if (vectors.length !== ids.length || vectors.length !== payloads.length) {
|
||||
throw new Error(
|
||||
`Baidu insert requires vectors, ids, and payloads of equal length (got ${vectors.length}/${ids.length}/${payloads.length}).`,
|
||||
);
|
||||
}
|
||||
|
||||
const rows = vectors.map((vector, index) => ({
|
||||
id: ids[index],
|
||||
data: memoryData(payloads[index] || {}),
|
||||
vector,
|
||||
textLemmatized: lemmatizedText(payloads[index] || {}),
|
||||
metadata: metadataPayload(payloads[index] || {}),
|
||||
}));
|
||||
|
||||
check(await client.upsert({ ...this.ns, rows }), "upsert");
|
||||
}
|
||||
|
||||
async search(
|
||||
query: number[],
|
||||
topK = 5,
|
||||
filters?: SearchFilters,
|
||||
): Promise<VectorStoreResult[]> {
|
||||
const { client, sdk } = await this.ready();
|
||||
const filter = this.filterOf(filters);
|
||||
|
||||
const request = new sdk.VectorTopkSearchRequest(
|
||||
"vector",
|
||||
new sdk.Vector(query),
|
||||
topK,
|
||||
)
|
||||
.Projections(PROJECTIONS)
|
||||
.Config(new sdk.VectorSearchConfig().Ef(200));
|
||||
if (filter) {
|
||||
request.Filter(filter);
|
||||
}
|
||||
|
||||
const response = (await client.vectorSearch({
|
||||
...this.ns,
|
||||
request,
|
||||
})) as SearchResponse;
|
||||
check(response, "vectorSearch");
|
||||
|
||||
return (response.rows ?? []).map((result) => ({
|
||||
id: String(result.row.id),
|
||||
payload: resultPayload(result.row),
|
||||
score: result.score,
|
||||
}));
|
||||
}
|
||||
|
||||
async keywordSearch(
|
||||
query: string,
|
||||
topK = 5,
|
||||
filters?: SearchFilters,
|
||||
): Promise<VectorStoreResult[] | null> {
|
||||
const { client, sdk } = await this.ready();
|
||||
|
||||
if (!this.supportsKeywordSearch) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const filter = this.filterOf(filters);
|
||||
const request = new sdk.BM25SearchRequest(BM25_INDEX, query)
|
||||
.Projections(PROJECTIONS)
|
||||
.Limit(topK);
|
||||
if (filter) {
|
||||
request.Filter(filter);
|
||||
}
|
||||
|
||||
const response = (await client.bm25Search({
|
||||
...this.ns,
|
||||
request,
|
||||
})) as SearchResponse;
|
||||
check(response, "bm25Search");
|
||||
|
||||
return (response.rows ?? []).map((result) => ({
|
||||
id: String(result.row.id),
|
||||
payload: resultPayload(result.row),
|
||||
score: result.score,
|
||||
}));
|
||||
}
|
||||
|
||||
async get(vectorId: string): Promise<VectorStoreResult | null> {
|
||||
const { client } = await this.ready();
|
||||
|
||||
const response: QueryResponse = await client.query({
|
||||
...this.ns,
|
||||
primaryKey: { id: vectorId },
|
||||
projections: PROJECTIONS,
|
||||
});
|
||||
check(response, `query '${vectorId}'`, ROW_KEY_NOT_FOUND);
|
||||
|
||||
if (!response.row || response.row.id === undefined) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return {
|
||||
id: String(response.row.id),
|
||||
payload: resultPayload(response.row),
|
||||
};
|
||||
}
|
||||
|
||||
async update(
|
||||
vectorId: string,
|
||||
vector: number[],
|
||||
payload: Record<string, any>,
|
||||
): Promise<void> {
|
||||
const { client } = await this.ready();
|
||||
|
||||
check(
|
||||
await client.upsert({
|
||||
...this.ns,
|
||||
rows: [
|
||||
{
|
||||
id: vectorId,
|
||||
data: memoryData(payload),
|
||||
vector,
|
||||
textLemmatized: lemmatizedText(payload),
|
||||
metadata: metadataPayload(payload),
|
||||
},
|
||||
],
|
||||
}),
|
||||
`upsert '${vectorId}'`,
|
||||
);
|
||||
}
|
||||
|
||||
async delete(vectorId: string): Promise<void> {
|
||||
const { client } = await this.ready();
|
||||
|
||||
check(
|
||||
await client.delete({ ...this.ns, primaryKey: { id: vectorId } }),
|
||||
`delete '${vectorId}'`,
|
||||
);
|
||||
}
|
||||
|
||||
async deleteCol(): Promise<void> {
|
||||
// The constructor starts initialize() without awaiting it. Let any in-flight run land
|
||||
// first, otherwise it recreates the table after dropTable() and reset() is a no-op.
|
||||
await this._initPromise?.catch(() => undefined);
|
||||
this._initPromise = undefined;
|
||||
this.supportsKeywordSearch = false;
|
||||
|
||||
const sdk = await this.loadSdk();
|
||||
const client = await this.ensureClient();
|
||||
const { ServerErrCode } = sdk;
|
||||
|
||||
const dropped = check(
|
||||
await client.dropTable(this.databaseName, this.tableName),
|
||||
`dropTable '${this.tableName}'`,
|
||||
ServerErrCode.TableNotExist,
|
||||
);
|
||||
if (dropped === ServerErrCode.TableNotExist) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Drops are asynchronous; recreating the table before it is gone fails.
|
||||
await this.pollTable(
|
||||
client,
|
||||
(response) =>
|
||||
check(
|
||||
response,
|
||||
`descTable '${this.tableName}'`,
|
||||
ServerErrCode.TableNotExist,
|
||||
) === ServerErrCode.TableNotExist,
|
||||
"dropped",
|
||||
);
|
||||
}
|
||||
|
||||
async reset(): Promise<void> {
|
||||
await this.deleteCol();
|
||||
await this.initialize();
|
||||
}
|
||||
|
||||
async list(
|
||||
filters?: SearchFilters,
|
||||
topK = 100,
|
||||
): Promise<[VectorStoreResult[], number]> {
|
||||
const { client } = await this.ready();
|
||||
|
||||
const response: SelectResponse = await client.select({
|
||||
...this.ns,
|
||||
filter: this.filterOf(filters),
|
||||
projections: PROJECTIONS,
|
||||
limit: topK,
|
||||
});
|
||||
check(response, "select");
|
||||
|
||||
const memories = (response.rows ?? []).map((row) => ({
|
||||
id: String(row.id),
|
||||
payload: resultPayload(row),
|
||||
}));
|
||||
return [memories, memories.length];
|
||||
}
|
||||
|
||||
async getUserId(): Promise<string> {
|
||||
return this.storeUserId;
|
||||
}
|
||||
|
||||
async setUserId(userId: string): Promise<void> {
|
||||
this.storeUserId = userId;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,617 @@
|
||||
import {
|
||||
AutoBuildPolicyType,
|
||||
FieldType,
|
||||
IndexType,
|
||||
InvertedIndexFieldAttribute,
|
||||
MetricType,
|
||||
PartitionType,
|
||||
ServerErrCode,
|
||||
TableState,
|
||||
} from "@mochow/mochow-sdk-node";
|
||||
import { BaiduDB } from "../src/vector_stores/baidu";
|
||||
|
||||
// No jest.mock() here on purpose: the real SDK supplies the enums and the search request
|
||||
// classes (which carry an internal `set` map the client reads, so they cannot be hand-rolled
|
||||
// as plain literals). Only the network-facing MochowClient is faked, via the `client` config.
|
||||
|
||||
const OK = { code: 0, msg: "" };
|
||||
const DIMS = 1536;
|
||||
|
||||
const normalTable = (schema: unknown = { fields: [], indexes: [] }) => ({
|
||||
...OK,
|
||||
table: { state: TableState.Normal, schema },
|
||||
});
|
||||
|
||||
const CORE_FIELDS = [
|
||||
{ fieldName: "id", fieldType: FieldType.String },
|
||||
{ fieldName: "data", fieldType: FieldType.Text },
|
||||
{ fieldName: "vector", fieldType: FieldType.FloatVector, dimension: DIMS },
|
||||
{ fieldName: "metadata", fieldType: "JSON" },
|
||||
];
|
||||
|
||||
const BM25_FIELDS = [
|
||||
...CORE_FIELDS,
|
||||
{ fieldName: "textLemmatized", fieldType: FieldType.Text },
|
||||
];
|
||||
|
||||
/** Records call order, so ordering regressions (deleteCol vs. in-flight init) are visible. */
|
||||
function fakeClient(overrides: Record<string, (...args: any[]) => any> = {}) {
|
||||
const calls: string[] = [];
|
||||
const track =
|
||||
(name: string, impl: (...args: any[]) => any) =>
|
||||
(...args: any[]) => {
|
||||
calls.push(name);
|
||||
return impl(...args);
|
||||
};
|
||||
|
||||
const client: any = {
|
||||
calls,
|
||||
createDatabase: jest.fn(track("createDatabase", async () => OK)),
|
||||
createTable: jest.fn(track("createTable", async () => OK)),
|
||||
dropTable: jest.fn(track("dropTable", async () => OK)),
|
||||
descTable: jest.fn(track("descTable", async () => normalTable())),
|
||||
upsert: jest.fn(async () => OK),
|
||||
delete: jest.fn(async () => OK),
|
||||
query: jest.fn(),
|
||||
select: jest.fn(),
|
||||
vectorSearch: jest.fn(),
|
||||
bm25Search: jest.fn(),
|
||||
};
|
||||
|
||||
for (const [name, impl] of Object.entries(overrides)) {
|
||||
client[name] = jest.fn(track(name, impl));
|
||||
}
|
||||
return client;
|
||||
}
|
||||
|
||||
const makeStore = (client: any, extra: Record<string, unknown> = {}) =>
|
||||
new BaiduDB({
|
||||
endpoint: "http://127.0.0.1:5287",
|
||||
account: "root",
|
||||
apiKey: "test-key",
|
||||
databaseName: "mem0_db",
|
||||
tableName: "mem0",
|
||||
embeddingModelDims: DIMS,
|
||||
client,
|
||||
...extra,
|
||||
} as any);
|
||||
|
||||
/** Run the poll loop's setTimeout inline so tests never wait the real 2s interval. */
|
||||
const runTimersInline = () =>
|
||||
jest.spyOn(global, "setTimeout").mockImplementation(((fn: () => void) => {
|
||||
fn();
|
||||
return 0;
|
||||
}) as any);
|
||||
|
||||
beforeEach(() => {
|
||||
jest.spyOn(console, "warn").mockImplementation(() => {});
|
||||
jest.spyOn(console, "error").mockImplementation(() => {});
|
||||
});
|
||||
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
describe("BaiduDB config", () => {
|
||||
it("rejects a missing required field", () => {
|
||||
expect(() => makeStore(fakeClient(), { tableName: "" })).toThrow(
|
||||
/non-empty 'tableName'/,
|
||||
);
|
||||
});
|
||||
|
||||
it("does not require endpoint credentials when a client is injected", () => {
|
||||
expect(() =>
|
||||
makeStore(fakeClient(), { endpoint: "", account: "", apiKey: "" }),
|
||||
).not.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe("BaiduDB table provisioning", () => {
|
||||
it("creates the table with the schema mem0 needs", async () => {
|
||||
const client = fakeClient();
|
||||
await makeStore(client).initialize();
|
||||
|
||||
const spec = client.createTable.mock.calls[0][0];
|
||||
expect(client.createDatabase).toHaveBeenCalledWith("mem0_db");
|
||||
expect(spec.database).toBe("mem0_db");
|
||||
expect(spec.table).toBe("mem0");
|
||||
expect(spec.enableDynamicField).toBe(false);
|
||||
// Mochow rejects a partition without partitionType.
|
||||
expect(spec.partition).toEqual({
|
||||
partitionType: PartitionType.HASH,
|
||||
partitionNum: 1,
|
||||
});
|
||||
|
||||
const fields = spec.schema.fields.map((f: any) => [
|
||||
f.fieldName,
|
||||
f.fieldType,
|
||||
]);
|
||||
expect(fields).toEqual([
|
||||
["id", FieldType.String],
|
||||
["data", FieldType.Text],
|
||||
["vector", FieldType.FloatVector],
|
||||
["textLemmatized", FieldType.Text],
|
||||
["metadata", "JSON"],
|
||||
]);
|
||||
expect(spec.schema.fields[0]).toMatchObject({
|
||||
primaryKey: true,
|
||||
partitionKey: true,
|
||||
notNull: true,
|
||||
});
|
||||
expect(spec.schema.fields[2].dimension).toBe(DIMS);
|
||||
});
|
||||
|
||||
it("builds a vector index with a genuine row-count-increment auto-build policy", async () => {
|
||||
const client = fakeClient();
|
||||
await makeStore(client).initialize();
|
||||
|
||||
const [vectorIndex] = client.createTable.mock.calls[0][0].schema.indexes;
|
||||
expect(vectorIndex).toMatchObject({
|
||||
indexName: "vector_idx",
|
||||
indexType: IndexType.HNSW,
|
||||
field: "vector",
|
||||
metricType: MetricType.L2,
|
||||
params: { M: 16, efConstruction: 200 },
|
||||
autoBuild: true,
|
||||
});
|
||||
// Regression guard: sdk.AutoBuildIncrement() stamps policyType "TIMING" in 2.1.5.
|
||||
expect(vectorIndex.autoBuildPolicy).toEqual({
|
||||
policyType: AutoBuildPolicyType.Increment,
|
||||
rowCountIncrement: 10000,
|
||||
});
|
||||
expect(AutoBuildPolicyType.Increment).not.toBe(AutoBuildPolicyType.Timing);
|
||||
});
|
||||
|
||||
it("declares the filtering and BM25 indexes with an indexType", async () => {
|
||||
const client = fakeClient();
|
||||
await makeStore(client).initialize();
|
||||
|
||||
const [, filtering, bm25] =
|
||||
client.createTable.mock.calls[0][0].schema.indexes;
|
||||
expect(filtering).toEqual({
|
||||
indexName: "metadata_filtering_idx",
|
||||
indexType: IndexType.FilteringIndex,
|
||||
fields: ["metadata"],
|
||||
});
|
||||
// Memory.search() hands keywordSearch() an already-lemmatized query, so raw `data` is
|
||||
// not worth indexing — only the lemmatized column is. The index is therefore named for
|
||||
// that column and must never be called "data_bm25_idx": that is the name Python's
|
||||
// keyword_search() queries with a raw, unstemmed query, and it must keep missing (and so
|
||||
// falling back to vector search) rather than half-matching this stemmed index.
|
||||
expect(bm25).toMatchObject({
|
||||
indexName: "text_lemmatized_bm25_idx",
|
||||
indexType: IndexType.InvertedIndex,
|
||||
fields: ["textLemmatized"],
|
||||
fieldAttributes: [InvertedIndexFieldAttribute.Analyzed],
|
||||
});
|
||||
});
|
||||
|
||||
it("honours a configured metric type", async () => {
|
||||
const client = fakeClient();
|
||||
await makeStore(client, { metricType: "COSINE" }).initialize();
|
||||
const [vectorIndex] = client.createTable.mock.calls[0][0].schema.indexes;
|
||||
expect(vectorIndex.metricType).toBe(MetricType.COSINE);
|
||||
});
|
||||
|
||||
it("tolerates an existing database and table", async () => {
|
||||
const client = fakeClient({
|
||||
createDatabase: async () => ({
|
||||
code: ServerErrCode.DBAlreadyExist,
|
||||
msg: "db exists",
|
||||
}),
|
||||
createTable: async () => ({
|
||||
code: ServerErrCode.TableAlreadyExist,
|
||||
msg: "table exists",
|
||||
}),
|
||||
descTable: async () => normalTable({ fields: BM25_FIELDS, indexes: [] }),
|
||||
});
|
||||
await expect(makeStore(client).initialize()).resolves.toBeUndefined();
|
||||
});
|
||||
|
||||
it("waits for a CREATING table to become NORMAL", async () => {
|
||||
runTimersInline();
|
||||
const states = [
|
||||
TableState.Creating,
|
||||
TableState.Creating,
|
||||
TableState.Normal,
|
||||
];
|
||||
const client = fakeClient({
|
||||
descTable: async () => ({
|
||||
...OK,
|
||||
table: { state: states.shift(), schema: { fields: [], indexes: [] } },
|
||||
}),
|
||||
});
|
||||
|
||||
await makeStore(client).initialize();
|
||||
expect(client.descTable).toHaveBeenCalledTimes(3);
|
||||
});
|
||||
|
||||
it("surfaces a non-zero envelope as an error rather than succeeding", async () => {
|
||||
const client = fakeClient({
|
||||
createTable: async () => ({
|
||||
code: ServerErrCode.InvalidTableSchema,
|
||||
msg: "bad schema",
|
||||
}),
|
||||
});
|
||||
await expect(makeStore(client).initialize()).rejects.toThrow(
|
||||
/createTable 'mem0' failed \(code 60\): bad schema/,
|
||||
);
|
||||
});
|
||||
|
||||
it("rejects an existing table whose vector dimension disagrees", async () => {
|
||||
const client = fakeClient({
|
||||
createTable: async () => ({
|
||||
code: ServerErrCode.TableAlreadyExist,
|
||||
msg: "",
|
||||
}),
|
||||
descTable: async () =>
|
||||
normalTable({
|
||||
fields: [
|
||||
CORE_FIELDS[0],
|
||||
CORE_FIELDS[1],
|
||||
{
|
||||
fieldName: "vector",
|
||||
fieldType: FieldType.FloatVector,
|
||||
dimension: 768,
|
||||
},
|
||||
CORE_FIELDS[3],
|
||||
],
|
||||
indexes: [],
|
||||
}),
|
||||
});
|
||||
await expect(makeStore(client).initialize()).rejects.toThrow(
|
||||
/stores 768-dimensional vectors, but 'embeddingModelDims' is 1536/,
|
||||
);
|
||||
});
|
||||
|
||||
it("rejects an existing table missing the core schema", async () => {
|
||||
const client = fakeClient({
|
||||
createTable: async () => ({
|
||||
code: ServerErrCode.TableAlreadyExist,
|
||||
msg: "",
|
||||
}),
|
||||
descTable: async () =>
|
||||
normalTable({ fields: [CORE_FIELDS[0]], indexes: [] }),
|
||||
});
|
||||
await expect(makeStore(client).initialize()).rejects.toThrow(
|
||||
/missing the id\/data\/vector\/metadata schema/,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("BaiduDB keyword search support detection", () => {
|
||||
it("fails closed when an existing table has no inverted index", async () => {
|
||||
const client = fakeClient({
|
||||
createTable: async () => ({
|
||||
code: ServerErrCode.TableAlreadyExist,
|
||||
msg: "",
|
||||
}),
|
||||
descTable: async () => normalTable({ fields: CORE_FIELDS, indexes: [] }),
|
||||
});
|
||||
const store = makeStore(client);
|
||||
|
||||
await expect(store.keywordSearch("hello")).resolves.toBeNull();
|
||||
expect(client.bm25Search).not.toHaveBeenCalled();
|
||||
expect(console.warn).toHaveBeenCalledWith(
|
||||
expect.stringContaining("text_lemmatized_bm25_idx"),
|
||||
);
|
||||
});
|
||||
|
||||
it("enables keyword search when the existing table carries the BM25 index", async () => {
|
||||
const client = fakeClient({
|
||||
createTable: async () => ({
|
||||
code: ServerErrCode.TableAlreadyExist,
|
||||
msg: "",
|
||||
}),
|
||||
descTable: async () =>
|
||||
normalTable({
|
||||
fields: BM25_FIELDS,
|
||||
indexes: [{ indexName: "text_lemmatized_bm25_idx" }],
|
||||
}),
|
||||
});
|
||||
client.bm25Search.mockResolvedValue({ ...OK, rows: [] });
|
||||
|
||||
await expect(makeStore(client).keywordSearch("hello")).resolves.toEqual([]);
|
||||
expect(client.bm25Search).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("queries the inverted index with the caller's already-lemmatized text", async () => {
|
||||
const client = fakeClient();
|
||||
client.bm25Search.mockResolvedValue({
|
||||
...OK,
|
||||
rows: [
|
||||
{ row: { id: "m1", data: "loves pizza", metadata: {} }, score: 3.5 },
|
||||
],
|
||||
});
|
||||
|
||||
const results = await makeStore(client).keywordSearch("love pizza", 7, {
|
||||
userId: "alice",
|
||||
});
|
||||
expect(results).toEqual([
|
||||
{ id: "m1", payload: { data: "loves pizza" }, score: 3.5 },
|
||||
]);
|
||||
|
||||
const { request, ...ns } = client.bm25Search.mock.calls[0][0];
|
||||
expect(ns).toEqual({ database: "mem0_db", table: "mem0" });
|
||||
expect(request.indexName).toBe("text_lemmatized_bm25_idx");
|
||||
expect(request.searchText).toBe("love pizza");
|
||||
expect(request.limit).toBe(7);
|
||||
expect(request.filter).toBe('metadata["userId"] = "alice"');
|
||||
});
|
||||
});
|
||||
|
||||
describe("BaiduDB writes", () => {
|
||||
it("upserts the whole batch in one call and mirrors textLemmatized out of the payload", async () => {
|
||||
const client = fakeClient();
|
||||
|
||||
await makeStore(client).insert(
|
||||
[
|
||||
[1, 2],
|
||||
[3, 4],
|
||||
],
|
||||
["a", "b"],
|
||||
[
|
||||
{ data: "loves pizza", textLemmatized: "love pizza" },
|
||||
{ data: "runs daily" },
|
||||
],
|
||||
);
|
||||
|
||||
expect(client.upsert).toHaveBeenCalledTimes(1);
|
||||
expect(client.upsert.mock.calls[0][0]).toEqual({
|
||||
database: "mem0_db",
|
||||
table: "mem0",
|
||||
rows: [
|
||||
{
|
||||
id: "a",
|
||||
data: "loves pizza",
|
||||
vector: [1, 2],
|
||||
textLemmatized: "love pizza",
|
||||
metadata: {},
|
||||
},
|
||||
// Falls back to `data` when the caller did not lemmatize.
|
||||
{
|
||||
id: "b",
|
||||
data: "runs daily",
|
||||
vector: [3, 4],
|
||||
textLemmatized: "runs daily",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
});
|
||||
});
|
||||
|
||||
it("refuses a ragged batch instead of silently truncating it", async () => {
|
||||
await expect(
|
||||
makeStore(fakeClient()).insert([[1]], ["a", "b"], [{}]),
|
||||
).rejects.toThrow(/equal length \(got 1\/2\/1\)/);
|
||||
});
|
||||
|
||||
it("updates and deletes by primary key", async () => {
|
||||
const client = fakeClient();
|
||||
const store = makeStore(client);
|
||||
|
||||
await store.update("m1", [9], { data: "new" });
|
||||
expect(client.upsert.mock.calls[0][0].rows).toEqual([
|
||||
{
|
||||
id: "m1",
|
||||
data: "new",
|
||||
vector: [9],
|
||||
textLemmatized: "new",
|
||||
metadata: {},
|
||||
},
|
||||
]);
|
||||
|
||||
await store.delete("m1");
|
||||
expect(client.delete).toHaveBeenCalledWith({
|
||||
database: "mem0_db",
|
||||
table: "mem0",
|
||||
primaryKey: { id: "m1" },
|
||||
});
|
||||
});
|
||||
|
||||
it("throws when the server rejects an upsert", async () => {
|
||||
const client = fakeClient();
|
||||
client.upsert.mockResolvedValue({ code: 100, msg: "duplicate key" });
|
||||
|
||||
await expect(makeStore(client).insert([[1]], ["a"], [{}])).rejects.toThrow(
|
||||
/upsert failed \(code 100\): duplicate key/,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("BaiduDB reads", () => {
|
||||
it("maps vector search hits out of the nested row envelope", async () => {
|
||||
const client = fakeClient();
|
||||
client.vectorSearch.mockResolvedValue({
|
||||
...OK,
|
||||
rows: [
|
||||
{
|
||||
row: { id: "m1", data: "x", metadata: {} },
|
||||
distance: 0.2,
|
||||
score: 0.8,
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const results = await makeStore(client).search([1, 2, 3], 5, {
|
||||
userId: "alice",
|
||||
});
|
||||
expect(results).toEqual([{ id: "m1", payload: { data: "x" }, score: 0.8 }]);
|
||||
|
||||
const { request } = client.vectorSearch.mock.calls[0][0];
|
||||
expect(request.vectorField).toBe("vector");
|
||||
expect(request.vector).toEqual({ vector: [1, 2, 3] });
|
||||
expect(request.limit).toBe(5);
|
||||
expect(request.filter).toBe('metadata["userId"] = "alice"');
|
||||
expect(request.projections).toEqual(["id", "data", "metadata"]);
|
||||
expect(request.config.params).toEqual({ ef: 200 });
|
||||
});
|
||||
|
||||
it("omits the filter when no filters are supplied", async () => {
|
||||
const client = fakeClient();
|
||||
client.vectorSearch.mockResolvedValue({ ...OK, rows: [] });
|
||||
await makeStore(client).search([1], 5);
|
||||
expect(client.vectorSearch.mock.calls[0][0].request.filter).toBeUndefined();
|
||||
});
|
||||
|
||||
it("escapes quotes and rejects unsafe filter keys and values", async () => {
|
||||
const client = fakeClient();
|
||||
client.vectorSearch.mockResolvedValue({ ...OK, rows: [] });
|
||||
const store = makeStore(client);
|
||||
|
||||
await store.search([1], 5, { userId: 'a"b', runId: 3, agentId: true });
|
||||
expect(client.vectorSearch.mock.calls[0][0].request.filter).toBe(
|
||||
'metadata["userId"] = "a\\"b" AND metadata["runId"] = 3 AND metadata["agentId"] = true',
|
||||
);
|
||||
|
||||
await expect(store.search([1], 5, { "bad key": "x" })).rejects.toThrow(
|
||||
/Invalid filter key/,
|
||||
);
|
||||
await expect(
|
||||
store.search([1], 5, { userId: ["a"] as any }),
|
||||
).rejects.toThrow(/must be str, int, float, or bool, got array/);
|
||||
});
|
||||
|
||||
it("returns null for a missing id and throws on a real query failure", async () => {
|
||||
const client = fakeClient();
|
||||
const store = makeStore(client);
|
||||
|
||||
client.query.mockResolvedValue({ ...OK, row: {} });
|
||||
await expect(store.get("nope")).resolves.toBeNull();
|
||||
|
||||
client.query.mockResolvedValue({
|
||||
...OK,
|
||||
row: { id: "m1", data: "stored text", metadata: { a: 1 } },
|
||||
});
|
||||
await expect(store.get("m1")).resolves.toEqual({
|
||||
id: "m1",
|
||||
payload: { a: 1, data: "stored text" },
|
||||
});
|
||||
|
||||
client.query.mockResolvedValue({ code: 2, msg: "invalid parameter" });
|
||||
await expect(store.get("m1")).rejects.toThrow(
|
||||
/query 'm1' failed \(code 2\): invalid parameter/,
|
||||
);
|
||||
});
|
||||
|
||||
// The server signals a missing primary key with code 101; pymochow and mochow-sdk-go both
|
||||
// name it (ROW_KEY_NOT_FOUND / RowKeyNotFound). The Node SDK's ServerErrCode stops at 100,
|
||||
// so it has to be spelled out. Memory.get()/update()/delete() all branch on a null here.
|
||||
it("returns null when the server reports the row key is missing", async () => {
|
||||
const client = fakeClient();
|
||||
const store = makeStore(client);
|
||||
|
||||
client.query.mockResolvedValue({ code: 101, msg: "row key not found" });
|
||||
await expect(store.get("nope")).resolves.toBeNull();
|
||||
});
|
||||
|
||||
it("lists flat select rows and reports how many came back", async () => {
|
||||
const client = fakeClient();
|
||||
client.select.mockResolvedValue({
|
||||
...OK,
|
||||
isTruncated: false,
|
||||
nextMarker: "",
|
||||
rows: [{ id: "m1", data: "x", metadata: {} }, { id: "m2" }],
|
||||
});
|
||||
|
||||
await expect(
|
||||
makeStore(client).list({ userId: "alice" }, 50),
|
||||
).resolves.toEqual([
|
||||
[
|
||||
{ id: "m1", payload: { data: "x" } },
|
||||
{ id: "m2", payload: {} },
|
||||
],
|
||||
2,
|
||||
]);
|
||||
expect(client.select).toHaveBeenCalledWith({
|
||||
database: "mem0_db",
|
||||
table: "mem0",
|
||||
filter: 'metadata["userId"] = "alice"',
|
||||
projections: ["id", "data", "metadata"],
|
||||
limit: 50,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("BaiduDB deleteCol", () => {
|
||||
it("waits for the drop to land before returning", async () => {
|
||||
runTimersInline();
|
||||
const client = fakeClient();
|
||||
const store = makeStore(client);
|
||||
await store.initialize();
|
||||
|
||||
client.descTable
|
||||
.mockResolvedValueOnce({ ...OK, table: { state: TableState.Deleting } })
|
||||
.mockResolvedValueOnce({
|
||||
code: ServerErrCode.TableNotExist,
|
||||
msg: "gone",
|
||||
});
|
||||
|
||||
await store.deleteCol();
|
||||
expect(client.dropTable).toHaveBeenCalledWith("mem0_db", "mem0");
|
||||
expect(client.descTable).toHaveBeenCalledTimes(3); // 1 from initialize + 2 polls
|
||||
});
|
||||
|
||||
it("is a no-op when the table is already gone", async () => {
|
||||
const client = fakeClient();
|
||||
const store = makeStore(client);
|
||||
await store.initialize();
|
||||
|
||||
client.dropTable.mockResolvedValue({
|
||||
code: ServerErrCode.TableNotExist,
|
||||
msg: "gone",
|
||||
});
|
||||
await expect(store.deleteCol()).resolves.toBeUndefined();
|
||||
});
|
||||
|
||||
// Regression: deleteCol() used to run alongside the fire-and-forget initialize() the
|
||||
// constructor starts, so the in-flight createTable landed *after* dropTable and the table
|
||||
// survived reset().
|
||||
it("does not race the initialize() the constructor kicks off", async () => {
|
||||
runTimersInline();
|
||||
let release: () => void = () => {};
|
||||
const gate = new Promise<void>((resolve) => {
|
||||
release = resolve;
|
||||
});
|
||||
|
||||
// The table exists after the first init, is gone once dropTable lands, then exists again.
|
||||
const descQueue: unknown[] = [
|
||||
normalTable(),
|
||||
{ code: ServerErrCode.TableNotExist, msg: "gone" },
|
||||
];
|
||||
const client = fakeClient({
|
||||
createDatabase: async () => {
|
||||
await gate;
|
||||
return OK;
|
||||
},
|
||||
descTable: async () => descQueue.shift() ?? normalTable(),
|
||||
});
|
||||
|
||||
const store = makeStore(client); // initialize() is now in flight, parked on `gate`
|
||||
const resetting = store.reset();
|
||||
release();
|
||||
await resetting;
|
||||
|
||||
expect(client.calls).toEqual([
|
||||
"createDatabase",
|
||||
"createTable",
|
||||
"descTable",
|
||||
"dropTable",
|
||||
"descTable",
|
||||
"createDatabase",
|
||||
"createTable",
|
||||
"descTable",
|
||||
]);
|
||||
expect(client.calls.indexOf("dropTable")).toBeGreaterThan(
|
||||
client.calls.indexOf("createTable"),
|
||||
);
|
||||
expect(client.createTable).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
});
|
||||
|
||||
describe("BaiduDB user id", () => {
|
||||
it("round-trips the store user id", async () => {
|
||||
const store = makeStore(fakeClient());
|
||||
await expect(store.getUserId()).resolves.toBe("anonymous-baidu-user");
|
||||
await store.setUserId("alice");
|
||||
await expect(store.getUserId()).resolves.toBe("alice");
|
||||
});
|
||||
});
|
||||
@@ -138,6 +138,11 @@ jest.mock("../src/vector_stores/qdrant", () => ({
|
||||
.fn()
|
||||
.mockImplementation((config) => ({ type: "qdrant", config })),
|
||||
}));
|
||||
jest.mock("../src/vector_stores/baidu", () => ({
|
||||
BaiduDB: jest
|
||||
.fn()
|
||||
.mockImplementation((config) => ({ type: "baidu", config })),
|
||||
}));
|
||||
jest.mock("../src/vector_stores/redis", () => ({
|
||||
RedisDB: jest
|
||||
.fn()
|
||||
@@ -320,6 +325,7 @@ describe("VectorStoreFactory", () => {
|
||||
});
|
||||
|
||||
test.each([
|
||||
["baidu"],
|
||||
["qdrant"],
|
||||
["redis"],
|
||||
["valkey"],
|
||||
@@ -335,9 +341,8 @@ describe("VectorStoreFactory", () => {
|
||||
["s3_vectors"],
|
||||
["weaviate"],
|
||||
])("creates vector store for provider '%s'", (provider) => {
|
||||
expect(() =>
|
||||
VectorStoreFactory.create(provider, dummyVSConfig),
|
||||
).not.toThrow();
|
||||
const result = VectorStoreFactory.create(provider, dummyVSConfig) as any;
|
||||
expect(result.config).toBe(dummyVSConfig);
|
||||
});
|
||||
|
||||
test("throws for unsupported provider", () => {
|
||||
|
||||
@@ -5,6 +5,7 @@ const external = [
|
||||
"openai",
|
||||
"@anthropic-ai/sdk",
|
||||
"@aws-sdk/client-s3vectors",
|
||||
"@mochow/mochow-sdk-node",
|
||||
"groq-sdk",
|
||||
"cohere-ai",
|
||||
"@huggingface/transformers",
|
||||
|
||||
Reference in New Issue
Block a user