use serde::{Deserialize, Serialize};
use crate::error::Result;
use crate::memory::{RecalledMemory, Session};
use crate::DakeraClient;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentSummary {
pub agent_id: String,
pub memory_count: i64,
pub session_count: i64,
pub active_sessions: i64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentStats {
pub agent_id: String,
pub total_memories: i64,
#[serde(default)]
pub memories_by_type: std::collections::HashMap<String, i64>,
pub total_sessions: i64,
pub active_sessions: i64,
#[serde(skip_serializing_if = "Option::is_none")]
pub avg_importance: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub oldest_memory_at: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub newest_memory_at: Option<String>,
}
impl DakeraClient {
pub async fn list_agents(&self) -> Result<Vec<AgentSummary>> {
let url = format!("{}/v1/agents", self.base_url);
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn agent_memories(
&self,
agent_id: &str,
memory_type: Option<&str>,
limit: Option<u32>,
) -> Result<Vec<RecalledMemory>> {
let mut url = format!("{}/v1/agents/{}/memories", self.base_url, agent_id);
let mut params = Vec::new();
if let Some(t) = memory_type {
params.push(format!("memory_type={}", t));
}
if let Some(l) = limit {
params.push(format!("limit={}", l));
}
if !params.is_empty() {
url.push('?');
url.push_str(¶ms.join("&"));
}
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn agent_stats(&self, agent_id: &str) -> Result<AgentStats> {
let url = format!("{}/v1/agents/{}/stats", self.base_url, agent_id);
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn subscribe_agent_events(
&self,
agent_id: &str,
tags: Option<Vec<String>>,
) -> crate::error::Result<
tokio::sync::mpsc::Receiver<crate::error::Result<crate::events::MemoryEvent>>,
> {
let (tx, rx) = tokio::sync::mpsc::channel(64);
let client = self.clone();
let agent_id = agent_id.to_owned();
tokio::spawn(async move {
loop {
match client.stream_memory_events().await {
Err(_) => {
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
continue;
}
Ok(mut inner_rx) => {
while let Some(result) = inner_rx.recv().await {
match result {
Err(e) => {
let _ = tx.send(Err(e)).await;
break;
}
Ok(event) => {
if event.event_type == "connected" {
continue;
}
if event.agent_id != agent_id {
continue;
}
if let Some(ref filter_tags) = tags {
let event_tags = event.tags.as_deref().unwrap_or(&[]);
if !filter_tags.iter().any(|t| event_tags.contains(t)) {
continue;
}
}
if tx.send(Ok(event)).await.is_err() {
return; }
}
}
}
}
}
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
}
});
Ok(rx)
}
pub async fn agent_sessions(
&self,
agent_id: &str,
active_only: Option<bool>,
limit: Option<u32>,
) -> Result<Vec<Session>> {
let mut url = format!("{}/v1/agents/{}/sessions", self.base_url, agent_id);
let mut params = Vec::new();
if let Some(active) = active_only {
params.push(format!("active_only={}", active));
}
if let Some(l) = limit {
params.push(format!("limit={}", l));
}
if !params.is_empty() {
url.push('?');
url.push_str(¶ms.join("&"));
}
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn wake_up(
&self,
agent_id: &str,
top_n: Option<u32>,
min_importance: Option<f32>,
) -> Result<WakeUpResponse> {
let mut url = format!("{}/v1/agents/{}/wake-up", self.base_url, agent_id);
let mut params = Vec::new();
if let Some(n) = top_n {
params.push(format!("top_n={}", n));
}
if let Some(mi) = min_importance {
params.push(format!("min_importance={}", mi));
}
if !params.is_empty() {
url.push('?');
url.push_str(¶ms.join("&"));
}
let response = self.client.get(&url).send().await?;
self.handle_response(response).await
}
pub async fn compress(&self, agent_id: &str) -> Result<CompressResponse> {
let url = format!("{}/v1/agents/{}/compress", self.base_url, agent_id);
let response = self.client.post(&url).send().await?;
self.handle_response(response).await
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Memory {
pub id: String,
pub content: String,
pub memory_type: String,
pub importance: f32,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub created_at: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub updated_at: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub access_count: Option<i64>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WakeUpResponse {
pub agent_id: String,
pub memories: Vec<Memory>,
pub total_available: i64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompressResponse {
pub agent_id: String,
#[serde(default)]
pub memories_scanned: i64,
#[serde(default, alias = "removed_count")]
pub originals_deprecated: i64,
#[serde(default)]
pub clusters_found: i64,
#[serde(default)]
pub summaries_created: i64,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub deprecated_ids: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub duration_ms: Option<f64>,
}