use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use crate::chunk::ChunkKind;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct MemoryQuery {
pub text: String,
#[serde(default)]
pub kinds: Option<Vec<ChunkKind>>,
#[serde(default)]
pub since: Option<DateTime<Utc>>,
#[serde(default)]
pub until: Option<DateTime<Utc>>,
#[serde(default)]
pub session_id: Option<String>,
#[serde(default)]
pub caller_scope: CallerScope,
#[serde(default)]
pub project_root_prefix: Option<String>,
#[serde(default = "default_k")]
pub k: usize,
#[serde(default = "default_true")]
pub include_rollups: bool,
#[serde(default)]
pub min_importance: Option<f32>,
#[serde(default)]
pub profile: RetrievalProfile,
pub caller_id: String,
}
fn default_k() -> usize {
8
}
fn default_true() -> bool {
true
}
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq, Hash)]
#[serde(rename_all = "snake_case")]
pub enum CallerScope {
#[default]
Own,
OwnPlusShared,
Global,
}
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq, Hash)]
#[serde(rename_all = "snake_case")]
pub enum RetrievalProfile {
#[default]
AgentChatTurn,
AgentDelegatedJob,
NLCompilerSimilarRules,
NLCompilerSimilarFires,
AuditTimeline,
UserSearch,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct MemoryPredicate {
#[serde(default)]
pub kinds: Option<Vec<ChunkKind>>,
#[serde(default)]
pub callers: Option<Vec<String>>,
#[serde(default)]
pub session_ids: Option<Vec<String>>,
#[serde(default)]
pub project_root_prefix: Option<String>,
#[serde(default)]
pub before: Option<DateTime<Utc>>,
#[serde(default)]
pub after: Option<DateTime<Utc>>,
#[serde(default)]
pub pinned: Option<bool>,
#[serde(default)]
pub importance_below: Option<f32>,
#[serde(default)]
pub content_contains: Option<String>,
}
impl MemoryPredicate {
pub fn is_empty(&self) -> bool {
self.kinds.is_none()
&& self.callers.is_none()
&& self.session_ids.is_none()
&& self.project_root_prefix.is_none()
&& self.before.is_none()
&& self.after.is_none()
&& self.pinned.is_none()
&& self.importance_below.is_none()
&& self.content_contains.is_none()
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn caller_scope_default_is_own() {
assert_eq!(CallerScope::default(), CallerScope::Own);
}
#[test]
fn retrieval_profile_default_is_agent_chat_turn() {
assert_eq!(RetrievalProfile::default(), RetrievalProfile::AgentChatTurn);
}
#[test]
fn empty_predicate_is_empty() {
assert!(MemoryPredicate::default().is_empty());
let p = MemoryPredicate {
content_contains: Some("hi".into()),
..Default::default()
};
assert!(!p.is_empty());
}
#[test]
fn query_round_trip() {
let q = MemoryQuery {
text: "Q4 report".into(),
kinds: Some(vec![ChunkKind::Chat, ChunkKind::Action]),
since: None,
until: None,
session_id: Some("sess_abc".into()),
caller_scope: CallerScope::OwnPlusShared,
project_root_prefix: Some("/Users/dim/Workspace".into()),
k: 8,
include_rollups: true,
min_importance: Some(0.3),
profile: RetrievalProfile::AgentChatTurn,
caller_id: "embedded".into(),
};
let s = serde_json::to_string(&q).unwrap();
let back: MemoryQuery = serde_json::from_str(&s).unwrap();
assert_eq!(q, back);
}
#[test]
fn query_defaults_apply_on_deserialize() {
let raw = json!({
"text": "Q4 report",
"caller_id": "embedded"
})
.to_string();
let q: MemoryQuery = serde_json::from_str(&raw).unwrap();
assert_eq!(q.k, 8);
assert!(q.include_rollups);
assert_eq!(q.profile, RetrievalProfile::AgentChatTurn);
assert_eq!(q.caller_scope, CallerScope::Own);
}
}