Skip to content

Commit c8a4665

Browse files
committed
feat: add search for the runtime's /v1/search endpoint
Adds Client::search, exposing vector, keyword, and hybrid search from Rust. Previously only spice.js could reach /v1/search; users of every other SDK had to hand-roll the HTTP call. SearchRequest follows the DatasetRefreshRequest builder pattern, and SearchError is a typed enum alongside DatasetError and QueryError. Response types are modelled on the runtime's actual wire shape: matches holds a list per column because one column can contribute several chunks to a match, and data / primary_key / metadata are omitted by the runtime when empty. Errors carry the runtime's plain-text explanation alongside the status code, and requests the runtime would reject with a 400 are validated before the request so the error names the field to fix.
1 parent d1cf846 commit c8a4665

5 files changed

Lines changed: 563 additions & 0 deletions

File tree

README.md

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,41 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
168168
}
169169
```
170170

171+
### Search
172+
173+
`search` finds documents similar to a piece of text using the runtime's `/v1/search` endpoint. It runs against datasets that have an embedding column and a loaded embedding model — see [Search & Retrieval](https://docs.spice.ai/features/search-and-retrieval) for how to configure them. Like dataset refresh, it uses the HTTP API, so `http_url()` must be configured.
174+
175+
```rust,no_run
176+
use spiceai::{ClientBuilder, SearchRequest};
177+
178+
#[tokio::main]
179+
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
180+
let client = ClientBuilder::new()
181+
.http_url("http://localhost:8090")
182+
.build()
183+
.await?;
184+
185+
let response = client
186+
.search(
187+
SearchRequest::new("tokyo plane tickets")
188+
.with_datasets(["app_messages"])
189+
.with_limit(3)
190+
.with_additional_columns(["timestamp"]),
191+
)
192+
.await?;
193+
194+
println!("{} matches in {}ms", response.len(), response.duration_ms);
195+
for m in response {
196+
println!("{} {} {:?}", m.score, m.dataset, m.matches);
197+
}
198+
Ok(())
199+
}
200+
```
201+
202+
Adding `with_keywords([...])` runs a lexical pass alongside the vector pass, which the runtime combines into a single hybrid ranking. `with_where("user_id = 42")` filters candidate rows with a SQL predicate.
203+
204+
Each `SearchMatch` carries `dataset`, `score` (higher is more similar), `matches` (matched values keyed by source column — a list per column, since one column can contribute several chunks to a match), `primary_key`, `data`, and `metadata`.
205+
171206
## Documentation
172207

173208
Check out our [Documentation](https://docs.spice.ai/sdks/rust-sdk) to learn more about how to use the Rust SDK.

src/client.rs

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use crate::{
66
config::{GenericError, SPICE_CLOUD_FLIGHT_ADDR, SPICE_LOCAL_FLIGHT_ADDR},
77
dataset::{DatasetError, DatasetRefreshRequest, DatasetRefreshResponse},
88
flight::{SqlFlightClient, is_connection_reset_generic_error},
9+
search::{SearchError, SearchRequest, SearchResponse},
910
tls::{FlightChannelBuilder, ensure_crypto_provider, new_tls_flight_channel},
1011
};
1112
use arrow::record_batch::RecordBatch;
@@ -575,6 +576,55 @@ impl SpiceClient {
575576

576577
http_client.refresh_dataset(dataset_name, &request).await
577578
}
579+
580+
/// Searches datasets for documents similar to the request's text.
581+
///
582+
/// Runs against datasets that have an embedding column and a loaded
583+
/// embedding model — see [Search & Retrieval](https://docs.spice.ai/features/search-and-retrieval)
584+
/// for how to configure them. Adding keywords to the request turns this
585+
/// into a hybrid search, combining a lexical pass with the vector scores.
586+
///
587+
/// **Note:** Requires [`http_url()`](SpiceClientBuilder::http_url) to be configured.
588+
///
589+
/// # Example
590+
///
591+
/// ```no_run
592+
/// # use spiceai::{ClientBuilder, SearchRequest};
593+
/// # #[tokio::main]
594+
/// # async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
595+
/// let client = ClientBuilder::new()
596+
/// .http_url("http://localhost:8090")
597+
/// .build()
598+
/// .await?;
599+
///
600+
/// let response = client
601+
/// .search(
602+
/// SearchRequest::new("tokyo plane tickets")
603+
/// .with_datasets(["app_messages"])
604+
/// .with_limit(3),
605+
/// )
606+
/// .await?;
607+
///
608+
/// for m in response {
609+
/// println!("{} {} {:?}", m.score, m.dataset, m.matches);
610+
/// }
611+
/// # Ok(())
612+
/// # }
613+
/// ```
614+
///
615+
/// # Errors
616+
///
617+
/// - [`SearchError::InvalidRequest`] if the request has empty text or a zero limit
618+
/// - [`SearchError::HttpError`] if the HTTP endpoint is not configured or unreachable
619+
/// - [`SearchError::SearchFailed`] if the runtime rejects the search
620+
pub async fn search(&self, request: SearchRequest) -> Result<SearchResponse, SearchError> {
621+
let http_client = self.http_client.as_ref().ok_or(SearchError::HttpError {
622+
message: "HTTP endpoint not configured. Use ClientBuilder::http_url() to set it."
623+
.to_string(),
624+
})?;
625+
626+
http_client.search(&request).await
627+
}
578628
}
579629

