1use serde::{Deserialize, Serialize};
7
8use crate::error::Result;
9use crate::DakeraClient;
10
11#[derive(Debug, Clone, Serialize, Deserialize, Default)]
17pub enum MemoryType {
18 #[default]
19 Episodic,
20 Semantic,
21 Procedural,
22 Working,
23}
24
25#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct StoreMemoryRequest {
28 pub agent_id: String,
29 pub content: String,
30 #[serde(default)]
31 pub memory_type: MemoryType,
32 #[serde(default = "default_importance")]
33 pub importance: f32,
34 #[serde(default)]
35 pub tags: Vec<String>,
36 #[serde(skip_serializing_if = "Option::is_none")]
37 pub session_id: Option<String>,
38 #[serde(skip_serializing_if = "Option::is_none")]
39 pub metadata: Option<serde_json::Value>,
40 #[serde(skip_serializing_if = "Option::is_none")]
43 pub ttl_seconds: Option<u64>,
44 #[serde(skip_serializing_if = "Option::is_none")]
48 pub expires_at: Option<u64>,
49}
50
51fn default_importance() -> f32 {
52 0.5
53}
54
55impl StoreMemoryRequest {
56 pub fn new(agent_id: impl Into<String>, content: impl Into<String>) -> Self {
58 Self {
59 agent_id: agent_id.into(),
60 content: content.into(),
61 memory_type: MemoryType::default(),
62 importance: 0.5,
63 tags: Vec::new(),
64 session_id: None,
65 metadata: None,
66 ttl_seconds: None,
67 expires_at: None,
68 }
69 }
70
71 pub fn with_type(mut self, memory_type: MemoryType) -> Self {
73 self.memory_type = memory_type;
74 self
75 }
76
77 pub fn with_importance(mut self, importance: f32) -> Self {
79 self.importance = importance.clamp(0.0, 1.0);
80 self
81 }
82
83 pub fn with_tags(mut self, tags: Vec<String>) -> Self {
85 self.tags = tags;
86 self
87 }
88
89 pub fn with_session(mut self, session_id: impl Into<String>) -> Self {
91 self.session_id = Some(session_id.into());
92 self
93 }
94
95 pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self {
97 self.metadata = Some(metadata);
98 self
99 }
100
101 pub fn with_ttl(mut self, ttl_seconds: u64) -> Self {
104 self.ttl_seconds = Some(ttl_seconds);
105 self
106 }
107
108 pub fn with_expires_at(mut self, expires_at: u64) -> Self {
111 self.expires_at = Some(expires_at);
112 self
113 }
114}
115
116#[derive(Debug, Clone, Serialize, Deserialize)]
118pub struct StoreMemoryResponse {
119 pub memory_id: String,
120 pub agent_id: String,
121 pub namespace: String,
122}
123
124#[derive(Debug, Clone, Serialize, Deserialize)]
126pub struct RecallRequest {
127 pub agent_id: String,
128 pub query: String,
129 #[serde(default = "default_top_k")]
130 pub top_k: usize,
131 #[serde(skip_serializing_if = "Option::is_none")]
132 pub memory_type: Option<MemoryType>,
133 #[serde(default)]
134 pub min_importance: f32,
135 #[serde(skip_serializing_if = "Option::is_none")]
136 pub session_id: Option<String>,
137 #[serde(default)]
138 pub tags: Vec<String>,
139}
140
141fn default_top_k() -> usize {
142 5
143}
144
145impl RecallRequest {
146 pub fn new(agent_id: impl Into<String>, query: impl Into<String>) -> Self {
148 Self {
149 agent_id: agent_id.into(),
150 query: query.into(),
151 top_k: 5,
152 memory_type: None,
153 min_importance: 0.0,
154 session_id: None,
155 tags: Vec::new(),
156 }
157 }
158
159 pub fn with_top_k(mut self, top_k: usize) -> Self {
161 self.top_k = top_k;
162 self
163 }
164
165 pub fn with_type(mut self, memory_type: MemoryType) -> Self {
167 self.memory_type = Some(memory_type);
168 self
169 }
170
171 pub fn with_min_importance(mut self, min: f32) -> Self {
173 self.min_importance = min;
174 self
175 }
176
177 pub fn with_session(mut self, session_id: impl Into<String>) -> Self {
179 self.session_id = Some(session_id.into());
180 self
181 }
182
183 pub fn with_tags(mut self, tags: Vec<String>) -> Self {
185 self.tags = tags;
186 self
187 }
188}
189
190#[derive(Debug, Clone, Serialize, Deserialize)]
192pub struct RecalledMemory {
193 pub id: String,
194 pub content: String,
195 pub memory_type: MemoryType,
196 pub importance: f32,
197 pub score: f32,
198 #[serde(default)]
199 pub tags: Vec<String>,
200 #[serde(skip_serializing_if = "Option::is_none")]
201 pub session_id: Option<String>,
202 #[serde(skip_serializing_if = "Option::is_none")]
203 pub metadata: Option<serde_json::Value>,
204 pub created_at: u64,
205 pub last_accessed_at: u64,
206 pub access_count: u32,
207}
208
209#[derive(Debug, Clone, Serialize, Deserialize)]
211pub struct RecallResponse {
212 pub memories: Vec<RecalledMemory>,
213 pub total_found: usize,
214}
215
216#[derive(Debug, Clone, Serialize, Deserialize)]
218pub struct ForgetRequest {
219 pub agent_id: String,
220 #[serde(default)]
221 pub memory_ids: Vec<String>,
222 #[serde(default)]
223 pub tags: Vec<String>,
224 #[serde(skip_serializing_if = "Option::is_none")]
225 pub session_id: Option<String>,
226 #[serde(skip_serializing_if = "Option::is_none")]
227 pub before_timestamp: Option<u64>,
228}
229
230impl ForgetRequest {
231 pub fn by_ids(agent_id: impl Into<String>, ids: Vec<String>) -> Self {
233 Self {
234 agent_id: agent_id.into(),
235 memory_ids: ids,
236 tags: Vec::new(),
237 session_id: None,
238 before_timestamp: None,
239 }
240 }
241
242 pub fn by_tags(agent_id: impl Into<String>, tags: Vec<String>) -> Self {
244 Self {
245 agent_id: agent_id.into(),
246 memory_ids: Vec::new(),
247 tags,
248 session_id: None,
249 before_timestamp: None,
250 }
251 }
252
253 pub fn by_session(agent_id: impl Into<String>, session_id: impl Into<String>) -> Self {
255 Self {
256 agent_id: agent_id.into(),
257 memory_ids: Vec::new(),
258 tags: Vec::new(),
259 session_id: Some(session_id.into()),
260 before_timestamp: None,
261 }
262 }
263}
264
265#[derive(Debug, Clone, Serialize, Deserialize)]
267pub struct ForgetResponse {
268 pub deleted_count: u64,
269}
270
271#[derive(Debug, Clone, Serialize, Deserialize)]
273pub struct SessionStartRequest {
274 pub agent_id: String,
275 #[serde(skip_serializing_if = "Option::is_none")]
276 pub metadata: Option<serde_json::Value>,
277}
278
279#[derive(Debug, Clone, Serialize, Deserialize)]
281pub struct Session {
282 pub id: String,
283 pub agent_id: String,
284 pub started_at: u64,
285 #[serde(skip_serializing_if = "Option::is_none")]
286 pub ended_at: Option<u64>,
287 #[serde(skip_serializing_if = "Option::is_none")]
288 pub summary: Option<String>,
289 #[serde(skip_serializing_if = "Option::is_none")]
290 pub metadata: Option<serde_json::Value>,
291}
292
293#[derive(Debug, Clone, Serialize, Deserialize)]
295pub struct SessionEndRequest {
296 #[serde(skip_serializing_if = "Option::is_none")]
297 pub summary: Option<String>,
298}
299
300#[derive(Debug, Clone, Serialize, Deserialize)]
302pub struct UpdateMemoryRequest {
303 #[serde(skip_serializing_if = "Option::is_none")]
304 pub content: Option<String>,
305 #[serde(skip_serializing_if = "Option::is_none")]
306 pub metadata: Option<serde_json::Value>,
307 #[serde(skip_serializing_if = "Option::is_none")]
308 pub memory_type: Option<MemoryType>,
309}
310
311#[derive(Debug, Clone, Serialize, Deserialize)]
313pub struct UpdateImportanceRequest {
314 pub memory_ids: Vec<String>,
315 pub importance: f32,
316}
317
318#[derive(Debug, Clone, Serialize, Deserialize)]
320pub struct ConsolidateRequest {
321 #[serde(skip_serializing_if = "Option::is_none")]
322 pub memory_type: Option<String>,
323 #[serde(skip_serializing_if = "Option::is_none")]
324 pub threshold: Option<f32>,
325 #[serde(default)]
326 pub dry_run: bool,
327}
328
329#[derive(Debug, Clone, Serialize, Deserialize)]
331pub struct ConsolidateResponse {
332 pub consolidated_count: usize,
333 pub removed_count: usize,
334 pub new_memories: Vec<String>,
335}
336
337#[derive(Debug, Clone, Serialize, Deserialize)]
339pub struct FeedbackRequest {
340 pub memory_id: String,
341 pub feedback: String,
342 #[serde(skip_serializing_if = "Option::is_none")]
343 pub relevance_score: Option<f32>,
344}
345
346#[derive(Debug, Clone, Serialize, Deserialize)]
348pub struct FeedbackResponse {
349 pub status: String,
350 pub updated_importance: Option<f32>,
351}
352
353#[derive(Debug, Clone, Serialize, Deserialize, Default)]
362pub struct BatchMemoryFilter {
363 #[serde(skip_serializing_if = "Option::is_none")]
365 pub tags: Option<Vec<String>>,
366 #[serde(skip_serializing_if = "Option::is_none")]
368 pub min_importance: Option<f32>,
369 #[serde(skip_serializing_if = "Option::is_none")]
371 pub max_importance: Option<f32>,
372 #[serde(skip_serializing_if = "Option::is_none")]
374 pub created_after: Option<u64>,
375 #[serde(skip_serializing_if = "Option::is_none")]
377 pub created_before: Option<u64>,
378 #[serde(skip_serializing_if = "Option::is_none")]
380 pub memory_type: Option<MemoryType>,
381 #[serde(skip_serializing_if = "Option::is_none")]
383 pub session_id: Option<String>,
384}
385
386impl BatchMemoryFilter {
387 pub fn with_tags(mut self, tags: Vec<String>) -> Self {
389 self.tags = Some(tags);
390 self
391 }
392
393 pub fn with_min_importance(mut self, min: f32) -> Self {
395 self.min_importance = Some(min);
396 self
397 }
398
399 pub fn with_max_importance(mut self, max: f32) -> Self {
401 self.max_importance = Some(max);
402 self
403 }
404
405 pub fn with_session(mut self, session_id: impl Into<String>) -> Self {
407 self.session_id = Some(session_id.into());
408 self
409 }
410}
411
412#[derive(Debug, Clone, Serialize, Deserialize)]
414pub struct BatchRecallRequest {
415 pub agent_id: String,
417 #[serde(default)]
419 pub filter: BatchMemoryFilter,
420 #[serde(default = "default_batch_limit")]
422 pub limit: usize,
423}
424
425fn default_batch_limit() -> usize {
426 100
427}
428
429impl BatchRecallRequest {
430 pub fn new(agent_id: impl Into<String>) -> Self {
432 Self {
433 agent_id: agent_id.into(),
434 filter: BatchMemoryFilter::default(),
435 limit: 100,
436 }
437 }
438
439 pub fn with_filter(mut self, filter: BatchMemoryFilter) -> Self {
441 self.filter = filter;
442 self
443 }
444
445 pub fn with_limit(mut self, limit: usize) -> Self {
447 self.limit = limit;
448 self
449 }
450}
451
452#[derive(Debug, Clone, Serialize, Deserialize)]
454pub struct BatchRecallResponse {
455 pub memories: Vec<RecalledMemory>,
456 pub total: usize,
458 pub filtered: usize,
460}
461
462#[derive(Debug, Clone, Serialize, Deserialize)]
464pub struct BatchForgetRequest {
465 pub agent_id: String,
467 pub filter: BatchMemoryFilter,
469}
470
471impl BatchForgetRequest {
472 pub fn new(agent_id: impl Into<String>, filter: BatchMemoryFilter) -> Self {
474 Self {
475 agent_id: agent_id.into(),
476 filter,
477 }
478 }
479}
480
481#[derive(Debug, Clone, Serialize, Deserialize)]
483pub struct BatchForgetResponse {
484 pub deleted_count: usize,
485}
486
487impl DakeraClient {
492 pub async fn store_memory(&self, request: StoreMemoryRequest) -> Result<StoreMemoryResponse> {
516 let url = format!("{}/v1/memory/store", self.base_url);
517 let response = self.client.post(&url).json(&request).send().await?;
518 self.handle_response(response).await
519 }
520
521 pub async fn recall(&self, request: RecallRequest) -> Result<RecallResponse> {
542 let url = format!("{}/v1/memory/recall", self.base_url);
543 let response = self.client.post(&url).json(&request).send().await?;
544 self.handle_response(response).await
545 }
546
547 pub async fn recall_simple(
549 &self,
550 agent_id: &str,
551 query: &str,
552 top_k: usize,
553 ) -> Result<RecallResponse> {
554 self.recall(RecallRequest::new(agent_id, query).with_top_k(top_k))
555 .await
556 }
557
558 pub async fn get_memory(&self, memory_id: &str) -> Result<RecalledMemory> {
560 let url = format!("{}/v1/memory/get/{}", self.base_url, memory_id);
561 let response = self.client.get(&url).send().await?;
562 self.handle_response(response).await
563 }
564
565 pub async fn forget(&self, request: ForgetRequest) -> Result<ForgetResponse> {
567 let url = format!("{}/v1/memory/forget", self.base_url);
568 let response = self.client.post(&url).json(&request).send().await?;
569 self.handle_response(response).await
570 }
571
572 pub async fn search_memories(&self, request: RecallRequest) -> Result<RecallResponse> {
574 let url = format!("{}/v1/memory/search", self.base_url);
575 let response = self.client.post(&url).json(&request).send().await?;
576 self.handle_response(response).await
577 }
578
579 pub async fn update_memory(
581 &self,
582 agent_id: &str,
583 memory_id: &str,
584 request: UpdateMemoryRequest,
585 ) -> Result<StoreMemoryResponse> {
586 let url = format!(
587 "{}/v1/agents/{}/memories/{}",
588 self.base_url, agent_id, memory_id
589 );
590 let response = self.client.put(&url).json(&request).send().await?;
591 self.handle_response(response).await
592 }
593
594 pub async fn update_importance(
596 &self,
597 agent_id: &str,
598 request: UpdateImportanceRequest,
599 ) -> Result<serde_json::Value> {
600 let url = format!(
601 "{}/v1/agents/{}/memories/importance",
602 self.base_url, agent_id
603 );
604 let response = self.client.put(&url).json(&request).send().await?;
605 self.handle_response(response).await
606 }
607
608 pub async fn consolidate(
610 &self,
611 agent_id: &str,
612 request: ConsolidateRequest,
613 ) -> Result<ConsolidateResponse> {
614 let url = format!(
615 "{}/v1/agents/{}/memories/consolidate",
616 self.base_url, agent_id
617 );
618 let response = self.client.post(&url).json(&request).send().await?;
619 self.handle_response(response).await
620 }
621
622 pub async fn memory_feedback(
624 &self,
625 agent_id: &str,
626 request: FeedbackRequest,
627 ) -> Result<FeedbackResponse> {
628 let url = format!("{}/v1/agents/{}/memories/feedback", self.base_url, agent_id);
629 let response = self.client.post(&url).json(&request).send().await?;
630 self.handle_response(response).await
631 }
632
633 pub async fn start_session(&self, agent_id: &str) -> Result<Session> {
639 let url = format!("{}/v1/sessions/start", self.base_url);
640 let request = SessionStartRequest {
641 agent_id: agent_id.to_string(),
642 metadata: None,
643 };
644 let response = self.client.post(&url).json(&request).send().await?;
645 self.handle_response(response).await
646 }
647
648 pub async fn start_session_with_metadata(
650 &self,
651 agent_id: &str,
652 metadata: serde_json::Value,
653 ) -> Result<Session> {
654 let url = format!("{}/v1/sessions/start", self.base_url);
655 let request = SessionStartRequest {
656 agent_id: agent_id.to_string(),
657 metadata: Some(metadata),
658 };
659 let response = self.client.post(&url).json(&request).send().await?;
660 self.handle_response(response).await
661 }
662
663 pub async fn end_session(&self, session_id: &str, summary: Option<String>) -> Result<Session> {
665 let url = format!("{}/v1/sessions/{}/end", self.base_url, session_id);
666 let request = SessionEndRequest { summary };
667 let response = self.client.post(&url).json(&request).send().await?;
668 self.handle_response(response).await
669 }
670
671 pub async fn get_session(&self, session_id: &str) -> Result<Session> {
673 let url = format!("{}/v1/sessions/{}", self.base_url, session_id);
674 let response = self.client.get(&url).send().await?;
675 self.handle_response(response).await
676 }
677
678 pub async fn list_sessions(&self, agent_id: &str) -> Result<Vec<Session>> {
680 let url = format!("{}/v1/sessions?agent_id={}", self.base_url, agent_id);
681 let response = self.client.get(&url).send().await?;
682 self.handle_response(response).await
683 }
684
685 pub async fn session_memories(&self, session_id: &str) -> Result<RecallResponse> {
687 let url = format!("{}/v1/sessions/{}/memories", self.base_url, session_id);
688 let response = self.client.get(&url).send().await?;
689 self.handle_response(response).await
690 }
691
692 pub async fn batch_recall(&self, request: BatchRecallRequest) -> Result<BatchRecallResponse> {
716 let url = format!("{}/v1/memories/recall/batch", self.base_url);
717 let response = self.client.post(&url).json(&request).send().await?;
718 self.handle_response(response).await
719 }
720
721 pub async fn batch_forget(&self, request: BatchForgetRequest) -> Result<BatchForgetResponse> {
741 let url = format!("{}/v1/memories/forget/batch", self.base_url);
742 let response = self.client.delete(&url).json(&request).send().await?;
743 self.handle_response(response).await
744 }
745}