mod common;
mod scroll;
mod text;
mod vector;
use std::{fmt, ops::Deref};
use serde::{Deserialize, Serialize};
use crate::DataRecord;
use super::ApplyDefaultLocator;
pub use common::{
Projection, QueryScope, ScopedQuery, TextQueryOptions, VectorQueryOptions, VectorSearchParams,
};
pub use scroll::ScrollRequest;
pub use text::TextSearchRequest;
pub use vector::NearestRequest;
#[derive(Clone, Debug, Default, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
#[serde(transparent)]
pub struct PageToken(String);
impl PageToken {
#[must_use]
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
#[must_use]
pub fn into_string(self) -> String {
self.0
}
pub(crate) fn effective_offset(
page_token: Option<&str>,
offset: Option<usize>,
) -> Result<Option<usize>, crate::StorageError> {
match page_token {
Some(token) => token.parse::<usize>().map(Some).map_err(|error| {
crate::StorageError::InvalidConfig(format!(
"invalid numeric page token `{token}`: {error}"
))
}),
None => Ok(offset),
}
}
#[must_use]
pub(crate) fn next_offset(
limit: Option<usize>,
effective_offset: Option<usize>,
count: usize,
) -> Option<Self> {
let limit = limit?;
if count < limit {
None
} else {
Some((effective_offset.unwrap_or(0) + count).to_string().into())
}
}
}
impl AsRef<str> for PageToken {
fn as_ref(&self) -> &str {
self.as_str()
}
}
impl Deref for PageToken {
type Target = str;
fn deref(&self) -> &Self::Target {
self.as_str()
}
}
impl fmt::Display for PageToken {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}
impl From<String> for PageToken {
fn from(value: String) -> Self {
Self(value)
}
}
impl From<&str> for PageToken {
fn from(value: &str) -> Self {
Self(value.to_owned())
}
}
impl From<&String> for PageToken {
fn from(value: &String) -> Self {
Self(value.clone())
}
}
impl From<PageToken> for String {
fn from(value: PageToken) -> Self {
value.into_string()
}
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub enum ScoreKind {
Similarity,
Distance,
Relevance,
Unknown,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct ScoredRecord {
pub record: DataRecord,
pub raw_score: Option<f64>,
pub score_kind: ScoreKind,
}
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
pub struct QueryResponse {
pub hits: Vec<ScoredRecord>,
pub next_page: Option<PageToken>,
}
impl QueryResponse {
pub(crate) fn from_hits(hits: Vec<ScoredRecord>, next_page: Option<PageToken>) -> Self {
Self { hits, next_page }
}
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub enum QueryRequest {
Nearest(NearestRequest),
Text(TextSearchRequest),
Scroll(ScrollRequest),
}
impl ApplyDefaultLocator for QueryRequest {
fn apply_default_locator(self, default: &crate::StorageLocator) -> Self {
match self {
Self::Nearest(request) => Self::Nearest(request.apply_default_locator(default)),
Self::Text(request) => Self::Text(request.apply_default_locator(default)),
Self::Scroll(request) => Self::Scroll(request.apply_default_locator(default)),
}
}
}