580630
/// Builder for creating a `SpiceClient`.
@@ -1403,4 +1453,116 @@ mod tests {
14031453
.expect("refresh dataset with explicit overrides");
14041454
assert_eq!(custom_refresh.message, "Refresh scheduled");
14051455
}
1456+
1457+
#[tokio::test]
1458+
async fn test_search_posts_request_and_parses_response() {
1459+
let server = MockServer::start().await;
1460+
1461+
Mock::given(method("POST"))
1462+
.and(path("/v1/search"))
1463+
.and(body_json(json!({
1464+
"text": "tokyo plane tickets",
1465+
"datasets": ["app_messages"],
1466+
"limit": 3,
1467+
"additional_columns": ["timestamp"],
1468+
})))
1469+
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
1470+
"results": [
1471+
{
1472+
"matches": {"message": ["I booked us some tickets", "direct to Narita"]},
1473+
"dataset": "app_messages",
1474+
"primary_key": {"id": "6fd5a215"},
1475+
"data": {"timestamp": 1_724_716_542_i64},
1476+
"_score": 0.914_321
1477+
},
1478+
{
1479+
"matches": {"message": ["we're sitting together"]},
1480+
"dataset": "app_messages",
1481+
"_score": 0.787_654
1482+
}
1483+
],
1484+
"duration_ms": 42
1485+
})))
1486+
.mount(&server)
1487+
.await;
1488+
1489+
let client = test_client(Some(&server.uri()));
1490+
1491+
let response = client
1492+
.search(
1493+
SearchRequest::new("tokyo plane tickets")
1494+
.with_datasets(["app_messages"])
1495+
.with_limit(3)
1496+
.with_additional_columns(["timestamp"]),
1497+
)
1498+
.await
1499+
.expect("search succeeds");
1500+
1501+
assert_eq!(response.duration_ms, 42);
1502+
assert_eq!(response.len(), 2);
1503+
1504+
let first = &response.results[0];
1505+
assert_eq!(first.dataset, "app_messages");
1506+
// One column can contribute several chunks to a single match.
1507+
assert_eq!(first.matches["message"].len(), 2);
1508+
assert_eq!(first.data["timestamp"], 1_724_716_542_i64);
1509+
1510+
// The runtime omits data and primary_key when they are empty.
1511+
assert!(response.results[1].data.is_empty());
1512+
assert!(response.results[1].primary_key.is_empty());
1513+
}
1514+
1515+
#[tokio::test]
1516+
async fn test_search_surfaces_runtime_error_body() {
1517+
let server = MockServer::start().await;
1518+
1519+
Mock::given(method("POST"))
1520+
.and(path("/v1/search"))
1521+
.respond_with(ResponseTemplate::new(400).set_body_string("No data sources provided"))
1522+
.mount(&server)
1523+
.await;
1524+
1525+
let client = test_client(Some(&server.uri()));
1526+
1527+
let err = client
1528+
.search(SearchRequest::new("tokyo"))
1529+
.await
1530+
.expect_err("runtime rejects the search");
1531+
1532+
let message = err.to_string();
1533+
assert!(message.contains("No data sources provided"), "{message}");
1534+
assert!(message.contains("400"), "{message}");
1535+
}
1536+
1537+
#[tokio::test]
1538+
async fn test_search_validates_before_sending() {
1539+
let server = MockServer::start().await;
1540+
// No mock is mounted: a request reaching the server fails the test.
1541+
let client = test_client(Some(&server.uri()));
1542+
1543+
let err = client
1544+
.search(SearchRequest::new(""))
1545+
.await
1546+
.expect_err("empty text is rejected");
1547+
assert!(err.to_string().contains("non-empty"), "{err}");
1548+
1549+
assert!(
1550+
server
1551+
.received_requests()
1552+
.await
1553+
.unwrap_or_default()
1554+
.is_empty()
1555+
);
1556+
}
1557+
1558+
#[tokio::test]
1559+
async fn test_search_requires_http_url() {
1560+
let client = test_client(None);
1561+
1562+
let err = client
1563+
.search(SearchRequest::new("tokyo"))
1564+
.await
1565+
.expect_err("search without an HTTP endpoint fails");
1566+
assert!(err.to_string().contains("http_url"), "{err}");
1567+
}
14061568
}

