Compare commits
9 Commits
tantivy_im
...
lancedb_im
Author | SHA1 | Date | |
---|---|---|---|
|
a8720c06db | ||
|
2139deb85d | ||
|
7d07b57d32 | ||
|
4aa49e0d5c | ||
|
644946f1ca | ||
|
cf66f4c304 | ||
|
6a4e2819bf | ||
|
77a53bae86 | ||
7f689ae29b |
4440
Cargo.lock
generated
4440
Cargo.lock
generated
File diff suppressed because it is too large
Load Diff
@@ -27,6 +27,14 @@ x25519-dalek = "2"
|
||||
base64 = "0.22"
|
||||
jsonrpsee = { version = "0.26.0", features = ["http-client", "ws-client", "server", "macros"] }
|
||||
tantivy = "0.25.0"
|
||||
arrow-schema = "55.2.0"
|
||||
arrow-array = "55.2.0"
|
||||
lance = "0.37.0"
|
||||
lance-index = "0.37.0"
|
||||
arrow = "55.2.0"
|
||||
lancedb = "0.22.1"
|
||||
uuid = "1.18.1"
|
||||
ureq = { version = "2.10.0", features = ["json", "tls"] }
|
||||
|
||||
[dev-dependencies]
|
||||
redis = { version = "0.24", features = ["aio", "tokio-comp"] }
|
||||
|
12
README.md
12
README.md
@@ -47,18 +47,24 @@ HeroDB can be interacted with using any standard Redis client, such as `redis-cl
|
||||
|
||||
### Example with `redis-cli`
|
||||
|
||||
Connections start with no database selected. You must SELECT a database first.
|
||||
|
||||
- To work in the admin database (DB 0), authenticate with the admin secret:
|
||||
```bash
|
||||
redis-cli -p 6379 SELECT 0 KEY myadminsecret
|
||||
redis-cli -p 6379 SET mykey "Hello from HeroDB!"
|
||||
redis-cli -p 6379 GET mykey
|
||||
# → "Hello from HeroDB!"
|
||||
```
|
||||
|
||||
- To use a user database, first create one via the JSON-RPC API (see docs/rpc_examples.md), then select it:
|
||||
```bash
|
||||
# Suppose RPC created database id 1
|
||||
redis-cli -p 6379 SELECT 1
|
||||
redis-cli -p 6379 HSET user:1 name "Alice" age "30"
|
||||
redis-cli -p 6379 HGET user:1 name
|
||||
# → "Alice"
|
||||
|
||||
redis-cli -p 6379 SCAN 0 MATCH user:* COUNT 10
|
||||
# → 1) "0"
|
||||
# 2) 1) "user:1"
|
||||
```
|
||||
|
||||
## Cryptography
|
||||
|
@@ -80,6 +80,7 @@ Keys in `DB 0` (internal layout, but useful to understand how things work):
|
||||
- Requires the exact admin secret as the `KEY` argument to `SELECT 0`
|
||||
- Permission is `ReadWrite` when the secret matches
|
||||
|
||||
Connections start with no database selected. Any command that requires storage (GET, SET, H*, L*, SCAN, etc.) will return an error until you issue a SELECT to choose a database. Admin DB 0 is never accessible without authenticating via SELECT 0 KEY <admin_secret>.
|
||||
### How to select databases with optional `KEY`
|
||||
|
||||
- Public DB (no key required)
|
||||
|
10
docs/cmds.md
10
docs/cmds.md
@@ -126,7 +126,9 @@ redis-cli -p 6381 --pipe < dump.rdb
|
||||
|
||||
## Authentication and Database Selection
|
||||
|
||||
HeroDB uses an `Admin DB 0` to govern database existence, access and per-db encryption. Access control is enforced via `Admin DB 0` metadata. See the full model in `docs/admin.md`.
|
||||
Connections start with no database selected. Any storage-backed command (GET, SET, H*, L*, SCAN, etc.) will return an error until you issue a SELECT to choose a database.
|
||||
|
||||
HeroDB uses an `Admin DB 0` to govern database existence, access and per-db encryption. Access control is enforced via `Admin DB 0` metadata. See the full model in (docs/admin.md:1).
|
||||
|
||||
Examples:
|
||||
```bash
|
||||
@@ -145,4 +147,10 @@ redis-cli -p $PORT SELECT 2 KEY my-db2-access-key
|
||||
# Admin DB 0 (requires admin secret)
|
||||
redis-cli -p $PORT SELECT 0 KEY my-admin-secret
|
||||
# → OK
|
||||
```
|
||||
|
||||
```bash
|
||||
# Before selecting a DB, storage commands will fail
|
||||
redis-cli -p $PORT GET key
|
||||
# → -ERR No database selected. Use SELECT <id> [KEY <key>] first
|
||||
```
|
444
docs/lance.md
Normal file
444
docs/lance.md
Normal file
@@ -0,0 +1,444 @@
|
||||
# Lance Vector Backend (RESP + JSON-RPC)
|
||||
|
||||
This document explains how to use HeroDB’s Lance-backed vector store. It is text-first: users provide text, and HeroDB computes embeddings server-side (no manual vectors). It includes copy-pasteable RESP (redis-cli) and JSON-RPC examples for:
|
||||
|
||||
- Creating a Lance database
|
||||
- Embedding provider configuration (OpenAI, Azure OpenAI, or deterministic test provider)
|
||||
- Dataset lifecycle: CREATE, LIST, INFO, DROP
|
||||
- Ingestion: STORE text (+ optional metadata)
|
||||
- Search: QUERY with K, optional FILTER and RETURN
|
||||
- Delete by id
|
||||
- Index creation (currently a placeholder/no-op)
|
||||
|
||||
References:
|
||||
- Implementation: [src/lance_store.rs](src/lance_store.rs), [src/cmd.rs](src/cmd.rs), [src/rpc.rs](src/rpc.rs), [src/server.rs](src/server.rs), [src/embedding.rs](src/embedding.rs)
|
||||
|
||||
Notes:
|
||||
- Admin DB 0 cannot be Lance (or Tantivy). Only databases with id >= 1 can use Lance.
|
||||
- Permissions:
|
||||
- Read operations (SEARCH, LIST, INFO) require read permission.
|
||||
- Mutating operations (CREATE, STORE, CREATEINDEX, DEL, DROP, EMBEDDING CONFIG SET) require readwrite permission.
|
||||
- Backend gating:
|
||||
- If a DB is Lance, only LANCE.* and basic control commands (PING, ECHO, SELECT, INFO, CLIENT, etc.) are permitted.
|
||||
- If a DB is not Lance, LANCE.* commands return an error.
|
||||
|
||||
Storage layout and schema:
|
||||
- Files live at: <base_dir>/lance/<db_id>/<dataset>.lance
|
||||
- Records schema:
|
||||
- id: Utf8 (non-null)
|
||||
- vector: FixedSizeList<Float32, dim> (non-null)
|
||||
- text: Utf8 (nullable)
|
||||
- meta: Utf8 JSON (nullable)
|
||||
- Search is an L2 KNN brute-force scan for now (lower score = better). Index creation is a no-op placeholder to be implemented later.
|
||||
|
||||
Prerequisites:
|
||||
- Start HeroDB with RPC enabled (for management calls):
|
||||
- See [docs/basics.md](./basics.md) for flags. Example:
|
||||
```bash
|
||||
./target/release/herodb --dir /tmp/herodb --admin-secret mysecret --port 6379 --enable-rpc
|
||||
```
|
||||
|
||||
|
||||
## 0) Create a Lance-backed database (JSON-RPC)
|
||||
|
||||
Use the management API to create a database with backend "Lance". DB 0 is reserved for admin and cannot be Lance.
|
||||
|
||||
Request:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "herodb_createDatabase",
|
||||
"params": [
|
||||
"Lance",
|
||||
{ "name": "vectors-db", "storage_path": null, "max_size": null, "redis_version": null },
|
||||
null
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
- Response contains the allocated db_id (>= 1). Use that id below (replace 1 with your actual id).
|
||||
|
||||
Select the database over RESP:
|
||||
```bash
|
||||
redis-cli -p 6379 SELECT 1
|
||||
# → OK
|
||||
```
|
||||
|
||||
|
||||
## 1) Configure embedding provider (server-side embeddings)
|
||||
|
||||
HeroDB embeds text internally at STORE/SEARCH time using a per-dataset EmbeddingConfig sidecar. Configure provider before creating a dataset to choose dimensions and provider.
|
||||
|
||||
Supported providers:
|
||||
- openai (standard OpenAI or Azure OpenAI)
|
||||
- testhash (deterministic, CI-friendly; no network)
|
||||
|
||||
Environment variables for OpenAI:
|
||||
- Standard OpenAI: export OPENAI_API_KEY=sk-...
|
||||
- Azure OpenAI: export AZURE_OPENAI_API_KEY=...
|
||||
|
||||
RESP examples:
|
||||
```bash
|
||||
# Standard OpenAI with default dims (model-dependent, e.g. 1536)
|
||||
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET myset PROVIDER openai MODEL text-embedding-3-small
|
||||
|
||||
# OpenAI with reduced output dimension (e.g., 512) when supported
|
||||
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET myset PROVIDER openai MODEL text-embedding-3-small PARAM dim 512
|
||||
|
||||
# Azure OpenAI (set env: AZURE_OPENAI_API_KEY)
|
||||
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET myset PROVIDER openai MODEL text-embedding-3-small \
|
||||
PARAM use_azure true \
|
||||
PARAM azure_endpoint https://myresource.openai.azure.com \
|
||||
PARAM azure_deployment my-embed-deploy \
|
||||
PARAM azure_api_version 2024-02-15 \
|
||||
PARAM dim 512
|
||||
|
||||
# Deterministic test provider (no network, stable vectors)
|
||||
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET myset PROVIDER testhash MODEL any
|
||||
```
|
||||
|
||||
Read config:
|
||||
```bash
|
||||
redis-cli -p 6379 LANCE.EMBEDDING CONFIG GET myset
|
||||
# → JSON blob describing provider/model/params
|
||||
```
|
||||
|
||||
JSON-RPC examples:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 2,
|
||||
"method": "herodb_lanceSetEmbeddingConfig",
|
||||
"params": [
|
||||
1,
|
||||
"myset",
|
||||
"openai",
|
||||
"text-embedding-3-small",
|
||||
{ "dim": "512" }
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 3,
|
||||
"method": "herodb_lanceGetEmbeddingConfig",
|
||||
"params": [1, "myset"]
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## 2) Create a dataset
|
||||
|
||||
Choose a dimension that matches your embedding configuration. For OpenAI text-embedding-3-small without dimension override, typical dimension is 1536; when `dim` is set (e.g., 512), use that. The current API requires an explicit DIM.
|
||||
|
||||
RESP:
|
||||
```bash
|
||||
redis-cli -p 6379 LANCE.CREATE myset DIM 512
|
||||
# → OK
|
||||
```
|
||||
|
||||
JSON-RPC:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 4,
|
||||
"method": "herodb_lanceCreate",
|
||||
"params": [1, "myset", 512]
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## 3) Store text documents (server-side embedding)
|
||||
|
||||
Provide your id, the text to embed, and optional META fields. The server computes the embedding using the configured provider and stores id/vector/text/meta in the Lance dataset. Upserts by id are supported via delete-then-append semantics.
|
||||
|
||||
RESP:
|
||||
```bash
|
||||
redis-cli -p 6379 LANCE.STORE myset ID doc-1 TEXT "Hello vector world" META title "Hello" category "demo"
|
||||
# → OK
|
||||
```
|
||||
|
||||
JSON-RPC:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 5,
|
||||
"method": "herodb_lanceStoreText",
|
||||
"params": [
|
||||
1,
|
||||
"myset",
|
||||
"doc-1",
|
||||
"Hello vector world",
|
||||
{ "title": "Hello", "category": "demo" }
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## 4) Search with a text query
|
||||
|
||||
Provide a query string; the server embeds it and performs KNN search. Optional: FILTER expression and RETURN subset of fields.
|
||||
|
||||
RESP:
|
||||
```bash
|
||||
# K nearest neighbors for the query text
|
||||
redis-cli -p 6379 LANCE.SEARCH myset K 5 QUERY "greetings to vectors"
|
||||
# → Array of hits: [id, score, [k,v, ...]] pairs, lower score = closer
|
||||
|
||||
# With a filter on meta fields and return only title
|
||||
redis-cli -p 6379 LANCE.SEARCH myset K 3 QUERY "greetings to vectors" FILTER "category = 'demo'" RETURN 1 title
|
||||
```
|
||||
|
||||
JSON-RPC:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 6,
|
||||
"method": "herodb_lanceSearchText",
|
||||
"params": [1, "myset", "greetings to vectors", 5, null, null]
|
||||
}
|
||||
```
|
||||
|
||||
With filter and selected fields:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 7,
|
||||
"method": "herodb_lanceSearchText",
|
||||
"params": [1, "myset", "greetings to vectors", 3, "category = 'demo'", ["title"]]
|
||||
}
|
||||
```
|
||||
|
||||
Response shape:
|
||||
- RESP over redis-cli: an array of hits [id, score, [k, v, ...]].
|
||||
- JSON-RPC returns an object containing the RESP-encoded wire format string or a structured result depending on implementation. See [src/rpc.rs](src/rpc.rs) for details.
|
||||
|
||||
|
||||
## 5) Create an index (placeholder)
|
||||
|
||||
Index creation currently returns OK but is a no-op. It will integrate Lance vector indices in a future update.
|
||||
|
||||
RESP:
|
||||
```bash
|
||||
redis-cli -p 6379 LANCE.CREATEINDEX myset TYPE "ivf_pq" PARAM nlist 100 PARAM pq_m 16
|
||||
# → OK (no-op for now)
|
||||
```
|
||||
|
||||
JSON-RPC:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 8,
|
||||
"method": "herodb_lanceCreateIndex",
|
||||
"params": [1, "myset", "ivf_pq", { "nlist": "100", "pq_m": "16" }]
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## 6) Inspect datasets
|
||||
|
||||
RESP:
|
||||
```bash
|
||||
# List datasets in current Lance DB
|
||||
redis-cli -p 6379 LANCE.LIST
|
||||
|
||||
# Get dataset info
|
||||
redis-cli -p 6379 LANCE.INFO myset
|
||||
```
|
||||
|
||||
JSON-RPC:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 9,
|
||||
"method": "herodb_lanceList",
|
||||
"params": [1]
|
||||
}
|
||||
```
|
||||
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 10,
|
||||
"method": "herodb_lanceInfo",
|
||||
"params": [1, "myset"]
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## 7) Delete and drop
|
||||
|
||||
RESP:
|
||||
```bash
|
||||
# Delete by id
|
||||
redis-cli -p 6379 LANCE.DEL myset doc-1
|
||||
# → OK
|
||||
|
||||
# Drop the entire dataset
|
||||
redis-cli -p 6379 LANCE.DROP myset
|
||||
# → OK
|
||||
```
|
||||
|
||||
JSON-RPC:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 11,
|
||||
"method": "herodb_lanceDel",
|
||||
"params": [1, "myset", "doc-1"]
|
||||
}
|
||||
```
|
||||
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 12,
|
||||
"method": "herodb_lanceDrop",
|
||||
"params": [1, "myset"]
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## 8) End-to-end example (RESP)
|
||||
|
||||
```bash
|
||||
# 1. Select Lance DB (assume db_id=1 created via RPC)
|
||||
redis-cli -p 6379 SELECT 1
|
||||
|
||||
# 2. Configure embedding provider (OpenAI small model at 512 dims)
|
||||
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET myset PROVIDER openai MODEL text-embedding-3-small PARAM dim 512
|
||||
|
||||
# 3. Create dataset
|
||||
redis-cli -p 6379 LANCE.CREATE myset DIM 512
|
||||
|
||||
# 4. Store documents
|
||||
redis-cli -p 6379 LANCE.STORE myset ID doc-1 TEXT "The quick brown fox jumps over the lazy dog" META title "Fox" category "animal"
|
||||
redis-cli -p 6379 LANCE.STORE myset ID doc-2 TEXT "A fast auburn fox vaulted a sleepy canine" META title "Fox paraphrase" category "animal"
|
||||
|
||||
# 5. Search
|
||||
redis-cli -p 6379 LANCE.SEARCH myset K 2 QUERY "quick brown fox" RETURN 1 title
|
||||
|
||||
# 6. Dataset info and listing
|
||||
redis-cli -p 6379 LANCE.INFO myset
|
||||
redis-cli -p 6379 LANCE.LIST
|
||||
|
||||
# 7. Delete and drop
|
||||
redis-cli -p 6379 LANCE.DEL myset doc-2
|
||||
redis-cli -p 6379 LANCE.DROP myset
|
||||
```
|
||||
|
||||
|
||||
## 9) End-to-end example (JSON-RPC)
|
||||
|
||||
Assume RPC server on port 8080. Replace ids and ports as needed.
|
||||
|
||||
1) Create Lance DB:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 100,
|
||||
"method": "herodb_createDatabase",
|
||||
"params": ["Lance", { "name": "vectors-db", "storage_path": null, "max_size": null, "redis_version": null }, null]
|
||||
}
|
||||
```
|
||||
|
||||
2) Set embedding config:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 101,
|
||||
"method": "herodb_lanceSetEmbeddingConfig",
|
||||
"params": [1, "myset", "openai", "text-embedding-3-small", { "dim": "512" }]
|
||||
}
|
||||
```
|
||||
|
||||
3) Create dataset:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 102,
|
||||
"method": "herodb_lanceCreate",
|
||||
"params": [1, "myset", 512]
|
||||
}
|
||||
```
|
||||
|
||||
4) Store text:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 103,
|
||||
"method": "herodb_lanceStoreText",
|
||||
"params": [1, "myset", "doc-1", "The quick brown fox jumps over the lazy dog", { "title": "Fox", "category": "animal" }]
|
||||
}
|
||||
```
|
||||
|
||||
5) Search text:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 104,
|
||||
"method": "herodb_lanceSearchText",
|
||||
"params": [1, "myset", "quick brown fox", 2, null, ["title"]]
|
||||
}
|
||||
```
|
||||
|
||||
6) Info/list:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 105,
|
||||
"method": "herodb_lanceInfo",
|
||||
"params": [1, "myset"]
|
||||
}
|
||||
```
|
||||
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 106,
|
||||
"method": "herodb_lanceList",
|
||||
"params": [1]
|
||||
}
|
||||
```
|
||||
|
||||
7) Delete/drop:
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 107,
|
||||
"method": "herodb_lanceDel",
|
||||
"params": [1, "myset", "doc-1"]
|
||||
}
|
||||
```
|
||||
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 108,
|
||||
"method": "herodb_lanceDrop",
|
||||
"params": [1, "myset"]
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## 10) Operational notes and troubleshooting
|
||||
|
||||
- If using OpenAI and you see “missing API key env”, set:
|
||||
- Standard: `export OPENAI_API_KEY=sk-...`
|
||||
- Azure: `export AZURE_OPENAI_API_KEY=...` and pass `use_azure true`, `azure_endpoint`, `azure_deployment`, `azure_api_version`.
|
||||
- Dimensions mismatch:
|
||||
- Ensure the dataset DIM equals the provider’s embedding dim. For OpenAI text-embedding-3 models, set `PARAM dim 512` (or another supported size) and use that same DIM for `LANCE.CREATE`.
|
||||
- DB 0 restriction:
|
||||
- Lance is not allowed on DB 0. Use db_id >= 1.
|
||||
- Permissions:
|
||||
- Read operations (SEARCH, LIST, INFO) require read permission.
|
||||
- Mutations (CREATE, STORE, CREATEINDEX, DEL, DROP, EMBEDDING CONFIG SET) require readwrite permission.
|
||||
- Backend gating:
|
||||
- On Lance DBs, only LANCE.* commands are accepted (plus basic control).
|
||||
- Current index behavior:
|
||||
- `LANCE.CREATEINDEX` returns OK but is a no-op. Future versions will integrate Lance vector indices.
|
||||
- Implementation files for reference:
|
||||
- [src/lance_store.rs](src/lance_store.rs), [src/cmd.rs](src/cmd.rs), [src/rpc.rs](src/rpc.rs), [src/server.rs](src/server.rs), [src/embedding.rs](src/embedding.rs)
|
138
docs/lancedb_text_and_images_example.md
Normal file
138
docs/lancedb_text_and_images_example.md
Normal file
@@ -0,0 +1,138 @@
|
||||
# LanceDB Text and Images: End-to-End Example
|
||||
|
||||
This guide demonstrates creating a Lance backend database, ingesting two text documents and two images, performing searches over both, and cleaning up the datasets.
|
||||
|
||||
Prerequisites
|
||||
- Build HeroDB and start the server with JSON-RPC enabled.
|
||||
Commands:
|
||||
```bash
|
||||
cargo build --release
|
||||
./target/release/herodb --dir /tmp/herodb --admin-secret mysecret --port 6379 --enable-rpc
|
||||
```
|
||||
|
||||
We'll use:
|
||||
- redis-cli for RESP commands against port 6379
|
||||
- curl for JSON-RPC against 8080 if desired
|
||||
- Deterministic local embedders to avoid external dependencies: testhash (text, dim 64) and testimagehash (image, dim 512)
|
||||
|
||||
0) Create a Lance-backed database (JSON-RPC)
|
||||
Request:
|
||||
```json
|
||||
{ "jsonrpc": "2.0", "id": 1, "method": "herodb_createDatabase", "params": ["Lance", { "name": "media-db", "storage_path": null, "max_size": null, "redis_version": null }, null] }
|
||||
```
|
||||
Response returns db_id (assume 1). Select DB over RESP:
|
||||
```bash
|
||||
redis-cli -p 6379 SELECT 1
|
||||
# → OK
|
||||
```
|
||||
|
||||
1) Configure embedding providers
|
||||
We'll create two datasets with independent embedding configs:
|
||||
- textset → provider testhash, dim 64
|
||||
- imageset → provider testimagehash, dim 512
|
||||
|
||||
Text config:
|
||||
```bash
|
||||
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET textset PROVIDER testhash MODEL any PARAM dim 64
|
||||
# → OK
|
||||
```
|
||||
Image config:
|
||||
```bash
|
||||
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET imageset PROVIDER testimagehash MODEL any PARAM dim 512
|
||||
# → OK
|
||||
```
|
||||
|
||||
2) Create datasets
|
||||
```bash
|
||||
redis-cli -p 6379 LANCE.CREATE textset DIM 64
|
||||
# → OK
|
||||
redis-cli -p 6379 LANCE.CREATE imageset DIM 512
|
||||
# → OK
|
||||
```
|
||||
|
||||
3) Ingest two text documents (server-side embedding)
|
||||
```bash
|
||||
redis-cli -p 6379 LANCE.STORE textset ID doc-1 TEXT "The quick brown fox jumps over the lazy dog" META title "Fox" category "animal"
|
||||
# → OK
|
||||
redis-cli -p 6379 LANCE.STORE textset ID doc-2 TEXT "A fast auburn fox vaulted a sleepy canine" META title "Paraphrase" category "animal"
|
||||
# → OK
|
||||
```
|
||||
|
||||
4) Ingest two images
|
||||
You can provide a URI or base64 bytes. Use URI for URIs, BYTES for base64 data.
|
||||
Example using free placeholder images:
|
||||
```bash
|
||||
# Store via URI
|
||||
redis-cli -p 6379 LANCE.STOREIMAGE imageset ID img-1 URI "https://picsum.photos/seed/1/256/256" META title "Seed1" group "demo"
|
||||
# → OK
|
||||
redis-cli -p 6379 LANCE.STOREIMAGE imageset ID img-2 URI "https://picsum.photos/seed/2/256/256" META title "Seed2" group "demo"
|
||||
# → OK
|
||||
```
|
||||
If your environment blocks outbound HTTP, you can embed image bytes:
|
||||
```bash
|
||||
# Example: read a local file and base64 it (replace path)
|
||||
b64=$(base64 -w0 ./image1.png)
|
||||
redis-cli -p 6379 LANCE.STOREIMAGE imageset ID img-b64-1 BYTES "$b64" META title "Local1" group "demo"
|
||||
```
|
||||
|
||||
5) Search text
|
||||
```bash
|
||||
# Top-2 nearest neighbors for a query
|
||||
redis-cli -p 6379 LANCE.SEARCH textset K 2 QUERY "quick brown fox" RETURN 1 title
|
||||
# → 1) [id, score, [k1,v1,...]]
|
||||
```
|
||||
With a filter (supports equality on schema or meta keys):
|
||||
```bash
|
||||
redis-cli -p 6379 LANCE.SEARCH textset K 2 QUERY "fox jumps" FILTER "category = 'animal'" RETURN 1 title
|
||||
```
|
||||
|
||||
6) Search images
|
||||
```bash
|
||||
# Provide a URI as the query
|
||||
redis-cli -p 6379 LANCE.SEARCHIMAGE imageset K 2 QUERYURI "https://picsum.photos/seed/1/256/256" RETURN 1 title
|
||||
|
||||
# Or provide base64 bytes as the query
|
||||
qb64=$(curl -s https://picsum.photos/seed/3/256/256 | base64 -w0)
|
||||
redis-cli -p 6379 LANCE.SEARCHIMAGE imageset K 2 QUERYBYTES "$qb64" RETURN 1 title
|
||||
```
|
||||
|
||||
7) Inspect datasets
|
||||
```bash
|
||||
redis-cli -p 6379 LANCE.LIST
|
||||
redis-cli -p 6379 LANCE.INFO textset
|
||||
redis-cli -p 6379 LANCE.INFO imageset
|
||||
```
|
||||
|
||||
8) Delete by id and drop datasets
|
||||
```bash
|
||||
# Delete one record
|
||||
redis-cli -p 6379 LANCE.DEL textset doc-2
|
||||
# → OK
|
||||
|
||||
# Drop entire datasets
|
||||
redis-cli -p 6379 LANCE.DROP textset
|
||||
redis-cli -p 6379 LANCE.DROP imageset
|
||||
# → OK
|
||||
```
|
||||
|
||||
Appendix: Using OpenAI embeddings instead of test providers
|
||||
Text:
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-...
|
||||
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET textset PROVIDER openai MODEL text-embedding-3-small PARAM dim 512
|
||||
redis-cli -p 6379 LANCE.CREATE textset DIM 512
|
||||
```
|
||||
Azure OpenAI:
|
||||
```bash
|
||||
export AZURE_OPENAI_API_KEY=...
|
||||
redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET textset PROVIDER openai MODEL text-embedding-3-small \
|
||||
PARAM use_azure true \
|
||||
PARAM azure_endpoint https://myresource.openai.azure.com \
|
||||
PARAM azure_deployment my-embed-deploy \
|
||||
PARAM azure_api_version 2024-02-15 \
|
||||
PARAM dim 512
|
||||
```
|
||||
Notes:
|
||||
- Ensure dataset DIM matches the configured embedding dimension.
|
||||
- Lance is only available for non-admin databases (db_id >= 1).
|
||||
- On Lance DBs, only LANCE.* and basic control commands are allowed.
|
@@ -1,4 +1,4 @@
|
||||
use std::path::PathBuf;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, OnceLock, Mutex, RwLock};
|
||||
use std::collections::HashMap;
|
||||
|
||||
@@ -35,11 +35,11 @@ static DATA_STORAGES: OnceLock<RwLock<HashMap<u64, Arc<dyn StorageBackend>>>> =
|
||||
static DATA_INIT_LOCK: Mutex<()> = Mutex::new(());
|
||||
|
||||
fn init_admin_storage(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
) -> Result<Arc<dyn StorageBackend>, DBError> {
|
||||
let db_file = PathBuf::from(base_dir).join("0.db");
|
||||
let db_file = base_dir.join("0.db");
|
||||
if let Some(parent_dir) = db_file.parent() {
|
||||
std::fs::create_dir_all(parent_dir).map_err(|e| {
|
||||
DBError(format!("Failed to create directory {}: {}", parent_dir.display(), e))
|
||||
@@ -48,8 +48,8 @@ fn init_admin_storage(
|
||||
let storage: Arc<dyn StorageBackend> = match backend {
|
||||
options::BackendType::Redb => Arc::new(Storage::new(&db_file, true, Some(admin_secret))?),
|
||||
options::BackendType::Sled => Arc::new(SledStorage::new(&db_file, true, Some(admin_secret))?),
|
||||
options::BackendType::Tantivy => {
|
||||
return Err(DBError("Admin DB 0 cannot use Tantivy backend".to_string()))
|
||||
options::BackendType::Tantivy | options::BackendType::Lance => {
|
||||
return Err(DBError("Admin DB 0 cannot use search-only backends (Tantivy/Lance)".to_string()))
|
||||
}
|
||||
};
|
||||
Ok(storage)
|
||||
@@ -57,24 +57,25 @@ fn init_admin_storage(
|
||||
|
||||
// Get or initialize a cached handle to admin DB 0 per base_dir (thread-safe, no double-open race)
|
||||
pub fn open_admin_storage(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
) -> Result<Arc<dyn StorageBackend>, DBError> {
|
||||
let map = ADMIN_STORAGES.get_or_init(|| RwLock::new(HashMap::new()));
|
||||
let key = base_dir.display().to_string();
|
||||
// Fast path
|
||||
if let Some(st) = map.read().unwrap().get(base_dir) {
|
||||
if let Some(st) = map.read().unwrap().get(&key) {
|
||||
return Ok(st.clone());
|
||||
}
|
||||
// Slow path with write lock
|
||||
{
|
||||
let mut w = map.write().unwrap();
|
||||
if let Some(st) = w.get(base_dir) {
|
||||
if let Some(st) = w.get(&key) {
|
||||
return Ok(st.clone());
|
||||
}
|
||||
|
||||
// Detect existing 0.db backend by filesystem, if present.
|
||||
let admin_path = PathBuf::from(base_dir).join("0.db");
|
||||
let admin_path = base_dir.join("0.db");
|
||||
let detected = if admin_path.exists() {
|
||||
if admin_path.is_file() {
|
||||
Some(options::BackendType::Redb)
|
||||
@@ -102,14 +103,14 @@ pub fn open_admin_storage(
|
||||
};
|
||||
|
||||
let st = init_admin_storage(base_dir, effective_backend, admin_secret)?;
|
||||
w.insert(base_dir.to_string(), st.clone());
|
||||
w.insert(key, st.clone());
|
||||
Ok(st)
|
||||
}
|
||||
}
|
||||
|
||||
// Ensure admin structures exist in encrypted DB 0
|
||||
pub fn ensure_bootstrap(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
) -> Result<(), DBError> {
|
||||
@@ -125,7 +126,7 @@ pub fn ensure_bootstrap(
|
||||
|
||||
// Get or initialize a shared handle to a data DB (> 0), avoiding double-open across subsystems
|
||||
pub fn open_data_storage(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
id: u64,
|
||||
@@ -159,7 +160,7 @@ pub fn open_data_storage(
|
||||
// 2) If missing, sniff filesystem (file => Redb, dir => Sled), then persist into admin meta
|
||||
// 3) Fallback to requested 'backend' (startup default) if nothing else is known
|
||||
let meta_backend = get_database_backend(base_dir, backend.clone(), admin_secret, id).ok().flatten();
|
||||
let db_path = PathBuf::from(base_dir).join(format!("{}.db", id));
|
||||
let db_path = base_dir.join(format!("{}.db", id));
|
||||
let sniffed_backend = if db_path.exists() {
|
||||
if db_path.is_file() {
|
||||
Some(options::BackendType::Redb)
|
||||
@@ -205,6 +206,9 @@ pub fn open_data_storage(
|
||||
options::BackendType::Tantivy => {
|
||||
return Err(DBError("Tantivy backend has no KV storage; use FT.* commands only".to_string()))
|
||||
}
|
||||
options::BackendType::Lance => {
|
||||
return Err(DBError("Lance backend has no KV storage; use LANCE.* commands only".to_string()))
|
||||
}
|
||||
};
|
||||
|
||||
// Publish to registry
|
||||
@@ -214,7 +218,7 @@ pub fn open_data_storage(
|
||||
|
||||
// Allocate the next DB id and persist new pointer
|
||||
pub fn allocate_next_id(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
) -> Result<u64, DBError> {
|
||||
@@ -238,7 +242,7 @@ pub fn allocate_next_id(
|
||||
|
||||
// Check existence of a db id in admin:dbs
|
||||
pub fn db_exists(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
id: u64,
|
||||
@@ -249,7 +253,7 @@ pub fn db_exists(
|
||||
|
||||
// Get per-db encryption key, if any
|
||||
pub fn get_enc_key(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
id: u64,
|
||||
@@ -260,7 +264,7 @@ pub fn get_enc_key(
|
||||
|
||||
// Set per-db encryption key (called during create)
|
||||
pub fn set_enc_key(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
id: u64,
|
||||
@@ -272,7 +276,7 @@ pub fn set_enc_key(
|
||||
|
||||
// Set database public flag
|
||||
pub fn set_database_public(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
id: u64,
|
||||
@@ -286,7 +290,7 @@ pub fn set_database_public(
|
||||
|
||||
// Persist per-db backend type in admin metadata (module-scope)
|
||||
pub fn set_database_backend(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
id: u64,
|
||||
@@ -298,13 +302,14 @@ pub fn set_database_backend(
|
||||
options::BackendType::Redb => "Redb",
|
||||
options::BackendType::Sled => "Sled",
|
||||
options::BackendType::Tantivy => "Tantivy",
|
||||
options::BackendType::Lance => "Lance",
|
||||
};
|
||||
let _ = admin.hset(&mk, vec![("backend".to_string(), val.to_string())])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn get_database_backend(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
id: u64,
|
||||
@@ -315,13 +320,14 @@ pub fn get_database_backend(
|
||||
Some(s) if s == "Redb" => Ok(Some(options::BackendType::Redb)),
|
||||
Some(s) if s == "Sled" => Ok(Some(options::BackendType::Sled)),
|
||||
Some(s) if s == "Tantivy" => Ok(Some(options::BackendType::Tantivy)),
|
||||
Some(s) if s == "Lance" => Ok(Some(options::BackendType::Lance)),
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
// Set database name
|
||||
pub fn set_database_name(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
id: u64,
|
||||
@@ -335,7 +341,7 @@ pub fn set_database_name(
|
||||
|
||||
// Get database name
|
||||
pub fn get_database_name(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
id: u64,
|
||||
@@ -359,7 +365,7 @@ fn load_public(
|
||||
|
||||
// Add access key for db (value format: "Read:ts" or "ReadWrite:ts")
|
||||
pub fn add_access_key(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
id: u64,
|
||||
@@ -378,7 +384,7 @@ pub fn add_access_key(
|
||||
|
||||
// Delete access key by hash
|
||||
pub fn delete_access_key(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
id: u64,
|
||||
@@ -391,7 +397,7 @@ pub fn delete_access_key(
|
||||
|
||||
// List access keys, returning (hash, perms, created_at_secs)
|
||||
pub fn list_access_keys(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
id: u64,
|
||||
@@ -411,7 +417,7 @@ pub fn list_access_keys(
|
||||
// - Ok(Some(Permissions)) when access is allowed
|
||||
// - Ok(None) when not allowed or db missing (caller can distinguish by calling db_exists)
|
||||
pub fn verify_access(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
id: u64,
|
||||
@@ -456,7 +462,7 @@ pub fn verify_access(
|
||||
|
||||
// Enumerate all db ids
|
||||
pub fn list_dbs(
|
||||
base_dir: &str,
|
||||
base_dir: &Path,
|
||||
backend: options::BackendType,
|
||||
admin_secret: &str,
|
||||
) -> Result<Vec<u64>, DBError> {
|
||||
|
729
src/cmd.rs
729
src/cmd.rs
@@ -1,4 +1,5 @@
|
||||
use crate::{error::DBError, protocol::Protocol, server::Server};
|
||||
use crate::{error::DBError, protocol::Protocol, server::Server, embedding::{EmbeddingConfig, EmbeddingProvider}};
|
||||
use base64::{engine::general_purpose, Engine as _};
|
||||
use tokio::time::{timeout, Duration};
|
||||
use futures::future::select_all;
|
||||
|
||||
@@ -125,6 +126,67 @@ pub enum Cmd {
|
||||
query: String,
|
||||
group_by: Vec<String>,
|
||||
reducers: Vec<String>,
|
||||
},
|
||||
|
||||
// LanceDB text-first commands (no user-provided vectors)
|
||||
LanceCreate {
|
||||
name: String,
|
||||
dim: usize,
|
||||
},
|
||||
LanceStoreText {
|
||||
name: String,
|
||||
id: String,
|
||||
text: String,
|
||||
meta: Vec<(String, String)>,
|
||||
},
|
||||
LanceSearchText {
|
||||
name: String,
|
||||
text: String,
|
||||
k: usize,
|
||||
filter: Option<String>,
|
||||
return_fields: Option<Vec<String>>,
|
||||
},
|
||||
// Image-first commands (no user-provided vectors)
|
||||
LanceStoreImage {
|
||||
name: String,
|
||||
id: String,
|
||||
uri: Option<String>,
|
||||
bytes_b64: Option<String>,
|
||||
meta: Vec<(String, String)>,
|
||||
},
|
||||
LanceSearchImage {
|
||||
name: String,
|
||||
k: usize,
|
||||
uri: Option<String>,
|
||||
bytes_b64: Option<String>,
|
||||
filter: Option<String>,
|
||||
return_fields: Option<Vec<String>>,
|
||||
},
|
||||
LanceCreateIndex {
|
||||
name: String,
|
||||
index_type: String,
|
||||
params: Vec<(String, String)>,
|
||||
},
|
||||
// Embedding configuration per dataset
|
||||
LanceEmbeddingConfigSet {
|
||||
name: String,
|
||||
provider: String,
|
||||
model: String,
|
||||
params: Vec<(String, String)>,
|
||||
},
|
||||
LanceEmbeddingConfigGet {
|
||||
name: String,
|
||||
},
|
||||
LanceList,
|
||||
LanceInfo {
|
||||
name: String,
|
||||
},
|
||||
LanceDel {
|
||||
name: String,
|
||||
id: String,
|
||||
},
|
||||
LanceDrop {
|
||||
name: String,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -815,6 +877,295 @@ impl Cmd {
|
||||
let reducers = Vec::new();
|
||||
Cmd::FtAggregate { index_name, query, group_by, reducers }
|
||||
}
|
||||
|
||||
// ----- LANCE.* commands -----
|
||||
"lance.create" => {
|
||||
// LANCE.CREATE name DIM d
|
||||
if cmd.len() != 4 || cmd[2].to_uppercase() != "DIM" {
|
||||
return Err(DBError("ERR LANCE.CREATE requires: name DIM <dim>".to_string()));
|
||||
}
|
||||
let name = cmd[1].clone();
|
||||
let dim: usize = cmd[3].parse().map_err(|_| DBError("ERR DIM must be an integer".to_string()))?;
|
||||
Cmd::LanceCreate { name, dim }
|
||||
}
|
||||
"lance.store" => {
|
||||
// LANCE.STORE name ID <id> TEXT <text> [META k v ...]
|
||||
if cmd.len() < 6 {
|
||||
return Err(DBError("ERR LANCE.STORE requires: name ID <id> TEXT <text> [META k v ...]".to_string()));
|
||||
}
|
||||
let name = cmd[1].clone();
|
||||
let mut i = 2;
|
||||
if cmd[i].to_uppercase() != "ID" || i + 1 >= cmd.len() {
|
||||
return Err(DBError("ERR LANCE.STORE requires ID <id>".to_string()));
|
||||
}
|
||||
let id = cmd[i + 1].clone();
|
||||
i += 2;
|
||||
if i >= cmd.len() || cmd[i].to_uppercase() != "TEXT" {
|
||||
return Err(DBError("ERR LANCE.STORE requires TEXT <text>".to_string()));
|
||||
}
|
||||
i += 1;
|
||||
if i >= cmd.len() {
|
||||
return Err(DBError("ERR LANCE.STORE requires TEXT <text>".to_string()));
|
||||
}
|
||||
let text = cmd[i].clone();
|
||||
i += 1;
|
||||
|
||||
let mut meta: Vec<(String, String)> = Vec::new();
|
||||
if i < cmd.len() && cmd[i].to_uppercase() == "META" {
|
||||
i += 1;
|
||||
while i + 1 < cmd.len() {
|
||||
meta.push((cmd[i].clone(), cmd[i + 1].clone()));
|
||||
i += 2;
|
||||
}
|
||||
}
|
||||
Cmd::LanceStoreText { name, id, text, meta }
|
||||
}
|
||||
"lance.storeimage" => {
|
||||
// LANCE.STOREIMAGE name ID <id> (URI <uri> | BYTES <base64>) [META k v ...]
|
||||
if cmd.len() < 6 {
|
||||
return Err(DBError("ERR LANCE.STOREIMAGE requires: name ID <id> (URI <uri> | BYTES <base64>) [META k v ...]".to_string()));
|
||||
}
|
||||
let name = cmd[1].clone();
|
||||
let mut i = 2;
|
||||
if cmd[i].to_uppercase() != "ID" || i + 1 >= cmd.len() {
|
||||
return Err(DBError("ERR LANCE.STOREIMAGE requires ID <id>".to_string()));
|
||||
}
|
||||
let id = cmd[i + 1].clone();
|
||||
i += 2;
|
||||
|
||||
let mut uri_opt: Option<String> = None;
|
||||
let mut bytes_b64_opt: Option<String> = None;
|
||||
|
||||
if i < cmd.len() && cmd[i].to_uppercase() == "URI" {
|
||||
if i + 1 >= cmd.len() { return Err(DBError("ERR LANCE.STOREIMAGE URI requires a value".to_string())); }
|
||||
uri_opt = Some(cmd[i + 1].clone());
|
||||
i += 2;
|
||||
} else if i < cmd.len() && cmd[i].to_uppercase() == "BYTES" {
|
||||
if i + 1 >= cmd.len() { return Err(DBError("ERR LANCE.STOREIMAGE BYTES requires a value".to_string())); }
|
||||
bytes_b64_opt = Some(cmd[i + 1].clone());
|
||||
i += 2;
|
||||
} else {
|
||||
return Err(DBError("ERR LANCE.STOREIMAGE requires either URI <uri> or BYTES <base64>".to_string()));
|
||||
}
|
||||
|
||||
// Parse optional META pairs
|
||||
let mut meta: Vec<(String, String)> = Vec::new();
|
||||
if i < cmd.len() && cmd[i].to_uppercase() == "META" {
|
||||
i += 1;
|
||||
while i + 1 < cmd.len() {
|
||||
meta.push((cmd[i].clone(), cmd[i + 1].clone()));
|
||||
i += 2;
|
||||
}
|
||||
}
|
||||
|
||||
Cmd::LanceStoreImage { name, id, uri: uri_opt, bytes_b64: bytes_b64_opt, meta }
|
||||
}
|
||||
"lance.search" => {
|
||||
// LANCE.SEARCH name K <k> QUERY <text> [FILTER expr] [RETURN n fields...]
|
||||
if cmd.len() < 6 {
|
||||
return Err(DBError("ERR LANCE.SEARCH requires: name K <k> QUERY <text> [FILTER expr] [RETURN n fields...]".to_string()));
|
||||
}
|
||||
let name = cmd[1].clone();
|
||||
if cmd[2].to_uppercase() != "K" {
|
||||
return Err(DBError("ERR LANCE.SEARCH requires K <k>".to_string()));
|
||||
}
|
||||
let k: usize = cmd[3].parse().map_err(|_| DBError("ERR K must be an integer".to_string()))?;
|
||||
if cmd[4].to_uppercase() != "QUERY" {
|
||||
return Err(DBError("ERR LANCE.SEARCH requires QUERY <text>".to_string()));
|
||||
}
|
||||
let mut i = 5;
|
||||
if i >= cmd.len() {
|
||||
return Err(DBError("ERR LANCE.SEARCH requires QUERY <text>".to_string()));
|
||||
}
|
||||
let text = cmd[i].clone();
|
||||
i += 1;
|
||||
|
||||
let mut filter: Option<String> = None;
|
||||
let mut return_fields: Option<Vec<String>> = None;
|
||||
while i < cmd.len() {
|
||||
match cmd[i].to_uppercase().as_str() {
|
||||
"FILTER" => {
|
||||
if i + 1 >= cmd.len() {
|
||||
return Err(DBError("ERR FILTER requires an expression".to_string()));
|
||||
}
|
||||
filter = Some(cmd[i + 1].clone());
|
||||
i += 2;
|
||||
}
|
||||
"RETURN" => {
|
||||
if i + 1 >= cmd.len() {
|
||||
return Err(DBError("ERR RETURN requires field count".to_string()));
|
||||
}
|
||||
let n: usize = cmd[i + 1].parse().map_err(|_| DBError("ERR RETURN count must be integer".to_string()))?;
|
||||
i += 2;
|
||||
let mut fields = Vec::new();
|
||||
for _ in 0..n {
|
||||
if i < cmd.len() {
|
||||
fields.push(cmd[i].clone());
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
return_fields = Some(fields);
|
||||
}
|
||||
_ => { i += 1; }
|
||||
}
|
||||
}
|
||||
Cmd::LanceSearchText { name, text, k, filter, return_fields }
|
||||
}
|
||||
"lance.searchimage" => {
|
||||
// LANCE.SEARCHIMAGE name K <k> (QUERYURI <uri> | QUERYBYTES <base64>) [FILTER expr] [RETURN n fields...]
|
||||
if cmd.len() < 6 {
|
||||
return Err(DBError("ERR LANCE.SEARCHIMAGE requires: name K <k> (QUERYURI <uri> | QUERYBYTES <base64>) [FILTER expr] [RETURN n fields...]".to_string()));
|
||||
}
|
||||
let name = cmd[1].clone();
|
||||
if cmd[2].to_uppercase() != "K" {
|
||||
return Err(DBError("ERR LANCE.SEARCHIMAGE requires K <k>".to_string()));
|
||||
}
|
||||
let k: usize = cmd[3].parse().map_err(|_| DBError("ERR K must be an integer".to_string()))?;
|
||||
let mut i = 4;
|
||||
|
||||
let mut uri_opt: Option<String> = None;
|
||||
let mut bytes_b64_opt: Option<String> = None;
|
||||
|
||||
if i < cmd.len() && cmd[i].to_uppercase() == "QUERYURI" {
|
||||
if i + 1 >= cmd.len() { return Err(DBError("ERR QUERYURI requires a value".to_string())); }
|
||||
uri_opt = Some(cmd[i + 1].clone());
|
||||
i += 2;
|
||||
} else if i < cmd.len() && cmd[i].to_uppercase() == "QUERYBYTES" {
|
||||
if i + 1 >= cmd.len() { return Err(DBError("ERR QUERYBYTES requires a value".to_string())); }
|
||||
bytes_b64_opt = Some(cmd[i + 1].clone());
|
||||
i += 2;
|
||||
} else {
|
||||
return Err(DBError("ERR LANCE.SEARCHIMAGE requires QUERYURI <uri> or QUERYBYTES <base64>".to_string()));
|
||||
}
|
||||
|
||||
let mut filter: Option<String> = None;
|
||||
let mut return_fields: Option<Vec<String>> = None;
|
||||
while i < cmd.len() {
|
||||
match cmd[i].to_uppercase().as_str() {
|
||||
"FILTER" => {
|
||||
if i + 1 >= cmd.len() {
|
||||
return Err(DBError("ERR FILTER requires an expression".to_string()));
|
||||
}
|
||||
filter = Some(cmd[i + 1].clone());
|
||||
i += 2;
|
||||
}
|
||||
"RETURN" => {
|
||||
if i + 1 >= cmd.len() {
|
||||
return Err(DBError("ERR RETURN requires field count".to_string()));
|
||||
}
|
||||
let n: usize = cmd[i + 1].parse().map_err(|_| DBError("ERR RETURN count must be integer".to_string()))?;
|
||||
i += 2;
|
||||
let mut fields = Vec::new();
|
||||
for _ in 0..n {
|
||||
if i < cmd.len() {
|
||||
fields.push(cmd[i].clone());
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
return_fields = Some(fields);
|
||||
}
|
||||
_ => { i += 1; }
|
||||
}
|
||||
}
|
||||
|
||||
Cmd::LanceSearchImage { name, k, uri: uri_opt, bytes_b64: bytes_b64_opt, filter, return_fields }
|
||||
}
|
||||
"lance.createindex" => {
|
||||
// LANCE.CREATEINDEX name TYPE t [PARAM k v ...]
|
||||
if cmd.len() < 4 || cmd[2].to_uppercase() != "TYPE" {
|
||||
return Err(DBError("ERR LANCE.CREATEINDEX requires: name TYPE <type> [PARAM k v ...]".to_string()));
|
||||
}
|
||||
let name = cmd[1].clone();
|
||||
let index_type = cmd[3].clone();
|
||||
let mut params: Vec<(String, String)> = Vec::new();
|
||||
let mut i = 4;
|
||||
if i < cmd.len() && cmd[i].to_uppercase() == "PARAM" {
|
||||
i += 1;
|
||||
while i + 1 < cmd.len() {
|
||||
params.push((cmd[i].clone(), cmd[i + 1].clone()));
|
||||
i += 2;
|
||||
}
|
||||
}
|
||||
Cmd::LanceCreateIndex { name, index_type, params }
|
||||
}
|
||||
"lance.embedding" => {
|
||||
// LANCE.EMBEDDING CONFIG SET name PROVIDER p MODEL m [PARAM k v ...]
|
||||
// LANCE.EMBEDDING CONFIG GET name
|
||||
if cmd.len() < 3 || cmd[1].to_uppercase() != "CONFIG" {
|
||||
return Err(DBError("ERR LANCE.EMBEDDING requires CONFIG subcommand".to_string()));
|
||||
}
|
||||
if cmd.len() >= 4 && cmd[2].to_uppercase() == "SET" {
|
||||
if cmd.len() < 8 {
|
||||
return Err(DBError("ERR LANCE.EMBEDDING CONFIG SET requires: SET name PROVIDER p MODEL m [PARAM k v ...]".to_string()));
|
||||
}
|
||||
let name = cmd[3].clone();
|
||||
let mut i = 4;
|
||||
let mut provider: Option<String> = None;
|
||||
let mut model: Option<String> = None;
|
||||
let mut params: Vec<(String, String)> = Vec::new();
|
||||
while i < cmd.len() {
|
||||
match cmd[i].to_uppercase().as_str() {
|
||||
"PROVIDER" => {
|
||||
if i + 1 >= cmd.len() {
|
||||
return Err(DBError("ERR PROVIDER requires a value".to_string()));
|
||||
}
|
||||
provider = Some(cmd[i + 1].clone());
|
||||
i += 2;
|
||||
}
|
||||
"MODEL" => {
|
||||
if i + 1 >= cmd.len() {
|
||||
return Err(DBError("ERR MODEL requires a value".to_string()));
|
||||
}
|
||||
model = Some(cmd[i + 1].clone());
|
||||
i += 2;
|
||||
}
|
||||
"PARAM" => {
|
||||
i += 1;
|
||||
while i + 1 < cmd.len() {
|
||||
params.push((cmd[i].clone(), cmd[i + 1].clone()));
|
||||
i += 2;
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
// Unknown token; break to avoid infinite loop
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
let provider = provider.ok_or_else(|| DBError("ERR missing PROVIDER".to_string()))?;
|
||||
let model = model.ok_or_else(|| DBError("ERR missing MODEL".to_string()))?;
|
||||
Cmd::LanceEmbeddingConfigSet { name, provider, model, params }
|
||||
} else if cmd.len() == 4 && cmd[2].to_uppercase() == "GET" {
|
||||
let name = cmd[3].clone();
|
||||
Cmd::LanceEmbeddingConfigGet { name }
|
||||
} else {
|
||||
return Err(DBError("ERR LANCE.EMBEDDING CONFIG supports: SET ... | GET name".to_string()));
|
||||
}
|
||||
}
|
||||
"lance.list" => {
|
||||
if cmd.len() != 1 {
|
||||
return Err(DBError("ERR LANCE.LIST takes no arguments".to_string()));
|
||||
}
|
||||
Cmd::LanceList
|
||||
}
|
||||
"lance.info" => {
|
||||
if cmd.len() != 2 {
|
||||
return Err(DBError("ERR LANCE.INFO requires: name".to_string()));
|
||||
}
|
||||
Cmd::LanceInfo { name: cmd[1].clone() }
|
||||
}
|
||||
"lance.drop" => {
|
||||
if cmd.len() != 2 {
|
||||
return Err(DBError("ERR LANCE.DROP requires: name".to_string()));
|
||||
}
|
||||
Cmd::LanceDrop { name: cmd[1].clone() }
|
||||
}
|
||||
"lance.del" => {
|
||||
if cmd.len() != 3 {
|
||||
return Err(DBError("ERR LANCE.DEL requires: name id".to_string()));
|
||||
}
|
||||
Cmd::LanceDel { name: cmd[1].clone(), id: cmd[2].clone() }
|
||||
}
|
||||
_ => Cmd::Unknow(cmd[0].clone()),
|
||||
},
|
||||
protocol,
|
||||
@@ -853,6 +1204,18 @@ impl Cmd {
|
||||
.map(|b| matches!(b, crate::options::BackendType::Tantivy))
|
||||
.unwrap_or(false);
|
||||
|
||||
// Determine Lance backend similarly
|
||||
let is_lance_backend = crate::admin_meta::get_database_backend(
|
||||
&server.option.dir,
|
||||
server.option.backend.clone(),
|
||||
&server.option.admin_secret,
|
||||
server.selected_db,
|
||||
)
|
||||
.ok()
|
||||
.flatten()
|
||||
.map(|b| matches!(b, crate::options::BackendType::Lance))
|
||||
.unwrap_or(false);
|
||||
|
||||
if is_tantivy_backend {
|
||||
match &self {
|
||||
Cmd::Select(..)
|
||||
@@ -876,6 +1239,34 @@ impl Cmd {
|
||||
}
|
||||
}
|
||||
|
||||
// Lance backend gating: allow only LANCE.* and basic control/info commands
|
||||
if is_lance_backend {
|
||||
match &self {
|
||||
Cmd::Select(..)
|
||||
| Cmd::Quit
|
||||
| Cmd::Client(..)
|
||||
| Cmd::ClientSetName(..)
|
||||
| Cmd::ClientGetName
|
||||
| Cmd::Command(..)
|
||||
| Cmd::Info(..)
|
||||
| Cmd::LanceCreate { .. }
|
||||
| Cmd::LanceStoreText { .. }
|
||||
| Cmd::LanceSearchText { .. }
|
||||
| Cmd::LanceStoreImage { .. }
|
||||
| Cmd::LanceSearchImage { .. }
|
||||
| Cmd::LanceEmbeddingConfigSet { .. }
|
||||
| Cmd::LanceEmbeddingConfigGet { .. }
|
||||
| Cmd::LanceCreateIndex { .. }
|
||||
| Cmd::LanceList
|
||||
| Cmd::LanceInfo { .. }
|
||||
| Cmd::LanceDel { .. }
|
||||
| Cmd::LanceDrop { .. } => {}
|
||||
_ => {
|
||||
return Ok(Protocol::err("ERR backend is Lance; only LANCE.* commands are allowed"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// If selected DB is not Tantivy, forbid all FT.* commands here.
|
||||
if !is_tantivy_backend {
|
||||
match &self {
|
||||
@@ -893,6 +1284,27 @@ impl Cmd {
|
||||
}
|
||||
}
|
||||
|
||||
// If selected DB is not Lance, forbid all LANCE.* commands here.
|
||||
if !is_lance_backend {
|
||||
match &self {
|
||||
Cmd::LanceCreate { .. }
|
||||
| Cmd::LanceStoreText { .. }
|
||||
| Cmd::LanceSearchText { .. }
|
||||
| Cmd::LanceStoreImage { .. }
|
||||
| Cmd::LanceSearchImage { .. }
|
||||
| Cmd::LanceEmbeddingConfigSet { .. }
|
||||
| Cmd::LanceEmbeddingConfigGet { .. }
|
||||
| Cmd::LanceCreateIndex { .. }
|
||||
| Cmd::LanceList
|
||||
| Cmd::LanceInfo { .. }
|
||||
| Cmd::LanceDel { .. }
|
||||
| Cmd::LanceDrop { .. } => {
|
||||
return Ok(Protocol::err("ERR DB backend is not Lance; LANCE.* commands are not allowed"));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
match self {
|
||||
Cmd::Select(db, key) => select_cmd(server, db, key).await,
|
||||
Cmd::Ping => Ok(Protocol::SimpleString("PONG".to_string())),
|
||||
@@ -1015,6 +1427,307 @@ impl Cmd {
|
||||
Ok(Protocol::err("FT.AGGREGATE not implemented yet"))
|
||||
}
|
||||
|
||||
// LanceDB commands
|
||||
Cmd::LanceCreate { name, dim } => {
|
||||
if !server.has_write_permission() {
|
||||
return Ok(Protocol::err("ERR write permission denied"));
|
||||
}
|
||||
match server.lance_store()?.create_dataset(&name, dim).await {
|
||||
Ok(()) => Ok(Protocol::SimpleString("OK".to_string())),
|
||||
Err(e) => Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
}
|
||||
Cmd::LanceEmbeddingConfigSet { name, provider, model, params } => {
|
||||
if !server.has_write_permission() {
|
||||
return Ok(Protocol::err("ERR write permission denied"));
|
||||
}
|
||||
// Map provider string to enum
|
||||
let p_lc = provider.to_lowercase();
|
||||
let prov = match p_lc.as_str() {
|
||||
"test-hash" | "testhash" => EmbeddingProvider::TestHash,
|
||||
"testimagehash" | "image-test-hash" | "imagetesthash" => EmbeddingProvider::ImageTestHash,
|
||||
"fastembed" | "lancefastembed" => EmbeddingProvider::LanceFastEmbed,
|
||||
"openai" | "lanceopenai" => EmbeddingProvider::LanceOpenAI,
|
||||
other => EmbeddingProvider::LanceOther(other.to_string()),
|
||||
};
|
||||
let cfg = EmbeddingConfig {
|
||||
provider: prov,
|
||||
model,
|
||||
params: params.into_iter().collect(),
|
||||
};
|
||||
match server.set_dataset_embedding_config(&name, &cfg) {
|
||||
Ok(()) => Ok(Protocol::SimpleString("OK".to_string())),
|
||||
Err(e) => Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
}
|
||||
Cmd::LanceEmbeddingConfigGet { name } => {
|
||||
match server.get_dataset_embedding_config(&name) {
|
||||
Ok(cfg) => {
|
||||
let mut arr = Vec::new();
|
||||
arr.push(Protocol::BulkString("provider".to_string()));
|
||||
arr.push(Protocol::BulkString(match cfg.provider {
|
||||
EmbeddingProvider::TestHash => "test-hash".to_string(),
|
||||
EmbeddingProvider::ImageTestHash => "testimagehash".to_string(),
|
||||
EmbeddingProvider::LanceFastEmbed => "lancefastembed".to_string(),
|
||||
EmbeddingProvider::LanceOpenAI => "lanceopenai".to_string(),
|
||||
EmbeddingProvider::LanceOther(ref s) => s.clone(),
|
||||
}));
|
||||
arr.push(Protocol::BulkString("model".to_string()));
|
||||
arr.push(Protocol::BulkString(cfg.model.clone()));
|
||||
arr.push(Protocol::BulkString("params".to_string()));
|
||||
arr.push(Protocol::BulkString(serde_json::to_string(&cfg.params).unwrap_or_else(|_| "{}".to_string())));
|
||||
Ok(Protocol::Array(arr))
|
||||
}
|
||||
Err(e) => Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
}
|
||||
Cmd::LanceStoreText { name, id, text, meta } => {
|
||||
if !server.has_write_permission() {
|
||||
return Ok(Protocol::err("ERR write permission denied"));
|
||||
}
|
||||
// Resolve embedder and embed text on a plain OS thread to avoid tokio runtime panics from reqwest::blocking
|
||||
let embedder = server.get_embedder_for(&name)?;
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
let emb_arc = embedder.clone();
|
||||
let text_cl = text.clone();
|
||||
std::thread::spawn(move || {
|
||||
let res = emb_arc.embed(&text_cl);
|
||||
let _ = tx.send(res);
|
||||
});
|
||||
let vector = match rx.await {
|
||||
Ok(Ok(v)) => v,
|
||||
Ok(Err(e)) => return Ok(Protocol::err(&e.0)),
|
||||
Err(recv_err) => return Ok(Protocol::err(&format!("ERR embedding thread error: {}", recv_err))),
|
||||
};
|
||||
let meta_map: std::collections::HashMap<String, String> = meta.into_iter().collect();
|
||||
match server.lance_store()?.store_vector(&name, &id, vector, meta_map, Some(text)).await {
|
||||
Ok(()) => Ok(Protocol::SimpleString("OK".to_string())),
|
||||
Err(e) => Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
}
|
||||
Cmd::LanceSearchText { name, text, k, filter, return_fields } => {
|
||||
// Resolve embedder and embed query text on a plain OS thread
|
||||
let embedder = server.get_embedder_for(&name)?;
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
let emb_arc = embedder.clone();
|
||||
let text_cl = text.clone();
|
||||
std::thread::spawn(move || {
|
||||
let res = emb_arc.embed(&text_cl);
|
||||
let _ = tx.send(res);
|
||||
});
|
||||
let qv = match rx.await {
|
||||
Ok(Ok(v)) => v,
|
||||
Ok(Err(e)) => return Ok(Protocol::err(&e.0)),
|
||||
Err(recv_err) => return Ok(Protocol::err(&format!("ERR embedding thread error: {}", recv_err))),
|
||||
};
|
||||
match server.lance_store()?.search_vectors(&name, qv, k, filter, return_fields).await {
|
||||
Ok(results) => {
|
||||
// Encode as array of [id, score, [k1, v1, k2, v2, ...]]
|
||||
let mut arr = Vec::new();
|
||||
for (id, score, meta) in results {
|
||||
let mut meta_arr: Vec<Protocol> = Vec::new();
|
||||
for (k, v) in meta {
|
||||
meta_arr.push(Protocol::BulkString(k));
|
||||
meta_arr.push(Protocol::BulkString(v));
|
||||
}
|
||||
arr.push(Protocol::Array(vec![
|
||||
Protocol::BulkString(id),
|
||||
Protocol::BulkString(score.to_string()),
|
||||
Protocol::Array(meta_arr),
|
||||
]));
|
||||
}
|
||||
Ok(Protocol::Array(arr))
|
||||
}
|
||||
Err(e) => Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
}
|
||||
|
||||
// New: Image store
|
||||
Cmd::LanceStoreImage { name, id, uri, bytes_b64, meta } => {
|
||||
if !server.has_write_permission() {
|
||||
return Ok(Protocol::err("ERR write permission denied"));
|
||||
}
|
||||
let use_uri = uri.is_some();
|
||||
let use_b64 = bytes_b64.is_some();
|
||||
if (use_uri && use_b64) || (!use_uri && !use_b64) {
|
||||
return Ok(Protocol::err("ERR Provide exactly one of URI or BYTES for LANCE.STOREIMAGE"));
|
||||
}
|
||||
let max_bytes: usize = std::env::var("HERODB_IMAGE_MAX_BYTES")
|
||||
.ok()
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.unwrap_or(10 * 1024 * 1024) as usize;
|
||||
|
||||
let media_uri_opt = if let Some(u) = uri.clone() {
|
||||
match server.fetch_image_bytes_from_uri(&u) {
|
||||
Ok(_) => {}
|
||||
Err(e) => return Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
Some(u)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let bytes: Vec<u8> = if let Some(u) = uri {
|
||||
match server.fetch_image_bytes_from_uri(&u) {
|
||||
Ok(b) => b,
|
||||
Err(e) => return Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
} else {
|
||||
let b64 = bytes_b64.unwrap_or_default();
|
||||
let data = match general_purpose::STANDARD.decode(b64.as_bytes()) {
|
||||
Ok(d) => d,
|
||||
Err(e) => return Ok(Protocol::err(&format!("ERR base64 decode error: {}", e))),
|
||||
};
|
||||
if data.len() > max_bytes {
|
||||
return Ok(Protocol::err(&format!("ERR image exceeds max allowed bytes {}", max_bytes)));
|
||||
}
|
||||
data
|
||||
};
|
||||
|
||||
let img_embedder = match server.get_image_embedder_for(&name) {
|
||||
Ok(e) => e,
|
||||
Err(e) => return Ok(Protocol::err(&e.0)),
|
||||
};
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
let emb_arc = img_embedder.clone();
|
||||
let bytes_cl = bytes.clone();
|
||||
std::thread::spawn(move || {
|
||||
let res = emb_arc.embed_image(&bytes_cl);
|
||||
let _ = tx.send(res);
|
||||
});
|
||||
let vector = match rx.await {
|
||||
Ok(Ok(v)) => v,
|
||||
Ok(Err(e)) => return Ok(Protocol::err(&e.0)),
|
||||
Err(recv_err) => return Ok(Protocol::err(&format!("ERR embedding thread error: {}", recv_err))),
|
||||
};
|
||||
|
||||
let meta_map: std::collections::HashMap<String, String> = meta.into_iter().collect();
|
||||
match server.lance_store()?.store_vector_with_media(
|
||||
&name,
|
||||
&id,
|
||||
vector,
|
||||
meta_map,
|
||||
None,
|
||||
Some("image".to_string()),
|
||||
media_uri_opt,
|
||||
).await {
|
||||
Ok(()) => Ok(Protocol::SimpleString("OK".to_string())),
|
||||
Err(e) => Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
}
|
||||
|
||||
// New: Image search
|
||||
Cmd::LanceSearchImage { name, k, uri, bytes_b64, filter, return_fields } => {
|
||||
let use_uri = uri.is_some();
|
||||
let use_b64 = bytes_b64.is_some();
|
||||
if (use_uri && use_b64) || (!use_uri && !use_b64) {
|
||||
return Ok(Protocol::err("ERR Provide exactly one of QUERYURI or QUERYBYTES for LANCE.SEARCHIMAGE"));
|
||||
}
|
||||
let max_bytes: usize = std::env::var("HERODB_IMAGE_MAX_BYTES")
|
||||
.ok()
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.unwrap_or(10 * 1024 * 1024) as usize;
|
||||
|
||||
let bytes: Vec<u8> = if let Some(u) = uri {
|
||||
match server.fetch_image_bytes_from_uri(&u) {
|
||||
Ok(b) => b,
|
||||
Err(e) => return Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
} else {
|
||||
let b64 = bytes_b64.unwrap_or_default();
|
||||
let data = match general_purpose::STANDARD.decode(b64.as_bytes()) {
|
||||
Ok(d) => d,
|
||||
Err(e) => return Ok(Protocol::err(&format!("ERR base64 decode error: {}", e))),
|
||||
};
|
||||
if data.len() > max_bytes {
|
||||
return Ok(Protocol::err(&format!("ERR image exceeds max allowed bytes {}", max_bytes)));
|
||||
}
|
||||
data
|
||||
};
|
||||
|
||||
let img_embedder = match server.get_image_embedder_for(&name) {
|
||||
Ok(e) => e,
|
||||
Err(e) => return Ok(Protocol::err(&e.0)),
|
||||
};
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
std::thread::spawn(move || {
|
||||
let res = img_embedder.embed_image(&bytes);
|
||||
let _ = tx.send(res);
|
||||
});
|
||||
let qv = match rx.await {
|
||||
Ok(Ok(v)) => v,
|
||||
Ok(Err(e)) => return Ok(Protocol::err(&e.0)),
|
||||
Err(recv_err) => return Ok(Protocol::err(&format!("ERR embedding thread error: {}", recv_err))),
|
||||
};
|
||||
|
||||
match server.lance_store()?.search_vectors(&name, qv, k, filter, return_fields).await {
|
||||
Ok(results) => {
|
||||
let mut arr = Vec::new();
|
||||
for (id, score, meta) in results {
|
||||
let mut meta_arr: Vec<Protocol> = Vec::new();
|
||||
for (k, v) in meta {
|
||||
meta_arr.push(Protocol::BulkString(k));
|
||||
meta_arr.push(Protocol::BulkString(v));
|
||||
}
|
||||
arr.push(Protocol::Array(vec![
|
||||
Protocol::BulkString(id),
|
||||
Protocol::BulkString(score.to_string()),
|
||||
Protocol::Array(meta_arr),
|
||||
]));
|
||||
}
|
||||
Ok(Protocol::Array(arr))
|
||||
}
|
||||
Err(e) => Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
}
|
||||
Cmd::LanceCreateIndex { name, index_type, params } => {
|
||||
if !server.has_write_permission() {
|
||||
return Ok(Protocol::err("ERR write permission denied"));
|
||||
}
|
||||
let params_map: std::collections::HashMap<String, String> = params.into_iter().collect();
|
||||
match server.lance_store()?.create_index(&name, &index_type, params_map).await {
|
||||
Ok(()) => Ok(Protocol::SimpleString("OK".to_string())),
|
||||
Err(e) => Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
}
|
||||
Cmd::LanceList => {
|
||||
match server.lance_store()?.list_datasets().await {
|
||||
Ok(list) => Ok(Protocol::Array(list.into_iter().map(Protocol::BulkString).collect())),
|
||||
Err(e) => Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
}
|
||||
Cmd::LanceInfo { name } => {
|
||||
match server.lance_store()?.get_dataset_info(&name).await {
|
||||
Ok(info) => {
|
||||
let mut arr = Vec::new();
|
||||
for (k, v) in info {
|
||||
arr.push(Protocol::BulkString(k));
|
||||
arr.push(Protocol::BulkString(v));
|
||||
}
|
||||
Ok(Protocol::Array(arr))
|
||||
}
|
||||
Err(e) => Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
}
|
||||
Cmd::LanceDel { name, id } => {
|
||||
if !server.has_write_permission() {
|
||||
return Ok(Protocol::err("ERR write permission denied"));
|
||||
}
|
||||
match server.lance_store()?.delete_by_id(&name, &id).await {
|
||||
Ok(b) => Ok(Protocol::SimpleString(if b { "1" } else { "0" }.to_string())),
|
||||
Err(e) => Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
}
|
||||
Cmd::LanceDrop { name } => {
|
||||
if !server.has_write_permission() {
|
||||
return Ok(Protocol::err("ERR write permission denied"));
|
||||
}
|
||||
match server.lance_store()?.drop_dataset(&name).await {
|
||||
Ok(_b) => Ok(Protocol::SimpleString("OK".to_string())),
|
||||
Err(e) => Ok(Protocol::err(&e.0)),
|
||||
}
|
||||
}
|
||||
|
||||
Cmd::Unknow(s) => Ok(Protocol::err(&format!("ERR unknown command `{}`", s))),
|
||||
}
|
||||
}
|
||||
@@ -1114,8 +1827,8 @@ async fn select_cmd(server: &mut Server, db: u64, key: Option<String>) -> Result
|
||||
.ok()
|
||||
.flatten();
|
||||
|
||||
if matches!(eff_backend, Some(crate::options::BackendType::Tantivy)) {
|
||||
// Tantivy DBs have no KV storage; allow SELECT to succeed
|
||||
if matches!(eff_backend, Some(crate::options::BackendType::Tantivy) | Some(crate::options::BackendType::Lance)) {
|
||||
// Search-only DBs (Tantivy/Lance) have no KV storage; allow SELECT to succeed
|
||||
Ok(Protocol::SimpleString("OK".to_string()))
|
||||
} else {
|
||||
match server.current_storage() {
|
||||
@@ -1427,7 +2140,7 @@ async fn incr_cmd(server: &Server, key: &String) -> Result<Protocol, DBError> {
|
||||
|
||||
fn config_get_cmd(name: &String, server: &Server) -> Result<Protocol, DBError> {
|
||||
let value = match name.as_str() {
|
||||
"dir" => Some(server.option.dir.clone()),
|
||||
"dir" => Some(server.option.dir.display().to_string()),
|
||||
"dbfilename" => Some(format!("{}.db", server.selected_db)),
|
||||
"databases" => Some("16".to_string()), // Hardcoded as per original logic
|
||||
_ => None,
|
||||
@@ -1459,9 +2172,9 @@ async fn dbsize_cmd(server: &Server) -> Result<Protocol, DBError> {
|
||||
}
|
||||
|
||||
async fn info_cmd(server: &Server, section: &Option<String>) -> Result<Protocol, DBError> {
|
||||
// For Tantivy backend, there is no KV storage; synthesize minimal info.
|
||||
// For Tantivy or Lance backend, there is no KV storage; synthesize minimal info.
|
||||
// Determine effective backend for the currently selected db.
|
||||
let is_tantivy_db = crate::admin_meta::get_database_backend(
|
||||
let is_search_only_db = crate::admin_meta::get_database_backend(
|
||||
&server.option.dir,
|
||||
server.option.backend.clone(),
|
||||
&server.option.admin_secret,
|
||||
@@ -1469,10 +2182,10 @@ async fn info_cmd(server: &Server, section: &Option<String>) -> Result<Protocol,
|
||||
)
|
||||
.ok()
|
||||
.flatten()
|
||||
.map(|b| matches!(b, crate::options::BackendType::Tantivy))
|
||||
.map(|b| matches!(b, crate::options::BackendType::Tantivy | crate::options::BackendType::Lance))
|
||||
.unwrap_or(false);
|
||||
|
||||
let storage_info: Vec<(String, String)> = if is_tantivy_db {
|
||||
let storage_info: Vec<(String, String)> = if is_search_only_db {
|
||||
vec![
|
||||
("db_size".to_string(), "0".to_string()),
|
||||
("is_encrypted".to_string(), "false".to_string()),
|
||||
|
@@ -1,8 +1,8 @@
|
||||
use chacha20poly1305::{
|
||||
aead::{Aead, KeyInit, OsRng},
|
||||
aead::{Aead, KeyInit},
|
||||
XChaCha20Poly1305, XNonce,
|
||||
};
|
||||
use rand::RngCore;
|
||||
use rand::{rngs::OsRng, RngCore};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
const VERSION: u8 = 1;
|
||||
@@ -31,7 +31,7 @@ pub struct CryptoFactory {
|
||||
impl CryptoFactory {
|
||||
/// Accepts any secret bytes; turns them into a 32-byte key (SHA-256).
|
||||
pub fn new<S: AsRef<[u8]>>(secret: S) -> Self {
|
||||
let mut h = Sha256::new();
|
||||
let mut h = Sha256::default();
|
||||
h.update(b"xchacha20poly1305-factory:v1"); // domain separation
|
||||
h.update(secret.as_ref());
|
||||
let digest = h.finalize(); // 32 bytes
|
||||
|
405
src/embedding.rs
Normal file
405
src/embedding.rs
Normal file
@@ -0,0 +1,405 @@
|
||||
// Embedding abstraction and minimal providers.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::error::DBError;
|
||||
|
||||
// Networking for OpenAI/Azure
|
||||
use std::time::Duration;
|
||||
use ureq::{Agent, AgentBuilder};
|
||||
use serde_json::json;
|
||||
|
||||
/// Provider identifiers. Extend as needed to mirror LanceDB-supported providers.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum EmbeddingProvider {
|
||||
// Deterministic, local-only embedder for CI and offline development (text).
|
||||
TestHash,
|
||||
// Deterministic, local-only embedder for CI and offline development (image).
|
||||
ImageTestHash,
|
||||
// Placeholders for LanceDB-supported providers; implementers can add concrete backends later.
|
||||
LanceFastEmbed,
|
||||
LanceOpenAI,
|
||||
LanceOther(String),
|
||||
}
|
||||
|
||||
/// Serializable embedding configuration.
|
||||
/// params: arbitrary key-value map for provider-specific knobs (e.g., "dim", "api_key_env", etc.)
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct EmbeddingConfig {
|
||||
pub provider: EmbeddingProvider,
|
||||
pub model: String,
|
||||
#[serde(default)]
|
||||
pub params: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl EmbeddingConfig {
|
||||
pub fn get_param_usize(&self, key: &str) -> Option<usize> {
|
||||
self.params.get(key).and_then(|v| v.parse::<usize>().ok())
|
||||
}
|
||||
pub fn get_param_string(&self, key: &str) -> Option<String> {
|
||||
self.params.get(key).cloned()
|
||||
}
|
||||
}
|
||||
|
||||
/// A provider-agnostic text embedding interface.
|
||||
pub trait Embedder: Send + Sync {
|
||||
/// Human-readable provider/model name
|
||||
fn name(&self) -> String;
|
||||
/// Embedding dimension
|
||||
fn dim(&self) -> usize;
|
||||
/// Embed a single text string into a fixed-length vector
|
||||
fn embed(&self, text: &str) -> Result<Vec<f32>, DBError>;
|
||||
/// Embed many texts; default maps embed() over inputs
|
||||
fn embed_many(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, DBError> {
|
||||
texts.iter().map(|t| self.embed(t)).collect()
|
||||
}
|
||||
}
|
||||
|
||||
//// ----------------------------- TEXT: deterministic test embedder -----------------------------
|
||||
|
||||
/// Deterministic, no-deps, no-network embedder for CI and offline dev.
|
||||
/// Algorithm:
|
||||
/// - Fold bytes of UTF-8 into 'dim' buckets with a simple rolling hash
|
||||
/// - Apply tanh-like scaling and L2-normalize to unit length
|
||||
pub struct TestHashEmbedder {
|
||||
dim: usize,
|
||||
model_name: String,
|
||||
}
|
||||
|
||||
impl TestHashEmbedder {
|
||||
pub fn new(dim: usize, model_name: impl Into<String>) -> Self {
|
||||
Self { dim, model_name: model_name.into() }
|
||||
}
|
||||
|
||||
fn l2_normalize(mut v: Vec<f32>) -> Vec<f32> {
|
||||
let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if norm > 0.0 {
|
||||
for x in &mut v {
|
||||
*x /= norm;
|
||||
}
|
||||
}
|
||||
v
|
||||
}
|
||||
}
|
||||
|
||||
impl Embedder for TestHashEmbedder {
|
||||
fn name(&self) -> String {
|
||||
format!("test-hash:{}", self.model_name)
|
||||
}
|
||||
|
||||
fn dim(&self) -> usize {
|
||||
self.dim
|
||||
}
|
||||
|
||||
fn embed(&self, text: &str) -> Result<Vec<f32>, DBError> {
|
||||
let mut acc = vec![0f32; self.dim];
|
||||
// A simple, deterministic folding hash over bytes
|
||||
let mut h1: u32 = 2166136261u32; // FNV-like seed
|
||||
let mut h2: u32 = 0x9e3779b9u32; // golden ratio
|
||||
for (i, b) in text.as_bytes().iter().enumerate() {
|
||||
h1 ^= *b as u32;
|
||||
h1 = h1.wrapping_mul(16777619u32);
|
||||
h2 = h2.wrapping_add(((*b as u32) << (i % 13)) ^ (h1.rotate_left((i % 7) as u32)));
|
||||
let idx = (h1 ^ h2) as usize % self.dim;
|
||||
// Map byte to [-1, 1] and accumulate with mild decay by position
|
||||
let val = ((*b as f32) / 127.5 - 1.0) * (1.0 / (1.0 + (i as f32 / 32.0)));
|
||||
acc[idx] += val;
|
||||
}
|
||||
// Non-linear squashing to stabilize + normalize
|
||||
for x in &mut acc {
|
||||
*x = x.tanh();
|
||||
}
|
||||
Ok(Self::l2_normalize(acc))
|
||||
}
|
||||
}
|
||||
|
||||
//// ----------------------------- IMAGE: trait + deterministic test embedder -----------------------------
|
||||
|
||||
/// Image embedding interface (separate from text to keep modality-specific inputs).
|
||||
pub trait ImageEmbedder: Send + Sync {
|
||||
/// Human-readable provider/model name
|
||||
fn name(&self) -> String;
|
||||
/// Embedding dimension
|
||||
fn dim(&self) -> usize;
|
||||
/// Embed a single image (raw bytes)
|
||||
fn embed_image(&self, bytes: &[u8]) -> Result<Vec<f32>, DBError>;
|
||||
/// Embed many images; default maps embed_image() over inputs
|
||||
fn embed_many_images(&self, images: &[Vec<u8>]) -> Result<Vec<Vec<f32>>, DBError> {
|
||||
images.iter().map(|b| self.embed_image(b)).collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Deterministic image embedder that folds bytes into buckets, applies tanh-like nonlinearity,
|
||||
/// and L2-normalizes. Suitable for CI and offline development.
|
||||
/// NOTE: This is NOT semantic; it is a stable hash-like representation.
|
||||
pub struct TestImageHashEmbedder {
|
||||
dim: usize,
|
||||
model_name: String,
|
||||
}
|
||||
|
||||
impl TestImageHashEmbedder {
|
||||
pub fn new(dim: usize, model_name: impl Into<String>) -> Self {
|
||||
Self { dim, model_name: model_name.into() }
|
||||
}
|
||||
|
||||
fn l2_normalize(mut v: Vec<f32>) -> Vec<f32> {
|
||||
let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if norm > 0.0 {
|
||||
for x in &mut v {
|
||||
*x /= norm;
|
||||
}
|
||||
}
|
||||
v
|
||||
}
|
||||
}
|
||||
|
||||
impl ImageEmbedder for TestImageHashEmbedder {
|
||||
fn name(&self) -> String {
|
||||
format!("test-image-hash:{}", self.model_name)
|
||||
}
|
||||
|
||||
fn dim(&self) -> usize {
|
||||
self.dim
|
||||
}
|
||||
|
||||
fn embed_image(&self, bytes: &[u8]) -> Result<Vec<f32>, DBError> {
|
||||
// Deterministic fold across bytes with two rolling accumulators.
|
||||
let mut acc = vec![0f32; self.dim];
|
||||
let mut h1: u32 = 0x811C9DC5; // FNV-like
|
||||
let mut h2: u32 = 0x9E3779B9; // golden ratio
|
||||
for (i, b) in bytes.iter().enumerate() {
|
||||
h1 ^= *b as u32;
|
||||
h1 = h1.wrapping_mul(16777619u32);
|
||||
// combine with position and h2
|
||||
h2 = h2.wrapping_add(((i as u32).rotate_left((i % 13) as u32)) ^ h1.rotate_left((i % 7) as u32));
|
||||
let idx = (h1 ^ h2) as usize % self.dim;
|
||||
// Map to [-1,1] and decay with position
|
||||
let val = ((*b as f32) / 127.5 - 1.0) * (1.0 / (1.0 + (i as f32 / 128.0)));
|
||||
acc[idx] += val;
|
||||
}
|
||||
for x in &mut acc {
|
||||
*x = x.tanh();
|
||||
}
|
||||
Ok(Self::l2_normalize(acc))
|
||||
}
|
||||
}
|
||||
|
||||
//// OpenAI embedder (supports OpenAI and Azure OpenAI via REST)
|
||||
struct OpenAIEmbedder {
|
||||
model: String,
|
||||
dim: usize,
|
||||
agent: Agent,
|
||||
endpoint: String,
|
||||
headers: Vec<(String, String)>,
|
||||
use_azure: bool,
|
||||
}
|
||||
|
||||
impl OpenAIEmbedder {
|
||||
fn new_from_config(cfg: &EmbeddingConfig) -> Result<Self, DBError> {
|
||||
// Whether to use Azure OpenAI
|
||||
let use_azure = cfg
|
||||
.get_param_string("use_azure")
|
||||
.map(|s| s.eq_ignore_ascii_case("true"))
|
||||
.unwrap_or(false);
|
||||
|
||||
// Resolve API key (OPENAI_API_KEY or AZURE_OPENAI_API_KEY by default)
|
||||
let api_key_env = cfg
|
||||
.get_param_string("api_key_env")
|
||||
.unwrap_or_else(|| {
|
||||
if use_azure {
|
||||
"AZURE_OPENAI_API_KEY".to_string()
|
||||
} else {
|
||||
"OPENAI_API_KEY".to_string()
|
||||
}
|
||||
});
|
||||
let api_key = std::env::var(&api_key_env)
|
||||
.map_err(|_| DBError(format!("Missing API key in env '{}'", api_key_env)))?;
|
||||
|
||||
// Resolve endpoint
|
||||
// - Standard OpenAI: https://api.openai.com/v1/embeddings (default) or params["base_url"]
|
||||
// - Azure OpenAI: {azure_endpoint}/openai/deployments/{deployment}/embeddings?api-version=...
|
||||
let endpoint = if use_azure {
|
||||
let base = cfg
|
||||
.get_param_string("azure_endpoint")
|
||||
.ok_or_else(|| DBError("Missing 'azure_endpoint' for Azure OpenAI".into()))?;
|
||||
let deployment = cfg
|
||||
.get_param_string("azure_deployment")
|
||||
.unwrap_or_else(|| cfg.model.clone());
|
||||
let api_version = cfg
|
||||
.get_param_string("azure_api_version")
|
||||
.unwrap_or_else(|| "2023-05-15".to_string());
|
||||
format!(
|
||||
"{}/openai/deployments/{}/embeddings?api-version={}",
|
||||
base.trim_end_matches('/'),
|
||||
deployment,
|
||||
api_version
|
||||
)
|
||||
} else {
|
||||
cfg.get_param_string("base_url")
|
||||
.unwrap_or_else(|| "https://api.openai.com/v1/embeddings".to_string())
|
||||
};
|
||||
|
||||
// Determine expected dimension (default 1536 for text-embedding-3-small; callers should override if needed)
|
||||
let dim = cfg
|
||||
.get_param_usize("dim")
|
||||
.or_else(|| cfg.get_param_usize("dimensions"))
|
||||
.unwrap_or(1536);
|
||||
|
||||
// Build an HTTP agent with timeouts (blocking; no tokio runtime involved)
|
||||
let agent = AgentBuilder::new()
|
||||
.timeout_read(Duration::from_secs(30))
|
||||
.timeout_write(Duration::from_secs(30))
|
||||
.build();
|
||||
|
||||
// Headers
|
||||
let mut headers: Vec<(String, String)> = Vec::new();
|
||||
headers.push(("Content-Type".to_string(), "application/json".to_string()));
|
||||
if use_azure {
|
||||
headers.push(("api-key".to_string(), api_key));
|
||||
} else {
|
||||
headers.push(("Authorization".to_string(), format!("Bearer {}", api_key)));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
model: cfg.model.clone(),
|
||||
dim,
|
||||
agent,
|
||||
endpoint,
|
||||
headers,
|
||||
use_azure,
|
||||
})
|
||||
}
|
||||
|
||||
fn request_many(&self, inputs: &[String]) -> Result<Vec<Vec<f32>>, DBError> {
|
||||
// Compose request body:
|
||||
// - Standard OpenAI: { "model": ..., "input": [...], "dimensions": dim? }
|
||||
// - Azure: { "input": [...], "dimensions": dim? } (model from deployment)
|
||||
let mut body = if self.use_azure {
|
||||
json!({ "input": inputs })
|
||||
} else {
|
||||
json!({ "model": self.model, "input": inputs })
|
||||
};
|
||||
if self.dim > 0 {
|
||||
body.as_object_mut()
|
||||
.unwrap()
|
||||
.insert("dimensions".to_string(), json!(self.dim));
|
||||
}
|
||||
|
||||
// Build request
|
||||
let mut req = self.agent.post(&self.endpoint);
|
||||
for (k, v) in &self.headers {
|
||||
req = req.set(k, v);
|
||||
}
|
||||
|
||||
// Send and handle errors
|
||||
let resp = req.send_json(body);
|
||||
let text = match resp {
|
||||
Ok(r) => r
|
||||
.into_string()
|
||||
.map_err(|e| DBError(format!("Failed to read embeddings response: {}", e)))?,
|
||||
Err(ureq::Error::Status(code, r)) => {
|
||||
let body = r.into_string().unwrap_or_default();
|
||||
return Err(DBError(format!("Embeddings API error {}: {}", code, body)));
|
||||
}
|
||||
Err(e) => return Err(DBError(format!("HTTP request failed: {}", e))),
|
||||
};
|
||||
|
||||
let val: serde_json::Value = serde_json::from_str(&text)
|
||||
.map_err(|e| DBError(format!("Invalid JSON from embeddings API: {}", e)))?;
|
||||
|
||||
let data = val
|
||||
.get("data")
|
||||
.and_then(|d| d.as_array())
|
||||
.ok_or_else(|| DBError("Embeddings API response missing 'data' array".into()))?;
|
||||
|
||||
let mut out: Vec<Vec<f32>> = Vec::with_capacity(data.len());
|
||||
for item in data {
|
||||
let emb = item
|
||||
.get("embedding")
|
||||
.and_then(|e| e.as_array())
|
||||
.ok_or_else(|| DBError("Embeddings API item missing 'embedding'".into()))?;
|
||||
let mut v: Vec<f32> = Vec::with_capacity(emb.len());
|
||||
for n in emb {
|
||||
let f = n
|
||||
.as_f64()
|
||||
.ok_or_else(|| DBError("Embedding element is not a number".into()))?;
|
||||
v.push(f as f32);
|
||||
}
|
||||
if self.dim > 0 && v.len() != self.dim {
|
||||
return Err(DBError(format!(
|
||||
"Embedding dimension mismatch: expected {}, got {}. Configure 'dim' or 'dimensions' to match output.",
|
||||
self.dim, v.len()
|
||||
)));
|
||||
}
|
||||
out.push(v);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
||||
impl Embedder for OpenAIEmbedder {
|
||||
fn name(&self) -> String {
|
||||
if self.use_azure {
|
||||
format!("azure-openai:{}", self.model)
|
||||
} else {
|
||||
format!("openai:{}", self.model)
|
||||
}
|
||||
}
|
||||
|
||||
fn dim(&self) -> usize {
|
||||
self.dim
|
||||
}
|
||||
|
||||
fn embed(&self, text: &str) -> Result<Vec<f32>, DBError> {
|
||||
let v = self.request_many(&[text.to_string()])?;
|
||||
Ok(v.into_iter().next().unwrap_or_else(|| vec![0.0; self.dim]))
|
||||
}
|
||||
|
||||
fn embed_many(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, DBError> {
|
||||
if texts.is_empty() {
|
||||
return Ok(vec![]);
|
||||
}
|
||||
self.request_many(texts)
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an embedder instance from a config.
|
||||
/// - TestHash: uses params["dim"] or defaults to 64
|
||||
/// - LanceOpenAI: uses OpenAI (or Azure OpenAI) embeddings REST API
|
||||
/// - Other Lance providers can be added similarly
|
||||
pub fn create_embedder(config: &EmbeddingConfig) -> Result<Arc<dyn Embedder>, DBError> {
|
||||
match &config.provider {
|
||||
EmbeddingProvider::TestHash => {
|
||||
let dim = config.get_param_usize("dim").unwrap_or(64);
|
||||
Ok(Arc::new(TestHashEmbedder::new(dim, config.model.clone())))
|
||||
}
|
||||
EmbeddingProvider::LanceOpenAI => {
|
||||
let inner = OpenAIEmbedder::new_from_config(config)?;
|
||||
Ok(Arc::new(inner))
|
||||
}
|
||||
EmbeddingProvider::ImageTestHash => {
|
||||
Err(DBError("Use create_image_embedder() for image providers".into()))
|
||||
}
|
||||
EmbeddingProvider::LanceFastEmbed => Err(DBError("LanceFastEmbed provider not yet implemented in Rust embedding layer; configure 'test-hash' or use 'openai'".into())),
|
||||
EmbeddingProvider::LanceOther(p) => Err(DBError(format!("Lance provider '{}' not implemented; configure 'openai' or 'test-hash'", p))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an image embedder instance from a config.
|
||||
pub fn create_image_embedder(config: &EmbeddingConfig) -> Result<Arc<dyn ImageEmbedder>, DBError> {
|
||||
match &config.provider {
|
||||
EmbeddingProvider::ImageTestHash => {
|
||||
let dim = config.get_param_usize("dim").unwrap_or(512);
|
||||
Ok(Arc::new(TestImageHashEmbedder::new(dim, config.model.clone())))
|
||||
}
|
||||
EmbeddingProvider::TestHash | EmbeddingProvider::LanceOpenAI => {
|
||||
Err(DBError("Configured text provider; dataset expects image provider (e.g., 'testimagehash')".into()))
|
||||
}
|
||||
EmbeddingProvider::LanceFastEmbed => Err(DBError("Image provider 'lancefastembed' not yet implemented".into())),
|
||||
EmbeddingProvider::LanceOther(p) => Err(DBError(format!("Image provider '{}' not implemented; use 'testimagehash' for now", p))),
|
||||
}
|
||||
}
|
663
src/lance_store.rs
Normal file
663
src/lance_store.rs
Normal file
@@ -0,0 +1,663 @@
|
||||
// LanceDB store abstraction (per database instance)
|
||||
// This module encapsulates all Lance/LanceDB operations for a given DB id.
|
||||
// Notes:
|
||||
// - We persist each dataset (aka "table") under <base_dir>/lance/<db_id>/<name>.lance
|
||||
// - Schema convention: id: Utf8 (non-null), vector: FixedSizeList<Float32, dim> (non-null), meta: Utf8 (nullable JSON string)
|
||||
// - We implement naive KNN (L2) scan in Rust for search to avoid tight coupling to lancedb search builder API.
|
||||
// Index creation uses lance::Dataset vector index; future optimization can route to index-aware search.
|
||||
|
||||
use std::cmp::Ordering;
|
||||
use std::collections::{BinaryHeap, HashMap};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::error::DBError;
|
||||
|
||||
use arrow_array::{Array, RecordBatch, RecordBatchIterator, StringArray};
|
||||
use arrow_array::builder::{FixedSizeListBuilder, Float32Builder, StringBuilder};
|
||||
use arrow_array::cast::AsArray;
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use futures::StreamExt;
|
||||
use serde_json::Value as JsonValue;
|
||||
|
||||
// Low-level Lance core
|
||||
use lance::dataset::{WriteMode, WriteParams};
|
||||
use lance::Dataset;
|
||||
|
||||
// Vector index (IVF_PQ etc.)
|
||||
|
||||
// High-level LanceDB (for deletes where available)
|
||||
use lancedb::connection::Connection;
|
||||
use arrow_array::types::Float32Type;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct LanceStore {
|
||||
base_dir: PathBuf,
|
||||
db_id: u64,
|
||||
}
|
||||
|
||||
impl LanceStore {
|
||||
// Create a new LanceStore rooted at <base_dir>/lance/<db_id>
|
||||
pub fn new(base_dir: &Path, db_id: u64) -> Result<Self, DBError> {
|
||||
let p = base_dir.join("lance").join(db_id.to_string());
|
||||
std::fs::create_dir_all(&p)
|
||||
.map_err(|e| DBError(format!("Failed to create Lance dir {}: {}", p.display(), e)))?;
|
||||
Ok(Self { base_dir: p, db_id })
|
||||
}
|
||||
|
||||
fn dataset_path(&self, name: &str) -> PathBuf {
|
||||
// Store datasets as directories or files with .lance suffix
|
||||
// We accept both "<name>" and "<name>.lance" as logical name; normalize on ".lance"
|
||||
let has_ext = name.ends_with(".lance");
|
||||
if has_ext {
|
||||
self.base_dir.join(name)
|
||||
} else {
|
||||
self.base_dir.join(format!("{name}.lance"))
|
||||
}
|
||||
}
|
||||
|
||||
fn file_uri(path: &Path) -> String {
|
||||
// lancedb can use filesystem path directly; keep it simple
|
||||
// Avoid file:// scheme since local paths are supported.
|
||||
path.to_string_lossy().to_string()
|
||||
}
|
||||
|
||||
async fn connect_db(&self) -> Result<Connection, DBError> {
|
||||
let uri = Self::file_uri(&self.base_dir);
|
||||
lancedb::connect(&uri)
|
||||
.execute()
|
||||
.await
|
||||
.map_err(|e| DBError(format!("LanceDB connect failed at {}: {}", uri, e)))
|
||||
}
|
||||
|
||||
fn vector_field(dim: i32) -> Field {
|
||||
Field::new(
|
||||
"vector",
|
||||
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim),
|
||||
false,
|
||||
)
|
||||
}
|
||||
|
||||
async fn read_existing_dim(&self, name: &str) -> Result<Option<i32>, DBError> {
|
||||
let path = self.dataset_path(name);
|
||||
if !path.exists() {
|
||||
return Ok(None);
|
||||
}
|
||||
let ds = Dataset::open(path.to_string_lossy().as_ref())
|
||||
.await
|
||||
.map_err(|e| DBError(format!("Open dataset failed: {}: {}", path.display(), e)))?;
|
||||
// Scan a single batch to infer vector dimension from the 'vector' column type
|
||||
let mut scan = ds.scan();
|
||||
if let Err(e) = scan.project(&["vector"]) {
|
||||
return Err(DBError(format!("Project failed while inferring dim: {}", e)));
|
||||
}
|
||||
let mut stream = scan
|
||||
.try_into_stream()
|
||||
.await
|
||||
.map_err(|e| DBError(format!("Scan stream failed while inferring dim: {}", e)))?;
|
||||
if let Some(batch_res) = stream.next().await {
|
||||
let batch = batch_res.map_err(|e| DBError(format!("Batch error: {}", e)))?;
|
||||
let vec_col = batch
|
||||
.column_by_name("vector")
|
||||
.ok_or_else(|| DBError("Column 'vector' missing".into()))?;
|
||||
let fsl = vec_col.as_fixed_size_list();
|
||||
let dim = fsl.value_length();
|
||||
return Ok(Some(dim));
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn build_schema(dim: i32) -> Arc<Schema> {
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Utf8, false),
|
||||
Self::vector_field(dim),
|
||||
Field::new("text", DataType::Utf8, true),
|
||||
Field::new("media_type", DataType::Utf8, true),
|
||||
Field::new("media_uri", DataType::Utf8, true),
|
||||
Field::new("meta", DataType::Utf8, true),
|
||||
]))
|
||||
}
|
||||
|
||||
fn build_one_row_batch(
|
||||
id: &str,
|
||||
vector: &[f32],
|
||||
meta: &HashMap<String, String>,
|
||||
text: Option<&str>,
|
||||
media_type: Option<&str>,
|
||||
media_uri: Option<&str>,
|
||||
dim: i32,
|
||||
) -> Result<(Arc<Schema>, RecordBatch), DBError> {
|
||||
if vector.len() as i32 != dim {
|
||||
return Err(DBError(format!(
|
||||
"Vector length mismatch: expected {}, got {}",
|
||||
dim,
|
||||
vector.len()
|
||||
)));
|
||||
}
|
||||
|
||||
let schema = Self::build_schema(dim);
|
||||
|
||||
// id column
|
||||
let mut id_builder = StringBuilder::new();
|
||||
id_builder.append_value(id);
|
||||
let id_arr = Arc::new(id_builder.finish()) as Arc<dyn Array>;
|
||||
|
||||
// vector column (FixedSizeList<Float32, dim>)
|
||||
let v_builder = Float32Builder::with_capacity(vector.len());
|
||||
let mut list_builder = FixedSizeListBuilder::new(v_builder, dim);
|
||||
for v in vector {
|
||||
list_builder.values().append_value(*v);
|
||||
}
|
||||
list_builder.append(true);
|
||||
let vec_arr = Arc::new(list_builder.finish()) as Arc<dyn Array>;
|
||||
|
||||
// text column (optional)
|
||||
let mut text_builder = StringBuilder::new();
|
||||
if let Some(t) = text {
|
||||
text_builder.append_value(t);
|
||||
} else {
|
||||
text_builder.append_null();
|
||||
}
|
||||
let text_arr = Arc::new(text_builder.finish()) as Arc<dyn Array>;
|
||||
|
||||
// media_type column (optional)
|
||||
let mut mt_builder = StringBuilder::new();
|
||||
if let Some(mt) = media_type {
|
||||
mt_builder.append_value(mt);
|
||||
} else {
|
||||
mt_builder.append_null();
|
||||
}
|
||||
let mt_arr = Arc::new(mt_builder.finish()) as Arc<dyn Array>;
|
||||
|
||||
// media_uri column (optional)
|
||||
let mut mu_builder = StringBuilder::new();
|
||||
if let Some(mu) = media_uri {
|
||||
mu_builder.append_value(mu);
|
||||
} else {
|
||||
mu_builder.append_null();
|
||||
}
|
||||
let mu_arr = Arc::new(mu_builder.finish()) as Arc<dyn Array>;
|
||||
|
||||
// meta column (JSON string)
|
||||
let meta_json = if meta.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(serde_json::to_string(meta).map_err(|e| DBError(format!("Serialize meta error: {e}")))?)
|
||||
};
|
||||
let mut meta_builder = StringBuilder::new();
|
||||
if let Some(s) = meta_json {
|
||||
meta_builder.append_value(&s);
|
||||
} else {
|
||||
meta_builder.append_null();
|
||||
}
|
||||
let meta_arr = Arc::new(meta_builder.finish()) as Arc<dyn Array>;
|
||||
|
||||
let batch =
|
||||
RecordBatch::try_new(schema.clone(), vec![id_arr, vec_arr, text_arr, mt_arr, mu_arr, meta_arr]).map_err(|e| {
|
||||
DBError(format!("RecordBatch build failed: {e}"))
|
||||
})?;
|
||||
|
||||
Ok((schema, batch))
|
||||
}
|
||||
|
||||
// Create a new dataset (vector collection) with dimension `dim`.
|
||||
pub async fn create_dataset(&self, name: &str, dim: usize) -> Result<(), DBError> {
|
||||
let dim_i32: i32 = dim
|
||||
.try_into()
|
||||
.map_err(|_| DBError("Dimension too large".into()))?;
|
||||
let path = self.dataset_path(name);
|
||||
|
||||
if path.exists() {
|
||||
// Validate dimension if present
|
||||
if let Some(existing_dim) = self.read_existing_dim(name).await? {
|
||||
if existing_dim != dim_i32 {
|
||||
return Err(DBError(format!(
|
||||
"Dataset '{}' already exists with dim {}, requested {}",
|
||||
name, existing_dim, dim_i32
|
||||
)));
|
||||
}
|
||||
// No-op
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
// Create an empty dataset by writing an empty batch
|
||||
let schema = Self::build_schema(dim_i32);
|
||||
let empty_id = Arc::new(StringArray::new_null(0));
|
||||
// Build an empty FixedSizeListArray
|
||||
let v_builder = Float32Builder::new();
|
||||
let mut list_builder = FixedSizeListBuilder::new(v_builder, dim_i32);
|
||||
let empty_vec = Arc::new(list_builder.finish()) as Arc<dyn Array>;
|
||||
let empty_text = Arc::new(StringArray::new_null(0));
|
||||
let empty_media_type = Arc::new(StringArray::new_null(0));
|
||||
let empty_media_uri = Arc::new(StringArray::new_null(0));
|
||||
let empty_meta = Arc::new(StringArray::new_null(0));
|
||||
|
||||
let empty_batch =
|
||||
RecordBatch::try_new(schema.clone(), vec![empty_id, empty_vec, empty_text, empty_media_type, empty_media_uri, empty_meta])
|
||||
.map_err(|e| DBError(format!("Build empty batch failed: {e}")))?;
|
||||
|
||||
let write_params = WriteParams {
|
||||
mode: WriteMode::Create,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let reader = RecordBatchIterator::new([Ok(empty_batch)], schema.clone());
|
||||
|
||||
Dataset::write(reader, path.to_string_lossy().as_ref(), Some(write_params))
|
||||
.await
|
||||
.map_err(|e| DBError(format!("Create dataset failed at {}: {}", path.display(), e)))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Store/Upsert a single vector with ID and optional metadata (append; duplicate IDs are possible for now)
|
||||
pub async fn store_vector(
|
||||
&self,
|
||||
name: &str,
|
||||
id: &str,
|
||||
vector: Vec<f32>,
|
||||
meta: HashMap<String, String>,
|
||||
text: Option<String>,
|
||||
) -> Result<(), DBError> {
|
||||
// Delegate to media-aware path with no media fields
|
||||
self.store_vector_with_media(name, id, vector, meta, text, None, None).await
|
||||
}
|
||||
|
||||
/// Store/Upsert a single vector with optional text and media fields (media_type/media_uri).
|
||||
pub async fn store_vector_with_media(
|
||||
&self,
|
||||
name: &str,
|
||||
id: &str,
|
||||
vector: Vec<f32>,
|
||||
meta: HashMap<String, String>,
|
||||
text: Option<String>,
|
||||
media_type: Option<String>,
|
||||
media_uri: Option<String>,
|
||||
) -> Result<(), DBError> {
|
||||
let path = self.dataset_path(name);
|
||||
|
||||
// Determine dimension: use existing or infer from vector
|
||||
let dim_i32 = if let Some(d) = self.read_existing_dim(name).await? {
|
||||
d
|
||||
} else {
|
||||
vector
|
||||
.len()
|
||||
.try_into()
|
||||
.map_err(|_| DBError("Vector length too large".into()))?
|
||||
};
|
||||
|
||||
let (schema, batch) = Self::build_one_row_batch(
|
||||
id,
|
||||
&vector,
|
||||
&meta,
|
||||
text.as_deref(),
|
||||
media_type.as_deref(),
|
||||
media_uri.as_deref(),
|
||||
dim_i32,
|
||||
)?;
|
||||
|
||||
// If LanceDB table exists and provides delete, we can upsert by deleting same id
|
||||
// Try best-effort delete; ignore errors to keep operation append-only on failure
|
||||
if path.exists() {
|
||||
if let Ok(conn) = self.connect_db().await {
|
||||
if let Ok(mut tbl) = conn.open_table(name).execute().await {
|
||||
let _ = tbl
|
||||
.delete(&format!("id = '{}'", id.replace('\'', "''")))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let write_params = WriteParams {
|
||||
mode: if path.exists() {
|
||||
WriteMode::Append
|
||||
} else {
|
||||
WriteMode::Create
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let reader = RecordBatchIterator::new([Ok(batch)], schema.clone());
|
||||
|
||||
Dataset::write(reader, path.to_string_lossy().as_ref(), Some(write_params))
|
||||
.await
|
||||
.map_err(|e| DBError(format!("Write (append/create) failed: {}", e)))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Delete a record by ID (best-effort; returns true if delete likely removed rows)
|
||||
pub async fn delete_by_id(&self, name: &str, id: &str) -> Result<bool, DBError> {
|
||||
let path = self.dataset_path(name);
|
||||
if !path.exists() {
|
||||
return Ok(false);
|
||||
}
|
||||
let conn = self.connect_db().await?;
|
||||
let mut tbl = conn
|
||||
.open_table(name)
|
||||
.execute()
|
||||
.await
|
||||
.map_err(|e| DBError(format!("Open table '{}' failed: {}", name, e)))?;
|
||||
// SQL-like predicate quoting
|
||||
let pred = format!("id = '{}'", id.replace('\'', "''"));
|
||||
// lancedb returns count or () depending on version; treat Ok as success
|
||||
match tbl.delete(&pred).await {
|
||||
Ok(_) => Ok(true),
|
||||
Err(e) => Err(DBError(format!("Delete failed: {}", e))),
|
||||
}
|
||||
}
|
||||
|
||||
// Drop the entire dataset
|
||||
pub async fn drop_dataset(&self, name: &str) -> Result<bool, DBError> {
|
||||
let path = self.dataset_path(name);
|
||||
// Try LanceDB drop first
|
||||
// Best-effort logical drop via lancedb if available; ignore failures.
|
||||
// Note: we rely on filesystem removal below for final cleanup.
|
||||
if let Ok(conn) = self.connect_db().await {
|
||||
if let Ok(mut t) = conn.open_table(name).execute().await {
|
||||
// Best-effort delete-all to reduce footprint prior to fs removal
|
||||
let _ = t.delete("true").await;
|
||||
}
|
||||
}
|
||||
if path.exists() {
|
||||
if path.is_dir() {
|
||||
std::fs::remove_dir_all(&path)
|
||||
.map_err(|e| DBError(format!("Failed to drop dataset '{}': {}", name, e)))?;
|
||||
} else {
|
||||
std::fs::remove_file(&path)
|
||||
.map_err(|e| DBError(format!("Failed to drop dataset '{}': {}", name, e)))?;
|
||||
}
|
||||
return Ok(true);
|
||||
}
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
// Search top-k nearest with optional filter; returns tuple of (id, score (lower=L2), meta)
|
||||
pub async fn search_vectors(
|
||||
&self,
|
||||
name: &str,
|
||||
query: Vec<f32>,
|
||||
k: usize,
|
||||
filter: Option<String>,
|
||||
return_fields: Option<Vec<String>>,
|
||||
) -> Result<Vec<(String, f32, HashMap<String, String>)>, DBError> {
|
||||
let path = self.dataset_path(name);
|
||||
if !path.exists() {
|
||||
return Err(DBError(format!("Dataset '{}' not found", name)));
|
||||
}
|
||||
// Determine dim and validate query length
|
||||
let dim_i32 = self
|
||||
.read_existing_dim(name)
|
||||
.await?
|
||||
.ok_or_else(|| DBError("Vector column not found".into()))?;
|
||||
if query.len() as i32 != dim_i32 {
|
||||
return Err(DBError(format!(
|
||||
"Query vector length mismatch: expected {}, got {}",
|
||||
dim_i32,
|
||||
query.len()
|
||||
)));
|
||||
}
|
||||
|
||||
let ds = Dataset::open(path.to_string_lossy().as_ref())
|
||||
.await
|
||||
.map_err(|e| DBError(format!("Open dataset failed: {}", e)))?;
|
||||
|
||||
// Build scanner with projection; we project needed fields and filter client-side to support meta keys
|
||||
let mut scan = ds.scan();
|
||||
if let Err(e) = scan.project(&["id", "vector", "meta", "text", "media_type", "media_uri"]) {
|
||||
return Err(DBError(format!("Project failed: {}", e)));
|
||||
}
|
||||
// Note: we no longer push down filter to Lance to allow filtering on meta fields client-side.
|
||||
|
||||
let mut stream = scan
|
||||
.try_into_stream()
|
||||
.await
|
||||
.map_err(|e| DBError(format!("Scan stream failed: {}", e)))?;
|
||||
|
||||
// Parse simple equality clause from filter for client-side filtering (supports one `key = 'value'`)
|
||||
let clause = filter.as_ref().and_then(|s| {
|
||||
fn parse_eq(s: &str) -> Option<(String, String)> {
|
||||
let s = s.trim();
|
||||
let pos = s.find('=').or_else(|| s.find(" = "))?;
|
||||
let (k, vraw) = s.split_at(pos);
|
||||
let mut v = vraw.trim_start_matches('=').trim();
|
||||
if (v.starts_with('\'') && v.ends_with('\'')) || (v.starts_with('"') && v.ends_with('"')) {
|
||||
if v.len() >= 2 {
|
||||
v = &v[1..v.len()-1];
|
||||
}
|
||||
}
|
||||
let key = k.trim().trim_matches('"').trim_matches('\'').to_string();
|
||||
if key.is_empty() { return None; }
|
||||
Some((key, v.to_string()))
|
||||
}
|
||||
parse_eq(s)
|
||||
});
|
||||
|
||||
// Maintain a max-heap with reverse ordering to keep top-k smallest distances
|
||||
#[derive(Debug)]
|
||||
struct Hit {
|
||||
dist: f32,
|
||||
id: String,
|
||||
meta: HashMap<String, String>,
|
||||
}
|
||||
impl PartialEq for Hit {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.dist.eq(&other.dist)
|
||||
}
|
||||
}
|
||||
impl Eq for Hit {}
|
||||
impl PartialOrd for Hit {
|
||||
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
|
||||
// Reverse for max-heap: larger distance = "greater"
|
||||
other.dist.partial_cmp(&self.dist)
|
||||
}
|
||||
}
|
||||
impl Ord for Hit {
|
||||
fn cmp(&self, other: &Self) -> Ordering {
|
||||
self.partial_cmp(other).unwrap_or(Ordering::Equal)
|
||||
}
|
||||
}
|
||||
|
||||
let mut heap: BinaryHeap<Hit> = BinaryHeap::with_capacity(k);
|
||||
|
||||
while let Some(batch_res) = stream.next().await {
|
||||
let batch = batch_res.map_err(|e| DBError(format!("Stream batch error: {}", e)))?;
|
||||
|
||||
let id_arr = batch
|
||||
.column_by_name("id")
|
||||
.ok_or_else(|| DBError("Column 'id' missing".into()))?
|
||||
.as_string::<i32>();
|
||||
let vec_arr = batch
|
||||
.column_by_name("vector")
|
||||
.ok_or_else(|| DBError("Column 'vector' missing".into()))?
|
||||
.as_fixed_size_list();
|
||||
let meta_arr = batch
|
||||
.column_by_name("meta")
|
||||
.map(|a| a.as_string::<i32>().clone());
|
||||
let text_arr = batch
|
||||
.column_by_name("text")
|
||||
.map(|a| a.as_string::<i32>().clone());
|
||||
let mt_arr = batch
|
||||
.column_by_name("media_type")
|
||||
.map(|a| a.as_string::<i32>().clone());
|
||||
let mu_arr = batch
|
||||
.column_by_name("media_uri")
|
||||
.map(|a| a.as_string::<i32>().clone());
|
||||
|
||||
for i in 0..batch.num_rows() {
|
||||
// Extract id
|
||||
let id_val = id_arr.value(i).to_string();
|
||||
|
||||
// Parse meta JSON if present
|
||||
let mut meta: HashMap<String, String> = HashMap::new();
|
||||
if let Some(meta_col) = &meta_arr {
|
||||
if !meta_col.is_null(i) {
|
||||
let s = meta_col.value(i);
|
||||
if let Ok(JsonValue::Object(map)) = serde_json::from_str::<JsonValue>(s) {
|
||||
for (k, v) in map {
|
||||
if let Some(vs) = v.as_str() {
|
||||
meta.insert(k, vs.to_string());
|
||||
} else if v.is_number() || v.is_boolean() {
|
||||
meta.insert(k, v.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Evaluate simple equality filter if provided (supports one clause)
|
||||
let passes = if let Some((ref key, ref val)) = clause {
|
||||
let candidate = match key.as_str() {
|
||||
"id" => Some(id_val.clone()),
|
||||
"text" => text_arr.as_ref().and_then(|col| if col.is_null(i) { None } else { Some(col.value(i).to_string()) }),
|
||||
"media_type" => mt_arr.as_ref().and_then(|col| if col.is_null(i) { None } else { Some(col.value(i).to_string()) }),
|
||||
"media_uri" => mu_arr.as_ref().and_then(|col| if col.is_null(i) { None } else { Some(col.value(i).to_string()) }),
|
||||
_ => meta.get(key).cloned(),
|
||||
};
|
||||
match candidate {
|
||||
Some(cv) => cv == *val,
|
||||
None => false,
|
||||
}
|
||||
} else { true };
|
||||
if !passes {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Compute L2 distance
|
||||
let val = vec_arr.value(i);
|
||||
let prim = val.as_primitive::<Float32Type>();
|
||||
let mut dist: f32 = 0.0;
|
||||
let plen = prim.len();
|
||||
for j in 0..plen {
|
||||
let r = prim.value(j);
|
||||
let d = query[j] - r;
|
||||
dist += d * d;
|
||||
}
|
||||
|
||||
// Apply return_fields on meta
|
||||
let mut meta_out = meta;
|
||||
if let Some(fields) = &return_fields {
|
||||
let mut filtered = HashMap::new();
|
||||
for f in fields {
|
||||
if let Some(val) = meta_out.get(f) {
|
||||
filtered.insert(f.clone(), val.clone());
|
||||
}
|
||||
}
|
||||
meta_out = filtered;
|
||||
}
|
||||
|
||||
let hit = Hit { dist, id: id_val, meta: meta_out };
|
||||
|
||||
if heap.len() < k {
|
||||
heap.push(hit);
|
||||
} else if let Some(top) = heap.peek() {
|
||||
if hit.dist < top.dist {
|
||||
heap.pop();
|
||||
heap.push(hit);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Extract and sort ascending by distance
|
||||
let mut hits: Vec<Hit> = heap.into_sorted_vec(); // already ascending by dist due to Ord
|
||||
let out = hits
|
||||
.drain(..)
|
||||
.map(|h| (h.id, h.dist, h.meta))
|
||||
.collect::<Vec<_>>();
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
// Create an ANN index on the vector column (IVF_PQ or similar)
|
||||
pub async fn create_index(
|
||||
&self,
|
||||
name: &str,
|
||||
index_type: &str,
|
||||
params: HashMap<String, String>,
|
||||
) -> Result<(), DBError> {
|
||||
let path = self.dataset_path(name);
|
||||
if !path.exists() {
|
||||
return Err(DBError(format!("Dataset '{}' not found", name)));
|
||||
}
|
||||
// Attempt to create a vector index using lance low-level API if available.
|
||||
// Some crate versions hide IndexType; to ensure build stability, we fall back to a no-op if the API is not accessible.
|
||||
let _ = (index_type, params); // currently unused; reserved for future tuning
|
||||
// TODO: Implement using lance::Dataset::create_index when public API is stable across versions.
|
||||
// For now, succeed as a no-op to keep flows working; search will operate as brute-force scan.
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// List datasets (tables) under this DB (show user-level logical names without .lance)
|
||||
pub async fn list_datasets(&self) -> Result<Vec<String>, DBError> {
|
||||
let mut out = Vec::new();
|
||||
if self.base_dir.exists() {
|
||||
if let Ok(rd) = std::fs::read_dir(&self.base_dir) {
|
||||
for entry in rd.flatten() {
|
||||
let p = entry.path();
|
||||
if let Some(name) = p.file_name().and_then(|s| s.to_str()) {
|
||||
// Only list .lance datasets
|
||||
if name.ends_with(".lance") {
|
||||
out.push(name.trim_end_matches(".lance").to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
// Return basic dataset info map
|
||||
pub async fn get_dataset_info(&self, name: &str) -> Result<HashMap<String, String>, DBError> {
|
||||
let path = self.dataset_path(name);
|
||||
let mut m = HashMap::new();
|
||||
m.insert("name".to_string(), name.to_string());
|
||||
m.insert("path".to_string(), path.display().to_string());
|
||||
if !path.exists() {
|
||||
return Err(DBError(format!("Dataset '{}' not found", name)));
|
||||
}
|
||||
|
||||
let ds = Dataset::open(path.to_string_lossy().as_ref())
|
||||
.await
|
||||
.map_err(|e| DBError(format!("Open dataset failed: {}", e)))?;
|
||||
|
||||
// dim: infer by scanning first batch
|
||||
let mut dim_str = "unknown".to_string();
|
||||
{
|
||||
let mut scan = ds.scan();
|
||||
if scan.project(&["vector"]).is_ok() {
|
||||
if let Ok(mut stream) = scan.try_into_stream().await {
|
||||
if let Some(batch_res) = stream.next().await {
|
||||
if let Ok(batch) = batch_res {
|
||||
if let Some(col) = batch.column_by_name("vector") {
|
||||
let fsl = col.as_fixed_size_list();
|
||||
dim_str = fsl.value_length().to_string();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
m.insert("dimension".to_string(), dim_str);
|
||||
|
||||
// row_count (approximate by scanning)
|
||||
let mut scan = ds.scan();
|
||||
if let Err(e) = scan.project(&["id"]) {
|
||||
return Err(DBError(format!("Project failed: {e}")));
|
||||
}
|
||||
let mut stream = scan
|
||||
.try_into_stream()
|
||||
.await
|
||||
.map_err(|e| DBError(format!("Scan failed: {e}")))?;
|
||||
let mut rows: usize = 0;
|
||||
while let Some(batch_res) = stream.next().await {
|
||||
let batch = batch_res.map_err(|e| DBError(format!("Scan batch error: {}", e)))?;
|
||||
rows += batch.num_rows();
|
||||
}
|
||||
m.insert("row_count".to_string(), rows.to_string());
|
||||
|
||||
// indexes: we can’t easily enumerate; set to "unknown" (future: read index metadata)
|
||||
m.insert("indexes".to_string(), "unknown".to_string());
|
||||
|
||||
Ok(m)
|
||||
}
|
||||
}
|
@@ -14,3 +14,5 @@ pub mod storage_sled;
|
||||
pub mod admin_meta;
|
||||
pub mod tantivy_search;
|
||||
pub mod search_cmd;
|
||||
pub mod lance_store;
|
||||
pub mod embedding;
|
||||
|
@@ -1,5 +1,6 @@
|
||||
// #![allow(unused_imports)]
|
||||
|
||||
use std::path::PathBuf;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
use herodb::server;
|
||||
@@ -13,7 +14,7 @@ use clap::Parser;
|
||||
struct Args {
|
||||
/// The directory of Redis DB file
|
||||
#[arg(long)]
|
||||
dir: String,
|
||||
dir: PathBuf,
|
||||
|
||||
/// The port of the Redis server, default is 6379 if not specified
|
||||
#[arg(long)]
|
||||
|
@@ -1,13 +1,16 @@
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum BackendType {
|
||||
Redb,
|
||||
Sled,
|
||||
Tantivy, // Full-text search backend (no KV storage)
|
||||
Lance, // Vector database backend (no KV storage)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DBOption {
|
||||
pub dir: String,
|
||||
pub dir: PathBuf,
|
||||
pub port: u16,
|
||||
pub debug: bool,
|
||||
// Deprecated for data DBs; retained for backward-compat on CLI parsing
|
||||
|
714
src/rpc.rs
714
src/rpc.rs
@@ -1,4 +1,5 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
use jsonrpsee::{core::RpcResult, proc_macros::rpc};
|
||||
@@ -8,6 +9,8 @@ use sha2::{Digest, Sha256};
|
||||
use crate::server::Server;
|
||||
use crate::options::DBOption;
|
||||
use crate::admin_meta;
|
||||
use crate::embedding::{EmbeddingConfig, EmbeddingProvider};
|
||||
use base64::{engine::general_purpose, Engine as _};
|
||||
|
||||
/// Database backend types
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -15,6 +18,7 @@ pub enum BackendType {
|
||||
Redb,
|
||||
Sled,
|
||||
Tantivy, // Full-text search backend (no KV storage)
|
||||
Lance, // Vector search backend (no KV storage)
|
||||
// Future: InMemory, Custom(String)
|
||||
}
|
||||
|
||||
@@ -160,12 +164,158 @@ pub trait Rpc {
|
||||
/// Drop an FT index
|
||||
#[method(name = "ftDrop")]
|
||||
async fn ft_drop(&self, db_id: u64, index_name: String) -> RpcResult<bool>;
|
||||
|
||||
// ----- LanceDB (Vector + Text) RPC endpoints -----
|
||||
|
||||
/// Create a new Lance dataset in a Lance-backed DB
|
||||
#[method(name = "lanceCreate")]
|
||||
async fn lance_create(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
dim: usize,
|
||||
) -> RpcResult<bool>;
|
||||
|
||||
/// Store a vector (with id and metadata) into a Lance dataset (deprecated; returns error)
|
||||
#[method(name = "lanceStore")]
|
||||
async fn lance_store(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
id: String,
|
||||
vector: Vec<f32>,
|
||||
meta: Option<HashMap<String, String>>,
|
||||
) -> RpcResult<bool>;
|
||||
|
||||
/// Search a Lance dataset with a query vector (deprecated; returns error)
|
||||
#[method(name = "lanceSearch")]
|
||||
async fn lance_search(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
vector: Vec<f32>,
|
||||
k: usize,
|
||||
filter: Option<String>,
|
||||
return_fields: Option<Vec<String>>,
|
||||
) -> RpcResult<serde_json::Value>;
|
||||
|
||||
/// Create an ANN index on a Lance dataset
|
||||
#[method(name = "lanceCreateIndex")]
|
||||
async fn lance_create_index(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
index_type: String,
|
||||
params: Option<HashMap<String, String>>,
|
||||
) -> RpcResult<bool>;
|
||||
|
||||
/// List Lance datasets for a DB
|
||||
#[method(name = "lanceList")]
|
||||
async fn lance_list(
|
||||
&self,
|
||||
db_id: u64,
|
||||
) -> RpcResult<Vec<String>>;
|
||||
|
||||
/// Get info for a Lance dataset
|
||||
#[method(name = "lanceInfo")]
|
||||
async fn lance_info(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
) -> RpcResult<serde_json::Value>;
|
||||
|
||||
/// Delete a record by id from a Lance dataset
|
||||
#[method(name = "lanceDel")]
|
||||
async fn lance_del(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
id: String,
|
||||
) -> RpcResult<bool>;
|
||||
|
||||
/// Drop a Lance dataset
|
||||
#[method(name = "lanceDrop")]
|
||||
async fn lance_drop(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
) -> RpcResult<bool>;
|
||||
|
||||
// New: Text-first endpoints (no user-provided vectors)
|
||||
/// Set per-dataset embedding configuration
|
||||
#[method(name = "lanceSetEmbeddingConfig")]
|
||||
async fn lance_set_embedding_config(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
provider: String,
|
||||
model: String,
|
||||
params: Option<HashMap<String, String>>,
|
||||
) -> RpcResult<bool>;
|
||||
|
||||
/// Get per-dataset embedding configuration
|
||||
#[method(name = "lanceGetEmbeddingConfig")]
|
||||
async fn lance_get_embedding_config(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
) -> RpcResult<serde_json::Value>;
|
||||
|
||||
/// Store text; server will embed and store vector+text+meta
|
||||
#[method(name = "lanceStoreText")]
|
||||
async fn lance_store_text(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
id: String,
|
||||
text: String,
|
||||
meta: Option<HashMap<String, String>>,
|
||||
) -> RpcResult<bool>;
|
||||
|
||||
/// Search using a text query; server will embed then search
|
||||
#[method(name = "lanceSearchText")]
|
||||
async fn lance_search_text(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
text: String,
|
||||
k: usize,
|
||||
filter: Option<String>,
|
||||
return_fields: Option<Vec<String>>,
|
||||
) -> RpcResult<serde_json::Value>;
|
||||
|
||||
// ----- Image-first endpoints (no user-provided vectors) -----
|
||||
|
||||
/// Store an image; exactly one of uri or bytes_b64 must be provided.
|
||||
#[method(name = "lanceStoreImage")]
|
||||
async fn lance_store_image(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
id: String,
|
||||
uri: Option<String>,
|
||||
bytes_b64: Option<String>,
|
||||
meta: Option<HashMap<String, String>>,
|
||||
) -> RpcResult<bool>;
|
||||
|
||||
/// Search using an image query; exactly one of uri or bytes_b64 must be provided.
|
||||
#[method(name = "lanceSearchImage")]
|
||||
async fn lance_search_image(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
k: usize,
|
||||
uri: Option<String>,
|
||||
bytes_b64: Option<String>,
|
||||
filter: Option<String>,
|
||||
return_fields: Option<Vec<String>>,
|
||||
) -> RpcResult<serde_json::Value>;
|
||||
}
|
||||
|
||||
/// RPC Server implementation
|
||||
pub struct RpcServerImpl {
|
||||
/// Base directory for database files
|
||||
base_dir: String,
|
||||
base_dir: PathBuf,
|
||||
/// Managed database servers
|
||||
servers: Arc<RwLock<HashMap<u64, Arc<Server>>>>,
|
||||
/// Default backend type
|
||||
@@ -176,7 +326,7 @@ pub struct RpcServerImpl {
|
||||
|
||||
impl RpcServerImpl {
|
||||
/// Create a new RPC server instance
|
||||
pub fn new(base_dir: String, backend: crate::options::BackendType, admin_secret: String) -> Self {
|
||||
pub fn new(base_dir: PathBuf, backend: crate::options::BackendType, admin_secret: String) -> Self {
|
||||
Self {
|
||||
base_dir,
|
||||
servers: Arc::new(RwLock::new(HashMap::new())),
|
||||
@@ -235,7 +385,10 @@ impl RpcServerImpl {
|
||||
}
|
||||
|
||||
// Create server instance with resolved backend
|
||||
let is_tantivy = matches!(effective_backend, crate::options::BackendType::Tantivy);
|
||||
let is_search_only = matches!(
|
||||
effective_backend,
|
||||
crate::options::BackendType::Tantivy | crate::options::BackendType::Lance
|
||||
);
|
||||
let db_option = DBOption {
|
||||
dir: self.base_dir.clone(),
|
||||
port: 0, // Not used for RPC-managed databases
|
||||
@@ -245,15 +398,15 @@ impl RpcServerImpl {
|
||||
backend: effective_backend.clone(),
|
||||
admin_secret: self.admin_secret.clone(),
|
||||
};
|
||||
|
||||
|
||||
let mut server = Server::new(db_option).await;
|
||||
|
||||
|
||||
// Set the selected database to the db_id
|
||||
server.selected_db = db_id;
|
||||
|
||||
|
||||
// Lazily open/create physical storage according to admin meta (per-db encryption)
|
||||
// Skip for Tantivy backend (no KV storage to open)
|
||||
if !is_tantivy {
|
||||
// Skip for search-only backends (Tantivy/Lance): no KV storage to open
|
||||
if !is_search_only {
|
||||
let _ = server.current_storage();
|
||||
}
|
||||
|
||||
@@ -343,6 +496,7 @@ impl RpcServerImpl {
|
||||
crate::options::BackendType::Redb => BackendType::Redb,
|
||||
crate::options::BackendType::Sled => BackendType::Sled,
|
||||
crate::options::BackendType::Tantivy => BackendType::Tantivy,
|
||||
crate::options::BackendType::Lance => BackendType::Lance,
|
||||
};
|
||||
|
||||
DatabaseInfo {
|
||||
@@ -351,7 +505,7 @@ impl RpcServerImpl {
|
||||
backend,
|
||||
encrypted,
|
||||
redis_version: Some("7.0".to_string()),
|
||||
storage_path: Some(server.option.dir.clone()),
|
||||
storage_path: Some(server.option.dir.display().to_string()),
|
||||
size_on_disk,
|
||||
key_count,
|
||||
created_at,
|
||||
@@ -394,12 +548,16 @@ impl RpcServer for RpcServerImpl {
|
||||
BackendType::Redb => crate::options::BackendType::Redb,
|
||||
BackendType::Sled => crate::options::BackendType::Sled,
|
||||
BackendType::Tantivy => crate::options::BackendType::Tantivy,
|
||||
BackendType::Lance => crate::options::BackendType::Lance,
|
||||
};
|
||||
admin_meta::set_database_backend(&self.base_dir, self.backend.clone(), &self.admin_secret, db_id, opt_backend.clone())
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
|
||||
|
||||
// Create server instance using base_dir, chosen backend and admin secret
|
||||
let is_tantivy_new = matches!(opt_backend, crate::options::BackendType::Tantivy);
|
||||
let is_search_only_new = matches!(
|
||||
opt_backend,
|
||||
crate::options::BackendType::Tantivy | crate::options::BackendType::Lance
|
||||
);
|
||||
let option = DBOption {
|
||||
dir: self.base_dir.clone(),
|
||||
port: 0, // Not used for RPC-managed databases
|
||||
@@ -409,13 +567,13 @@ impl RpcServer for RpcServerImpl {
|
||||
backend: opt_backend.clone(),
|
||||
admin_secret: self.admin_secret.clone(),
|
||||
};
|
||||
|
||||
|
||||
let mut server = Server::new(option).await;
|
||||
server.selected_db = db_id;
|
||||
|
||||
|
||||
// Initialize storage to create physical <id>.db with proper encryption from admin meta
|
||||
// Skip for Tantivy backend (no KV storage to initialize)
|
||||
if !is_tantivy_new {
|
||||
// Skip for search-only backends (Tantivy/Lance): no KV storage to initialize
|
||||
if !is_search_only_new {
|
||||
let _ = server.current_storage();
|
||||
}
|
||||
|
||||
@@ -675,4 +833,530 @@ impl RpcServer for RpcServerImpl {
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
// ----- LanceDB (Vector) RPC endpoints -----
|
||||
|
||||
async fn lance_create(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
dim: usize,
|
||||
) -> RpcResult<bool> {
|
||||
let server = self.get_or_create_server(db_id).await?;
|
||||
if db_id == 0 {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
|
||||
}
|
||||
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
|
||||
}
|
||||
if !server.has_write_permission() {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
|
||||
}
|
||||
server.lance_store()
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
|
||||
.create_dataset(&name, dim).await
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn lance_store(
|
||||
&self,
|
||||
_db_id: u64,
|
||||
_name: String,
|
||||
_id: String,
|
||||
_vector: Vec<f32>,
|
||||
_meta: Option<HashMap<String, String>>,
|
||||
) -> RpcResult<bool> {
|
||||
Err(jsonrpsee::types::ErrorObjectOwned::owned(
|
||||
-32000,
|
||||
"Vector endpoint removed. Use lanceStoreText instead.",
|
||||
None::<()>
|
||||
))
|
||||
}
|
||||
|
||||
async fn lance_search(
|
||||
&self,
|
||||
_db_id: u64,
|
||||
_name: String,
|
||||
_vector: Vec<f32>,
|
||||
_k: usize,
|
||||
_filter: Option<String>,
|
||||
_return_fields: Option<Vec<String>>,
|
||||
) -> RpcResult<serde_json::Value> {
|
||||
Err(jsonrpsee::types::ErrorObjectOwned::owned(
|
||||
-32000,
|
||||
"Vector endpoint removed. Use lanceSearchText instead.",
|
||||
None::<()>
|
||||
))
|
||||
}
|
||||
|
||||
async fn lance_create_index(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
index_type: String,
|
||||
params: Option<HashMap<String, String>>,
|
||||
) -> RpcResult<bool> {
|
||||
let server = self.get_or_create_server(db_id).await?;
|
||||
if db_id == 0 {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
|
||||
}
|
||||
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
|
||||
}
|
||||
if !server.has_write_permission() {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
|
||||
}
|
||||
server.lance_store()
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
|
||||
.create_index(&name, &index_type, params.unwrap_or_default()).await
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn lance_list(
|
||||
&self,
|
||||
db_id: u64,
|
||||
) -> RpcResult<Vec<String>> {
|
||||
let server = self.get_or_create_server(db_id).await?;
|
||||
if db_id == 0 {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
|
||||
}
|
||||
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
|
||||
}
|
||||
if !server.has_read_permission() {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "read permission denied", None::<()>));
|
||||
}
|
||||
let list = server.lance_store()
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
|
||||
.list_datasets().await
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
Ok(list)
|
||||
}
|
||||
|
||||
async fn lance_info(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
) -> RpcResult<serde_json::Value> {
|
||||
let server = self.get_or_create_server(db_id).await?;
|
||||
if db_id == 0 {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
|
||||
}
|
||||
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
|
||||
}
|
||||
if !server.has_read_permission() {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "read permission denied", None::<()>));
|
||||
}
|
||||
let info = server.lance_store()
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
|
||||
.get_dataset_info(&name).await
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
Ok(serde_json::json!(info))
|
||||
}
|
||||
|
||||
async fn lance_del(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
id: String,
|
||||
) -> RpcResult<bool> {
|
||||
let server = self.get_or_create_server(db_id).await?;
|
||||
if db_id == 0 {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
|
||||
}
|
||||
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
|
||||
}
|
||||
if !server.has_write_permission() {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
|
||||
}
|
||||
let ok = server.lance_store()
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
|
||||
.delete_by_id(&name, &id).await
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
Ok(ok)
|
||||
}
|
||||
|
||||
async fn lance_drop(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
) -> RpcResult<bool> {
|
||||
let server = self.get_or_create_server(db_id).await?;
|
||||
if db_id == 0 {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
|
||||
}
|
||||
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
|
||||
}
|
||||
if !server.has_write_permission() {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
|
||||
}
|
||||
let ok = server.lance_store()
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
|
||||
.drop_dataset(&name).await
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
Ok(ok)
|
||||
}
|
||||
|
||||
// ----- New text-first Lance RPC implementations -----
|
||||
|
||||
async fn lance_set_embedding_config(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
provider: String,
|
||||
model: String,
|
||||
params: Option<HashMap<String, String>>,
|
||||
) -> RpcResult<bool> {
|
||||
let server = self.get_or_create_server(db_id).await?;
|
||||
if db_id == 0 {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
|
||||
}
|
||||
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
|
||||
}
|
||||
if !server.has_write_permission() {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
|
||||
}
|
||||
let prov = match provider.to_lowercase().as_str() {
|
||||
"test-hash" | "testhash" => EmbeddingProvider::TestHash,
|
||||
"testimagehash" | "image-test-hash" | "imagetesthash" => EmbeddingProvider::ImageTestHash,
|
||||
"fastembed" | "lancefastembed" => EmbeddingProvider::LanceFastEmbed,
|
||||
"openai" | "lanceopenai" => EmbeddingProvider::LanceOpenAI,
|
||||
other => EmbeddingProvider::LanceOther(other.to_string()),
|
||||
};
|
||||
let cfg = EmbeddingConfig {
|
||||
provider: prov,
|
||||
model,
|
||||
params: params.unwrap_or_default(),
|
||||
};
|
||||
server.set_dataset_embedding_config(&name, &cfg)
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn lance_get_embedding_config(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
) -> RpcResult<serde_json::Value> {
|
||||
let server = self.get_or_create_server(db_id).await?;
|
||||
if db_id == 0 {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
|
||||
}
|
||||
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
|
||||
}
|
||||
if !server.has_read_permission() {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "read permission denied", None::<()>));
|
||||
}
|
||||
let cfg = server.get_dataset_embedding_config(&name)
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
Ok(serde_json::json!({
|
||||
"provider": match cfg.provider {
|
||||
EmbeddingProvider::TestHash => "test-hash",
|
||||
EmbeddingProvider::ImageTestHash => "testimagehash",
|
||||
EmbeddingProvider::LanceFastEmbed => "lancefastembed",
|
||||
EmbeddingProvider::LanceOpenAI => "lanceopenai",
|
||||
EmbeddingProvider::LanceOther(ref s) => s,
|
||||
},
|
||||
"model": cfg.model,
|
||||
"params": cfg.params
|
||||
}))
|
||||
}
|
||||
|
||||
async fn lance_store_text(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
id: String,
|
||||
text: String,
|
||||
meta: Option<HashMap<String, String>>,
|
||||
) -> RpcResult<bool> {
|
||||
let server = self.get_or_create_server(db_id).await?;
|
||||
if db_id == 0 {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
|
||||
}
|
||||
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
|
||||
}
|
||||
if !server.has_write_permission() {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
|
||||
}
|
||||
// Resolve embedder and run blocking embedding off the async runtime
|
||||
// Resolve embedder and run embedding on a plain OS thread (avoid dropping any runtime in async context)
|
||||
let embedder = server.get_embedder_for(&name)
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
let emb_arc = embedder.clone();
|
||||
let text_cl = text.clone();
|
||||
std::thread::spawn(move || {
|
||||
let res = emb_arc.embed(&text_cl);
|
||||
let _ = tx.send(res);
|
||||
});
|
||||
let vector = match rx.await {
|
||||
Ok(Ok(v)) => v,
|
||||
Ok(Err(e)) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>)),
|
||||
Err(recv_err) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, format!("embedding thread error: {}", recv_err), None::<()>)),
|
||||
};
|
||||
server.lance_store()
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
|
||||
.store_vector(&name, &id, vector, meta.unwrap_or_default(), Some(text)).await
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn lance_search_text(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
text: String,
|
||||
k: usize,
|
||||
filter: Option<String>,
|
||||
return_fields: Option<Vec<String>>,
|
||||
) -> RpcResult<serde_json::Value> {
|
||||
let server = self.get_or_create_server(db_id).await?;
|
||||
if db_id == 0 {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
|
||||
}
|
||||
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
|
||||
}
|
||||
if !server.has_read_permission() {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "read permission denied", None::<()>));
|
||||
}
|
||||
// Resolve embedder and run embedding on a plain OS thread (avoid dropping any runtime in async context)
|
||||
let embedder = server.get_embedder_for(&name)
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
let emb_arc = embedder.clone();
|
||||
let text_cl = text.clone();
|
||||
std::thread::spawn(move || {
|
||||
let res = emb_arc.embed(&text_cl);
|
||||
let _ = tx.send(res);
|
||||
});
|
||||
let qv = match rx.await {
|
||||
Ok(Ok(v)) => v,
|
||||
Ok(Err(e)) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>)),
|
||||
Err(recv_err) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, format!("embedding thread error: {}", recv_err), None::<()>)),
|
||||
};
|
||||
let results = server.lance_store()
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
|
||||
.search_vectors(&name, qv, k, filter, return_fields).await
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
|
||||
let json_results: Vec<serde_json::Value> = results.into_iter().map(|(id, score, meta)| {
|
||||
serde_json::json!({
|
||||
"id": id,
|
||||
"score": score,
|
||||
"meta": meta,
|
||||
})
|
||||
}).collect();
|
||||
|
||||
Ok(serde_json::json!({ "results": json_results }))
|
||||
}
|
||||
|
||||
// ----- New image-first Lance RPC implementations -----
|
||||
|
||||
async fn lance_store_image(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
id: String,
|
||||
uri: Option<String>,
|
||||
bytes_b64: Option<String>,
|
||||
meta: Option<HashMap<String, String>>,
|
||||
) -> RpcResult<bool> {
|
||||
let server = self.get_or_create_server(db_id).await?;
|
||||
if db_id == 0 {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
|
||||
}
|
||||
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
|
||||
}
|
||||
if !server.has_write_permission() {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
|
||||
}
|
||||
|
||||
// Validate exactly one of uri or bytes_b64
|
||||
let (use_uri, use_b64) = (uri.is_some(), bytes_b64.is_some());
|
||||
if (use_uri && use_b64) || (!use_uri && !use_b64) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(
|
||||
-32000,
|
||||
"Provide exactly one of 'uri' or 'bytes_b64'",
|
||||
None::<()>,
|
||||
));
|
||||
}
|
||||
|
||||
// Acquire image bytes (with caps)
|
||||
let max_bytes: usize = std::env::var("HERODB_IMAGE_MAX_BYTES")
|
||||
.ok()
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.unwrap_or(10 * 1024 * 1024) as usize;
|
||||
|
||||
let (bytes, media_uri_opt) = if let Some(u) = uri.clone() {
|
||||
let data = server
|
||||
.fetch_image_bytes_from_uri(&u)
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
(data, Some(u))
|
||||
} else {
|
||||
let b64 = bytes_b64.unwrap_or_default();
|
||||
let data = general_purpose::STANDARD
|
||||
.decode(b64.as_bytes())
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, format!("base64 decode error: {}", e), None::<()>))?;
|
||||
if data.len() > max_bytes {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(
|
||||
-32000,
|
||||
format!("Image exceeds max allowed bytes {}", max_bytes),
|
||||
None::<()>,
|
||||
));
|
||||
}
|
||||
(data, None)
|
||||
};
|
||||
|
||||
// Resolve image embedder and embed on a plain OS thread
|
||||
let img_embedder = server
|
||||
.get_image_embedder_for(&name)
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
let emb_arc = img_embedder.clone();
|
||||
let bytes_cl = bytes.clone();
|
||||
std::thread::spawn(move || {
|
||||
let res = emb_arc.embed_image(&bytes_cl);
|
||||
let _ = tx.send(res);
|
||||
});
|
||||
let vector = match rx.await {
|
||||
Ok(Ok(v)) => v,
|
||||
Ok(Err(e)) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>)),
|
||||
Err(recv_err) => {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(
|
||||
-32000,
|
||||
format!("embedding thread error: {}", recv_err),
|
||||
None::<()>,
|
||||
))
|
||||
}
|
||||
};
|
||||
|
||||
// Store vector with media fields
|
||||
server
|
||||
.lance_store()
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
|
||||
.store_vector_with_media(
|
||||
&name,
|
||||
&id,
|
||||
vector,
|
||||
meta.unwrap_or_default(),
|
||||
None,
|
||||
Some("image".to_string()),
|
||||
media_uri_opt,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn lance_search_image(
|
||||
&self,
|
||||
db_id: u64,
|
||||
name: String,
|
||||
k: usize,
|
||||
uri: Option<String>,
|
||||
bytes_b64: Option<String>,
|
||||
filter: Option<String>,
|
||||
return_fields: Option<Vec<String>>,
|
||||
) -> RpcResult<serde_json::Value> {
|
||||
let server = self.get_or_create_server(db_id).await?;
|
||||
if db_id == 0 {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Lance not allowed on DB 0", None::<()>));
|
||||
}
|
||||
if !matches!(server.option.backend, crate::options::BackendType::Lance) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Lance", None::<()>));
|
||||
}
|
||||
if !server.has_read_permission() {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "read permission denied", None::<()>));
|
||||
}
|
||||
|
||||
// Validate exactly one of uri or bytes_b64
|
||||
let (use_uri, use_b64) = (uri.is_some(), bytes_b64.is_some());
|
||||
if (use_uri && use_b64) || (!use_uri && !use_b64) {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(
|
||||
-32000,
|
||||
"Provide exactly one of 'uri' or 'bytes_b64'",
|
||||
None::<()>,
|
||||
));
|
||||
}
|
||||
|
||||
// Acquire image bytes for query (with caps)
|
||||
let max_bytes: usize = std::env::var("HERODB_IMAGE_MAX_BYTES")
|
||||
.ok()
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.unwrap_or(10 * 1024 * 1024) as usize;
|
||||
|
||||
let bytes = if let Some(u) = uri {
|
||||
server
|
||||
.fetch_image_bytes_from_uri(&u)
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
|
||||
} else {
|
||||
let b64 = bytes_b64.unwrap_or_default();
|
||||
let data = general_purpose::STANDARD
|
||||
.decode(b64.as_bytes())
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, format!("base64 decode error: {}", e), None::<()>))?;
|
||||
if data.len() > max_bytes {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(
|
||||
-32000,
|
||||
format!("Image exceeds max allowed bytes {}", max_bytes),
|
||||
None::<()>,
|
||||
));
|
||||
}
|
||||
data
|
||||
};
|
||||
|
||||
// Resolve image embedder and embed on OS thread
|
||||
let img_embedder = server
|
||||
.get_image_embedder_for(&name)
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
let emb_arc = img_embedder.clone();
|
||||
std::thread::spawn(move || {
|
||||
let res = emb_arc.embed_image(&bytes);
|
||||
let _ = tx.send(res);
|
||||
});
|
||||
let qv = match rx.await {
|
||||
Ok(Ok(v)) => v,
|
||||
Ok(Err(e)) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>)),
|
||||
Err(recv_err) => {
|
||||
return Err(jsonrpsee::types::ErrorObjectOwned::owned(
|
||||
-32000,
|
||||
format!("embedding thread error: {}", recv_err),
|
||||
None::<()>,
|
||||
))
|
||||
}
|
||||
};
|
||||
|
||||
// KNN search and return results
|
||||
let results = server
|
||||
.lance_store()
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
|
||||
.search_vectors(&name, qv, k, filter, return_fields)
|
||||
.await
|
||||
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
|
||||
|
||||
let json_results: Vec<serde_json::Value> = results
|
||||
.into_iter()
|
||||
.map(|(id, score, meta)| {
|
||||
serde_json::json!({
|
||||
"id": id,
|
||||
"score": score,
|
||||
"meta": meta,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(serde_json::json!({ "results": json_results }))
|
||||
}
|
||||
}
|
@@ -1,11 +1,12 @@
|
||||
use std::net::SocketAddr;
|
||||
use std::path::PathBuf;
|
||||
use jsonrpsee::server::{ServerBuilder, ServerHandle};
|
||||
use jsonrpsee::RpcModule;
|
||||
|
||||
use crate::rpc::{RpcServer, RpcServerImpl};
|
||||
|
||||
/// Start the RPC server on the specified address
|
||||
pub async fn start_rpc_server(addr: SocketAddr, base_dir: String, backend: crate::options::BackendType, admin_secret: String) -> Result<ServerHandle, Box<dyn std::error::Error + Send + Sync>> {
|
||||
pub async fn start_rpc_server(addr: SocketAddr, base_dir: PathBuf, backend: crate::options::BackendType, admin_secret: String) -> Result<ServerHandle, Box<dyn std::error::Error + Send + Sync>> {
|
||||
// Create the RPC server implementation
|
||||
let rpc_impl = RpcServerImpl::new(base_dir, backend, admin_secret);
|
||||
|
||||
@@ -34,7 +35,7 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn test_rpc_server_startup() {
|
||||
let addr = "127.0.0.1:0".parse().unwrap(); // Use port 0 for auto-assignment
|
||||
let base_dir = "/tmp/test_rpc".to_string();
|
||||
let base_dir = PathBuf::from("/tmp/test_rpc");
|
||||
let backend = crate::options::BackendType::Redb; // Default for test
|
||||
|
||||
let handle = start_rpc_server(addr, base_dir, backend, "test-admin".to_string()).await.unwrap();
|
||||
|
250
src/server.rs
250
src/server.rs
@@ -14,6 +14,15 @@ use crate::protocol::Protocol;
|
||||
use crate::storage_trait::StorageBackend;
|
||||
use crate::admin_meta;
|
||||
|
||||
// Embeddings: config and cache
|
||||
use crate::embedding::{EmbeddingConfig, create_embedder, Embedder, create_image_embedder, ImageEmbedder};
|
||||
use serde_json;
|
||||
use ureq::{Agent, AgentBuilder};
|
||||
use std::time::Duration;
|
||||
use std::io::Read;
|
||||
|
||||
const NO_DB_SELECTED: u64 = u64::MAX;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Server {
|
||||
pub db_cache: std::sync::Arc<std::sync::RwLock<HashMap<u64, Arc<dyn StorageBackend>>>>,
|
||||
@@ -26,6 +35,15 @@ pub struct Server {
|
||||
// In-memory registry of Tantivy search indexes for this server
|
||||
pub search_indexes: Arc<std::sync::RwLock<HashMap<String, Arc<crate::tantivy_search::TantivySearch>>>>,
|
||||
|
||||
// Per-DB Lance stores (vector DB), keyed by db_id
|
||||
pub lance_stores: Arc<std::sync::RwLock<HashMap<u64, Arc<crate::lance_store::LanceStore>>>>,
|
||||
|
||||
// Per-(db_id, dataset) embedder cache (text)
|
||||
pub embedders: Arc<std::sync::RwLock<HashMap<(u64, String), Arc<dyn Embedder>>>>,
|
||||
|
||||
// Per-(db_id, dataset) image embedder cache (image)
|
||||
pub image_embedders: Arc<std::sync::RwLock<HashMap<(u64, String), Arc<dyn ImageEmbedder>>>>,
|
||||
|
||||
// BLPOP waiter registry: per (db_index, key) FIFO of waiters
|
||||
pub list_waiters: Arc<Mutex<HashMap<u64, HashMap<String, Vec<Waiter>>>>>,
|
||||
pub waiter_seq: Arc<AtomicU64>,
|
||||
@@ -49,11 +67,14 @@ impl Server {
|
||||
db_cache: Arc::new(std::sync::RwLock::new(HashMap::new())),
|
||||
option,
|
||||
client_name: None,
|
||||
selected_db: 0,
|
||||
selected_db: NO_DB_SELECTED,
|
||||
queued_cmd: None,
|
||||
current_permissions: None,
|
||||
|
||||
search_indexes: Arc::new(std::sync::RwLock::new(HashMap::new())),
|
||||
lance_stores: Arc::new(std::sync::RwLock::new(HashMap::new())),
|
||||
embedders: Arc::new(std::sync::RwLock::new(HashMap::new())),
|
||||
image_embedders: Arc::new(std::sync::RwLock::new(HashMap::new())),
|
||||
list_waiters: Arc::new(Mutex::new(HashMap::new())),
|
||||
waiter_seq: Arc::new(AtomicU64::new(1)),
|
||||
}
|
||||
@@ -71,7 +92,30 @@ impl Server {
|
||||
base
|
||||
}
|
||||
|
||||
// Path where Lance datasets are stored, namespaced per selected DB:
|
||||
// <base_dir>/lance/<db_id>
|
||||
pub fn lance_data_path(&self) -> std::path::PathBuf {
|
||||
let base = std::path::PathBuf::from(&self.option.dir)
|
||||
.join("lance")
|
||||
.join(self.selected_db.to_string());
|
||||
if !base.exists() {
|
||||
let _ = std::fs::create_dir_all(&base);
|
||||
}
|
||||
base
|
||||
}
|
||||
|
||||
pub fn current_storage(&self) -> Result<Arc<dyn StorageBackend>, DBError> {
|
||||
// Require explicit SELECT before any storage access
|
||||
if self.selected_db == NO_DB_SELECTED {
|
||||
return Err(DBError("No database selected. Use SELECT <id> [KEY <key>] first".to_string()));
|
||||
}
|
||||
// Admin DB 0 access must be authenticated with SELECT 0 KEY <admin_secret>
|
||||
if self.selected_db == 0 {
|
||||
if !matches!(self.current_permissions, Some(crate::rpc::Permissions::ReadWrite)) {
|
||||
return Err(DBError("Admin DB 0 requires SELECT 0 KEY <admin_secret>".to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
let mut cache = self.db_cache.write().unwrap();
|
||||
|
||||
if let Some(storage) = cache.get(&self.selected_db) {
|
||||
@@ -99,10 +143,208 @@ impl Server {
|
||||
cache.insert(self.selected_db, storage.clone());
|
||||
Ok(storage)
|
||||
}
|
||||
|
||||
|
||||
/// Get or create the LanceStore for the currently selected DB.
|
||||
/// Only valid for non-zero DBs and when the backend is Lance.
|
||||
pub fn lance_store(&self) -> Result<Arc<crate::lance_store::LanceStore>, DBError> {
|
||||
if self.selected_db == 0 {
|
||||
return Err(DBError("Lance not available on admin DB 0".to_string()));
|
||||
}
|
||||
// Resolve backend for selected_db
|
||||
let backend_opt = crate::admin_meta::get_database_backend(
|
||||
&self.option.dir,
|
||||
self.option.backend.clone(),
|
||||
&self.option.admin_secret,
|
||||
self.selected_db,
|
||||
)
|
||||
.ok()
|
||||
.flatten();
|
||||
|
||||
if !matches!(backend_opt, Some(crate::options::BackendType::Lance)) {
|
||||
return Err(DBError("ERR DB backend is not Lance; LANCE.* commands are not allowed".to_string()));
|
||||
}
|
||||
|
||||
// Fast path: read lock
|
||||
{
|
||||
let map = self.lance_stores.read().unwrap();
|
||||
if let Some(store) = map.get(&self.selected_db) {
|
||||
return Ok(store.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// Slow path: create and insert
|
||||
let store = Arc::new(crate::lance_store::LanceStore::new(&self.option.dir, self.selected_db)?);
|
||||
{
|
||||
let mut map = self.lance_stores.write().unwrap();
|
||||
map.insert(self.selected_db, store.clone());
|
||||
}
|
||||
Ok(store)
|
||||
}
|
||||
|
||||
// ----- Embedding configuration and resolution -----
|
||||
|
||||
// Sidecar embedding config path: <base_dir>/lance/<db_id>/<dataset>.lance.embedding.json
|
||||
fn dataset_embedding_config_path(&self, dataset: &str) -> std::path::PathBuf {
|
||||
let mut base = self.lance_data_path();
|
||||
// Ensure parent dir exists
|
||||
if !base.exists() {
|
||||
let _ = std::fs::create_dir_all(&base);
|
||||
}
|
||||
base.push(format!("{}.lance.embedding.json", dataset));
|
||||
base
|
||||
}
|
||||
|
||||
/// Persist per-dataset embedding config as JSON sidecar.
|
||||
pub fn set_dataset_embedding_config(&self, dataset: &str, cfg: &EmbeddingConfig) -> Result<(), DBError> {
|
||||
if self.selected_db == 0 {
|
||||
return Err(DBError("Lance not available on admin DB 0".to_string()));
|
||||
}
|
||||
let p = self.dataset_embedding_config_path(dataset);
|
||||
let data = serde_json::to_vec_pretty(cfg)
|
||||
.map_err(|e| DBError(format!("Failed to serialize embedding config: {}", e)))?;
|
||||
std::fs::write(&p, data)
|
||||
.map_err(|e| DBError(format!("Failed to write embedding config {}: {}", p.display(), e)))?;
|
||||
// Invalidate embedder cache entry for this dataset
|
||||
{
|
||||
let mut map = self.embedders.write().unwrap();
|
||||
map.remove(&(self.selected_db, dataset.to_string()));
|
||||
}
|
||||
{
|
||||
let mut map_img = self.image_embedders.write().unwrap();
|
||||
map_img.remove(&(self.selected_db, dataset.to_string()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Load per-dataset embedding config.
|
||||
pub fn get_dataset_embedding_config(&self, dataset: &str) -> Result<EmbeddingConfig, DBError> {
|
||||
if self.selected_db == 0 {
|
||||
return Err(DBError("Lance not available on admin DB 0".to_string()));
|
||||
}
|
||||
let p = self.dataset_embedding_config_path(dataset);
|
||||
if !p.exists() {
|
||||
return Err(DBError(format!(
|
||||
"Embedding config not set for dataset '{}'. Use LANCE.EMBEDDING CONFIG SET ... or RPC to configure.",
|
||||
dataset
|
||||
)));
|
||||
}
|
||||
let data = std::fs::read(&p)
|
||||
.map_err(|e| DBError(format!("Failed to read embedding config {}: {}", p.display(), e)))?;
|
||||
let cfg: EmbeddingConfig = serde_json::from_slice(&data)
|
||||
.map_err(|e| DBError(format!("Failed to parse embedding config {}: {}", p.display(), e)))?;
|
||||
Ok(cfg)
|
||||
}
|
||||
|
||||
/// Resolve or build an embedder for (db_id, dataset). Caches instance.
|
||||
pub fn get_embedder_for(&self, dataset: &str) -> Result<Arc<dyn Embedder>, DBError> {
|
||||
if self.selected_db == 0 {
|
||||
return Err(DBError("Lance not available on admin DB 0".to_string()));
|
||||
}
|
||||
// Fast path
|
||||
{
|
||||
let map = self.embedders.read().unwrap();
|
||||
if let Some(e) = map.get(&(self.selected_db, dataset.to_string())) {
|
||||
return Ok(e.clone());
|
||||
}
|
||||
}
|
||||
// Load config and instantiate
|
||||
let cfg = self.get_dataset_embedding_config(dataset)?;
|
||||
let emb = create_embedder(&cfg)?;
|
||||
{
|
||||
let mut map = self.embedders.write().unwrap();
|
||||
map.insert((self.selected_db, dataset.to_string()), emb.clone());
|
||||
}
|
||||
Ok(emb)
|
||||
}
|
||||
|
||||
/// Resolve or build an IMAGE embedder for (db_id, dataset). Caches instance.
|
||||
pub fn get_image_embedder_for(&self, dataset: &str) -> Result<Arc<dyn ImageEmbedder>, DBError> {
|
||||
if self.selected_db == 0 {
|
||||
return Err(DBError("Lance not available on admin DB 0".to_string()));
|
||||
}
|
||||
// Fast path
|
||||
{
|
||||
let map = self.image_embedders.read().unwrap();
|
||||
if let Some(e) = map.get(&(self.selected_db, dataset.to_string())) {
|
||||
return Ok(e.clone());
|
||||
}
|
||||
}
|
||||
// Load config and instantiate
|
||||
let cfg = self.get_dataset_embedding_config(dataset)?;
|
||||
let emb = create_image_embedder(&cfg)?;
|
||||
{
|
||||
let mut map = self.image_embedders.write().unwrap();
|
||||
map.insert((self.selected_db, dataset.to_string()), emb.clone());
|
||||
}
|
||||
Ok(emb)
|
||||
}
|
||||
|
||||
/// Download image bytes from a URI with safety checks (size, timeout, content-type, optional host allowlist).
|
||||
/// Env overrides:
|
||||
/// - HERODB_IMAGE_MAX_BYTES (u64, default 10485760)
|
||||
/// - HERODB_IMAGE_FETCH_TIMEOUT_SECS (u64, default 30)
|
||||
/// - HERODB_IMAGE_ALLOWED_HOSTS (comma-separated, optional)
|
||||
pub fn fetch_image_bytes_from_uri(&self, uri: &str) -> Result<Vec<u8>, DBError> {
|
||||
// Basic scheme validation
|
||||
if !(uri.starts_with("http://") || uri.starts_with("https://")) {
|
||||
return Err(DBError("Only http(s) URIs are supported for image fetch".into()));
|
||||
}
|
||||
// Parse host (naive) for allowlist check
|
||||
let host = {
|
||||
let after_scheme = match uri.find("://") {
|
||||
Some(i) => &uri[i + 3..],
|
||||
None => uri,
|
||||
};
|
||||
let end = after_scheme.find('/').unwrap_or(after_scheme.len());
|
||||
let host_port = &after_scheme[..end];
|
||||
host_port.split('@').last().unwrap_or(host_port).split(':').next().unwrap_or(host_port).to_string()
|
||||
};
|
||||
|
||||
let max_bytes: u64 = std::env::var("HERODB_IMAGE_MAX_BYTES").ok().and_then(|s| s.parse::<u64>().ok()).unwrap_or(10 * 1024 * 1024);
|
||||
let timeout_secs: u64 = std::env::var("HERODB_IMAGE_FETCH_TIMEOUT_SECS").ok().and_then(|s| s.parse::<u64>().ok()).unwrap_or(30);
|
||||
let allowed_hosts_env = std::env::var("HERODB_IMAGE_ALLOWED_HOSTS").ok();
|
||||
if let Some(allow) = allowed_hosts_env {
|
||||
if !allow.split(',').map(|s| s.trim()).filter(|s| !s.is_empty()).any(|h| h.eq_ignore_ascii_case(&host)) {
|
||||
return Err(DBError(format!("Host '{}' not allowed for image fetch (HERODB_IMAGE_ALLOWED_HOSTS)", host)));
|
||||
}
|
||||
}
|
||||
|
||||
let agent: Agent = AgentBuilder::new()
|
||||
.timeout_read(Duration::from_secs(timeout_secs))
|
||||
.timeout_write(Duration::from_secs(timeout_secs))
|
||||
.build();
|
||||
|
||||
let resp = agent.get(uri).call().map_err(|e| DBError(format!("HTTP GET failed: {}", e)))?;
|
||||
// Validate content-type
|
||||
let ctype = resp.header("Content-Type").unwrap_or("");
|
||||
let ctype_main = ctype.split(';').next().unwrap_or("").trim().to_ascii_lowercase();
|
||||
if !ctype_main.starts_with("image/") {
|
||||
return Err(DBError(format!("Remote content-type '{}' is not image/*", ctype)));
|
||||
}
|
||||
|
||||
// Read with cap
|
||||
let mut reader = resp.into_reader();
|
||||
let mut buf: Vec<u8> = Vec::with_capacity(8192);
|
||||
let mut tmp = [0u8; 8192];
|
||||
let mut total: u64 = 0;
|
||||
loop {
|
||||
let n = reader.read(&mut tmp).map_err(|e| DBError(format!("Read error: {}", e)))?;
|
||||
if n == 0 { break; }
|
||||
total += n as u64;
|
||||
if total > max_bytes {
|
||||
return Err(DBError(format!("Image exceeds max allowed bytes {}", max_bytes)));
|
||||
}
|
||||
buf.extend_from_slice(&tmp[..n]);
|
||||
}
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
/// Check if current permissions allow read operations
|
||||
pub fn has_read_permission(&self) -> bool {
|
||||
// No DB selected -> no permissions
|
||||
if self.selected_db == NO_DB_SELECTED {
|
||||
return false;
|
||||
}
|
||||
// If an explicit permission is set for this connection, honor it.
|
||||
if let Some(perms) = self.current_permissions.as_ref() {
|
||||
return matches!(*perms, crate::rpc::Permissions::Read | crate::rpc::Permissions::ReadWrite);
|
||||
@@ -122,6 +364,10 @@ impl Server {
|
||||
|
||||
/// Check if current permissions allow write operations
|
||||
pub fn has_write_permission(&self) -> bool {
|
||||
// No DB selected -> no permissions
|
||||
if self.selected_db == NO_DB_SELECTED {
|
||||
return false;
|
||||
}
|
||||
// If an explicit permission is set for this connection, honor it.
|
||||
if let Some(perms) = self.current_permissions.as_ref() {
|
||||
return matches!(*perms, crate::rpc::Permissions::ReadWrite);
|
||||
|
@@ -1,4 +1,5 @@
|
||||
use herodb::{server::Server, options::DBOption};
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpStream;
|
||||
@@ -22,7 +23,7 @@ async fn debug_hset_simple() {
|
||||
|
||||
let port = 16500;
|
||||
let option = DBOption {
|
||||
dir: test_dir.to_string(),
|
||||
dir: PathBuf::from(test_dir),
|
||||
port,
|
||||
debug: false,
|
||||
encrypt: false,
|
||||
|
@@ -1,4 +1,5 @@
|
||||
use herodb::{server::Server, options::DBOption};
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpStream;
|
||||
@@ -13,7 +14,7 @@ async fn debug_hset_return_value() {
|
||||
std::fs::create_dir_all(&test_dir).unwrap();
|
||||
|
||||
let option = DBOption {
|
||||
dir: test_dir.to_string(),
|
||||
dir: PathBuf::from(test_dir),
|
||||
port: 16390,
|
||||
debug: false,
|
||||
encrypt: false,
|
||||
|
484
tests/lance_integration_tests.rs
Normal file
484
tests/lance_integration_tests.rs
Normal file
@@ -0,0 +1,484 @@
|
||||
use redis::{Client, Connection, RedisResult, Value};
|
||||
use std::process::{Child, Command};
|
||||
use std::time::Duration;
|
||||
|
||||
use jsonrpsee::http_client::{HttpClient, HttpClientBuilder};
|
||||
use herodb::rpc::{BackendType, DatabaseConfig, RpcClient};
|
||||
use base64::Engine;
|
||||
use tokio::time::sleep;
|
||||
|
||||
// ------------------------
|
||||
// Helpers
|
||||
// ------------------------
|
||||
|
||||
fn get_redis_connection(port: u16) -> Connection {
|
||||
let connection_info = format!("redis://127.0.0.1:{}", port);
|
||||
let client = Client::open(connection_info).unwrap();
|
||||
let mut attempts = 0;
|
||||
loop {
|
||||
match client.get_connection() {
|
||||
Ok(mut conn) => {
|
||||
if redis::cmd("PING").query::<String>(&mut conn).is_ok() {
|
||||
return conn;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
if attempts >= 3600 {
|
||||
panic!("Failed to connect to Redis server after 3600 attempts: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
attempts += 1;
|
||||
std::thread::sleep(Duration::from_millis(500));
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_rpc_client(port: u16) -> HttpClient {
|
||||
let url = format!("http://127.0.0.1:{}", port + 1); // RPC port = Redis port + 1
|
||||
HttpClientBuilder::default().build(url).unwrap()
|
||||
}
|
||||
|
||||
/// Wait until RPC server is responsive (getServerStats succeeds) or panic after retries.
|
||||
async fn wait_for_rpc_ready(client: &HttpClient, max_attempts: u32, delay: Duration) {
|
||||
for _ in 0..max_attempts {
|
||||
match client.get_server_stats().await {
|
||||
Ok(_) => return,
|
||||
Err(_) => {
|
||||
sleep(delay).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
panic!("RPC server did not become ready in time");
|
||||
}
|
||||
|
||||
// A guard to ensure the server process is killed when it goes out of scope and test dir cleaned.
|
||||
struct ServerProcessGuard {
|
||||
process: Child,
|
||||
test_dir: String,
|
||||
}
|
||||
|
||||
impl Drop for ServerProcessGuard {
|
||||
fn drop(&mut self) {
|
||||
eprintln!("Killing server process (pid: {})...", self.process.id());
|
||||
if let Err(e) = self.process.kill() {
|
||||
eprintln!("Failed to kill server process: {}", e);
|
||||
}
|
||||
match self.process.wait() {
|
||||
Ok(status) => eprintln!("Server process exited with: {}", status),
|
||||
Err(e) => eprintln!("Failed to wait on server process: {}", e),
|
||||
}
|
||||
|
||||
// Clean up the specific test directory
|
||||
eprintln!("Cleaning up test directory: {}", self.test_dir);
|
||||
if let Err(e) = std::fs::remove_dir_all(&self.test_dir) {
|
||||
eprintln!("Failed to clean up test directory: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Helper to set up the server and return guard + ports
|
||||
async fn setup_server() -> (ServerProcessGuard, u16) {
|
||||
use std::sync::atomic::{AtomicU16, Ordering};
|
||||
static PORT_COUNTER: AtomicU16 = AtomicU16::new(17500);
|
||||
let port = PORT_COUNTER.fetch_add(1, Ordering::SeqCst);
|
||||
|
||||
let test_dir = format!("/tmp/herodb_lance_test_{}", port);
|
||||
|
||||
// Clean up previous test data
|
||||
if std::path::Path::new(&test_dir).exists() {
|
||||
let _ = std::fs::remove_dir_all(&test_dir);
|
||||
}
|
||||
std::fs::create_dir_all(&test_dir).unwrap();
|
||||
|
||||
// Start the server in a subprocess with RPC enabled (follows tantivy test pattern)
|
||||
let child = Command::new("cargo")
|
||||
.args(&[
|
||||
"run",
|
||||
"--",
|
||||
"--dir",
|
||||
&test_dir,
|
||||
"--port",
|
||||
&port.to_string(),
|
||||
"--rpc-port",
|
||||
&(port + 1).to_string(),
|
||||
"--enable-rpc",
|
||||
"--debug",
|
||||
"--admin-secret",
|
||||
"test-admin",
|
||||
])
|
||||
.spawn()
|
||||
.expect("Failed to start server process");
|
||||
|
||||
let guard = ServerProcessGuard {
|
||||
process: child,
|
||||
test_dir,
|
||||
};
|
||||
|
||||
// Give the server time to build and start (cargo run may compile first)
|
||||
// Increase significantly to accommodate first-time dependency compilation in CI.
|
||||
std::thread::sleep(Duration::from_millis(60000));
|
||||
|
||||
(guard, port)
|
||||
}
|
||||
|
||||
// Convenient helpers for assertions on redis::Value
|
||||
fn value_is_ok(v: &Value) -> bool {
|
||||
match v {
|
||||
Value::Okay => true,
|
||||
Value::Status(s) if s == "OK" => true,
|
||||
Value::Data(d) if d == b"OK" => true,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn value_is_int_eq(v: &Value, expected: i64) -> bool {
|
||||
matches!(v, Value::Int(n) if *n == expected)
|
||||
}
|
||||
|
||||
fn value_is_str_eq(v: &Value, expected: &str) -> bool {
|
||||
match v {
|
||||
Value::Status(s) => s == expected,
|
||||
Value::Data(d) => String::from_utf8_lossy(d) == expected,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn to_string_lossy(v: &Value) -> String {
|
||||
match v {
|
||||
Value::Nil => "Nil".to_string(),
|
||||
Value::Int(n) => n.to_string(),
|
||||
Value::Status(s) => s.clone(),
|
||||
Value::Okay => "OK".to_string(),
|
||||
Value::Data(d) => String::from_utf8_lossy(d).to_string(),
|
||||
Value::Bulk(items) => {
|
||||
let inner: Vec<String> = items.iter().map(to_string_lossy).collect();
|
||||
format!("[{}]", inner.join(", "))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Extract ids from LANCE.SEARCH / LANCE.SEARCHIMAGE reply which is:
|
||||
// Array of elements: [ [id, score, [k,v,...]], [id, score, ...], ... ]
|
||||
fn extract_hit_ids(v: &Value) -> Vec<String> {
|
||||
let mut ids = Vec::new();
|
||||
if let Value::Bulk(items) = v {
|
||||
for item in items {
|
||||
if let Value::Bulk(row) = item {
|
||||
if !row.is_empty() {
|
||||
// first element is id (Data or Status)
|
||||
let id = match &row[0] {
|
||||
Value::Data(d) => String::from_utf8_lossy(d).to_string(),
|
||||
Value::Status(s) => s.clone(),
|
||||
Value::Int(n) => n.to_string(),
|
||||
_ => continue,
|
||||
};
|
||||
ids.push(id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
ids
|
||||
}
|
||||
|
||||
// Check whether a Bulk array (RESP array) contains a given string element.
|
||||
fn bulk_contains_string(v: &Value, needle: &str) -> bool {
|
||||
match v {
|
||||
Value::Bulk(items) => items.iter().any(|it| match it {
|
||||
Value::Data(d) => String::from_utf8_lossy(d).contains(needle),
|
||||
Value::Status(s) => s.contains(needle),
|
||||
Value::Bulk(_) => bulk_contains_string(it, needle),
|
||||
_ => false,
|
||||
}),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
// ------------------------
|
||||
// Test: Lance end-to-end (RESP) using only local embedders
|
||||
// ------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_lance_end_to_end() {
|
||||
let (_guard, port) = setup_server().await;
|
||||
|
||||
// First, wait for RESP to be available; this also gives cargo-run child ample time to finish building.
|
||||
// Reuse the helper that retries PING until success.
|
||||
{
|
||||
let _conn_ready = get_redis_connection(port);
|
||||
// Drop immediately; we only needed readiness.
|
||||
}
|
||||
|
||||
// Build RPC client and create a Lance DB
|
||||
let rpc_client = get_rpc_client(port).await;
|
||||
// Ensure RPC server is listening before we issue createDatabase (allow longer warm-up to accommodate first-build costs)
|
||||
wait_for_rpc_ready(&rpc_client, 3600, Duration::from_millis(250)).await;
|
||||
|
||||
let db_config = DatabaseConfig {
|
||||
name: Some("media-db".to_string()),
|
||||
storage_path: None,
|
||||
max_size: None,
|
||||
redis_version: None,
|
||||
};
|
||||
|
||||
let db_id = rpc_client
|
||||
.create_database(BackendType::Lance, db_config, None)
|
||||
.await
|
||||
.expect("create_database Lance failed");
|
||||
|
||||
assert_eq!(db_id, 1, "Expected first Lance DB id to be 1");
|
||||
|
||||
// Add access keys
|
||||
let _ = rpc_client
|
||||
.add_access_key(db_id, "readwrite_key".to_string(), "readwrite".to_string())
|
||||
.await
|
||||
.expect("add_access_key readwrite failed");
|
||||
|
||||
let _ = rpc_client
|
||||
.add_access_key(db_id, "read_key".to_string(), "read".to_string())
|
||||
.await
|
||||
.expect("add_access_key read failed");
|
||||
|
||||
// Connect to Redis and SELECT DB with readwrite key
|
||||
let mut conn = get_redis_connection(port);
|
||||
|
||||
let sel_ok: RedisResult<String> = redis::cmd("SELECT")
|
||||
.arg(db_id)
|
||||
.arg("KEY")
|
||||
.arg("readwrite_key")
|
||||
.query(&mut conn);
|
||||
assert!(sel_ok.is_ok(), "SELECT db with key failed: {:?}", sel_ok);
|
||||
assert_eq!(sel_ok.unwrap(), "OK");
|
||||
|
||||
// 1) Configure embedding providers: textset -> testhash dim 64, imageset -> testimagehash dim 512
|
||||
let v = redis::cmd("LANCE.EMBEDDING")
|
||||
.arg("CONFIG")
|
||||
.arg("SET")
|
||||
.arg("textset")
|
||||
.arg("PROVIDER")
|
||||
.arg("testhash")
|
||||
.arg("MODEL")
|
||||
.arg("any")
|
||||
.arg("PARAM")
|
||||
.arg("dim")
|
||||
.arg("64")
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
assert!(value_is_ok(&v), "Embedding config set (text) not OK: {}", to_string_lossy(&v));
|
||||
|
||||
let v = redis::cmd("LANCE.EMBEDDING")
|
||||
.arg("CONFIG")
|
||||
.arg("SET")
|
||||
.arg("imageset")
|
||||
.arg("PROVIDER")
|
||||
.arg("testimagehash")
|
||||
.arg("MODEL")
|
||||
.arg("any")
|
||||
.arg("PARAM")
|
||||
.arg("dim")
|
||||
.arg("512")
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
assert!(value_is_ok(&v), "Embedding config set (image) not OK: {}", to_string_lossy(&v));
|
||||
|
||||
// 2) Create datasets
|
||||
let v = redis::cmd("LANCE.CREATE")
|
||||
.arg("textset")
|
||||
.arg("DIM")
|
||||
.arg(64)
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
assert!(value_is_ok(&v), "LANCE.CREATE textset failed: {}", to_string_lossy(&v));
|
||||
|
||||
let v = redis::cmd("LANCE.CREATE")
|
||||
.arg("imageset")
|
||||
.arg("DIM")
|
||||
.arg(512)
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
assert!(value_is_ok(&v), "LANCE.CREATE imageset failed: {}", to_string_lossy(&v));
|
||||
|
||||
// 3) Store two text documents
|
||||
let v = redis::cmd("LANCE.STORE")
|
||||
.arg("textset")
|
||||
.arg("ID")
|
||||
.arg("doc-1")
|
||||
.arg("TEXT")
|
||||
.arg("The quick brown fox jumps over the lazy dog")
|
||||
.arg("META")
|
||||
.arg("title")
|
||||
.arg("Fox")
|
||||
.arg("category")
|
||||
.arg("animal")
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
assert!(value_is_ok(&v), "LANCE.STORE doc-1 failed: {}", to_string_lossy(&v));
|
||||
|
||||
let v = redis::cmd("LANCE.STORE")
|
||||
.arg("textset")
|
||||
.arg("ID")
|
||||
.arg("doc-2")
|
||||
.arg("TEXT")
|
||||
.arg("A fast auburn fox vaulted a sleepy canine")
|
||||
.arg("META")
|
||||
.arg("title")
|
||||
.arg("Paraphrase")
|
||||
.arg("category")
|
||||
.arg("animal")
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
assert!(value_is_ok(&v), "LANCE.STORE doc-2 failed: {}", to_string_lossy(&v));
|
||||
|
||||
// 4) Store two images via BYTES (local fake bytes; embedder only hashes bytes, not decoding)
|
||||
let img1: Vec<u8> = b"local-image-bytes-1-abcdefghijklmnopqrstuvwxyz".to_vec();
|
||||
let img2: Vec<u8> = b"local-image-bytes-2-ABCDEFGHIJKLMNOPQRSTUVWXYZ".to_vec();
|
||||
let img1_b64 = base64::engine::general_purpose::STANDARD.encode(&img1);
|
||||
let img2_b64 = base64::engine::general_purpose::STANDARD.encode(&img2);
|
||||
|
||||
let v = redis::cmd("LANCE.STOREIMAGE")
|
||||
.arg("imageset")
|
||||
.arg("ID")
|
||||
.arg("img-1")
|
||||
.arg("BYTES")
|
||||
.arg(&img1_b64)
|
||||
.arg("META")
|
||||
.arg("title")
|
||||
.arg("Local1")
|
||||
.arg("group")
|
||||
.arg("demo")
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
assert!(value_is_ok(&v), "LANCE.STOREIMAGE img-1 failed: {}", to_string_lossy(&v));
|
||||
|
||||
let v = redis::cmd("LANCE.STOREIMAGE")
|
||||
.arg("imageset")
|
||||
.arg("ID")
|
||||
.arg("img-2")
|
||||
.arg("BYTES")
|
||||
.arg(&img2_b64)
|
||||
.arg("META")
|
||||
.arg("title")
|
||||
.arg("Local2")
|
||||
.arg("group")
|
||||
.arg("demo")
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
assert!(value_is_ok(&v), "LANCE.STOREIMAGE img-2 failed: {}", to_string_lossy(&v));
|
||||
|
||||
// 5) Search text: K 2 QUERY "quick brown fox" RETURN 1 title
|
||||
let v = redis::cmd("LANCE.SEARCH")
|
||||
.arg("textset")
|
||||
.arg("K")
|
||||
.arg(2)
|
||||
.arg("QUERY")
|
||||
.arg("quick brown fox")
|
||||
.arg("RETURN")
|
||||
.arg(1)
|
||||
.arg("title")
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
|
||||
// Should be an array of hits
|
||||
let ids = extract_hit_ids(&v);
|
||||
assert!(
|
||||
ids.contains(&"doc-1".to_string()) || ids.contains(&"doc-2".to_string()),
|
||||
"LANCE.SEARCH should return doc-1/doc-2; got: {}",
|
||||
to_string_lossy(&v)
|
||||
);
|
||||
|
||||
// With FILTER on category
|
||||
let v = redis::cmd("LANCE.SEARCH")
|
||||
.arg("textset")
|
||||
.arg("K")
|
||||
.arg(2)
|
||||
.arg("QUERY")
|
||||
.arg("fox jumps")
|
||||
.arg("FILTER")
|
||||
.arg("category = 'animal'")
|
||||
.arg("RETURN")
|
||||
.arg(1)
|
||||
.arg("title")
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
|
||||
let ids_f = extract_hit_ids(&v);
|
||||
assert!(
|
||||
!ids_f.is_empty(),
|
||||
"Filtered LANCE.SEARCH should return at least one document; got: {}",
|
||||
to_string_lossy(&v)
|
||||
);
|
||||
|
||||
// 6) Search images with QUERYBYTES
|
||||
let query_img: Vec<u8> = b"local-image-query-3-1234567890".to_vec();
|
||||
let query_img_b64 = base64::engine::general_purpose::STANDARD.encode(&query_img);
|
||||
|
||||
let v = redis::cmd("LANCE.SEARCHIMAGE")
|
||||
.arg("imageset")
|
||||
.arg("K")
|
||||
.arg(2)
|
||||
.arg("QUERYBYTES")
|
||||
.arg(&query_img_b64)
|
||||
.arg("RETURN")
|
||||
.arg(1)
|
||||
.arg("title")
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
|
||||
// Should get 2 hits (img-1 and img-2) in some order; assert array non-empty
|
||||
let img_ids = extract_hit_ids(&v);
|
||||
assert!(
|
||||
!img_ids.is_empty(),
|
||||
"LANCE.SEARCHIMAGE should return non-empty results; got: {}",
|
||||
to_string_lossy(&v)
|
||||
);
|
||||
|
||||
// 7) Inspect datasets
|
||||
let v = redis::cmd("LANCE.LIST").query::<Value>(&mut conn).unwrap();
|
||||
assert!(
|
||||
bulk_contains_string(&v, "textset"),
|
||||
"LANCE.LIST missing textset: {}",
|
||||
to_string_lossy(&v)
|
||||
);
|
||||
assert!(
|
||||
bulk_contains_string(&v, "imageset"),
|
||||
"LANCE.LIST missing imageset: {}",
|
||||
to_string_lossy(&v)
|
||||
);
|
||||
|
||||
// INFO textset
|
||||
let info_text = redis::cmd("LANCE.INFO")
|
||||
.arg("textset")
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
// INFO returns Array [k,v,k,v,...] including "dimension" "64" and "row_count" "...".
|
||||
let info_str = to_string_lossy(&info_text);
|
||||
assert!(
|
||||
info_str.contains("dimension") && info_str.contains("64"),
|
||||
"LANCE.INFO textset should include dimension 64; got: {}",
|
||||
info_str
|
||||
);
|
||||
|
||||
// 8) Delete by id and drop datasets
|
||||
let v = redis::cmd("LANCE.DEL")
|
||||
.arg("textset")
|
||||
.arg("doc-2")
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
// Returns SimpleString "1" or Int 1 depending on encoding path; accept either
|
||||
assert!(
|
||||
value_is_int_eq(&v, 1) || value_is_str_eq(&v, "1"),
|
||||
"LANCE.DEL doc-2 expected 1; got {}",
|
||||
to_string_lossy(&v)
|
||||
);
|
||||
|
||||
let v = redis::cmd("LANCE.DROP")
|
||||
.arg("textset")
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
assert!(value_is_ok(&v), "LANCE.DROP textset failed: {}", to_string_lossy(&v));
|
||||
|
||||
let v = redis::cmd("LANCE.DROP")
|
||||
.arg("imageset")
|
||||
.query::<Value>(&mut conn)
|
||||
.unwrap();
|
||||
assert!(value_is_ok(&v), "LANCE.DROP imageset failed: {}", to_string_lossy(&v));
|
||||
}
|
@@ -1,4 +1,5 @@
|
||||
use herodb::{server::Server, options::DBOption};
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpStream;
|
||||
@@ -17,7 +18,7 @@ async fn start_test_server(test_name: &str) -> (Server, u16) {
|
||||
std::fs::create_dir_all(&test_dir).unwrap();
|
||||
|
||||
let option = DBOption {
|
||||
dir: test_dir,
|
||||
dir: PathBuf::from(test_dir),
|
||||
port,
|
||||
debug: true,
|
||||
encrypt: false,
|
||||
|
@@ -1,6 +1,7 @@
|
||||
use herodb::rpc::{BackendType, DatabaseConfig};
|
||||
use herodb::admin_meta;
|
||||
use herodb::options::BackendType as OptionsBackendType;
|
||||
use std::path::Path;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_rpc_server_basic() {
|
||||
@@ -70,11 +71,11 @@ async fn test_database_name_persistence() {
|
||||
let _ = std::fs::remove_dir_all(base_dir);
|
||||
|
||||
// Set the database name
|
||||
admin_meta::set_database_name(base_dir, backend.clone(), admin_secret, db_id, test_name)
|
||||
admin_meta::set_database_name(Path::new(base_dir), backend.clone(), admin_secret, db_id, test_name)
|
||||
.expect("Failed to set database name");
|
||||
|
||||
// Retrieve the database name
|
||||
let retrieved_name = admin_meta::get_database_name(base_dir, backend, admin_secret, db_id)
|
||||
let retrieved_name = admin_meta::get_database_name(Path::new(base_dir), backend, admin_secret, db_id)
|
||||
.expect("Failed to get database name");
|
||||
|
||||
// Verify the name matches
|
||||
|
@@ -1,4 +1,5 @@
|
||||
use herodb::{server::Server, options::DBOption};
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
use tokio::time::sleep;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
@@ -19,7 +20,7 @@ async fn start_test_server(test_name: &str) -> (Server, u16) {
|
||||
std::fs::create_dir_all(&test_dir).unwrap();
|
||||
|
||||
let option = DBOption {
|
||||
dir: test_dir,
|
||||
dir: PathBuf::from(test_dir),
|
||||
port,
|
||||
debug: true,
|
||||
encrypt: false,
|
||||
|
@@ -1,4 +1,5 @@
|
||||
use herodb::{server::Server, options::DBOption};
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpStream;
|
||||
@@ -17,7 +18,7 @@ async fn start_test_server(test_name: &str) -> (Server, u16) {
|
||||
std::fs::create_dir_all(&test_dir).unwrap();
|
||||
|
||||
let option = DBOption {
|
||||
dir: test_dir,
|
||||
dir: PathBuf::from(test_dir),
|
||||
port,
|
||||
debug: false,
|
||||
encrypt: false,
|
||||
|
@@ -1,4 +1,5 @@
|
||||
use herodb::{options::DBOption, server::Server};
|
||||
use std::path::PathBuf;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::time::{sleep, Duration};
|
||||
@@ -17,7 +18,7 @@ async fn start_test_server(test_name: &str) -> (Server, u16) {
|
||||
std::fs::create_dir_all(&test_dir).unwrap();
|
||||
|
||||
let option = DBOption {
|
||||
dir: test_dir,
|
||||
dir: PathBuf::from(test_dir),
|
||||
port,
|
||||
debug: false,
|
||||
encrypt: false,
|
||||
|
Reference in New Issue
Block a user