mockforge_intelligence/ai_studio/
conversation_store.rs

1//! Conversation storage for AI Studio chat sessions
2//!
3//! This module provides persistent storage for conversation history, allowing
4//! multi-turn conversations in the AI Studio. It supports both in-memory
5//! (for development) and file-based (for production) storage.
6
7use crate::ai_studio::chat_orchestrator::{ChatContext, ChatMessage};
8use chrono::{DateTime, Utc};
9use mockforge_foundation::Result;
10use serde::{Deserialize, Serialize};
11use std::collections::HashMap;
12use std::path::{Path, PathBuf};
13use std::sync::Arc;
14use tokio::fs;
15use tokio::sync::RwLock;
16use uuid::Uuid;
17
18/// Conversation storage backend
19pub struct ConversationStore {
20    /// In-memory cache of conversations
21    cache: Arc<RwLock<HashMap<String, Conversation>>>,
22    /// Storage path (if using file-based persistence)
23    storage_path: Option<PathBuf>,
24}
25
26/// A conversation with its history
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct Conversation {
29    /// Unique conversation ID
30    pub id: String,
31    /// Workspace ID this conversation belongs to
32    pub workspace_id: Option<String>,
33    /// Conversation history
34    pub messages: Vec<ChatMessage>,
35    /// When the conversation was created
36    pub created_at: DateTime<Utc>,
37    /// When the conversation was last updated
38    pub updated_at: DateTime<Utc>,
39    /// Optional metadata
40    #[serde(default)]
41    pub metadata: HashMap<String, serde_json::Value>,
42}
43
44impl Conversation {
45    /// Create a new conversation
46    pub fn new(workspace_id: Option<String>) -> Self {
47        let now = Utc::now();
48        Self {
49            id: Uuid::new_v4().to_string(),
50            workspace_id,
51            messages: Vec::new(),
52            created_at: now,
53            updated_at: now,
54            metadata: HashMap::new(),
55        }
56    }
57
58    /// Add a message to the conversation
59    pub fn add_message(&mut self, message: ChatMessage) {
60        self.messages.push(message);
61        self.updated_at = Utc::now();
62    }
63
64    /// Convert to ChatContext
65    pub fn to_context(&self) -> ChatContext {
66        ChatContext {
67            history: self.messages.clone(),
68            workspace_id: self.workspace_id.clone(),
69        }
70    }
71}
72
73impl ConversationStore {
74    /// Create a new in-memory conversation store
75    pub fn new() -> Self {
76        Self {
77            cache: Arc::new(RwLock::new(HashMap::new())),
78            storage_path: None,
79        }
80    }
81
82    /// Create a new conversation store with file-based persistence
83    pub fn with_persistence<P: AsRef<Path>>(storage_path: P) -> Self {
84        Self {
85            cache: Arc::new(RwLock::new(HashMap::new())),
86            storage_path: Some(storage_path.as_ref().to_path_buf()),
87        }
88    }
89
90    /// Initialize the store (load from disk if using persistence)
91    pub async fn initialize(&self) -> Result<()> {
92        if let Some(ref path) = self.storage_path {
93            // Ensure directory exists
94            if let Some(parent) = path.parent() {
95                fs::create_dir_all(parent).await.map_err(|e| {
96                    mockforge_foundation::Error::io_with_context(
97                        "create storage directory",
98                        e.to_string(),
99                    )
100                })?;
101            }
102
103            // Load existing conversations if file exists
104            if path.exists() {
105                let content = fs::read_to_string(path).await.map_err(|e| {
106                    mockforge_foundation::Error::io_with_context(
107                        "read conversation store",
108                        e.to_string(),
109                    )
110                })?;
111
112                let conversations: Vec<Conversation> =
113                    serde_json::from_str(&content).map_err(|e| {
114                        mockforge_foundation::Error::io_with_context(
115                            "parse conversation store",
116                            e.to_string(),
117                        )
118                    })?;
119
120                let mut cache = self.cache.write().await;
121                for conv in conversations {
122                    cache.insert(conv.id.clone(), conv);
123                }
124            }
125        }
126
127        Ok(())
128    }
129
130    /// Save conversations to disk (if using persistence)
131    async fn persist(&self) -> Result<()> {
132        if let Some(ref path) = self.storage_path {
133            let cache = self.cache.read().await;
134            let conversations: Vec<&Conversation> = cache.values().collect();
135
136            let content = serde_json::to_string_pretty(&conversations).map_err(|e| {
137                mockforge_foundation::Error::io_with_context(
138                    "serialize conversations",
139                    e.to_string(),
140                )
141            })?;
142
143            fs::write(path, content).await.map_err(|e| {
144                mockforge_foundation::Error::io_with_context(
145                    "write conversation store",
146                    e.to_string(),
147                )
148            })?;
149        }
150
151        Ok(())
152    }
153
154    /// Create a new conversation
155    pub async fn create_conversation(&self, workspace_id: Option<String>) -> Result<String> {
156        let conversation = Conversation::new(workspace_id);
157        let id = conversation.id.clone();
158
159        {
160            let mut cache = self.cache.write().await;
161            cache.insert(id.clone(), conversation);
162        }
163
164        self.persist().await?;
165        Ok(id)
166    }
167
168    /// Get a conversation by ID
169    pub async fn get_conversation(&self, id: &str) -> Result<Option<Conversation>> {
170        let cache = self.cache.read().await;
171        Ok(cache.get(id).cloned())
172    }
173
174    /// Add a message to a conversation
175    pub async fn add_message(&self, conversation_id: &str, message: ChatMessage) -> Result<()> {
176        let mut cache = self.cache.write().await;
177        if let Some(conversation) = cache.get_mut(conversation_id) {
178            conversation.add_message(message);
179            self.persist().await?;
180            Ok(())
181        } else {
182            Err(mockforge_foundation::Error::not_found("Conversation", conversation_id))
183        }
184    }
185
186    /// Get conversation context for chat
187    pub async fn get_context(&self, conversation_id: &str) -> Result<Option<ChatContext>> {
188        let conversation = self.get_conversation(conversation_id).await?;
189        Ok(conversation.map(|c| c.to_context()))
190    }
191
192    /// List conversations for a workspace
193    pub async fn list_conversations(
194        &self,
195        workspace_id: Option<&str>,
196    ) -> Result<Vec<Conversation>> {
197        let cache = self.cache.read().await;
198        let conversations: Vec<Conversation> = cache
199            .values()
200            .filter(|conv| {
201                if let Some(wid) = workspace_id {
202                    conv.workspace_id.as_deref() == Some(wid)
203                } else {
204                    true
205                }
206            })
207            .cloned()
208            .collect();
209
210        Ok(conversations)
211    }
212
213    /// Delete a conversation
214    pub async fn delete_conversation(&self, conversation_id: &str) -> Result<()> {
215        let mut cache = self.cache.write().await;
216        cache.remove(conversation_id);
217        self.persist().await?;
218        Ok(())
219    }
220
221    /// Clear old conversations (older than specified days)
222    pub async fn clear_old_conversations(&self, days: u64) -> Result<usize> {
223        let cutoff = Utc::now() - chrono::Duration::days(days as i64);
224        let mut cache = self.cache.write().await;
225        let mut removed = 0;
226
227        cache.retain(|_, conv| {
228            if conv.updated_at < cutoff {
229                removed += 1;
230                false
231            } else {
232                true
233            }
234        });
235
236        if removed > 0 {
237            self.persist().await?;
238        }
239
240        Ok(removed)
241    }
242}
243
244impl Default for ConversationStore {
245    fn default() -> Self {
246        Self::new()
247    }
248}
249
250/// Global conversation store instance
251static CONVERSATION_STORE: once_cell::sync::Lazy<Arc<ConversationStore>> =
252    once_cell::sync::Lazy::new(|| {
253        // Use file-based storage in .mockforge directory
254        let storage_path = dirs::home_dir()
255            .map(|home| home.join(".mockforge").join("conversations.json"))
256            .or_else(|| Some(PathBuf::from(".mockforge/conversations.json")));
257
258        if let Some(path) = storage_path {
259            Arc::new(ConversationStore::with_persistence(path))
260        } else {
261            Arc::new(ConversationStore::new())
262        }
263    });
264
265/// Get the global conversation store
266pub fn get_conversation_store() -> Arc<ConversationStore> {
267    CONVERSATION_STORE.clone()
268}
269
270/// Initialize the global conversation store
271pub async fn initialize_conversation_store() -> Result<()> {
272    CONVERSATION_STORE.initialize().await
273}
mockforge_intelligence/ai_studio/conversation_store.rs

mockforge_intelligence/ai_studio/
conversation_store.rs
