Skip to main content

dakera_client/
session.rs

1//! High-level session helper for LLM chat comparison patterns.
2//!
3//! [`ChatMemorySession`] wraps the low-level session/memory API into the
4//! three-step pattern used by the playground LLM chat comparison feature:
5//!
6//! 1. Create a session bound to an agent.
7//! 2. Store conversation turns with [`ChatMemorySession::store`].
8//! 3. Recall relevant context before generating the next response.
9//!
10//! # Example
11//!
12//! ```no_run
13//! use std::sync::Arc;
14//! use dakera_client::{DakeraClient, ChatMemorySession};
15//!
16//! # async fn run() -> dakera_client::Result<()> {
17//! let client = Arc::new(
18//!     DakeraClient::builder("http://localhost:3000")
19//!         .api_key("dk-mykey")
20//!         .build()?,
21//! );
22//!
23//! let session = ChatMemorySession::create(Arc::clone(&client), "chat-agent").await?;
24//! session.store("user", "My name is Alice and I like Rust.").await?;
25//! let context = session.recall("user preferences").await?;
26//! // Pass context to your LLM — or skip it for the baseline comparison arm.
27//! session.close().await?;
28//! # Ok(())
29//! # }
30//! ```
31
32use std::sync::Arc;
33
34use crate::memory::{RecallRequest, RecalledMemory, StoreMemoryRequest, StoreMemoryResponse};
35use crate::{DakeraClient, Result};
36
37/// High-level session helper for LLM chat comparison patterns.
38///
39/// Groups conversation turns under a single Dakera session so that:
40///
41/// * Every stored message is associated with `session_id` for scoped retrieval.
42/// * [`recall`][ChatMemorySession::recall] queries the agent's **full** memory —
43///   not just this session — so prior conversations inform the current exchange.
44///
45/// Create via [`ChatMemorySession::create`]; close via [`ChatMemorySession::close`].
46pub struct ChatMemorySession {
47    client: Arc<DakeraClient>,
48    agent_id: String,
49    session_id: String,
50}
51
52impl ChatMemorySession {
53    // -------------------------------------------------------------------------
54    // Factory
55    // -------------------------------------------------------------------------
56
57    /// Create a new Dakera session and return a [`ChatMemorySession`].
58    ///
59    /// # Arguments
60    ///
61    /// * `client`   – Shared [`DakeraClient`] instance.
62    /// * `agent_id` – Identifier for the agent whose memory to use.
63    pub async fn create(
64        client: Arc<DakeraClient>,
65        agent_id: impl Into<String>,
66    ) -> Result<ChatMemorySession> {
67        let agent_id = agent_id.into();
68        let session = client.start_session(&agent_id).await?;
69        Ok(ChatMemorySession {
70            client,
71            agent_id,
72            session_id: session.id,
73        })
74    }
75
76    /// Create a session with attached metadata.
77    pub async fn create_with_metadata(
78        client: Arc<DakeraClient>,
79        agent_id: impl Into<String>,
80        metadata: serde_json::Value,
81    ) -> Result<ChatMemorySession> {
82        let agent_id = agent_id.into();
83        let session = client
84            .start_session_with_metadata(&agent_id, metadata)
85            .await?;
86        Ok(ChatMemorySession {
87            client,
88            agent_id,
89            session_id: session.id,
90        })
91    }
92
93    // -------------------------------------------------------------------------
94    // Core operations
95    // -------------------------------------------------------------------------
96
97    /// Store a conversation turn in the session with default importance (0.6).
98    ///
99    /// The `role` (e.g. `"user"` or `"assistant"`) is appended to the memory's
100    /// tags automatically.
101    pub async fn store(&self, role: &str, content: &str) -> Result<StoreMemoryResponse> {
102        self.store_with_opts(role, content, 0.6, &[]).await
103    }
104
105    /// Store a conversation turn with custom importance and additional tags.
106    pub async fn store_with_opts(
107        &self,
108        role: &str,
109        content: &str,
110        importance: f32,
111        extra_tags: &[&str],
112    ) -> Result<StoreMemoryResponse> {
113        let mut tags: Vec<String> = extra_tags.iter().map(|&t| t.to_owned()).collect();
114        if !tags.iter().any(|t| t == role) {
115            tags.push(role.to_owned());
116        }
117        let request = StoreMemoryRequest::new(&self.agent_id, content)
118            .with_importance(importance)
119            .with_tags(tags)
120            .with_session(self.session_id.clone());
121        self.client.store_memory(request).await
122    }
123
124    /// Recall up to 5 memories relevant to `query` for this agent.
125    ///
126    /// Searches the agent's **full** memory, not just the current session, so
127    /// context from prior conversations is surfaced when relevant.
128    pub async fn recall(&self, query: &str) -> Result<Vec<RecalledMemory>> {
129        self.recall_top_k(query, 5).await
130    }
131
132    /// Recall up to `top_k` memories relevant to `query`.
133    pub async fn recall_top_k(&self, query: &str, top_k: usize) -> Result<Vec<RecalledMemory>> {
134        let request = RecallRequest::new(&self.agent_id, query).with_top_k(top_k);
135        let response = self.client.recall(request).await?;
136        Ok(response.memories)
137    }
138
139    /// End the Dakera session.
140    pub async fn close(self) -> Result<()> {
141        self.client.end_session(&self.session_id, None).await?;
142        Ok(())
143    }
144
145    // -------------------------------------------------------------------------
146    // Properties
147    // -------------------------------------------------------------------------
148
149    /// The underlying Dakera session ID.
150    pub fn session_id(&self) -> &str {
151        &self.session_id
152    }
153
154    /// The agent ID this session is bound to.
155    pub fn agent_id(&self) -> &str {
156        &self.agent_id
157    }
158}