src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ mod dataset;
66
mod flight;
77
mod params;
88
pub mod query;
9+
pub mod search;
910
pub mod tls;
1011
mod util;
1112

@@ -21,6 +22,7 @@ pub use query::{
2122
QueryError, QueryInfo, QueryJob, QueryListResponse, QueryResult, QueryResultStream,
2223
QueryStatus, QuerySubmitOptions, QuerySummary,
2324
};
25+
pub use search::{SearchError, SearchMatch, SearchRequest, SearchResponse};
2426

2527
// Further public exports and integrations
2628
pub use futures::StreamExt;

src/query.rs

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
3636
use crate::dataset::{DatasetError, DatasetRefreshRequest, DatasetRefreshResponse};
3737
use crate::params::{QueryParameterError, QueryParameters};
38+
use crate::search::{SearchError, SearchRequest, SearchResponse};
3839
use arrow::array::RecordBatch;
3940
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
4041
use futures::Stream;
@@ -758,6 +759,36 @@ impl QueryHttpClient {
758759
}
759760
}
760761

762+
pub async fn search(&self, request: &SearchRequest) -> Result<SearchResponse, SearchError> {
763+
request.validate()?;
764+
765+
let url = format!("{}/v1/search", self.base_url);
766+
767+
let response = self
768+
.add_auth(self.client.post(&url))
769+
.json(request)
770+
.send()
771+
.await
772+
.map_err(|e| SearchError::HttpError {
773+
message: e.to_string(),
774+
})?;
775+
776+
let status_code = response.status().as_u16();
777+
if status_code != 200 {
778+
// The runtime explains search failures in a plain-text body ("No
779+
// data sources provided"). Surface it, not just the status code.
780+
let response_body = response.text().await.unwrap_or_default();
781+
return Err(SearchError::SearchFailed {
782+
status_code,
783+
response_body: response_body.trim().to_string(),
784+
});
785+
}
786+
787+
response.json().await.map_err(|e| SearchError::ParseError {
788+
message: e.to_string(),
789+
})
790+
}
791+
761792
/// List queries with optional status filter and limit.
762793
pub async fn list_queries(
763794
&self,

0 commit comments

Comments
 (0)