WIP6: implementing image embedding as first step towards multi-model support

src/rpc.rs

@@ -10,6 +10,7 @@ use crate::server::Server;
 use crate::options::DBOption;
 use crate::admin_meta;
 use crate::embedding::{EmbeddingConfig, EmbeddingProvider};
+use base64::{engine::general_purpose, Engine as _};
 
 /// Database backend types
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -282,6 +283,33 @@ pub trait Rpc {
         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
@@ -1131,4 +1159,204 @@ impl RpcServer for RpcServerImpl {
 
         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 }))
+    }
 }