rs_adk/session/
vertex_ai.rs

1//! Vertex AI session service — managed session storage via Vertex AI REST API.
2//!
3//! Provides session persistence using the Vertex AI session management
4//! endpoint. Sessions are stored and managed by Google Cloud, with
5//! optional TTL-based expiration.
6
7use async_trait::async_trait;
8
9use super::{Session, SessionError, SessionId, SessionService};
10use crate::events::Event;
11
12/// Configuration for the Vertex AI session service.
13#[derive(Debug, Clone)]
14pub struct VertexAiSessionConfig {
15    /// Google Cloud project ID.
16    pub project: String,
17    /// Google Cloud region (e.g., `us-central1`).
18    pub location: String,
19    /// Optional time-to-live for sessions, in seconds.
20    /// If set, sessions expire after this duration of inactivity.
21    pub ttl_seconds: Option<u64>,
22}
23
24impl VertexAiSessionConfig {
25    /// Create a new Vertex AI session config.
26    pub fn new(project: impl Into<String>, location: impl Into<String>) -> Self {
27        Self {
28            project: project.into(),
29            location: location.into(),
30            ttl_seconds: None,
31        }
32    }
33
34    /// Set the session TTL in seconds.
35    pub fn ttl_seconds(mut self, ttl: u64) -> Self {
36        self.ttl_seconds = Some(ttl);
37        self
38    }
39
40    /// Construct the base URL for the Vertex AI session endpoint.
41    ///
42    /// Format: `https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project}/locations/{location}/reasoningEngines`
43    fn base_url(&self) -> String {
44        format!(
45            "https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project}/locations/{location}",
46            project = self.project,
47            location = self.location,
48        )
49    }
50
51    /// Construct the sessions endpoint URL for a specific reasoning engine.
52    fn sessions_url(&self, engine_id: &str) -> String {
53        format!(
54            "{}/reasoningEngines/{}/sessions",
55            self.base_url(),
56            engine_id,
57        )
58    }
59
60    /// Construct the URL for a specific session.
61    fn session_url(&self, engine_id: &str, session_id: &str) -> String {
62        format!("{}/{}", self.sessions_url(engine_id), session_id)
63    }
64
65    /// Construct the events endpoint URL for a specific session.
66    fn events_url(&self, engine_id: &str, session_id: &str) -> String {
67        format!("{}/events", self.session_url(engine_id, session_id))
68    }
69}
70
71/// Session service backed by the Vertex AI managed session endpoint.
72///
73/// Uses the Vertex AI REST API for session CRUD and event storage.
74/// Requires a valid Google Cloud project with the AI Platform API enabled.
75///
76/// Sessions are stored server-side by Google Cloud, providing managed
77/// persistence without requiring a separate database.
78pub struct VertexAiSessionService {
79    config: VertexAiSessionConfig,
80    // In a real implementation, this would hold:
81    // - An HTTP client (e.g., `reqwest::Client`)
82    // - An auth token provider (e.g., `rs_genai::VertexAIAuth`)
83    // - An optional reasoning engine ID
84}
85
86impl VertexAiSessionService {
87    /// Create a new Vertex AI session service.
88    pub fn new(config: VertexAiSessionConfig) -> Self {
89        Self { config }
90    }
91
92    /// Returns the configured project ID.
93    pub fn project(&self) -> &str {
94        &self.config.project
95    }
96
97    /// Returns the configured location.
98    pub fn location(&self) -> &str {
99        &self.config.location
100    }
101
102    /// Returns the configured TTL in seconds, if any.
103    pub fn ttl_seconds(&self) -> Option<u64> {
104        self.config.ttl_seconds
105    }
106}
107
108#[async_trait]
109impl SessionService for VertexAiSessionService {
110    async fn create_session(&self, app_name: &str, user_id: &str) -> Result<Session, SessionError> {
111        let _url = self.config.sessions_url(app_name);
112        let _user = user_id;
113
114        // Real implementation would:
115        // POST {sessions_url}
116        // Authorization: Bearer {token}
117        // Content-Type: application/json
118        //
119        // {
120        //   "userId": "{user_id}",
121        //   "ttl": "{ttl_seconds}s"   // if configured
122        // }
123        //
124        // Response: { "name": "...sessions/{id}", "userId": "...", ... }
125        let _ttl_body = self
126            .config
127            .ttl_seconds
128            .map(|t| format!("\"ttl\": \"{t}s\""));
129
130        todo!("POST to {_url} to create Vertex AI session for user={_user}")
131    }
132
133    async fn get_session(&self, id: &SessionId) -> Result<Option<Session>, SessionError> {
134        let _url = self.config.session_url("default", id.as_str());
135
136        // Real implementation would:
137        // GET {session_url}
138        // Authorization: Bearer {token}
139        //
140        // Response: { "name": "...sessions/{id}", "userId": "...", "state": {...}, ... }
141        // Returns None if 404.
142        todo!("GET {_url} to fetch Vertex AI session")
143    }
144
145    async fn list_sessions(
146        &self,
147        app_name: &str,
148        user_id: &str,
149    ) -> Result<Vec<Session>, SessionError> {
150        let _url = self.config.sessions_url(app_name);
151        let _user = user_id;
152
153        // Real implementation would:
154        // GET {sessions_url}?filter=userId={user_id}
155        // Authorization: Bearer {token}
156        //
157        // Response: { "sessions": [{ "name": "...", ... }, ...] }
158        todo!("GET {_url} to list Vertex AI sessions for user={_user}")
159    }
160
161    async fn delete_session(&self, id: &SessionId) -> Result<(), SessionError> {
162        let _url = self.config.session_url("default", id.as_str());
163
164        // Real implementation would:
165        // DELETE {session_url}
166        // Authorization: Bearer {token}
167        //
168        // Response: {} (empty on success)
169        todo!("DELETE {_url} to remove Vertex AI session")
170    }
171
172    async fn append_event(&self, id: &SessionId, event: Event) -> Result<(), SessionError> {
173        let _url = self.config.events_url("default", id.as_str());
174        let _event_json =
175            serde_json::to_value(&event).map_err(|e| SessionError::Storage(e.to_string()))?;
176
177        // Real implementation would:
178        // POST {events_url}
179        // Authorization: Bearer {token}
180        // Content-Type: application/json
181        //
182        // {
183        //   "author": "{event.author}",
184        //   "invocationId": "{event.invocation_id}",
185        //   "content": { "parts": [{ "text": "{event.content}" }] },
186        //   "actions": {event.actions}
187        // }
188        todo!("POST to {_url} to append event to Vertex AI session")
189    }
190
191    async fn get_events(&self, id: &SessionId) -> Result<Vec<Event>, SessionError> {
192        let _url = self.config.events_url("default", id.as_str());
193
194        // Real implementation would:
195        // GET {events_url}
196        // Authorization: Bearer {token}
197        //
198        // Response: { "sessionEvents": [{ "author": "...", ... }, ...] }
199        // Transform each Vertex AI event into our Event type.
200        todo!("GET {_url} to fetch events for Vertex AI session")
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207
208    #[test]
209    fn config_new() {
210        let config = VertexAiSessionConfig::new("my-project", "us-central1");
211        assert_eq!(config.project, "my-project");
212        assert_eq!(config.location, "us-central1");
213        assert!(config.ttl_seconds.is_none());
214    }
215
216    #[test]
217    fn config_with_ttl() {
218        let config = VertexAiSessionConfig::new("proj", "us-east1").ttl_seconds(3600);
219        assert_eq!(config.ttl_seconds, Some(3600));
220    }
221
222    #[test]
223    fn url_construction() {
224        let config = VertexAiSessionConfig::new("my-project", "us-central1");
225        assert_eq!(
226            config.base_url(),
227            "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/my-project/locations/us-central1"
228        );
229        assert!(config
230            .sessions_url("engine-1")
231            .contains("reasoningEngines/engine-1/sessions"));
232        assert!(config
233            .session_url("engine-1", "sess-1")
234            .contains("sessions/sess-1"));
235        assert!(config
236            .events_url("engine-1", "sess-1")
237            .contains("sessions/sess-1/events"));
238    }
239
240    #[test]
241    fn service_accessors() {
242        let svc = VertexAiSessionService::new(
243            VertexAiSessionConfig::new("proj", "us-west1").ttl_seconds(7200),
244        );
245        assert_eq!(svc.project(), "proj");
246        assert_eq!(svc.location(), "us-west1");
247        assert_eq!(svc.ttl_seconds(), Some(7200));
248    }
249}
rs_adk/session/vertex_ai.rs

rs_adk/session/
vertex_ai.rs
