| 1 | //! The seam between Ed and a host's chat storage. |
| 2 | //! |
| 3 | //! Ed owns no storage. A host that keeps several conversations implements |
| 4 | //! [`SessionStore`] over whatever it already has (files, SQLite, a keychain) |
| 5 | //! and hands the messages to [`Ed::restore_history`](crate::Ed::restore_history) |
| 6 | //! when the user switches conversation. |
| 7 | |
| 8 | use async_trait::async_trait; |
| 9 | use onde::inference::types::ChatMessage; |
| 10 | use serde::{Deserialize, Serialize}; |
| 11 | |
| 12 | /// What a session list needs to render a row, without loading the messages. |
| 13 | #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] |
| 14 | pub struct SessionMeta { |
| 15 | pub id: String, |
| 16 | pub title: String, |
| 17 | /// Unix seconds. |
| 18 | pub created_at: i64, |
| 19 | /// Unix seconds. |
| 20 | pub updated_at: i64, |
| 21 | pub message_count: usize, |
| 22 | } |
| 23 | |
| 24 | /// A storage failure. The store decides the wording; Ed only carries it. |
| 25 | #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] |
| 26 | #[error("{0}")] |
| 27 | pub struct SessionError(pub String); |
| 28 | |
| 29 | /// Persistent conversations, keyed by an id the host chooses. |
| 30 | /// |
| 31 | /// Contract: |
| 32 | /// - `load` of an unknown id is `Ok(None)`, not an error. |
| 33 | /// - `save` creates the session if it doesn't exist and replaces its |
| 34 | /// messages if it does. |
| 35 | /// - `delete` of an unknown id succeeds. |
| 36 | /// - `list` is ordered most recently updated first. |
| 37 | #[async_trait] |
| 38 | pub trait SessionStore: Send + Sync + 'static { |
| 39 | async fn list(&self) -> Result<Vec<SessionMeta>, SessionError>; |
| 40 | async fn load(&self, id: &str) -> Result<Option<Vec<ChatMessage>>, SessionError>; |
| 41 | async fn save(&self, id: &str, messages: &[ChatMessage]) -> Result<(), SessionError>; |
| 42 | async fn delete(&self, id: &str) -> Result<(), SessionError>; |
| 43 | } |
| 44 | |
| 45 | #[cfg(test)] |
| 46 | mod tests { |
| 47 | use std::collections::HashMap; |
| 48 | use std::sync::Mutex; |
| 49 | |
| 50 | use super::*; |
| 51 | |
| 52 | struct Stored { |
| 53 | messages: Vec<ChatMessage>, |
| 54 | created: i64, |
| 55 | updated: i64, |
| 56 | } |
| 57 | |
| 58 | /// The reference implementation of the contract above. |
| 59 | #[derive(Default)] |
| 60 | struct MemoryStore { |
| 61 | sessions: Mutex<HashMap<String, Stored>>, |
| 62 | clock: Mutex<i64>, |
| 63 | } |
| 64 | |
| 65 | impl MemoryStore { |
| 66 | fn tick(&self) -> i64 { |
| 67 | let mut clock = self.clock.lock().unwrap(); |
| 68 | *clock += 1; |
| 69 | *clock |
| 70 | } |
| 71 | } |
| 72 | |
| 73 | #[async_trait] |
| 74 | impl SessionStore for MemoryStore { |
| 75 | async fn list(&self) -> Result<Vec<SessionMeta>, SessionError> { |
| 76 | let mut metas: Vec<_> = self |
| 77 | .sessions |
| 78 | .lock() |
| 79 | .unwrap() |
| 80 | .iter() |
| 81 | .map(|(id, stored)| SessionMeta { |
| 82 | id: id.clone(), |
| 83 | title: stored |
| 84 | .messages |
| 85 | .first() |
| 86 | .map(|m| m.content.clone()) |
| 87 | .unwrap_or_default(), |
| 88 | created_at: stored.created, |
| 89 | updated_at: stored.updated, |
| 90 | message_count: stored.messages.len(), |
| 91 | }) |
| 92 | .collect(); |
| 93 | metas.sort_by(|a, b| b.updated_at.cmp(&a.updated_at)); |
| 94 | Ok(metas) |
| 95 | } |
| 96 | |
| 97 | async fn load(&self, id: &str) -> Result<Option<Vec<ChatMessage>>, SessionError> { |
| 98 | Ok(self |
| 99 | .sessions |
| 100 | .lock() |
| 101 | .unwrap() |
| 102 | .get(id) |
| 103 | .map(|s| s.messages.clone())) |
| 104 | } |
| 105 | |
| 106 | async fn save(&self, id: &str, messages: &[ChatMessage]) -> Result<(), SessionError> { |
| 107 | let now = self.tick(); |
| 108 | let mut sessions = self.sessions.lock().unwrap(); |
| 109 | let created = sessions.get(id).map_or(now, |s| s.created); |
| 110 | sessions.insert( |
| 111 | id.to_owned(), |
| 112 | Stored { |
| 113 | messages: messages.to_vec(), |
| 114 | created, |
| 115 | updated: now, |
| 116 | }, |
| 117 | ); |
| 118 | Ok(()) |
| 119 | } |
| 120 | |
| 121 | async fn delete(&self, id: &str) -> Result<(), SessionError> { |
| 122 | self.sessions.lock().unwrap().remove(id); |
| 123 | Ok(()) |
| 124 | } |
| 125 | } |
| 126 | |
| 127 | #[tokio::test] |
| 128 | async fn loading_an_unknown_session_is_none() { |
| 129 | let store = MemoryStore::default(); |
| 130 | assert!(store.load("nope").await.unwrap().is_none()); |
| 131 | } |
| 132 | |
| 133 | #[tokio::test] |
| 134 | async fn save_creates_then_replaces() { |
| 135 | let store = MemoryStore::default(); |
| 136 | store.save("a", &[ChatMessage::user("hi")]).await.unwrap(); |
| 137 | store |
| 138 | .save( |
| 139 | "a", |
| 140 | &[ChatMessage::user("hi"), ChatMessage::assistant("yo")], |
| 141 | ) |
| 142 | .await |
| 143 | .unwrap(); |
| 144 | |
| 145 | let loaded = store.load("a").await.unwrap().unwrap(); |
| 146 | assert_eq!(loaded.len(), 2); |
| 147 | let listed = store.list().await.unwrap(); |
| 148 | assert_eq!(listed.len(), 1); |
| 149 | assert_eq!(listed[0].message_count, 2); |
| 150 | assert!(listed[0].updated_at > listed[0].created_at); |
| 151 | } |
| 152 | |
| 153 | #[tokio::test] |
| 154 | async fn list_is_most_recently_updated_first() { |
| 155 | let store = MemoryStore::default(); |
| 156 | store.save("old", &[ChatMessage::user("1")]).await.unwrap(); |
| 157 | store.save("new", &[ChatMessage::user("2")]).await.unwrap(); |
| 158 | store.save("old", &[ChatMessage::user("3")]).await.unwrap(); |
| 159 | |
| 160 | let ids: Vec<_> = store |
| 161 | .list() |
| 162 | .await |
| 163 | .unwrap() |
| 164 | .into_iter() |
| 165 | .map(|m| m.id) |
| 166 | .collect(); |
| 167 | assert_eq!(ids, ["old", "new"]); |
| 168 | } |
| 169 | |
| 170 | #[tokio::test] |
| 171 | async fn delete_is_idempotent() { |
| 172 | let store = MemoryStore::default(); |
| 173 | store.save("a", &[]).await.unwrap(); |
| 174 | store.delete("a").await.unwrap(); |
| 175 | store.delete("a").await.unwrap(); |
| 176 | assert!(store.list().await.unwrap().is_empty()); |
| 177 | } |
| 178 | } |