use serde::{Deserialize, Serialize};
use crate::error::Result;
use crate::types::{
AgentFeedbackSummary, EdgeType, FeedbackHealthResponse, FeedbackHistoryResponse,
FeedbackResponse, FeedbackSignal, GraphExport, GraphLinkRequest, GraphLinkResponse,
GraphOptions, GraphPath, MemoryFeedbackBody, MemoryGraph, MemoryImportancePatch,
};
use crate::DakeraClient;
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub enum MemoryType {
#[default]
Episodic,
Semantic,
Procedural,
Working,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoreMemoryRequest {
pub agent_id: String,
pub content: String,
#[serde(default)]
pub memory_type: MemoryType,
#[serde(default = "default_importance")]
pub importance: f32,
#[serde(default)]
pub tags: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub session_id: Option<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>,
}
fn default_importance() -> f32 {
0.5
}
impl StoreMemoryRequest {
pub fn new(agent_id: impl Into<String>, content: impl Into<String>) -> Self {
Self {
agent_id: agent_id.into(),
content: content.into(),
memory_type: MemoryType::default(),
importance: 0.5,
tags: Vec::new(),
session_id: None,
metadata: None,
ttl_seconds: None,
expires_at: None,
}
}
pub fn with_type(mut self, memory_type: MemoryType) -> Self {
self.memory_type = memory_type;
self
}
pub fn with_importance(mut self, importance: f32) -> Self {
self.importance = importance.clamp(0.0, 1.0);
self
}
pub fn with_tags(mut self, tags: Vec<String>) -> Self {
self.tags = tags;
self
}
pub fn with_session(mut self, session_id: impl Into<String>) -> Self {
self.session_id = Some(session_id.into());
self
}
pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self {
self.metadata = Some(metadata);
self
}
pub fn with_ttl(mut self, ttl_seconds: u64) -> Self {
self.ttl_seconds = Some(ttl_seconds);
self
}
pub fn with_expires_at(mut self, expires_at: u64) -> Self {
self.expires_at = Some(expires_at);
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoreMemoryResponse {
pub memory_id: String,
pub agent_id: String,
pub namespace: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecallRequest {
pub agent_id: String,
pub query: String,
#[serde(default = "default_top_k")]
pub top_k: usize,
#[serde(skip_serializing_if = "Option::is_none")]
pub memory_type: Option<MemoryType>,
#[serde(default)]
pub min_importance: f32,
#[serde(skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
#[serde(default)]
pub tags: Vec<String>,
}
fn default_top_k() -> usize {
5
}
impl RecallRequest {
pub fn new(agent_id: impl Into<String>, query: impl Into<String>) -> Self {
Self {
agent_id: agent_id.into(),
query: query.into(),
top_k: 5,
memory_type: None,
min_importance: 0.0,
session_id: None,
tags: Vec::new(),
}
}
pub fn with_top_k(mut self, top_k: usize) -> Self {
self.top_k = top_k;
self
}
pub fn with_type(mut self, memory_type: MemoryType) -> Self {
self.memory_type = Some(memory_type);
self
}
pub fn with_min_importance(mut self, min: f32) -> Self {
self.min_importance = min;
self
}
pub fn with_session(mut self, session_id: impl Into<String>) -> Self {
self.session_id = Some(session_id.into());
self
}
pub fn with_tags(mut self, tags: Vec<String>) -> Self {
self.tags = tags;
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecalledMemory {
pub id: String,
pub content: String,
pub memory_type: MemoryType,
pub importance: f32,
pub score: f32,
#[serde(default)]
pub tags: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
pub created_at: u64,
pub last_accessed_at: u64,
pub access_count: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecallResponse {
pub memories: Vec<RecalledMemory>,
pub total_found: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ForgetRequest {
pub agent_id: String,
#[serde(default)]
pub memory_ids: Vec<String>,
#[serde(default)]
pub tags: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub before_timestamp: Option<u64>,
}
impl ForgetRequest {
pub fn by_ids(agent_id: impl Into<String>, ids: Vec<String>) -> Self {
Self {
agent_id: agent_id.into(),
memory_ids: ids,
tags: Vec::new(),
session_id: None,
before_timestamp: None,
}
}
pub fn by_tags(agent_id: impl Into<String>, tags: Vec<String>) -> Self {
Self {
agent_id: agent_id.into(),
memory_ids: Vec::new(),
tags,
session_id: None,
before_timestamp: None,
}
}
pub fn by_session(agent_id: impl Into<String>, session_id: impl Into<String>) -> Self {
Self {
agent_id: agent_id.into(),
memory_ids: Vec::new(),
tags: Vec::new(),
session_id: Some(session_id.into()),
before_timestamp: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ForgetResponse {
pub deleted_count: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionStartRequest {
pub agent_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
}
#[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>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionEndRequest {
#[serde(skip_serializing_if = "Option::is_none")]
pub summary: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdateMemoryRequest {
#[serde(skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub memory_type: Option<MemoryType>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdateImportanceRequest {
pub memory_ids: Vec<String>,
pub importance: f32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConsolidateRequest {
#[serde(skip_serializing_if = "Option::is_none")]
pub memory_type: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub threshold: Option<f32>,
#[serde(default)]
pub dry_run: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConsolidateResponse {
pub consolidated_count: usize,
pub removed_count: usize,
pub new_memories: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeedbackRequest {
pub memory_id: String,
pub feedback: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub relevance_score: Option<f32>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LegacyFeedbackResponse {
pub status: String,
pub updated_importance: Option<f32>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct BatchMemoryFilter {
#[serde(skip_serializing_if = "Option::is_none")]
pub tags: Option<Vec<String>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub min_importance: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_importance: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub created_after: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub created_before: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub memory_type: Option<MemoryType>,
#[serde(skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
}
impl BatchMemoryFilter {
pub fn with_tags(mut self, tags: Vec<String>) -> Self {
self.tags = Some(tags);
self
}
pub fn with_min_importance(mut self, min: f32) -> Self {
self.min_importance = Some(min);
self
}
pub fn with_max_importance(mut self, max: f32) -> Self {
self.max_importance = Some(max);
self
}
pub fn with_session(mut self, session_id: impl Into<String>) -> Self {
self.session_id = Some(session_id.into());
self
}
}
#[derive(Debug, Clone, Serialize, 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
}
impl BatchRecallRequest {
pub fn new(agent_id: impl Into<String>) -> Self {
Self {
agent_id: agent_id.into(),
filter: BatchMemoryFilter::default(),
limit: 100,
}
}
pub fn with_filter(mut self, filter: BatchMemoryFilter) -> Self {
self.filter = filter;
self
}
pub fn with_limit(mut self, limit: usize) -> Self {
self.limit = limit;
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BatchRecallResponse {
pub memories: Vec<RecalledMemory>,
pub total: usize,
pub filtered: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BatchForgetRequest {
pub agent_id: String,
pub filter: BatchMemoryFilter,
}
impl BatchForgetRequest {
pub fn new(agent_id: impl Into<String>, filter: BatchMemoryFilter) -> Self {
Self {
agent_id: agent_id.into(),
filter,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BatchForgetResponse {
pub deleted_count: usize,
}
impl DakeraClient {
pub async fn store_memory(&self, request: StoreMemoryRequest) -> Result<StoreMemoryResponse> {
let url = format!("{}/v1/memory/store", self.base_url);
let response = self.client.post(&url).json(&request).send().await?;
self.handle_response(response).await
}
pub async fn recall(&self, request: RecallRequest) -> Result<RecallResponse> {
let url = format!("{}/v1/memory/recall", self.base_url);
let response = self.client.post(&url).json(&request).send().await?;
self.handle_response(response).await
}
pub async fn recall_simple(
&self,
agent_id: &str,
query: &str,
top_k: usize,
) -> Result<RecallResponse> {
self.recall(RecallRequest::new(agent_id, query).with_top_k(top_k))
.await
}
pub async fn get_memory(&self, memory_id: &str) -> Result<RecalledMemory> {
let url = format!("{}/v1/memory/get/{}", self.base_url, memory_id);
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn forget(&self, request: ForgetRequest) -> Result<ForgetResponse> {
let url = format!("{}/v1/memory/forget", self.base_url);
let response = self.client.post(&url).json(&request).send().await?;
self.handle_response(response).await
}
pub async fn search_memories(&self, request: RecallRequest) -> Result<RecallResponse> {
let url = format!("{}/v1/memory/search", self.base_url);
let response = self.client.post(&url).json(&request).send().await?;
self.handle_response(response).await
}
pub async fn update_memory(
&self,
agent_id: &str,
memory_id: &str,
request: UpdateMemoryRequest,
) -> Result<StoreMemoryResponse> {
let url = format!(
"{}/v1/agents/{}/memories/{}",
self.base_url, agent_id, memory_id
);
let response = self.client.put(&url).json(&request).send().await?;
self.handle_response(response).await
}
pub async fn update_importance(
&self,
agent_id: &str,
request: UpdateImportanceRequest,
) -> Result<serde_json::Value> {
let url = format!(
"{}/v1/agents/{}/memories/importance",
self.base_url, agent_id
);
let response = self.client.put(&url).json(&request).send().await?;
self.handle_response(response).await
}
pub async fn consolidate(
&self,
agent_id: &str,
request: ConsolidateRequest,
) -> Result<ConsolidateResponse> {
let url = format!(
"{}/v1/agents/{}/memories/consolidate",
self.base_url, agent_id
);
let response = self.client.post(&url).json(&request).send().await?;
self.handle_response(response).await
}
pub async fn memory_feedback(
&self,
agent_id: &str,
request: FeedbackRequest,
) -> Result<LegacyFeedbackResponse> {
let url = format!("{}/v1/agents/{}/memories/feedback", self.base_url, agent_id);
let response = self.client.post(&url).json(&request).send().await?;
self.handle_response(response).await
}
pub async fn feedback_memory(
&self,
memory_id: &str,
agent_id: &str,
signal: FeedbackSignal,
) -> Result<FeedbackResponse> {
let url = format!("{}/v1/memories/{}/feedback", self.base_url, memory_id);
let body = MemoryFeedbackBody {
agent_id: agent_id.to_string(),
signal,
};
let response = self.client.post(&url).json(&body).send().await?;
self.handle_response(response).await
}
pub async fn get_memory_feedback_history(
&self,
memory_id: &str,
) -> Result<FeedbackHistoryResponse> {
let url = format!("{}/v1/memories/{}/feedback", self.base_url, memory_id);
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn get_agent_feedback_summary(&self, agent_id: &str) -> Result<AgentFeedbackSummary> {
let url = format!("{}/v1/agents/{}/feedback/summary", self.base_url, agent_id);
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn patch_memory_importance(
&self,
memory_id: &str,
agent_id: &str,
importance: f32,
) -> Result<FeedbackResponse> {
let url = format!("{}/v1/memories/{}/importance", self.base_url, memory_id);
let body = MemoryImportancePatch {
agent_id: agent_id.to_string(),
importance,
};
let response = self.client.patch(&url).json(&body).send().await?;
self.handle_response(response).await
}
pub async fn get_feedback_health(&self, agent_id: &str) -> Result<FeedbackHealthResponse> {
let url = format!("{}/v1/feedback/health?agent_id={}", self.base_url, agent_id);
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn memory_graph(
&self,
memory_id: &str,
options: GraphOptions,
) -> Result<MemoryGraph> {
let mut url = format!("{}/v1/memories/{}/graph", self.base_url, memory_id);
let depth = options.depth.unwrap_or(1);
url.push_str(&format!("?depth={}", depth));
if let Some(types) = &options.types {
let type_strs: Vec<String> = types
.iter()
.map(|t| {
serde_json::to_value(t)
.unwrap()
.as_str()
.unwrap_or("")
.to_string()
})
.collect();
if !type_strs.is_empty() {
url.push_str(&format!("&types={}", type_strs.join(",")));
}
}
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn memory_path(&self, source_id: &str, target_id: &str) -> Result<GraphPath> {
let url = format!(
"{}/v1/memories/{}/path?target={}",
self.base_url,
source_id,
urlencoding::encode(target_id)
);
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn memory_link(
&self,
source_id: &str,
target_id: &str,
edge_type: EdgeType,
) -> Result<GraphLinkResponse> {
let url = format!("{}/v1/memories/{}/links", self.base_url, source_id);
let request = GraphLinkRequest {
target_id: target_id.to_string(),
edge_type,
};
let response = self.client.post(&url).json(&request).send().await?;
self.handle_response(response).await
}
pub async fn agent_graph_export(&self, agent_id: &str, format: &str) -> Result<GraphExport> {
let url = format!(
"{}/v1/agents/{}/graph/export?format={}",
self.base_url, agent_id, format
);
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn start_session(&self, agent_id: &str) -> Result<Session> {
let url = format!("{}/v1/sessions/start", self.base_url);
let request = SessionStartRequest {
agent_id: agent_id.to_string(),
metadata: None,
};
let response = self.client.post(&url).json(&request).send().await?;
self.handle_response(response).await
}
pub async fn start_session_with_metadata(
&self,
agent_id: &str,
metadata: serde_json::Value,
) -> Result<Session> {
let url = format!("{}/v1/sessions/start", self.base_url);
let request = SessionStartRequest {
agent_id: agent_id.to_string(),
metadata: Some(metadata),
};
let response = self.client.post(&url).json(&request).send().await?;
self.handle_response(response).await
}
pub async fn end_session(&self, session_id: &str, summary: Option<String>) -> Result<Session> {
let url = format!("{}/v1/sessions/{}/end", self.base_url, session_id);
let request = SessionEndRequest { summary };
let response = self.client.post(&url).json(&request).send().await?;
self.handle_response(response).await
}
pub async fn get_session(&self, session_id: &str) -> Result<Session> {
let url = format!("{}/v1/sessions/{}", self.base_url, session_id);
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn list_sessions(&self, agent_id: &str) -> Result<Vec<Session>> {
let url = format!("{}/v1/sessions?agent_id={}", self.base_url, agent_id);
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn session_memories(&self, session_id: &str) -> Result<RecallResponse> {
let url = format!("{}/v1/sessions/{}/memories", self.base_url, session_id);
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn batch_recall(&self, request: BatchRecallRequest) -> Result<BatchRecallResponse> {
let url = format!("{}/v1/memories/recall/batch", self.base_url);
let response = self.client.post(&url).json(&request).send().await?;
self.handle_response(response).await
}
pub async fn batch_forget(&self, request: BatchForgetRequest) -> Result<BatchForgetResponse> {
let url = format!("{}/v1/memories/forget/batch", self.base_url);
let response = self.client.delete(&url).json(&request).send().await?;
self.handle_response(response).await
}
}