use serde::{Deserialize, Serialize};
pub type VectorId = String;
pub type NamespaceId = String;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Vector {
pub id: VectorId,
pub values: Vec<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub ttl_seconds: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub expires_at: Option<u64>,
}
impl Vector {
pub fn is_expired(&self) -> bool {
if let Some(expires_at) = self.expires_at {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
now >= expires_at
} else {
false
}
}
#[inline]
pub fn is_expired_at(&self, now_secs: u64) -> bool {
self.expires_at.is_some_and(|exp| now_secs >= exp)
}
pub fn apply_ttl(&mut self) {
if let Some(ttl) = self.ttl_seconds {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
self.expires_at = Some(now + ttl);
}
}
pub fn remaining_ttl(&self) -> Option<u64> {
self.expires_at.and_then(|expires_at| {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
if now < expires_at {
Some(expires_at - now)
} else {
None
}
})
}
}
#[derive(Debug, Deserialize)]
pub struct UpsertRequest {
pub vectors: Vec<Vector>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct UpsertResponse {
pub upserted_count: usize,
}
#[derive(Debug, Deserialize)]
pub struct ColumnUpsertRequest {
pub ids: Vec<VectorId>,
pub vectors: Vec<Vec<f32>>,
#[serde(default)]
pub attributes: std::collections::HashMap<String, Vec<serde_json::Value>>,
#[serde(default)]
pub ttl_seconds: Option<u64>,
#[serde(default)]
pub dimension: Option<usize>,
}
impl ColumnUpsertRequest {
pub fn to_vectors(&self) -> Result<Vec<Vector>, String> {
let count = self.ids.len();
if self.vectors.len() != count {
return Err(format!(
"vectors array length ({}) doesn't match ids array length ({})",
self.vectors.len(),
count
));
}
for (attr_name, attr_values) in &self.attributes {
if attr_values.len() != count {
return Err(format!(
"attribute '{}' array length ({}) doesn't match ids array length ({})",
attr_name,
attr_values.len(),
count
));
}
}
let expected_dim = if let Some(dim) = self.dimension {
Some(dim)
} else {
self.vectors.first().map(|v| v.len())
};
if let Some(expected) = expected_dim {
for (i, vec) in self.vectors.iter().enumerate() {
if vec.len() != expected {
return Err(format!(
"vectors[{}] has dimension {} but expected {}",
i,
vec.len(),
expected
));
}
}
}
let mut vectors = Vec::with_capacity(count);
for i in 0..count {
let metadata = if self.attributes.is_empty() {
None
} else {
let mut meta = serde_json::Map::new();
for (attr_name, attr_values) in &self.attributes {
let value = &attr_values[i];
if !value.is_null() {
meta.insert(attr_name.clone(), value.clone());
}
}
if meta.is_empty() {
None
} else {
Some(serde_json::Value::Object(meta))
}
};
let mut vector = Vector {
id: self.ids[i].clone(),
values: self.vectors[i].clone(),
metadata,
ttl_seconds: self.ttl_seconds,
expires_at: None,
};
vector.apply_ttl();
vectors.push(vector);
}
Ok(vectors)
}
}
#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum DistanceMetric {
#[default]
Cosine,
Euclidean,
DotProduct,
}
#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ReadConsistency {
Strong,
#[default]
Eventual,
#[serde(rename = "bounded_staleness")]
BoundedStaleness,
}
#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, PartialEq, Eq)]
pub struct StalenessConfig {
#[serde(default = "default_max_staleness_ms")]
pub max_staleness_ms: u64,
}
fn default_max_staleness_ms() -> u64 {
5000 }
#[derive(Debug, Deserialize)]
pub struct QueryRequest {
pub vector: Vec<f32>,
#[serde(default = "default_top_k")]
pub top_k: usize,
#[serde(default)]
pub distance_metric: DistanceMetric,
#[serde(default = "default_true")]
pub include_metadata: bool,
#[serde(default)]
pub include_vectors: bool,
#[serde(default)]
pub filter: Option<FilterExpression>,
#[serde(default)]
pub cursor: Option<String>,
#[serde(default)]
pub consistency: ReadConsistency,
#[serde(default)]
pub staleness_config: Option<StalenessConfig>,
}
fn default_top_k() -> usize {
10
}
fn default_true() -> bool {
true
}
#[derive(Debug, Serialize, Deserialize)]
pub struct SearchResult {
pub id: VectorId,
pub score: f32,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub vector: Option<Vec<f32>>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct QueryResponse {
pub results: Vec<SearchResult>,
#[serde(skip_serializing_if = "Option::is_none")]
pub next_cursor: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub has_more: Option<bool>,
#[serde(default)]
pub search_time_ms: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PaginationCursor {
pub last_score: f32,
pub last_id: String,
}
impl PaginationCursor {
pub fn new(last_score: f32, last_id: String) -> Self {
Self {
last_score,
last_id,
}
}
pub fn encode(&self) -> String {
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
let json = serde_json::to_string(self).unwrap_or_default();
URL_SAFE_NO_PAD.encode(json.as_bytes())
}
pub fn decode(cursor: &str) -> Option<Self> {
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
let bytes = URL_SAFE_NO_PAD.decode(cursor).ok()?;
let json = String::from_utf8(bytes).ok()?;
serde_json::from_str(&json).ok()
}
}
#[derive(Debug, Deserialize)]
pub struct DeleteRequest {
pub ids: Vec<VectorId>,
}
#[derive(Debug, Serialize)]
pub struct DeleteResponse {
pub deleted_count: usize,
}
#[derive(Debug, Clone, Deserialize)]
pub struct BatchQueryItem {
#[serde(default)]
pub id: Option<String>,
pub vector: Vec<f32>,
#[serde(default = "default_batch_top_k")]
pub top_k: u32,
#[serde(default)]
pub filter: Option<FilterExpression>,
#[serde(default)]
pub include_metadata: bool,
#[serde(default)]
pub consistency: ReadConsistency,
#[serde(default)]
pub staleness_config: Option<StalenessConfig>,
}
fn default_batch_top_k() -> u32 {
10
}
#[derive(Debug, Deserialize)]
pub struct BatchQueryRequest {
pub queries: Vec<BatchQueryItem>,
}
#[derive(Debug, Serialize)]
pub struct BatchQueryResult {
#[serde(skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
pub results: Vec<SearchResult>,
pub latency_ms: f64,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
#[derive(Debug, Serialize)]
pub struct BatchQueryResponse {
pub results: Vec<BatchQueryResult>,
pub total_latency_ms: f64,
pub query_count: usize,
}
#[derive(Debug, Deserialize)]
pub struct MultiVectorSearchRequest {
pub positive_vectors: Vec<Vec<f32>>,
#[serde(default)]
pub positive_weights: Option<Vec<f32>>,
#[serde(default)]
pub negative_vectors: Option<Vec<Vec<f32>>>,
#[serde(default)]
pub negative_weights: Option<Vec<f32>>,
#[serde(default = "default_top_k")]
pub top_k: usize,
#[serde(default)]
pub distance_metric: DistanceMetric,
#[serde(default)]
pub score_threshold: Option<f32>,
#[serde(default)]
pub enable_mmr: bool,
#[serde(default = "default_mmr_lambda")]
pub mmr_lambda: f32,
#[serde(default = "default_true")]
pub include_metadata: bool,
#[serde(default)]
pub include_vectors: bool,
#[serde(default)]
pub filter: Option<FilterExpression>,
#[serde(default)]
pub consistency: ReadConsistency,
#[serde(default)]
pub staleness_config: Option<StalenessConfig>,
}
fn default_mmr_lambda() -> f32 {
0.5
}
#[derive(Debug, Serialize, Deserialize)]
pub struct MultiVectorSearchResult {
pub id: VectorId,
pub score: f32,
#[serde(skip_serializing_if = "Option::is_none")]
pub mmr_score: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub original_rank: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub vector: Option<Vec<f32>>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct MultiVectorSearchResponse {
pub results: Vec<MultiVectorSearchResult>,
#[serde(skip_serializing_if = "Option::is_none")]
pub computed_query_vector: Option<Vec<f32>>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct IndexDocumentRequest {
pub id: String,
pub text: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
}
#[derive(Debug, Deserialize)]
pub struct IndexDocumentsRequest {
pub documents: Vec<IndexDocumentRequest>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct IndexDocumentsResponse {
pub indexed_count: usize,
}
#[derive(Debug, Deserialize)]
pub struct FullTextSearchRequest {
pub query: String,
#[serde(default = "default_top_k")]
pub top_k: usize,
#[serde(default)]
pub filter: Option<FilterExpression>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct FullTextSearchResult {
pub id: String,
pub score: f32,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct FullTextSearchResponse {
pub results: Vec<FullTextSearchResult>,
#[serde(default)]
pub search_time_ms: u64,
}
#[derive(Debug, Deserialize)]
pub struct DeleteDocumentsRequest {
pub ids: Vec<String>,
}
#[derive(Debug, Serialize)]
pub struct DeleteDocumentsResponse {
pub deleted_count: usize,
}
#[derive(Debug, Serialize)]
pub struct FullTextIndexStats {
pub document_count: u32,
pub unique_terms: usize,
pub avg_doc_length: f32,
}
#[derive(Debug, Deserialize)]
pub struct HybridSearchRequest {
#[serde(default)]
pub vector: Option<Vec<f32>>,
pub text: String,
#[serde(default = "default_top_k")]
pub top_k: usize,
#[serde(default = "default_vector_weight")]
pub vector_weight: f32,
#[serde(default)]
pub distance_metric: DistanceMetric,
#[serde(default = "default_true")]
pub include_metadata: bool,
#[serde(default)]
pub include_vectors: bool,
#[serde(default)]
pub filter: Option<FilterExpression>,
}
fn default_vector_weight() -> f32 {
0.5 }
#[derive(Debug, Serialize, Deserialize)]
pub struct HybridSearchResult {
pub id: String,
pub score: f32,
pub vector_score: f32,
pub text_score: f32,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub vector: Option<Vec<f32>>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct HybridSearchResponse {
pub results: Vec<HybridSearchResult>,
#[serde(default)]
pub search_time_ms: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum FilterValue {
String(String),
Number(f64),
Integer(i64),
Boolean(bool),
StringArray(Vec<String>),
NumberArray(Vec<f64>),
}
impl FilterValue {
pub fn as_f64(&self) -> Option<f64> {
match self {
FilterValue::Number(n) => Some(*n),
FilterValue::Integer(i) => Some(*i as f64),
_ => None,
}
}
pub fn as_str(&self) -> Option<&str> {
match self {
FilterValue::String(s) => Some(s.as_str()),
_ => None,
}
}
pub fn as_bool(&self) -> Option<bool> {
match self {
FilterValue::Boolean(b) => Some(*b),
_ => None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum FilterCondition {
#[serde(rename = "$eq")]
Eq(FilterValue),
#[serde(rename = "$ne")]
Ne(FilterValue),
#[serde(rename = "$gt")]
Gt(FilterValue),
#[serde(rename = "$gte")]
Gte(FilterValue),
#[serde(rename = "$lt")]
Lt(FilterValue),
#[serde(rename = "$lte")]
Lte(FilterValue),
#[serde(rename = "$in")]
In(Vec<FilterValue>),
#[serde(rename = "$nin")]
NotIn(Vec<FilterValue>),
#[serde(rename = "$exists")]
Exists(bool),
#[serde(rename = "$contains")]
Contains(String),
#[serde(rename = "$icontains")]
IContains(String),
#[serde(rename = "$startsWith")]
StartsWith(String),
#[serde(rename = "$endsWith")]
EndsWith(String),
#[serde(rename = "$glob")]
Glob(String),
#[serde(rename = "$regex")]
Regex(String),
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum FilterExpression {
And {
#[serde(rename = "$and")]
conditions: Vec<FilterExpression>,
},
Or {
#[serde(rename = "$or")]
conditions: Vec<FilterExpression>,
},
Field {
#[serde(flatten)]
field: std::collections::HashMap<String, FilterCondition>,
},
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct QuotaConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub max_vectors: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_storage_bytes: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_dimensions: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_metadata_bytes: Option<usize>,
#[serde(default)]
pub enforcement: QuotaEnforcement,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum QuotaEnforcement {
None,
Soft,
#[default]
Hard,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct QuotaUsage {
pub vector_count: u64,
pub storage_bytes: u64,
pub avg_dimensions: Option<usize>,
pub avg_metadata_bytes: Option<usize>,
pub last_updated: u64,
}
impl QuotaUsage {
pub fn new(vector_count: u64, storage_bytes: u64) -> Self {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
Self {
vector_count,
storage_bytes,
avg_dimensions: None,
avg_metadata_bytes: None,
last_updated: now,
}
}
pub fn touch(&mut self) {
self.last_updated = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QuotaStatus {
pub namespace: String,
pub config: QuotaConfig,
pub usage: QuotaUsage,
#[serde(skip_serializing_if = "Option::is_none")]
pub vector_usage_percent: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub storage_usage_percent: Option<f32>,
pub is_exceeded: bool,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub exceeded_quotas: Vec<String>,
}
impl QuotaStatus {
pub fn new(namespace: String, config: QuotaConfig, usage: QuotaUsage) -> Self {
let vector_usage_percent = config
.max_vectors
.map(|max| (usage.vector_count as f32 / max as f32) * 100.0);
let storage_usage_percent = config
.max_storage_bytes
.map(|max| (usage.storage_bytes as f32 / max as f32) * 100.0);
let mut exceeded_quotas = Vec::new();
if let Some(max) = config.max_vectors {
if usage.vector_count > max {
exceeded_quotas.push("max_vectors".to_string());
}
}
if let Some(max) = config.max_storage_bytes {
if usage.storage_bytes > max {
exceeded_quotas.push("max_storage_bytes".to_string());
}
}
let is_exceeded = !exceeded_quotas.is_empty();
Self {
namespace,
config,
usage,
vector_usage_percent,
storage_usage_percent,
is_exceeded,
exceeded_quotas,
}
}
}
#[derive(Debug, Deserialize)]
pub struct SetQuotaRequest {
pub config: QuotaConfig,
}
#[derive(Debug, Serialize)]
pub struct SetQuotaResponse {
pub success: bool,
pub namespace: String,
pub config: QuotaConfig,
pub message: String,
}
#[derive(Debug, Clone, Serialize)]
pub struct QuotaCheckResult {
pub allowed: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub reason: Option<String>,
pub usage: QuotaUsage,
#[serde(skip_serializing_if = "Option::is_none")]
pub exceeded_quota: Option<String>,
}
#[derive(Debug, Serialize)]
pub struct QuotaListResponse {
pub quotas: Vec<QuotaStatus>,
pub total: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub default_config: Option<QuotaConfig>,
}
#[derive(Debug, Serialize)]
pub struct DefaultQuotaResponse {
pub config: Option<QuotaConfig>,
}
#[derive(Debug, Deserialize)]
pub struct SetDefaultQuotaRequest {
pub config: Option<QuotaConfig>,
}
#[derive(Debug, Deserialize)]
pub struct QuotaCheckRequest {
pub vector_ids: Vec<String>,
#[serde(default)]
pub dimensions: Option<usize>,
#[serde(default)]
pub metadata_bytes: Option<usize>,
}
#[derive(Debug, Deserialize)]
pub struct ExportRequest {
#[serde(default = "default_export_top_k")]
pub top_k: usize,
#[serde(skip_serializing_if = "Option::is_none")]
pub cursor: Option<String>,
#[serde(default = "default_true")]
pub include_vectors: bool,
#[serde(default = "default_true")]
pub include_metadata: bool,
}
fn default_export_top_k() -> usize {
1000
}
impl Default for ExportRequest {
fn default() -> Self {
Self {
top_k: 1000,
cursor: None,
include_vectors: true,
include_metadata: true,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExportedVector {
pub id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub values: Option<Vec<f32>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub ttl_seconds: Option<u64>,
}
impl From<&Vector> for ExportedVector {
fn from(v: &Vector) -> Self {
Self {
id: v.id.clone(),
values: Some(v.values.clone()),
metadata: v.metadata.clone(),
ttl_seconds: v.ttl_seconds,
}
}
}
#[derive(Debug, Serialize)]
pub struct ExportResponse {
pub vectors: Vec<ExportedVector>,
#[serde(skip_serializing_if = "Option::is_none")]
pub next_cursor: Option<String>,
pub total_count: usize,
pub returned_count: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum RankBy {
VectorSearch {
field: String,
method: VectorSearchMethod,
query_vector: Vec<f32>,
},
FullTextSearch {
field: String,
method: String, query: String,
},
AttributeOrder {
field: String,
direction: SortDirection,
},
Sum(Vec<RankBy>),
Max(Vec<RankBy>),
Product { weight: f32, ranking: Box<RankBy> },
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
pub enum VectorSearchMethod {
#[default]
ANN,
#[serde(rename = "kNN")]
KNN,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
#[derive(Default)]
pub enum SortDirection {
Asc,
#[default]
Desc,
}
#[derive(Debug, Deserialize)]
pub struct UnifiedQueryRequest {
pub rank_by: RankByInput,
#[serde(default = "default_top_k")]
pub top_k: usize,
#[serde(default)]
pub filter: Option<FilterExpression>,
#[serde(default = "default_true")]
pub include_metadata: bool,
#[serde(default)]
pub include_vectors: bool,
#[serde(default)]
pub distance_metric: DistanceMetric,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(from = "serde_json::Value")]
pub struct RankByInput(pub RankBy);
impl From<serde_json::Value> for RankByInput {
fn from(value: serde_json::Value) -> Self {
RankByInput(parse_rank_by(&value).unwrap_or_else(|| {
RankBy::AttributeOrder {
field: "id".to_string(),
direction: SortDirection::Asc,
}
}))
}
}
fn parse_rank_by(value: &serde_json::Value) -> Option<RankBy> {
let arr = value.as_array()?;
if arr.is_empty() {
return None;
}
let first = arr.first()?.as_str()?;
match first {
"Sum" => {
let rankings = arr.get(1)?.as_array()?;
let parsed: Option<Vec<RankBy>> = rankings.iter().map(parse_rank_by).collect();
Some(RankBy::Sum(parsed?))
}
"Max" => {
let rankings = arr.get(1)?.as_array()?;
let parsed: Option<Vec<RankBy>> = rankings.iter().map(parse_rank_by).collect();
Some(RankBy::Max(parsed?))
}
"Product" => {
let weight = arr.get(1)?.as_f64()? as f32;
let ranking = parse_rank_by(arr.get(2)?)?;
Some(RankBy::Product {
weight,
ranking: Box::new(ranking),
})
}
"ANN" => {
let query_vector = parse_vector_array(arr.get(1)?)?;
Some(RankBy::VectorSearch {
field: "vector".to_string(),
method: VectorSearchMethod::ANN,
query_vector,
})
}
"kNN" => {
let query_vector = parse_vector_array(arr.get(1)?)?;
Some(RankBy::VectorSearch {
field: "vector".to_string(),
method: VectorSearchMethod::KNN,
query_vector,
})
}
field => {
let second = arr.get(1)?;
if let Some(method_str) = second.as_str() {
match method_str {
"ANN" => {
let query_vector = parse_vector_array(arr.get(2)?)?;
Some(RankBy::VectorSearch {
field: field.to_string(),
method: VectorSearchMethod::ANN,
query_vector,
})
}
"kNN" => {
let query_vector = parse_vector_array(arr.get(2)?)?;
Some(RankBy::VectorSearch {
field: field.to_string(),
method: VectorSearchMethod::KNN,
query_vector,
})
}
"BM25" => {
let query = arr.get(2)?.as_str()?;
Some(RankBy::FullTextSearch {
field: field.to_string(),
method: "BM25".to_string(),
query: query.to_string(),
})
}
"asc" => Some(RankBy::AttributeOrder {
field: field.to_string(),
direction: SortDirection::Asc,
}),
"desc" => Some(RankBy::AttributeOrder {
field: field.to_string(),
direction: SortDirection::Desc,
}),
_ => None,
}
} else {
None
}
}
}
}
fn parse_vector_array(value: &serde_json::Value) -> Option<Vec<f32>> {
let arr = value.as_array()?;
arr.iter().map(|v| v.as_f64().map(|n| n as f32)).collect()
}
#[derive(Debug, Serialize, Deserialize)]
pub struct UnifiedQueryResponse {
pub results: Vec<UnifiedSearchResult>,
#[serde(skip_serializing_if = "Option::is_none")]
pub next_cursor: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct UnifiedSearchResult {
pub id: String,
#[serde(rename = "$dist", skip_serializing_if = "Option::is_none")]
pub dist: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub vector: Option<Vec<f32>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AggregateFunction {
Count,
Sum { field: String },
Avg { field: String },
Min { field: String },
Max { field: String },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(from = "serde_json::Value")]
pub struct AggregateFunctionInput(pub AggregateFunction);
impl From<serde_json::Value> for AggregateFunctionInput {
fn from(value: serde_json::Value) -> Self {
parse_aggregate_function(&value)
.map(AggregateFunctionInput)
.unwrap_or_else(|| {
AggregateFunctionInput(AggregateFunction::Count)
})
}
}
fn parse_aggregate_function(value: &serde_json::Value) -> Option<AggregateFunction> {
let arr = value.as_array()?;
if arr.is_empty() {
return None;
}
let func_name = arr.first()?.as_str()?;
match func_name {
"Count" => Some(AggregateFunction::Count),
"Sum" => {
let field = arr.get(1)?.as_str()?;
Some(AggregateFunction::Sum {
field: field.to_string(),
})
}
"Avg" => {
let field = arr.get(1)?.as_str()?;
Some(AggregateFunction::Avg {
field: field.to_string(),
})
}
"Min" => {
let field = arr.get(1)?.as_str()?;
Some(AggregateFunction::Min {
field: field.to_string(),
})
}
"Max" => {
let field = arr.get(1)?.as_str()?;
Some(AggregateFunction::Max {
field: field.to_string(),
})
}
_ => None,
}
}
#[derive(Debug, Deserialize)]
pub struct AggregationRequest {
pub aggregate_by: std::collections::HashMap<String, AggregateFunctionInput>,
#[serde(default)]
pub group_by: Vec<String>,
#[serde(default)]
pub filter: Option<FilterExpression>,
#[serde(default = "default_agg_limit")]
pub limit: usize,
}
fn default_agg_limit() -> usize {
100
}
#[derive(Debug, Serialize, Deserialize)]
pub struct AggregationResponse {
#[serde(skip_serializing_if = "Option::is_none")]
pub aggregations: Option<std::collections::HashMap<String, serde_json::Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub aggregation_groups: Option<Vec<AggregationGroup>>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct AggregationGroup {
#[serde(flatten)]
pub group_key: std::collections::HashMap<String, serde_json::Value>,
#[serde(flatten)]
pub aggregations: std::collections::HashMap<String, serde_json::Value>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TextDocument {
pub id: VectorId,
pub text: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub ttl_seconds: Option<u64>,
}
#[derive(Debug, Deserialize)]
pub struct TextUpsertRequest {
pub documents: Vec<TextDocument>,
#[serde(default)]
pub model: Option<EmbeddingModelType>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct TextUpsertResponse {
pub upserted_count: usize,
pub tokens_processed: usize,
pub model: EmbeddingModelType,
pub embedding_time_ms: u64,
}
#[derive(Debug, Deserialize)]
pub struct TextQueryRequest {
pub text: String,
#[serde(default = "default_top_k")]
pub top_k: usize,
#[serde(default)]
pub filter: Option<FilterExpression>,
#[serde(default)]
pub include_vectors: bool,
#[serde(default = "default_true")]
pub include_text: bool,
#[serde(default)]
pub model: Option<EmbeddingModelType>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct TextQueryResponse {
pub results: Vec<TextSearchResult>,
pub model: EmbeddingModelType,
pub embedding_time_ms: u64,
pub search_time_ms: u64,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct TextSearchResult {
pub id: VectorId,
pub score: f32,
#[serde(skip_serializing_if = "Option::is_none")]
pub text: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub vector: Option<Vec<f32>>,
}
#[derive(Debug, Deserialize)]
pub struct BatchTextQueryRequest {
pub queries: Vec<String>,
#[serde(default = "default_top_k")]
pub top_k: usize,
#[serde(default)]
pub filter: Option<FilterExpression>,
#[serde(default)]
pub include_vectors: bool,
#[serde(default)]
pub model: Option<EmbeddingModelType>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct BatchTextQueryResponse {
pub results: Vec<Vec<TextSearchResult>>,
pub model: EmbeddingModelType,
pub embedding_time_ms: u64,
pub search_time_ms: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum EmbeddingModelType {
#[default]
#[serde(rename = "bge-large")]
BgeLarge,
#[serde(rename = "minilm")]
MiniLM,
#[serde(rename = "bge-small")]
BgeSmall,
#[serde(rename = "e5-small")]
E5Small,
}
impl EmbeddingModelType {
pub fn dimension(&self) -> usize {
match self {
EmbeddingModelType::BgeLarge => 1024,
EmbeddingModelType::MiniLM => 384,
EmbeddingModelType::BgeSmall => 384,
EmbeddingModelType::E5Small => 384,
}
}
}
impl std::fmt::Display for EmbeddingModelType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
EmbeddingModelType::BgeLarge => write!(f, "bge-large"),
EmbeddingModelType::MiniLM => write!(f, "minilm"),
EmbeddingModelType::BgeSmall => write!(f, "bge-small"),
EmbeddingModelType::E5Small => write!(f, "e5-small"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[derive(Default)]
pub enum MemoryType {
#[default]
Episodic,
Semantic,
Procedural,
Working,
}
impl std::fmt::Display for MemoryType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
MemoryType::Episodic => write!(f, "episodic"),
MemoryType::Semantic => write!(f, "semantic"),
MemoryType::Procedural => write!(f, "procedural"),
MemoryType::Working => write!(f, "working"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Memory {
pub id: String,
#[serde(default)]
pub memory_type: MemoryType,
pub content: String,
pub agent_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
#[serde(default = "default_importance")]
pub importance: f32,
#[serde(default)]
pub tags: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
pub created_at: u64,
pub last_accessed_at: u64,
#[serde(default)]
pub access_count: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub ttl_seconds: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub expires_at: Option<u64>,
}
fn default_importance() -> f32 {
0.5
}
impl Memory {
pub fn new(id: String, content: String, agent_id: String, memory_type: MemoryType) -> Self {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
Self {
id,
memory_type,
content,
agent_id,
session_id: None,
importance: 0.5,
tags: Vec::new(),
metadata: None,
created_at: now,
last_accessed_at: now,
access_count: 0,
ttl_seconds: None,
expires_at: None,
}
}
pub fn is_expired(&self) -> bool {
if let Some(expires_at) = self.expires_at {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
now >= expires_at
} else {
false
}
}
pub fn to_vector_metadata(&self) -> serde_json::Value {
let mut meta = serde_json::Map::new();
meta.insert("_dakera_type".to_string(), serde_json::json!("memory"));
meta.insert(
"memory_type".to_string(),
serde_json::json!(self.memory_type),
);
meta.insert("content".to_string(), serde_json::json!(self.content));
meta.insert("agent_id".to_string(), serde_json::json!(self.agent_id));
if let Some(ref sid) = self.session_id {
meta.insert("session_id".to_string(), serde_json::json!(sid));
}
meta.insert("importance".to_string(), serde_json::json!(self.importance));
meta.insert("tags".to_string(), serde_json::json!(self.tags));
meta.insert("created_at".to_string(), serde_json::json!(self.created_at));
meta.insert(
"last_accessed_at".to_string(),
serde_json::json!(self.last_accessed_at),
);
meta.insert(
"access_count".to_string(),
serde_json::json!(self.access_count),
);
if let Some(ref ttl) = self.ttl_seconds {
meta.insert("ttl_seconds".to_string(), serde_json::json!(ttl));
}
if let Some(ref expires) = self.expires_at {
meta.insert("expires_at".to_string(), serde_json::json!(expires));
}
if let Some(ref user_meta) = self.metadata {
meta.insert("user_metadata".to_string(), user_meta.clone());
}
serde_json::Value::Object(meta)
}
pub fn to_vector(&self, embedding: Vec<f32>) -> Vector {
let mut v = Vector {
id: self.id.clone(),
values: embedding,
metadata: Some(self.to_vector_metadata()),
ttl_seconds: self.ttl_seconds,
expires_at: self.expires_at,
};
v.apply_ttl();
v
}
pub fn from_vector(vector: &Vector) -> Option<Self> {
let meta = vector.metadata.as_ref()?.as_object()?;
let entry_type = meta.get("_dakera_type")?.as_str()?;
if entry_type != "memory" {
return None;
}
Some(Memory {
id: vector.id.clone(),
memory_type: serde_json::from_value(meta.get("memory_type")?.clone())
.unwrap_or_default(),
content: meta.get("content")?.as_str()?.to_string(),
agent_id: meta.get("agent_id")?.as_str()?.to_string(),
session_id: meta
.get("session_id")
.and_then(|v| v.as_str())
.map(String::from),
importance: meta
.get("importance")
.and_then(|v| v.as_f64())
.unwrap_or(0.5) as f32,
tags: meta
.get("tags")
.and_then(|v| serde_json::from_value(v.clone()).ok())
.unwrap_or_default(),
metadata: meta.get("user_metadata").cloned(),
created_at: meta.get("created_at").and_then(|v| v.as_u64()).unwrap_or(0),
last_accessed_at: meta
.get("last_accessed_at")
.and_then(|v| v.as_u64())
.unwrap_or(0),
access_count: meta
.get("access_count")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32,
ttl_seconds: vector.ttl_seconds,
expires_at: vector.expires_at,
})
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Session {
pub id: String,
pub agent_id: String,
pub started_at: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub ended_at: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub summary: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(default)]
pub memory_count: usize,
}
impl Session {
pub fn new(id: String, agent_id: String) -> Self {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
Self {
id,
agent_id,
started_at: now,
ended_at: None,
summary: None,
metadata: None,
memory_count: 0,
}
}
pub fn to_vector_metadata(&self) -> serde_json::Value {
let mut meta = serde_json::Map::new();
meta.insert("_dakera_type".to_string(), serde_json::json!("session"));
meta.insert("agent_id".to_string(), serde_json::json!(self.agent_id));
meta.insert("started_at".to_string(), serde_json::json!(self.started_at));
if let Some(ref ended) = self.ended_at {
meta.insert("ended_at".to_string(), serde_json::json!(ended));
}
if let Some(ref summary) = self.summary {
meta.insert("summary".to_string(), serde_json::json!(summary));
}
if let Some(ref user_meta) = self.metadata {
meta.insert("user_metadata".to_string(), user_meta.clone());
}
meta.insert(
"memory_count".to_string(),
serde_json::json!(self.memory_count),
);
serde_json::Value::Object(meta)
}
pub fn to_vector(&self, embedding: Vec<f32>) -> Vector {
Vector {
id: self.id.clone(),
values: embedding,
metadata: Some(self.to_vector_metadata()),
ttl_seconds: None,
expires_at: None,
}
}
pub fn from_vector(vector: &Vector) -> Option<Self> {
let meta = vector.metadata.as_ref()?.as_object()?;
let entry_type = meta.get("_dakera_type")?.as_str()?;
if entry_type != "session" {
return None;
}
Some(Session {
id: vector.id.clone(),
agent_id: meta.get("agent_id")?.as_str()?.to_string(),
started_at: meta.get("started_at").and_then(|v| v.as_u64()).unwrap_or(0),
ended_at: meta.get("ended_at").and_then(|v| v.as_u64()),
summary: meta
.get("summary")
.and_then(|v| v.as_str())
.map(String::from),
metadata: meta.get("user_metadata").cloned(),
memory_count: meta
.get("memory_count")
.and_then(|v| v.as_u64())
.unwrap_or(0) as usize,
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[derive(Default)]
pub enum DecayStrategy {
#[default]
Exponential,
Linear,
StepFunction,
PowerLaw,
Logarithmic,
Flat,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DecayConfig {
#[serde(default)]
pub strategy: DecayStrategy,
#[serde(default = "default_half_life_hours")]
pub half_life_hours: f64,
#[serde(default = "default_min_importance")]
pub min_importance: f32,
}
fn default_half_life_hours() -> f64 {
168.0 }
fn default_min_importance() -> f32 {
0.01
}
impl Default for DecayConfig {
fn default() -> Self {
Self {
strategy: DecayStrategy::default(),
half_life_hours: default_half_life_hours(),
min_importance: default_min_importance(),
}
}
}
#[derive(Debug, Deserialize)]
pub struct StoreMemoryRequest {
pub content: String,
pub agent_id: String,
#[serde(default)]
pub memory_type: MemoryType,
#[serde(skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
#[serde(default = "default_importance")]
pub importance: f32,
#[serde(default)]
pub tags: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub ttl_seconds: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub expires_at: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
}
#[derive(Debug, Serialize)]
pub struct StoreMemoryResponse {
pub memory: Memory,
pub embedding_time_ms: u64,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
#[serde(rename_all = "lowercase")]
pub enum RoutingMode {
#[default]
Auto,
Vector,
Bm25,
Hybrid,
}
#[derive(Debug, Deserialize)]
pub struct RecallRequest {
pub query: String,
pub agent_id: String,
#[serde(default = "default_top_k")]
pub top_k: usize,
#[serde(default)]
pub memory_type: Option<MemoryType>,
#[serde(default)]
pub session_id: Option<String>,
#[serde(default)]
pub tags: Option<Vec<String>>,
#[serde(default)]
pub min_importance: Option<f32>,
#[serde(default = "default_true")]
pub importance_weighted: bool,
#[serde(default)]
pub include_associated: bool,
#[serde(default)]
pub associated_memories_cap: Option<usize>,
#[serde(default)]
pub since: Option<String>,
#[serde(default)]
pub until: Option<String>,
#[serde(default)]
pub associated_memories_depth: Option<u8>,
#[serde(default)]
pub associated_memories_min_weight: Option<f32>,
#[serde(default)]
pub routing: RoutingMode,
#[serde(default = "default_true")]
pub rerank: bool,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct RecallResult {
pub memory: Memory,
pub score: f32,
#[serde(skip_serializing_if = "Option::is_none")]
pub weighted_score: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub smart_score: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub depth: Option<u8>,
}
#[derive(Debug, Serialize)]
pub struct RecallResponse {
pub memories: Vec<RecallResult>,
pub query_embedding_time_ms: u64,
pub search_time_ms: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub associated_memories: Option<Vec<RecallResult>>,
}
#[derive(Debug, Deserialize)]
pub struct ForgetRequest {
pub agent_id: String,
#[serde(default)]
pub memory_ids: Option<Vec<String>>,
#[serde(default)]
pub memory_type: Option<MemoryType>,
#[serde(default)]
pub session_id: Option<String>,
#[serde(default)]
pub tags: Option<Vec<String>>,
#[serde(default)]
pub below_importance: Option<f32>,
}
#[derive(Debug, Serialize)]
pub struct ForgetResponse {
pub deleted_count: usize,
}
#[derive(Debug, Deserialize)]
pub struct UpdateMemoryRequest {
#[serde(default)]
pub content: Option<String>,
#[serde(default)]
pub importance: Option<f32>,
#[serde(default)]
pub tags: Option<Vec<String>>,
#[serde(default)]
pub metadata: Option<serde_json::Value>,
#[serde(default)]
pub memory_type: Option<MemoryType>,
}
#[derive(Debug, Deserialize)]
pub struct UpdateImportanceRequest {
pub memory_id: String,
pub importance: f32,
pub agent_id: String,
}
#[derive(Debug, Deserialize)]
pub struct ConsolidateRequest {
pub agent_id: String,
#[serde(default)]
pub memory_ids: Option<Vec<String>>,
#[serde(default = "default_consolidation_threshold")]
pub threshold: f32,
#[serde(default)]
pub target_type: Option<MemoryType>,
}
fn default_consolidation_threshold() -> f32 {
0.85
}
#[derive(Debug, Serialize)]
pub struct ConsolidateResponse {
pub consolidated_memory: Memory,
pub source_memory_ids: Vec<String>,
pub memories_removed: usize,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum FeedbackSignal {
Upvote,
Downvote,
Flag,
Positive,
Negative,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeedbackHistoryEntry {
pub signal: FeedbackSignal,
pub timestamp: u64,
pub old_importance: f32,
pub new_importance: f32,
}
#[derive(Debug, Deserialize)]
pub struct FeedbackRequest {
pub agent_id: String,
pub memory_id: String,
pub signal: FeedbackSignal,
}
#[derive(Debug, Deserialize)]
pub struct MemoryFeedbackRequest {
pub agent_id: String,
pub signal: FeedbackSignal,
}
#[derive(Debug, Serialize)]
pub struct FeedbackResponse {
pub memory_id: String,
pub new_importance: f32,
pub signal: FeedbackSignal,
}
#[derive(Debug, Serialize)]
pub struct FeedbackHistoryResponse {
pub memory_id: String,
pub entries: Vec<FeedbackHistoryEntry>,
}
#[derive(Debug, Serialize)]
pub struct AgentFeedbackSummary {
pub agent_id: String,
pub upvotes: u64,
pub downvotes: u64,
pub flags: u64,
pub total_feedback: u64,
pub health_score: f32,
}
#[derive(Debug, Deserialize)]
pub struct MemoryImportancePatchRequest {
pub agent_id: String,
pub importance: f32,
}
#[derive(Debug, Deserialize)]
pub struct FeedbackHealthQuery {
pub agent_id: String,
}
#[derive(Debug, Serialize)]
pub struct FeedbackHealthResponse {
pub agent_id: String,
pub health_score: f32,
pub memory_count: usize,
pub avg_importance: f32,
}
#[derive(Debug, Deserialize)]
pub struct SearchMemoriesRequest {
pub agent_id: String,
#[serde(default)]
pub query: Option<String>,
#[serde(default)]
pub memory_type: Option<MemoryType>,
#[serde(default)]
pub session_id: Option<String>,
#[serde(default)]
pub tags: Option<Vec<String>>,
#[serde(default)]
pub min_importance: Option<f32>,
#[serde(default)]
pub max_importance: Option<f32>,
#[serde(default)]
pub created_after: Option<u64>,
#[serde(default)]
pub created_before: Option<u64>,
#[serde(default = "default_top_k")]
pub top_k: usize,
#[serde(default)]
pub sort_by: Option<MemorySortField>,
#[serde(default)]
pub routing: RoutingMode,
#[serde(default)]
pub rerank: bool,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MemorySortField {
CreatedAt,
LastAccessedAt,
Importance,
AccessCount,
}
#[derive(Debug, Serialize)]
pub struct SearchMemoriesResponse {
pub memories: Vec<RecallResult>,
pub total_count: usize,
}
#[derive(Debug, Deserialize)]
pub struct SessionStartRequest {
pub agent_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
}
#[derive(Debug, Serialize)]
pub struct SessionStartResponse {
pub session: Session,
}
#[derive(Debug, Deserialize)]
pub struct SessionEndRequest {
#[serde(default)]
pub summary: Option<String>,
#[serde(default)]
pub auto_summarize: bool,
}
#[derive(Debug, Serialize)]
pub struct SessionEndResponse {
pub session: Session,
pub memory_count: usize,
}
#[derive(Debug, Serialize)]
pub struct ListSessionsResponse {
pub sessions: Vec<Session>,
pub total: usize,
}
#[derive(Debug, Serialize)]
pub struct SessionMemoriesResponse {
pub session: Session,
pub memories: Vec<Memory>,
#[serde(skip_serializing_if = "Option::is_none")]
pub total: Option<usize>,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct AgentSummary {
pub agent_id: String,
pub memory_count: usize,
pub session_count: usize,
pub active_sessions: usize,
}
#[derive(Debug, Serialize)]
pub struct AgentStats {
pub agent_id: String,
pub total_memories: usize,
pub memories_by_type: std::collections::HashMap<String, usize>,
pub total_sessions: usize,
pub active_sessions: usize,
pub avg_importance: f32,
pub oldest_memory_at: Option<u64>,
pub newest_memory_at: Option<u64>,
}
#[derive(Debug, Serialize)]
pub struct WakeUpResponse {
pub agent_id: String,
pub memories: Vec<Memory>,
pub total_available: usize,
}
#[derive(Debug, Deserialize)]
pub struct KnowledgeGraphRequest {
pub agent_id: String,
pub memory_id: String,
#[serde(default = "default_graph_depth")]
pub depth: usize,
#[serde(default = "default_graph_min_similarity")]
pub min_similarity: f32,
}
fn default_graph_depth() -> usize {
2
}
fn default_graph_min_similarity() -> f32 {
0.7
}
#[derive(Debug, Serialize)]
pub struct KnowledgeGraphNode {
pub memory: Memory,
pub similarity: f32,
pub related: Vec<KnowledgeGraphEdge>,
}
#[derive(Debug, Serialize)]
pub struct KnowledgeGraphEdge {
pub memory_id: String,
pub similarity: f32,
pub shared_tags: Vec<String>,
}
#[derive(Debug, Serialize)]
pub struct KnowledgeGraphResponse {
pub root: KnowledgeGraphNode,
pub total_nodes: usize,
}
fn default_full_graph_max_nodes() -> usize {
200
}
fn default_full_graph_min_similarity() -> f32 {
0.50
}
fn default_full_graph_cluster_threshold() -> f32 {
0.60
}
fn default_full_graph_max_edges_per_node() -> usize {
8
}
#[derive(Debug, Deserialize)]
pub struct FullKnowledgeGraphRequest {
pub agent_id: String,
#[serde(default = "default_full_graph_max_nodes")]
pub max_nodes: usize,
#[serde(default = "default_full_graph_min_similarity")]
pub min_similarity: f32,
#[serde(default = "default_full_graph_cluster_threshold")]
pub cluster_threshold: f32,
#[serde(default = "default_full_graph_max_edges_per_node")]
pub max_edges_per_node: usize,
}
#[derive(Debug, Serialize)]
pub struct FullGraphNode {
pub id: String,
pub content: String,
pub memory_type: String,
pub importance: f32,
pub tags: Vec<String>,
pub created_at: Option<String>,
pub cluster_id: usize,
pub centrality: f32,
}
#[derive(Debug, Serialize)]
pub struct FullGraphEdge {
pub source: String,
pub target: String,
pub similarity: f32,
pub shared_tags: Vec<String>,
}
#[derive(Debug, Serialize)]
pub struct GraphCluster {
pub id: usize,
pub node_count: usize,
pub top_tags: Vec<String>,
pub avg_importance: f32,
}
#[derive(Debug, Serialize)]
pub struct GraphStats {
pub total_memories: usize,
pub included_memories: usize,
pub total_edges: usize,
pub cluster_count: usize,
pub density: f32,
pub hub_memory_id: Option<String>,
}
#[derive(Debug, Serialize)]
pub struct FullKnowledgeGraphResponse {
pub nodes: Vec<FullGraphNode>,
pub edges: Vec<FullGraphEdge>,
pub clusters: Vec<GraphCluster>,
pub stats: GraphStats,
}
#[derive(Debug, Deserialize)]
pub struct SummarizeRequest {
pub agent_id: String,
pub memory_ids: Vec<String>,
#[serde(default)]
pub target_type: Option<MemoryType>,
}
#[derive(Debug, Serialize)]
pub struct SummarizeResponse {
pub summary_memory: Memory,
pub source_count: usize,
}
#[derive(Debug, Deserialize)]
pub struct DeduplicateRequest {
pub agent_id: String,
#[serde(default = "default_dedup_threshold")]
pub threshold: f32,
#[serde(default)]
pub memory_type: Option<MemoryType>,
#[serde(default)]
pub dry_run: bool,
}
fn default_dedup_threshold() -> f32 {
0.92
}
#[derive(Debug, Serialize)]
pub struct DuplicateGroup {
pub canonical_id: String,
pub duplicate_ids: Vec<String>,
pub avg_similarity: f32,
}
#[derive(Debug, Serialize)]
pub struct DeduplicateResponse {
pub groups: Vec<DuplicateGroup>,
pub duplicates_found: usize,
pub duplicates_merged: usize,
}
fn default_cross_agent_min_similarity() -> f32 {
0.3
}
fn default_cross_agent_max_nodes_per_agent() -> usize {
50
}
fn default_cross_agent_max_cross_edges() -> usize {
200
}
#[derive(Debug, Deserialize)]
pub struct CrossAgentNetworkRequest {
#[serde(default)]
pub agent_ids: Option<Vec<String>>,
#[serde(default = "default_cross_agent_min_similarity")]
pub min_similarity: f32,
#[serde(default = "default_cross_agent_max_nodes_per_agent")]
pub max_nodes_per_agent: usize,
#[serde(default)]
pub min_importance: f32,
#[serde(default = "default_cross_agent_max_cross_edges")]
pub max_cross_edges: usize,
}
#[derive(Debug, Serialize)]
pub struct AgentNetworkInfo {
pub agent_id: String,
pub memory_count: usize,
pub avg_importance: f32,
}
#[derive(Debug, Serialize)]
pub struct AgentNetworkNode {
pub id: String,
pub agent_id: String,
pub content: String,
pub importance: f32,
pub tags: Vec<String>,
pub memory_type: String,
pub created_at: u64,
}
#[derive(Debug, Serialize)]
pub struct AgentNetworkEdge {
pub source: String,
pub target: String,
pub source_agent: String,
pub target_agent: String,
pub similarity: f32,
}
#[derive(Debug, Serialize)]
pub struct AgentNetworkStats {
pub total_agents: usize,
pub total_nodes: usize,
pub total_cross_edges: usize,
pub density: f32,
}
#[derive(Debug, Serialize)]
pub struct CrossAgentNetworkResponse {
pub node_count: usize,
pub agents: Vec<AgentNetworkInfo>,
pub nodes: Vec<AgentNetworkNode>,
pub edges: Vec<AgentNetworkEdge>,
pub stats: AgentNetworkStats,
}
#[derive(Debug, Deserialize, Default)]
pub struct BatchMemoryFilter {
#[serde(default)]
pub tags: Option<Vec<String>>,
#[serde(default)]
pub min_importance: Option<f32>,
#[serde(default)]
pub max_importance: Option<f32>,
#[serde(default)]
pub created_after: Option<u64>,
#[serde(default)]
pub created_before: Option<u64>,
#[serde(default)]
pub memory_type: Option<MemoryType>,
#[serde(default)]
pub session_id: Option<String>,
}
impl BatchMemoryFilter {
pub fn has_any(&self) -> bool {
self.tags.is_some()
|| self.min_importance.is_some()
|| self.max_importance.is_some()
|| self.created_after.is_some()
|| self.created_before.is_some()
|| self.memory_type.is_some()
|| self.session_id.is_some()
}
pub fn matches(&self, memory: &Memory) -> bool {
if let Some(ref tags) = self.tags {
if !tags.is_empty() && !tags.iter().all(|t| memory.tags.contains(t)) {
return false;
}
}
if let Some(min) = self.min_importance {
if memory.importance < min {
return false;
}
}
if let Some(max) = self.max_importance {
if memory.importance > max {
return false;
}
}
if let Some(after) = self.created_after {
if memory.created_at < after {
return false;
}
}
if let Some(before) = self.created_before {
if memory.created_at > before {
return false;
}
}
if let Some(ref mt) = self.memory_type {
if memory.memory_type != *mt {
return false;
}
}
if let Some(ref sid) = self.session_id {
if memory.session_id.as_ref() != Some(sid) {
return false;
}
}
true
}
}
#[derive(Debug, Deserialize)]
pub struct BatchRecallRequest {
pub agent_id: String,
#[serde(default)]
pub filter: BatchMemoryFilter,
#[serde(default = "default_batch_limit")]
pub limit: usize,
}
fn default_batch_limit() -> usize {
100
}
#[derive(Debug, Serialize)]
pub struct BatchRecallResponse {
pub memories: Vec<Memory>,
pub total: usize,
pub filtered: usize,
}
#[derive(Debug, Deserialize)]
pub struct BatchForgetRequest {
pub agent_id: String,
pub filter: BatchMemoryFilter,
}
#[derive(Debug, Serialize)]
pub struct BatchForgetResponse {
pub deleted_count: usize,
}
#[derive(Debug, Deserialize)]
pub struct NamespaceEntityConfigRequest {
pub extract_entities: bool,
#[serde(default)]
pub entity_types: Vec<String>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct NamespaceEntityConfigResponse {
pub namespace: String,
pub extract_entities: bool,
pub entity_types: Vec<String>,
}
#[derive(Debug, Deserialize)]
pub struct ExtractEntitiesRequest {
pub content: String,
#[serde(default)]
pub entity_types: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EntityResult {
pub entity_type: String,
pub value: String,
pub score: f32,
pub start: usize,
pub end: usize,
pub tag: String,
}
#[derive(Debug, Serialize)]
pub struct ExtractEntitiesResponse {
pub entities: Vec<EntityResult>,
pub count: usize,
}
#[derive(Debug, Deserialize)]
pub struct GraphTraverseQuery {
#[serde(default = "default_ce5_graph_depth")]
pub depth: u32,
}
fn default_ce5_graph_depth() -> u32 {
3
}
#[derive(Debug, Deserialize)]
pub struct GraphPathQuery {
pub to: String,
}
#[derive(Debug, Deserialize)]
pub struct MemoryLinkRequest {
pub target_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub label: Option<String>,
pub agent_id: String,
}
#[derive(Debug, Serialize)]
pub struct GraphTraverseResponse {
pub root_id: String,
pub depth: u32,
pub node_count: usize,
pub nodes: Vec<GraphNodeResponse>,
}
#[derive(Debug, Serialize)]
pub struct GraphNodeResponse {
pub memory_id: String,
pub depth: u32,
pub edges: Vec<GraphEdgeResponse>,
}
#[derive(Debug, Serialize)]
pub struct GraphEdgeResponse {
pub from_id: String,
pub to_id: String,
pub edge_type: String,
pub weight: f32,
pub created_at: u64,
}
#[derive(Debug, Serialize)]
pub struct GraphPathResponse {
pub from_id: String,
pub to_id: String,
pub path: Vec<String>,
pub hop_count: usize,
}
#[derive(Debug, Serialize)]
pub struct MemoryLinkResponse {
pub from_id: String,
pub to_id: String,
pub edge_type: String,
}
#[derive(Debug, Serialize)]
pub struct GraphExportResponse {
pub agent_id: String,
pub namespace: String,
pub node_count: usize,
pub edge_count: usize,
pub edges: Vec<GraphEdgeResponse>,
}
#[derive(Debug, Deserialize)]
pub struct KgQueryParams {
pub agent_id: String,
#[serde(default)]
pub root_id: Option<String>,
#[serde(default)]
pub edge_type: Option<String>,
#[serde(default)]
pub min_weight: Option<f32>,
#[serde(default = "default_kg_depth")]
pub max_depth: u32,
#[serde(default = "default_kg_limit")]
pub limit: usize,
}
fn default_kg_depth() -> u32 {
3
}
fn default_kg_limit() -> usize {
100
}
#[derive(Debug, Serialize)]
pub struct KgQueryResponse {
pub agent_id: String,
pub node_count: usize,
pub edge_count: usize,
pub edges: Vec<GraphEdgeResponse>,
}
#[derive(Debug, Deserialize)]
pub struct KgPathParams {
pub agent_id: String,
pub from: String,
pub to: String,
}
#[derive(Debug, Serialize)]
pub struct KgPathResponse {
pub agent_id: String,
pub from_id: String,
pub to_id: String,
pub hop_count: usize,
pub path: Vec<String>,
}
#[derive(Debug, Deserialize)]
pub struct KgExportParams {
pub agent_id: String,
#[serde(default = "default_kg_format")]
pub format: String,
}
fn default_kg_format() -> String {
"json".to_string()
}
#[derive(Debug, Serialize)]
pub struct KgExportJsonResponse {
pub agent_id: String,
pub format: String,
pub node_count: usize,
pub edge_count: usize,
pub edges: Vec<GraphEdgeResponse>,
}
fn default_working_ttl() -> Option<u64> {
Some(14_400) }
fn default_episodic_ttl() -> Option<u64> {
Some(2_592_000) }
fn default_semantic_ttl() -> Option<u64> {
Some(31_536_000) }
fn default_procedural_ttl() -> Option<u64> {
Some(63_072_000) }
fn default_working_decay() -> DecayStrategy {
DecayStrategy::Exponential
}
fn default_episodic_decay() -> DecayStrategy {
DecayStrategy::PowerLaw
}
fn default_semantic_decay() -> DecayStrategy {
DecayStrategy::Logarithmic
}
fn default_procedural_decay() -> DecayStrategy {
DecayStrategy::Flat
}
fn default_sr_factor() -> f64 {
1.0
}
fn default_sr_base_interval() -> u64 {
86_400 }
fn default_consolidation_enabled() -> bool {
false
}
fn default_policy_consolidation_threshold() -> f32 {
0.92
}
fn default_consolidation_interval_hours() -> u32 {
24
}
fn default_store_dedup_threshold() -> f32 {
0.95
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryPolicy {
#[serde(
default = "default_working_ttl",
skip_serializing_if = "Option::is_none"
)]
pub working_ttl_seconds: Option<u64>,
#[serde(
default = "default_episodic_ttl",
skip_serializing_if = "Option::is_none"
)]
pub episodic_ttl_seconds: Option<u64>,
#[serde(
default = "default_semantic_ttl",
skip_serializing_if = "Option::is_none"
)]
pub semantic_ttl_seconds: Option<u64>,
#[serde(
default = "default_procedural_ttl",
skip_serializing_if = "Option::is_none"
)]
pub procedural_ttl_seconds: Option<u64>,
#[serde(default = "default_working_decay")]
pub working_decay: DecayStrategy,
#[serde(default = "default_episodic_decay")]
pub episodic_decay: DecayStrategy,
#[serde(default = "default_semantic_decay")]
pub semantic_decay: DecayStrategy,
#[serde(default = "default_procedural_decay")]
pub procedural_decay: DecayStrategy,
#[serde(default = "default_sr_factor")]
pub spaced_repetition_factor: f64,
#[serde(default = "default_sr_base_interval")]
pub spaced_repetition_base_interval_seconds: u64,
#[serde(default = "default_consolidation_enabled")]
pub consolidation_enabled: bool,
#[serde(default = "default_policy_consolidation_threshold")]
pub consolidation_threshold: f32,
#[serde(default = "default_consolidation_interval_hours")]
pub consolidation_interval_hours: u32,
#[serde(default)]
pub consolidated_count: u64,
#[serde(default)]
pub rate_limit_enabled: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub rate_limit_stores_per_minute: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub rate_limit_recalls_per_minute: Option<u32>,
#[serde(default)]
pub dedup_on_store: bool,
#[serde(default = "default_store_dedup_threshold")]
pub dedup_threshold: f32,
}
impl Default for MemoryPolicy {
fn default() -> Self {
Self {
working_ttl_seconds: default_working_ttl(),
episodic_ttl_seconds: default_episodic_ttl(),
semantic_ttl_seconds: default_semantic_ttl(),
procedural_ttl_seconds: default_procedural_ttl(),
working_decay: default_working_decay(),
episodic_decay: default_episodic_decay(),
semantic_decay: default_semantic_decay(),
procedural_decay: default_procedural_decay(),
spaced_repetition_factor: default_sr_factor(),
spaced_repetition_base_interval_seconds: default_sr_base_interval(),
consolidation_enabled: default_consolidation_enabled(),
consolidation_threshold: default_policy_consolidation_threshold(),
consolidation_interval_hours: default_consolidation_interval_hours(),
consolidated_count: 0,
rate_limit_enabled: false,
rate_limit_stores_per_minute: None,
rate_limit_recalls_per_minute: None,
dedup_on_store: false,
dedup_threshold: default_store_dedup_threshold(),
}
}
}
impl MemoryPolicy {
pub fn ttl_for_type(&self, memory_type: &MemoryType) -> Option<u64> {
match memory_type {
MemoryType::Working => self.working_ttl_seconds,
MemoryType::Episodic => self.episodic_ttl_seconds,
MemoryType::Semantic => self.semantic_ttl_seconds,
MemoryType::Procedural => self.procedural_ttl_seconds,
}
}
pub fn decay_for_type(&self, memory_type: &MemoryType) -> DecayStrategy {
match memory_type {
MemoryType::Working => self.working_decay,
MemoryType::Episodic => self.episodic_decay,
MemoryType::Semantic => self.semantic_decay,
MemoryType::Procedural => self.procedural_decay,
}
}
pub fn spaced_repetition_extension(&self, access_count: u32) -> u64 {
if self.spaced_repetition_factor <= 0.0 {
return 0;
}
let ext = access_count as f64
* self.spaced_repetition_factor
* self.spaced_repetition_base_interval_seconds as f64;
ext.round() as u64
}
}