use bytes::BytesMut;
use qail_core::ast::Qail;
use crate::error::{QdrantError, QdrantResult};
use crate::grpc_transport::GrpcClient;
use crate::point::{Point, ScoredPoint};
use crate::proto_decoder;
use crate::proto_encoder;
pub struct GrpcDriver {
client: GrpcClient,
buffer: BytesMut,
}
impl GrpcDriver {
pub async fn connect(host: &str, port: u16) -> QdrantResult<Self> {
let client = GrpcClient::connect(host, port).await?;
Ok(Self {
client,
buffer: BytesMut::with_capacity(8192),
})
}
pub async fn connect_addr(addr: &str) -> QdrantResult<Self> {
let parts: Vec<&str> = addr.split(':').collect();
if parts.len() != 2 {
return Err(QdrantError::Connection(
"Invalid address format, expected host:port".to_string(),
));
}
let port: u16 = parts[1]
.parse()
.map_err(|_| QdrantError::Connection("Invalid port".to_string()))?;
Self::connect(parts[0], port).await
}
pub async fn search(
&mut self,
collection: &str,
vector: &[f32],
limit: u64,
score_threshold: Option<f32>,
) -> QdrantResult<Vec<ScoredPoint>> {
proto_encoder::encode_search_proto(
&mut self.buffer,
collection,
vector,
limit,
score_threshold,
None,
);
let response = self.client.search(self.buffer.clone().freeze()).await?;
proto_decoder::decode_search_response(&response)
}
pub async fn search_ast(&mut self, cmd: &Qail) -> QdrantResult<Vec<ScoredPoint>> {
let collection = if cmd.table.is_empty() {
return Err(QdrantError::Encode("Collection name required".to_string()));
} else {
&cmd.table
};
let vector = cmd.vector.as_ref().ok_or_else(|| {
QdrantError::Encode("Vector required for search".to_string())
})?;
let mut limit = 10u64;
for cage in &cmd.cages {
if let qail_core::ast::CageKind::Limit(n) = cage.kind {
limit = n as u64;
}
}
let score_threshold = cmd.score_threshold;
self.search(collection, vector, limit, score_threshold).await
}
pub async fn upsert(
&mut self,
collection: &str,
points: &[Point],
wait: bool,
) -> QdrantResult<()> {
proto_encoder::encode_upsert_proto(&mut self.buffer, collection, points, wait);
let _response = self.client.upsert(self.buffer.clone().freeze()).await?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_grpc_driver_struct() {
let buffer = BytesMut::with_capacity(1024);
assert!(buffer.capacity() >= 1024);
}
}