use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[non_exhaustive]
pub enum QuantizationMode {
F32,
Int8,
Dual,
#[default]
Auto,
}
impl QuantizationMode {
#[must_use]
pub fn parse(s: &str) -> Option<Self> {
match s.to_lowercase().as_str() {
"f32" | "full" | "exact" => Some(Self::F32),
"int8" | "sq8" | "quantized" => Some(Self::Int8),
"dual" | "hybrid" => Some(Self::Dual),
"auto" | "default" => Some(Self::Auto),
_ => None,
}
}
#[must_use]
pub const fn as_str(&self) -> &'static str {
match self {
Self::F32 => "f32",
Self::Int8 => "int8",
Self::Dual => "dual",
Self::Auto => "auto",
}
}
}
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
pub struct WithClause {
pub options: Vec<WithOption>,
}
impl WithClause {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_option(mut self, key: impl Into<String>, value: WithValue) -> Self {
self.options.push(WithOption {
key: key.into(),
value,
});
self
}
#[must_use]
pub fn get(&self, key: &str) -> Option<&WithValue> {
self.options
.iter()
.find(|opt| opt.key.eq_ignore_ascii_case(key))
.map(|opt| &opt.value)
}
#[must_use]
pub fn get_mode(&self) -> Option<&str> {
self.get("mode")
.or_else(|| self.get("quality"))
.and_then(|v| v.as_str())
}
#[must_use]
#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
pub fn get_ef_search(&self) -> Option<usize> {
self.get("ef_search")
.and_then(WithValue::as_integer)
.map(|v| v as usize)
}
#[must_use]
#[allow(clippy::cast_sign_loss)]
pub fn get_timeout_ms(&self) -> Option<u64> {
self.get("timeout_ms")
.and_then(WithValue::as_integer)
.map(|v| v as u64)
}
#[must_use]
pub fn get_rerank(&self) -> Option<bool> {
self.get("rerank").and_then(WithValue::as_bool)
}
#[must_use]
pub fn get_quantization(&self) -> Option<QuantizationMode> {
self.get("quantization")
.and_then(WithValue::as_str)
.and_then(QuantizationMode::parse)
}
#[must_use]
#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
pub fn get_oversampling(&self) -> Option<usize> {
self.get("oversampling")
.and_then(WithValue::as_integer)
.map(|v| v.max(1) as usize)
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct WithOption {
pub key: String,
pub value: WithValue,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum WithValue {
String(String),
Integer(i64),
Float(f64),
Boolean(bool),
Identifier(String),
}
impl WithValue {
#[must_use]
pub fn as_str(&self) -> Option<&str> {
match self {
Self::String(s) | Self::Identifier(s) => Some(s),
_ => None,
}
}
#[must_use]
pub fn as_integer(&self) -> Option<i64> {
match self {
Self::Integer(i) => Some(*i),
_ => None,
}
}
#[must_use]
pub fn as_float(&self) -> Option<f64> {
match self {
Self::Float(f) => Some(*f),
#[allow(clippy::cast_precision_loss)]
Self::Integer(i) => Some(*i as f64),
_ => None,
}
}
#[must_use]
pub fn as_bool(&self) -> Option<bool> {
match self {
Self::Boolean(b) => Some(*b),
_ => None,
}
}
}