lancedb_impl #15
							
								
								
									
										1
									
								
								.gitignore
									
									
									
									
										vendored
									
									
								
							
							
						
						
									
										1
									
								
								.gitignore
									
									
									
									
										vendored
									
									
								
							| @@ -3,6 +3,7 @@ | ||||
| debug/ | ||||
| target/ | ||||
| .vscode/ | ||||
| test_images/ | ||||
|  | ||||
| # These are backup files generated by rustfmt | ||||
| **/*.rs.bk | ||||
|   | ||||
							
								
								
									
										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"] } | ||||
|   | ||||
							
								
								
									
										27
									
								
								README.md
									
									
									
									
									
								
							
							
						
						
									
										27
									
								
								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 | ||||
| @@ -77,8 +83,13 @@ For examples, see [JSON-RPC Examples](docs/rpc_examples.md) and [Admin DB 0 Mode | ||||
|  | ||||
| For more detailed information on commands, features, and advanced usage, please refer to the documentation: | ||||
|  | ||||
| - [Basics](docs/basics.md) | ||||
| - [Supported Commands](docs/cmds.md) | ||||
| - [AGE Cryptography](docs/age.md) | ||||
| - [Admin DB 0 Model (access control, per-db encryption)](docs/admin.md) | ||||
| - [JSON-RPC Examples (management API)](docs/rpc_examples.md) | ||||
| - [Basics](docs/basics.md) - Launch options, symmetric encryption, and basic usage | ||||
| - [Supported Commands](docs/cmds.md) - Complete Redis command reference and backend comparison | ||||
| - [AGE Cryptography](docs/age.md) - Asymmetric encryption and digital signatures | ||||
| - [Admin DB 0 Model](docs/admin.md) - Database management, access control, and per-database encryption | ||||
| - [JSON-RPC Examples](docs/rpc_examples.md) - Management API examples | ||||
| - [Full-Text Search](docs/search.md) - Tantivy-powered search capabilities | ||||
| - [Tantivy Backend](docs/tantivy.md) - Tantivy as a dedicated database backend | ||||
| - [Lance Vector Store](docs/lance.md) - Vector embeddings and semantic search | ||||
| - [Lance Text and Images Example](docs/lancedb_text_and_images_example.md) - End-to-end vector search examples | ||||
| - [Local Embedder Tutorial](docs/local_embedder_full_example.md) - Complete embedding models tutorial | ||||
| @@ -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) | ||||
| @@ -176,6 +177,6 @@ Copyable JSON examples are provided in the [RPC examples documentation](./rpc_ex | ||||
|   - Encryption: `meta:db:<id>:enc` | ||||
|  | ||||
| For command examples and management payloads: | ||||
| - RESP command basics: `docs/basics.md` | ||||
| - Supported commands: `docs/cmds.md` | ||||
| - JSON-RPC examples: `docs/rpc_examples.md` | ||||
| - RESP command basics: [docs/basics.md](./basics.md) | ||||
| - Supported commands: [docs/cmds.md](./cmds.md) | ||||
| - JSON-RPC examples: [docs/rpc_examples.md](./rpc_examples.md) | ||||
| @@ -701,7 +701,7 @@ This expanded documentation includes all the list commands that were implemented | ||||
|  | ||||
| ## Updated Database Selection and Access Keys | ||||
|  | ||||
| HeroDB uses an `Admin DB 0` to control database existence, access, and encryption. Access to data DBs can be public (no key) or private (requires a key). See detailed model in `docs/admin.md`. | ||||
| HeroDB uses an `Admin DB 0` to control database existence, access, and encryption. Access to data DBs can be public (no key) or private (requires a key). See detailed model in [docs/admin.md](./admin.md). | ||||
|  | ||||
| Examples: | ||||
|  | ||||
|   | ||||
							
								
								
									
										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](./admin.md). | ||||
|  | ||||
| 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 | ||||
| ``` | ||||
							
								
								
									
										440
									
								
								docs/lance.md
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										440
									
								
								docs/lance.md
									
									
									
									
									
										Normal file
									
								
							| @@ -0,0 +1,440 @@ | ||||
| # 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 API or custom OpenAI-compatible endpoints) | ||||
| - testhash (deterministic, CI-friendly; no network) | ||||
|  | ||||
| Environment variable for OpenAI: | ||||
| - Standard OpenAI: export OPENAI_API_KEY=sk-... | ||||
|  | ||||
| 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 | ||||
|  | ||||
| # Custom OpenAI-compatible endpoint (e.g., self-hosted) | ||||
| redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET myset PROVIDER openai MODEL text-embedding-3-small \ | ||||
|   PARAM endpoint http://localhost:8081/v1/embeddings \ | ||||
|   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) | ||||
							
								
								
									
										134
									
								
								docs/lancedb_text_and_images_example.md
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										134
									
								
								docs/lancedb_text_and_images_example.md
									
									
									
									
									
										Normal file
									
								
							| @@ -0,0 +1,134 @@ | ||||
| # 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 | ||||
| ``` | ||||
| Custom OpenAI-compatible endpoint: | ||||
| ```bash | ||||
| redis-cli -p 6379 LANCE.EMBEDDING CONFIG SET textset PROVIDER openai MODEL text-embedding-3-small \ | ||||
|   PARAM endpoint http://localhost:8081/v1/embeddings \ | ||||
|   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. | ||||
							
								
								
									
										831
									
								
								docs/local_embedder_full_example.md
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										831
									
								
								docs/local_embedder_full_example.md
									
									
									
									
									
										Normal file
									
								
							| @@ -0,0 +1,831 @@ | ||||
| # HeroDB Embedding Models: Complete Tutorial | ||||
|  | ||||
| This tutorial demonstrates how to use embedding models with HeroDB for vector search, covering local self-hosted models, OpenAI's API, and deterministic test embedders. | ||||
|  | ||||
| ## Table of Contents | ||||
| - [Prerequisites](#prerequisites) | ||||
| - [Scenario 1: Local Embedding Model](#scenario-1-local-embedding-model-testing) | ||||
| - [Scenario 2: OpenAI API](#scenario-2-openai-api) | ||||
| - [Scenario 3: Deterministic Test Embedder](#scenario-3-deterministic-test-embedder-no-network) | ||||
| - [Troubleshooting](#troubleshooting) | ||||
|  | ||||
| --- | ||||
|  | ||||
| ## Prerequisites | ||||
|  | ||||
| ### Start HeroDB Server | ||||
|  | ||||
| Build and start HeroDB with RPC enabled: | ||||
|  | ||||
| ```bash | ||||
| cargo build --release | ||||
| ./target/release/herodb --dir ./data --admin-secret my-admin-secret --enable-rpc --rpc-port 8080 | ||||
| ``` | ||||
|  | ||||
| This starts: | ||||
| - Redis-compatible server on port 6379 | ||||
| - JSON-RPC server on port 8080 | ||||
|  | ||||
| ### Client Tools | ||||
|  | ||||
| For Redis-like commands: | ||||
| ```bash | ||||
| redis-cli -p 6379 | ||||
| ``` | ||||
|  | ||||
| For JSON-RPC calls, use `curl`: | ||||
| ```bash | ||||
| curl -X POST http://localhost:8080 \ | ||||
|   -H "Content-Type: application/json" \ | ||||
|   -d '{"jsonrpc":"2.0","id":1,"method":"herodb_METHOD","params":[...]}' | ||||
| ``` | ||||
|  | ||||
| --- | ||||
|  | ||||
| ## Scenario 1: Local Embedding Model (Testing) | ||||
|  | ||||
| Run your own embedding service locally for development, testing, or privacy. | ||||
|  | ||||
| ### Option A: Python Mock Server (Simplest) | ||||
|  | ||||
| This creates a minimal OpenAI-compatible embedding server for testing. | ||||
|  | ||||
| **1. Create `mock_embedder.py`:** | ||||
|  | ||||
| ```python | ||||
| from flask import Flask, request, jsonify | ||||
| import numpy as np | ||||
|  | ||||
| app = Flask(__name__) | ||||
|  | ||||
| @app.route('/v1/embeddings', methods=['POST']) | ||||
| def embeddings(): | ||||
|     """OpenAI-compatible embeddings endpoint""" | ||||
|     data = request.json | ||||
|     inputs = data.get('input', []) | ||||
|      | ||||
|     # Handle both single string and array | ||||
|     if isinstance(inputs, str): | ||||
|         inputs = [inputs] | ||||
|      | ||||
|     # Generate deterministic 768-dim embeddings (hash-based) | ||||
|     embeddings = [] | ||||
|     for text in inputs: | ||||
|         # Simple hash to vector (deterministic) | ||||
|         vec = np.zeros(768) | ||||
|         for i, char in enumerate(text[:768]): | ||||
|             vec[i % 768] += ord(char) / 255.0 | ||||
|          | ||||
|         # L2 normalize | ||||
|         norm = np.linalg.norm(vec) | ||||
|         if norm > 0: | ||||
|             vec = vec / norm | ||||
|          | ||||
|         embeddings.append(vec.tolist()) | ||||
|      | ||||
|     return jsonify({ | ||||
|         "data": [{"embedding": emb, "index": i} for i, emb in enumerate(embeddings)], | ||||
|         "model": data.get('model', 'mock-local'), | ||||
|         "usage": {"total_tokens": sum(len(t) for t in inputs)} | ||||
|     }) | ||||
|  | ||||
| if __name__ == '__main__': | ||||
|     print("Starting mock embedding server on http://127.0.0.1:8081") | ||||
|     app.run(host='127.0.0.1', port=8081, debug=False) | ||||
| ``` | ||||
|  | ||||
| **2. Install dependencies and run:** | ||||
|  | ||||
| ```bash | ||||
| pip install flask numpy | ||||
| python mock_embedder.py | ||||
| ``` | ||||
|  | ||||
| Output: `Starting mock embedding server on http://127.0.0.1:8081` | ||||
|  | ||||
| **3. Test the server (optional):** | ||||
|  | ||||
| ```bash | ||||
| curl -X POST http://127.0.0.1:8081/v1/embeddings \ | ||||
|   -H "Content-Type: application/json" \ | ||||
|   -d '{"input":["hello world"],"model":"test"}' | ||||
| ``` | ||||
|  | ||||
| You should see a JSON response with a 768-dimensional embedding. | ||||
|  | ||||
| ### End-to-End Example with Local Model | ||||
|  | ||||
| **Step 1: Create a Lance database** | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 1, | ||||
|   "method": "herodb_createDatabase", | ||||
|   "params": [ | ||||
|     "Lance", | ||||
|     { "name": "local-vectors", "storage_path": null, "max_size": null, "redis_version": null }, | ||||
|     null | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Expected response: | ||||
| ```json | ||||
| {"jsonrpc":"2.0","id":1,"result":1} | ||||
| ``` | ||||
|  | ||||
| The database ID is `1`. | ||||
|  | ||||
| **Step 2: Configure embedding for the dataset** | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 2, | ||||
|   "method": "herodb_lanceSetEmbeddingConfig", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "products", | ||||
|     { | ||||
|       "provider": "openai", | ||||
|       "model": "mock-local", | ||||
|       "dim": 768, | ||||
|       "endpoint": "http://127.0.0.1:8081/v1/embeddings", | ||||
|       "headers": { | ||||
|         "Authorization": "Bearer dummy" | ||||
|       }, | ||||
|       "timeout_ms": 30000 | ||||
|     } | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Redis-like: | ||||
| ```bash | ||||
| redis-cli -p 6379 | ||||
| SELECT 1 | ||||
| LANCE.EMBEDDING CONFIG SET products PROVIDER openai MODEL mock-local DIM 768 ENDPOINT http://127.0.0.1:8081/v1/embeddings HEADER Authorization "Bearer dummy" TIMEOUTMS 30000 | ||||
| ``` | ||||
|  | ||||
| Expected response: | ||||
| ```json | ||||
| {"jsonrpc":"2.0","id":2,"result":true} | ||||
| ``` | ||||
|  | ||||
| **Step 3: Verify configuration** | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 3, | ||||
|   "method": "herodb_lanceGetEmbeddingConfig", | ||||
|   "params": [1, "products"] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Redis-like: | ||||
| ```bash | ||||
| LANCE.EMBEDDING CONFIG GET products | ||||
| ``` | ||||
|  | ||||
| Expected: Returns your configuration with provider, model, dim, endpoint, etc. | ||||
|  | ||||
| **Step 4: Insert product data** | ||||
|  | ||||
| JSON-RPC (item 1): | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 4, | ||||
|   "method": "herodb_lanceStoreText", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "products", | ||||
|     "item-1", | ||||
|     "Waterproof hiking boots with ankle support and aggressive tread", | ||||
|     { "brand": "TrailMax", "category": "footwear", "price": "129.99" } | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Redis-like: | ||||
| ```bash | ||||
| LANCE.STORE products ID item-1 TEXT "Waterproof hiking boots with ankle support and aggressive tread" META brand TrailMax category footwear price 129.99 | ||||
| ``` | ||||
|  | ||||
| JSON-RPC (item 2): | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 5, | ||||
|   "method": "herodb_lanceStoreText", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "products", | ||||
|     "item-2", | ||||
|     "Lightweight running shoes with breathable mesh upper", | ||||
|     { "brand": "SpeedFit", "category": "footwear", "price": "89.99" } | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| JSON-RPC (item 3): | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 6, | ||||
|   "method": "herodb_lanceStoreText", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "products", | ||||
|     "item-3", | ||||
|     "Insulated winter jacket with removable hood and multiple pockets", | ||||
|     { "brand": "WarmTech", "category": "outerwear", "price": "199.99" } | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| JSON-RPC (item 4): | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 7, | ||||
|   "method": "herodb_lanceStoreText", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "products", | ||||
|     "item-4", | ||||
|     "Camping tent for 4 people with waterproof rainfly", | ||||
|     { "brand": "OutdoorPro", "category": "camping", "price": "249.99" } | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Expected response for each: `{"jsonrpc":"2.0","id":N,"result":true}` | ||||
|  | ||||
| **Step 5: Search by text query** | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 8, | ||||
|   "method": "herodb_lanceSearchText", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "products", | ||||
|     "boots for hiking in wet conditions", | ||||
|     3, | ||||
|     null, | ||||
|     ["brand", "category", "price"] | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Redis-like: | ||||
| ```bash | ||||
| LANCE.SEARCH products K 3 QUERY "boots for hiking in wet conditions" RETURN 3 brand category price | ||||
| ``` | ||||
|  | ||||
| Expected response: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 8, | ||||
|   "result": { | ||||
|     "results": [ | ||||
|       { | ||||
|         "id": "item-1", | ||||
|         "score": 0.234, | ||||
|         "meta": { | ||||
|           "brand": "TrailMax", | ||||
|           "category": "footwear", | ||||
|           "price": "129.99" | ||||
|         } | ||||
|       }, | ||||
|       ... | ||||
|     ] | ||||
|   } | ||||
| } | ||||
| ``` | ||||
|  | ||||
| **Step 6: Search with metadata filter** | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 9, | ||||
|   "method": "herodb_lanceSearchText", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "products", | ||||
|     "comfortable shoes for running", | ||||
|     5, | ||||
|     "category = 'footwear'", | ||||
|     null | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Redis-like: | ||||
| ```bash | ||||
| LANCE.SEARCH products K 5 QUERY "comfortable shoes for running" FILTER "category = 'footwear'" | ||||
| ``` | ||||
|  | ||||
| This returns only items where `category` equals `'footwear'`. | ||||
|  | ||||
| **Step 7: List datasets** | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 10, | ||||
|   "method": "herodb_lanceList", | ||||
|   "params": [1] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Redis-like: | ||||
| ```bash | ||||
| LANCE.LIST | ||||
| ``` | ||||
|  | ||||
| **Step 8: Get dataset info** | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 11, | ||||
|   "method": "herodb_lanceInfo", | ||||
|   "params": [1, "products"] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Redis-like: | ||||
| ```bash | ||||
| LANCE.INFO products | ||||
| ``` | ||||
|  | ||||
| Returns dimension, row count, and other metadata. | ||||
|  | ||||
| --- | ||||
|  | ||||
| ## Scenario 2: OpenAI API | ||||
|  | ||||
| Use OpenAI's production embedding service for semantic search. | ||||
|  | ||||
| ### Setup | ||||
|  | ||||
| **1. Set your API key:** | ||||
|  | ||||
| ```bash | ||||
| export OPENAI_API_KEY="sk-your-actual-openai-key-here" | ||||
| ``` | ||||
|  | ||||
| **2. Start HeroDB** (same as before): | ||||
|  | ||||
| ```bash | ||||
| ./target/release/herodb --dir ./data --admin-secret my-admin-secret --enable-rpc --rpc-port 8080 | ||||
| ``` | ||||
|  | ||||
| ### End-to-End Example with OpenAI | ||||
|  | ||||
| **Step 1: Create a Lance database** | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 1, | ||||
|   "method": "herodb_createDatabase", | ||||
|   "params": [ | ||||
|     "Lance", | ||||
|     { "name": "openai-vectors", "storage_path": null, "max_size": null, "redis_version": null }, | ||||
|     null | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Expected: `{"jsonrpc":"2.0","id":1,"result":1}` (database ID = 1) | ||||
|  | ||||
| **Step 2: Configure OpenAI embeddings** | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 2, | ||||
|   "method": "herodb_lanceSetEmbeddingConfig", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "documents", | ||||
|     { | ||||
|       "provider": "openai", | ||||
|       "model": "text-embedding-3-small", | ||||
|       "dim": 1536, | ||||
|       "endpoint": null, | ||||
|       "headers": {}, | ||||
|       "timeout_ms": 30000 | ||||
|     } | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Redis-like: | ||||
| ```bash | ||||
| redis-cli -p 6379 | ||||
| SELECT 1 | ||||
| LANCE.EMBEDDING CONFIG SET documents PROVIDER openai MODEL text-embedding-3-small DIM 1536 TIMEOUTMS 30000 | ||||
| ``` | ||||
|  | ||||
| Notes: | ||||
| - `endpoint` is `null` (defaults to OpenAI API: https://api.openai.com/v1/embeddings) | ||||
| - `headers` is empty (Authorization auto-added from OPENAI_API_KEY env var) | ||||
| - `dim` is 1536 for text-embedding-3-small | ||||
|  | ||||
| Expected: `{"jsonrpc":"2.0","id":2,"result":true}` | ||||
|  | ||||
| **Step 3: Insert documents** | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 3, | ||||
|   "method": "herodb_lanceStoreText", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "documents", | ||||
|     "doc-1", | ||||
|     "The quick brown fox jumps over the lazy dog", | ||||
|     { "source": "example", "lang": "en", "topic": "animals" } | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 4, | ||||
|   "method": "herodb_lanceStoreText", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "documents", | ||||
|     "doc-2", | ||||
|     "Machine learning models require large datasets for training and validation", | ||||
|     { "source": "tech", "lang": "en", "topic": "ai" } | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 5, | ||||
|   "method": "herodb_lanceStoreText", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "documents", | ||||
|     "doc-3", | ||||
|     "Python is a popular programming language for data science and web development", | ||||
|     { "source": "tech", "lang": "en", "topic": "programming" } | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Redis-like: | ||||
| ```bash | ||||
| LANCE.STORE documents ID doc-1 TEXT "The quick brown fox jumps over the lazy dog" META source example lang en topic animals | ||||
| LANCE.STORE documents ID doc-2 TEXT "Machine learning models require large datasets for training and validation" META source tech lang en topic ai | ||||
| LANCE.STORE documents ID doc-3 TEXT "Python is a popular programming language for data science and web development" META source tech lang en topic programming | ||||
| ``` | ||||
|  | ||||
| Expected for each: `{"jsonrpc":"2.0","id":N,"result":true}` | ||||
|  | ||||
| **Step 4: Semantic search** | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 6, | ||||
|   "method": "herodb_lanceSearchText", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "documents", | ||||
|     "artificial intelligence and neural networks", | ||||
|     3, | ||||
|     null, | ||||
|     ["source", "topic"] | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Redis-like: | ||||
| ```bash | ||||
| LANCE.SEARCH documents K 3 QUERY "artificial intelligence and neural networks" RETURN 2 source topic | ||||
| ``` | ||||
|  | ||||
| Expected response (doc-2 should rank highest due to semantic similarity): | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 6, | ||||
|   "result": { | ||||
|     "results": [ | ||||
|       { | ||||
|         "id": "doc-2", | ||||
|         "score": 0.123, | ||||
|         "meta": { | ||||
|           "source": "tech", | ||||
|           "topic": "ai" | ||||
|         } | ||||
|       }, | ||||
|       { | ||||
|         "id": "doc-3", | ||||
|         "score": 0.456, | ||||
|         "meta": { | ||||
|           "source": "tech", | ||||
|           "topic": "programming" | ||||
|         } | ||||
|       }, | ||||
|       { | ||||
|         "id": "doc-1", | ||||
|         "score": 0.789, | ||||
|         "meta": { | ||||
|           "source": "example", | ||||
|           "topic": "animals" | ||||
|         } | ||||
|       } | ||||
|     ] | ||||
|   } | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Note: Lower score = better match (L2 distance). | ||||
|  | ||||
| **Step 5: Search with filter** | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 7, | ||||
|   "method": "herodb_lanceSearchText", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "documents", | ||||
|     "programming and software", | ||||
|     5, | ||||
|     "topic = 'programming'", | ||||
|     null | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Redis-like: | ||||
| ```bash | ||||
| LANCE.SEARCH documents K 5 QUERY "programming and software" FILTER "topic = 'programming'" | ||||
| ``` | ||||
|  | ||||
| This returns only documents where `topic` equals `'programming'`. | ||||
|  | ||||
| --- | ||||
|  | ||||
|  | ||||
| --- | ||||
|  | ||||
| ## Scenario 3: Deterministic Test Embedder (No Network) | ||||
|  | ||||
| For CI/offline development, use the built-in test embedder that requires no external service. | ||||
|  | ||||
| ### Configuration | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 1, | ||||
|   "method": "herodb_lanceSetEmbeddingConfig", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "testdata", | ||||
|     { | ||||
|       "provider": "test", | ||||
|       "model": "dev", | ||||
|       "dim": 64, | ||||
|       "endpoint": null, | ||||
|       "headers": {}, | ||||
|       "timeout_ms": null | ||||
|     } | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| Redis-like: | ||||
| ```bash | ||||
| SELECT 1 | ||||
| LANCE.EMBEDDING CONFIG SET testdata PROVIDER test MODEL dev DIM 64 | ||||
| ``` | ||||
|  | ||||
| ### Usage | ||||
|  | ||||
| Use `lanceStoreText` and `lanceSearchText` as in previous scenarios. The embeddings are: | ||||
| - Deterministic (same text → same vector) | ||||
| - Fast (no network) | ||||
| - Not semantic (hash-based, not ML) | ||||
|  | ||||
| Perfect for testing the vector storage/search mechanics without external dependencies. | ||||
|  | ||||
| --- | ||||
|  | ||||
| ## Advanced: Custom Headers and Timeouts | ||||
|  | ||||
| ### Example: Local model with custom auth | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 1, | ||||
|   "method": "herodb_lanceSetEmbeddingConfig", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "secure-data", | ||||
|     { | ||||
|       "provider": "openai", | ||||
|       "model": "custom-model", | ||||
|       "dim": 512, | ||||
|       "endpoint": "http://192.168.1.100:9000/embeddings", | ||||
|       "headers": { | ||||
|         "Authorization": "Bearer my-local-token", | ||||
|         "X-Custom-Header": "value" | ||||
|       }, | ||||
|       "timeout_ms": 60000 | ||||
|     } | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| ### Example: OpenAI with explicit API key (not from env) | ||||
|  | ||||
| JSON-RPC: | ||||
| ```json | ||||
| { | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 1, | ||||
|   "method": "herodb_lanceSetEmbeddingConfig", | ||||
|   "params": [ | ||||
|     1, | ||||
|     "dataset", | ||||
|     { | ||||
|       "provider": "openai", | ||||
|       "model": "text-embedding-3-small", | ||||
|       "dim": 1536, | ||||
|       "endpoint": null, | ||||
|       "headers": { | ||||
|         "Authorization": "Bearer sk-your-key-here" | ||||
|       }, | ||||
|       "timeout_ms": 30000 | ||||
|     } | ||||
|   ] | ||||
| } | ||||
| ``` | ||||
|  | ||||
| --- | ||||
|  | ||||
| ## Troubleshooting | ||||
|  | ||||
| ### Error: "Embedding config not set for dataset" | ||||
|  | ||||
| **Cause:** You tried to use `lanceStoreText` or `lanceSearchText` without configuring an embedder. | ||||
|  | ||||
| **Solution:** Run `lanceSetEmbeddingConfig` first. | ||||
|  | ||||
| ### Error: "Embedding dimension mismatch: expected X, got Y" | ||||
|  | ||||
| **Cause:** The embedding service returned vectors of a different size than configured. | ||||
|  | ||||
| **Solution:**  | ||||
| - For OpenAI text-embedding-3-small, use `dim: 1536` | ||||
| - For your local mock (from this tutorial), use `dim: 768` | ||||
| - Check your embedding service's actual output dimension | ||||
|  | ||||
| ### Error: "Missing API key in env 'OPENAI_API_KEY'" | ||||
|  | ||||
| **Cause:** Using OpenAI provider without setting the API key. | ||||
|  | ||||
| **Solution:** | ||||
| - Set `export OPENAI_API_KEY="sk-..."` before starting HeroDB, OR | ||||
| - Pass the key explicitly in headers: `"Authorization": "Bearer sk-..."` | ||||
|  | ||||
| ### Error: "HTTP request failed" or "Embeddings API error 404" | ||||
|  | ||||
| **Cause:** Cannot reach the embedding endpoint. | ||||
|  | ||||
| **Solution:** | ||||
| - Verify your local server is running: `curl http://127.0.0.1:8081/v1/embeddings` | ||||
| - Check the endpoint URL in your config | ||||
| - Ensure firewall allows the connection | ||||
|  | ||||
| ### Error: "ERR DB backend is not Lance" | ||||
|  | ||||
| **Cause:** Trying to use LANCE.* commands on a non-Lance database. | ||||
|  | ||||
| **Solution:** Create the database with backend "Lance" (see Step 1). | ||||
|  | ||||
| ### Error: "write permission denied" | ||||
|  | ||||
| **Cause:** Database is private and you haven't authenticated. | ||||
|  | ||||
| **Solution:** Use `SELECT <db_id> KEY <access-key>` or make the database public via RPC. | ||||
|  | ||||
| --- | ||||
|  | ||||
| ## Complete Example Script (Bash + curl) | ||||
|  | ||||
| Save as `test_embeddings.sh`: | ||||
|  | ||||
| ```bash | ||||
| #!/bin/bash | ||||
|  | ||||
| RPC_URL="http://localhost:8080" | ||||
|  | ||||
| # 1. Create Lance database | ||||
| curl -X POST $RPC_URL -H "Content-Type: application/json" -d '{ | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 1, | ||||
|   "method": "herodb_createDatabase", | ||||
|   "params": ["Lance", {"name": "test-vectors", "storage_path": null, "max_size": null, "redis_version": null}, null] | ||||
| }' | ||||
|  | ||||
| echo -e "\n" | ||||
|  | ||||
| # 2. Configure local embedder | ||||
| curl -X POST $RPC_URL -H "Content-Type: application/json" -d '{ | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 2, | ||||
|   "method": "herodb_lanceSetEmbeddingConfig", | ||||
|   "params": [1, "products", { | ||||
|     "provider": "openai", | ||||
|     "model": "mock", | ||||
|     "dim": 768, | ||||
|     "endpoint": "http://127.0.0.1:8081/v1/embeddings", | ||||
|     "headers": {"Authorization": "Bearer dummy"}, | ||||
|     "timeout_ms": 30000 | ||||
|   }] | ||||
| }' | ||||
|  | ||||
| echo -e "\n" | ||||
|  | ||||
| # 3. Insert data | ||||
| curl -X POST $RPC_URL -H "Content-Type: application/json" -d '{ | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 3, | ||||
|   "method": "herodb_lanceStoreText", | ||||
|   "params": [1, "products", "item-1", "Hiking boots", {"brand": "TrailMax"}] | ||||
| }' | ||||
|  | ||||
| echo -e "\n" | ||||
|  | ||||
| # 4. Search | ||||
| curl -X POST $RPC_URL -H "Content-Type: application/json" -d '{ | ||||
|   "jsonrpc": "2.0", | ||||
|   "id": 4, | ||||
|   "method": "herodb_lanceSearchText", | ||||
|   "params": [1, "products", "outdoor footwear", 5, null, null] | ||||
| }' | ||||
|  | ||||
| echo -e "\n" | ||||
| ``` | ||||
|  | ||||
| Run: | ||||
| ```bash | ||||
| chmod +x test_embeddings.sh | ||||
| ./test_embeddings.sh | ||||
| ``` | ||||
|  | ||||
| --- | ||||
|  | ||||
| ## Summary | ||||
|  | ||||
| | Provider | Use Case | Endpoint | API Key | | ||||
| |----------|----------|----------|---------| | ||||
| | `openai` | Production semantic search | Default (OpenAI) or custom URL | OPENAI_API_KEY env or headers | | ||||
| | `openai` | Local self-hosted gateway | http://127.0.0.1:8081/... | Optional (depends on your service) | | ||||
| | `test` | CI/offline development | N/A (local hash) | None | | ||||
| | `image_test` | Image testing | N/A (local hash) | None | | ||||
|  | ||||
| **Notes:** | ||||
| - The `provider` field is always `"openai"` for OpenAI-compatible services (whether cloud or local). This is because it uses the OpenAI-compatible API shape. | ||||
| - Use `endpoint` to point to your local service | ||||
| - Use `headers` for custom authentication | ||||
| - `dim` must match your embedding service's output dimension | ||||
| - Once configured, `lanceStoreText` and `lanceSearchText` handle embedding automatically | ||||
| @@ -249,5 +249,5 @@ Troubleshooting | ||||
|  | ||||
| Related docs | ||||
|  | ||||
| - Command‑level search overview: [docs/search.md](docs/search.md:1) | ||||
| - RPC definitions: [src/rpc.rs](src/rpc.rs:1) | ||||
| - Command‑level search overview: [docs/search.md](./search.md) | ||||
| - RPC definitions: [src/rpc.rs](../src/rpc.rs) | ||||
							
								
								
									
										34
									
								
								mock_embedder.py
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										34
									
								
								mock_embedder.py
									
									
									
									
									
										Normal file
									
								
							| @@ -0,0 +1,34 @@ | ||||
| from flask import Flask, request, jsonify | ||||
| import numpy as np | ||||
|  | ||||
| app = Flask(__name__) | ||||
|  | ||||
| @app.route('/v1/embeddings', methods=['POST']) | ||||
| def embeddings(): | ||||
|     data = request.json | ||||
|     inputs = data.get('input', []) | ||||
|     if isinstance(inputs, str): | ||||
|         inputs = [inputs] | ||||
|      | ||||
|     # Generate deterministic 768-dim embeddings (hash-based) | ||||
|     embeddings = [] | ||||
|     for text in inputs: | ||||
|         # Simple hash to vector | ||||
|         vec = np.zeros(768) | ||||
|         for i, char in enumerate(text[:768]): | ||||
|             vec[i % 768] += ord(char) / 255.0 | ||||
|         # Normalize | ||||
|         norm = np.linalg.norm(vec) | ||||
|         if norm > 0: | ||||
|             vec = vec / norm | ||||
|         embeddings.append(vec.tolist()) | ||||
|      | ||||
|     return jsonify({ | ||||
|         "data": [{"embedding": emb} for emb in embeddings], | ||||
|         "model": data.get('model', 'mock'), | ||||
|         "usage": {"total_tokens": sum(len(t) for t in inputs)} | ||||
|     }) | ||||
|  | ||||
| if __name__ == '__main__': | ||||
|     app.run(host='127.0.0.1', port=8081) | ||||
|  | ||||
| @@ -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) | ||||
| @@ -206,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 | ||||
| @@ -299,6 +302,7 @@ 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(()) | ||||
| @@ -316,6 +320,7 @@ 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), | ||||
|     } | ||||
| } | ||||
|   | ||||
							
								
								
									
										763
									
								
								src/cmd.rs
									
									
									
									
									
								
							
							
						
						
									
										763
									
								
								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}; | ||||
| use base64::{engine::general_purpose, Engine as _}; | ||||
| use tokio::time::{timeout, Duration}; | ||||
| use futures::future::select_all; | ||||
|  | ||||
| @@ -125,6 +126,65 @@ 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, | ||||
|         config: EmbeddingConfig, | ||||
|     }, | ||||
|     LanceEmbeddingConfigGet { | ||||
|         name: String, | ||||
|     }, | ||||
|     LanceList, | ||||
|     LanceInfo { | ||||
|         name: String, | ||||
|     }, | ||||
|     LanceDel { | ||||
|         name: String, | ||||
|         id: String, | ||||
|     }, | ||||
|     LanceDrop { | ||||
|         name: String, | ||||
|     } | ||||
| } | ||||
|  | ||||
| @@ -815,6 +875,340 @@ 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 DIM d [ENDPOINT url] [HEADER k v]... [TIMEOUTMS t] | ||||
|                             // 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 DIM d [ENDPOINT url] [HEADER k v]... [TIMEOUTMS t]".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 dim: Option<usize> = None; | ||||
|                                 let mut endpoint: Option<String> = None; | ||||
|                                 let mut headers: std::collections::HashMap<String, String> = std::collections::HashMap::new(); | ||||
|                                 let mut timeout_ms: Option<u64> = None; | ||||
|  | ||||
|                                 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; | ||||
|                                         } | ||||
|                                         "DIM" => { | ||||
|                                             if i + 1 >= cmd.len() { | ||||
|                                                 return Err(DBError("ERR DIM requires a value".to_string())); | ||||
|                                             } | ||||
|                                             let d: usize = cmd[i + 1].parse().map_err(|_| DBError("ERR DIM must be an integer".to_string()))?; | ||||
|                                             dim = Some(d); | ||||
|                                             i += 2; | ||||
|                                         } | ||||
|                                         "ENDPOINT" => { | ||||
|                                             if i + 1 >= cmd.len() { | ||||
|                                                 return Err(DBError("ERR ENDPOINT requires a value".to_string())); | ||||
|                                             } | ||||
|                                             endpoint = Some(cmd[i + 1].clone()); | ||||
|                                             i += 2; | ||||
|                                         } | ||||
|                                         "HEADER" => { | ||||
|                                             if i + 2 >= cmd.len() { | ||||
|                                                 return Err(DBError("ERR HEADER requires key and value".to_string())); | ||||
|                                             } | ||||
|                                             headers.insert(cmd[i + 1].clone(), cmd[i + 2].clone()); | ||||
|                                             i += 3; | ||||
|                                         } | ||||
|                                         "TIMEOUTMS" => { | ||||
|                                             if i + 1 >= cmd.len() { | ||||
|                                                 return Err(DBError("ERR TIMEOUTMS requires a value".to_string())); | ||||
|                                             } | ||||
|                                             let t: u64 = cmd[i + 1].parse().map_err(|_| DBError("ERR TIMEOUTMS must be an integer".to_string()))?; | ||||
|                                             timeout_ms = Some(t); | ||||
|                                             i += 2; | ||||
|                                         } | ||||
|                                         _ => { | ||||
|                                             i += 1; | ||||
|                                         } | ||||
|                                     } | ||||
|                                 } | ||||
|  | ||||
|                                 let provider_str = provider.ok_or_else(|| DBError("ERR missing PROVIDER".to_string()))?; | ||||
|                                 let provider_enum = match provider_str.to_lowercase().as_str() { | ||||
|                                     "openai" => crate::embedding::EmbeddingProvider::openai, | ||||
|                                     "test" => crate::embedding::EmbeddingProvider::test, | ||||
|                                     "image_test" | "imagetest" | "image-test" => crate::embedding::EmbeddingProvider::image_test, | ||||
|                                     other => return Err(DBError(format!("ERR unsupported provider '{}'", other))), | ||||
|                                 }; | ||||
|                                 let model = model.ok_or_else(|| DBError("ERR missing MODEL".to_string()))?; | ||||
|                                 let dim = dim.ok_or_else(|| DBError("ERR missing DIM".to_string()))?; | ||||
|  | ||||
|                                 let config = EmbeddingConfig { | ||||
|                                     provider: provider_enum, | ||||
|                                     model, | ||||
|                                     dim, | ||||
|                                     endpoint, | ||||
|                                     headers, | ||||
|                                     timeout_ms, | ||||
|                                 }; | ||||
|  | ||||
|                                 Cmd::LanceEmbeddingConfigSet { name, config } | ||||
|                             } 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 +1247,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 +1282,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 +1327,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 +1470,300 @@ 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, config } => { | ||||
|                 if !server.has_write_permission() { | ||||
|                     return Ok(Protocol::err("ERR write permission denied")); | ||||
|                 } | ||||
|                 if config.dim == 0 { | ||||
|                     return Ok(Protocol::err("ERR embedding DIM must be > 0")); | ||||
|                 } | ||||
|                 match server.set_dataset_embedding_config(&name, &config) { | ||||
|                     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 { | ||||
|                             crate::embedding::EmbeddingProvider::openai => "openai".to_string(), | ||||
|                             crate::embedding::EmbeddingProvider::test => "test".to_string(), | ||||
|                             crate::embedding::EmbeddingProvider::image_test => "image_test".to_string(), | ||||
|                         })); | ||||
|                         arr.push(Protocol::BulkString("model".to_string())); | ||||
|                         arr.push(Protocol::BulkString(cfg.model.clone())); | ||||
|                         arr.push(Protocol::BulkString("dim".to_string())); | ||||
|                         arr.push(Protocol::BulkString(cfg.dim.to_string())); | ||||
|                         arr.push(Protocol::BulkString("endpoint".to_string())); | ||||
|                         arr.push(Protocol::BulkString(cfg.endpoint.clone().unwrap_or_default())); | ||||
|                         arr.push(Protocol::BulkString("timeout_ms".to_string())); | ||||
|                         arr.push(Protocol::BulkString(cfg.timeout_ms.map(|v| v.to_string()).unwrap_or_default())); | ||||
|                         arr.push(Protocol::BulkString("headers".to_string())); | ||||
|                         arr.push(Protocol::BulkString(serde_json::to_string(&cfg.headers).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 +1863,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() { | ||||
| @@ -1459,9 +2208,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 +2218,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 | ||||
|   | ||||
							
								
								
									
										353
									
								
								src/embedding.rs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										353
									
								
								src/embedding.rs
									
									
									
									
									
										Normal file
									
								
							| @@ -0,0 +1,353 @@ | ||||
| // Embedding abstraction with a single external provider (OpenAI-compatible) and local test providers. | ||||
|  | ||||
| use std::collections::HashMap; | ||||
| use std::sync::Arc; | ||||
|  | ||||
| use serde::{Deserialize, Serialize}; | ||||
|  | ||||
| use crate::error::DBError; | ||||
|  | ||||
| // Networking for OpenAI-compatible endpoints | ||||
| use std::time::Duration; | ||||
| use ureq::{Agent, AgentBuilder}; | ||||
| use serde_json::json; | ||||
|  | ||||
| /// Provider identifiers (minimal set). | ||||
| #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] | ||||
| #[serde(rename_all = "snake_case")] | ||||
| pub enum EmbeddingProvider { | ||||
|     /// External HTTP provider compatible with OpenAI's embeddings API. | ||||
|     openai, | ||||
|     /// Deterministic, local-only embedder for CI and offline development (text). | ||||
|     test, | ||||
|     /// Deterministic, local-only embedder for CI and offline development (image). | ||||
|     image_test, | ||||
| } | ||||
|  | ||||
| /// Serializable embedding configuration. | ||||
| /// - provider: "openai" | "test" | "image_test" | ||||
| /// - model: provider/model id (e.g., "text-embedding-3-small"), may be ignored by local gateways | ||||
| /// - dim: required output dimension (used to create Lance datasets and validate outputs) | ||||
| /// - endpoint: optional HTTP endpoint (defaults to OpenAI API when provider == openai) | ||||
| /// - headers: optional HTTP headers (e.g., Authorization). If empty and OPENAI_API_KEY is present, Authorization will be inferred. | ||||
| /// - timeout_ms: optional HTTP timeout in milliseconds (for both read and write) | ||||
| #[derive(Debug, Clone, Serialize, Deserialize)] | ||||
| pub struct EmbeddingConfig { | ||||
|     pub provider: EmbeddingProvider, | ||||
|     pub model: String, | ||||
|     pub dim: usize, | ||||
|     #[serde(default)] | ||||
|     pub endpoint: Option<String>, | ||||
|     #[serde(default)] | ||||
|     pub headers: HashMap<String, String>, | ||||
|     #[serde(default)] | ||||
|     pub timeout_ms: Option<u64>, | ||||
| } | ||||
|  | ||||
| /// 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() | ||||
|     } | ||||
| } | ||||
|  | ||||
| /// 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() | ||||
|     } | ||||
| } | ||||
|  | ||||
| //// ----------------------------- 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:{}", 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: deterministic test embedder ----------------------------- | ||||
|  | ||||
| /// 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!("image_test:{}", 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-compatible HTTP embedder ----------------------------- | ||||
|  | ||||
| struct OpenAIEmbedder { | ||||
|     model: String, | ||||
|     dim: usize, | ||||
|     agent: Agent, | ||||
|     endpoint: String, | ||||
|     headers: Vec<(String, String)>, | ||||
| } | ||||
|  | ||||
| impl OpenAIEmbedder { | ||||
|     fn new_from_config(cfg: &EmbeddingConfig) -> Result<Self, DBError> { | ||||
|         // Resolve endpoint | ||||
|         let endpoint = cfg.endpoint.clone().unwrap_or_else(|| { | ||||
|             "https://api.openai.com/v1/embeddings".to_string() | ||||
|         }); | ||||
|  | ||||
|         // Determine expected dimension (required by config) | ||||
|         let dim = cfg.dim; | ||||
|  | ||||
|         // Build an HTTP agent with timeouts (blocking; no tokio runtime involved) | ||||
|         let to_ms = cfg.timeout_ms.unwrap_or(30_000); | ||||
|         let agent = AgentBuilder::new() | ||||
|             .timeout_read(Duration::from_millis(to_ms)) | ||||
|             .timeout_write(Duration::from_millis(to_ms)) | ||||
|             .build(); | ||||
|  | ||||
|         // Headers: start from cfg.headers, and add Authorization from env if absent and available | ||||
|         let mut headers: Vec<(String, String)> = | ||||
|             cfg.headers.iter().map(|(k, v)| (k.clone(), v.clone())).collect(); | ||||
|  | ||||
|         if !headers.iter().any(|(k, _)| k.eq_ignore_ascii_case("content-type")) { | ||||
|             headers.push(("Content-Type".to_string(), "application/json".to_string())); | ||||
|         } | ||||
|  | ||||
|         if !headers.iter().any(|(k, _)| k.eq_ignore_ascii_case("authorization")) { | ||||
|             if let Ok(key) = std::env::var("OPENAI_API_KEY") { | ||||
|                 headers.push(("Authorization".to_string(), format!("Bearer {}", key))); | ||||
|             } | ||||
|         } | ||||
|  | ||||
|         Ok(Self { | ||||
|             model: cfg.model.clone(), | ||||
|             dim, | ||||
|             agent, | ||||
|             endpoint, | ||||
|             headers, | ||||
|         }) | ||||
|     } | ||||
|  | ||||
|     fn request_many(&self, inputs: &[String]) -> Result<Vec<Vec<f32>>, DBError> { | ||||
|         // Compose request body (OpenAI-compatible) | ||||
|         let mut body = 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' to match output.", | ||||
|                     self.dim, v.len() | ||||
|                 ))); | ||||
|             } | ||||
|             out.push(v); | ||||
|         } | ||||
|         Ok(out) | ||||
|     } | ||||
| } | ||||
|  | ||||
| impl Embedder for OpenAIEmbedder { | ||||
|     fn name(&self) -> String { | ||||
|         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. | ||||
| /// - openai: uses OpenAI-compatible embeddings REST API (endpoint override supported) | ||||
| /// - test: deterministic local text embedder (no network) | ||||
| /// - image_test: not valid for text (use create_image_embedder) | ||||
| pub fn create_embedder(config: &EmbeddingConfig) -> Result<Arc<dyn Embedder>, DBError> { | ||||
|     match &config.provider { | ||||
|         EmbeddingProvider::openai => { | ||||
|             let inner = OpenAIEmbedder::new_from_config(config)?; | ||||
|             Ok(Arc::new(inner)) | ||||
|         } | ||||
|         EmbeddingProvider::test => { | ||||
|             Ok(Arc::new(TestHashEmbedder::new(config.dim, config.model.clone()))) | ||||
|         } | ||||
|         EmbeddingProvider::image_test => { | ||||
|             Err(DBError("Use create_image_embedder() for image providers".into())) | ||||
|         } | ||||
|     } | ||||
| } | ||||
|  | ||||
| /// Create an image embedder instance from a config. | ||||
| /// - image_test: deterministic local image embedder | ||||
| pub fn create_image_embedder(config: &EmbeddingConfig) -> Result<Arc<dyn ImageEmbedder>, DBError> { | ||||
|     match &config.provider { | ||||
|         EmbeddingProvider::image_test => { | ||||
|             Ok(Arc::new(TestImageHashEmbedder::new(config.dim, config.model.clone()))) | ||||
|         } | ||||
|         EmbeddingProvider::test | EmbeddingProvider::openai => { | ||||
|             Err(DBError("Configured text provider; dataset expects image provider (e.g., 'image_test')".into())) | ||||
|         } | ||||
|     } | ||||
| } | ||||
							
								
								
									
										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; | ||||
|   | ||||
| @@ -5,6 +5,7 @@ pub enum BackendType { | ||||
|     Redb, | ||||
|     Sled, | ||||
|     Tantivy, // Full-text search backend (no KV storage) | ||||
|     Lance,   // Vector database backend (no KV storage) | ||||
| } | ||||
|  | ||||
| #[derive(Debug, Clone)] | ||||
|   | ||||
							
								
								
									
										693
									
								
								src/rpc.rs
									
									
									
									
									
								
							
							
						
						
									
										693
									
								
								src/rpc.rs
									
									
									
									
									
								
							| @@ -9,6 +9,8 @@ use sha2::{Digest, Sha256}; | ||||
| use crate::server::Server; | ||||
| use crate::options::DBOption; | ||||
| use crate::admin_meta; | ||||
| use crate::embedding::EmbeddingConfig; | ||||
| use base64::{engine::general_purpose, Engine as _}; | ||||
|  | ||||
| /// Database backend types | ||||
| #[derive(Debug, Clone, Serialize, Deserialize)] | ||||
| @@ -16,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) | ||||
| } | ||||
|  | ||||
| @@ -68,7 +71,7 @@ pub fn hash_key(key: &str) -> String { | ||||
| } | ||||
|  | ||||
| /// RPC trait for HeroDB management | ||||
| #[rpc(server, client, namespace = "herodb")] | ||||
| #[rpc(server, client, namespace = "hero")] | ||||
| pub trait Rpc { | ||||
|     /// Create a new database with specified configuration | ||||
|     #[method(name = "createDatabase")] | ||||
| @@ -161,6 +164,150 @@ 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, | ||||
|         config: EmbeddingConfig, | ||||
|     ) -> 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 | ||||
| @@ -236,7 +383,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 | ||||
| @@ -246,15 +396,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(); | ||||
|         } | ||||
|  | ||||
| @@ -344,6 +494,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 { | ||||
| @@ -395,12 +546,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 | ||||
| @@ -410,13 +565,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(); | ||||
|         } | ||||
|  | ||||
| @@ -676,4 +831,516 @@ 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, | ||||
|         config: EmbeddingConfig, | ||||
|     ) -> 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 provider and dimension (only a minimal set is allowed for now) | ||||
|         match config.provider { | ||||
|             crate::embedding::EmbeddingProvider::openai | ||||
|             | crate::embedding::EmbeddingProvider::test | ||||
|             | crate::embedding::EmbeddingProvider::image_test => {} | ||||
|         } | ||||
|         if config.dim == 0 { | ||||
|             return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "Invalid embedding config: dim must be > 0", None::<()>)); | ||||
|         } | ||||
|  | ||||
|         server.set_dataset_embedding_config(&name, &config) | ||||
|             .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::to_value(&cfg).unwrap_or(serde_json::json!({}))) | ||||
|     } | ||||
|  | ||||
|     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 })) | ||||
|     } | ||||
| } | ||||
							
								
								
									
										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); | ||||
|   | ||||
							
								
								
									
										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(5000)); | ||||
|  | ||||
|     (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)); | ||||
| } | ||||
		Reference in New Issue
	
	Block a user