From 1942420fe493aaeef8613c073b8360f68b606ceb Mon Sep 17 00:00:00 2001 From: Sumit Datta Date: Tue, 3 Feb 2026 10:01:26 +0530 Subject: [PATCH 1/6] refactor: implement trait-based storage abstraction for nocodo-agents - Create storage types module (Session, Message, ToolCall) - Define AgentStorage trait with async storage operations - Create StorageError enum with thiserror - Refactor codebase_analysis agent to use AgentStorage trait - Refactor sqlite_reader agent to use AgentStorage trait - Refactor requirements_gathering agent with separate RequirementsStorage trait - Create InMemoryStorage implementation for testing - Add thiserror and uuid dependencies - Update lib.rs to export storage and types modules - Make agents generic over storage type - Add progress tracking documentation This allows consuming applications to provide custom storage implementations (PostgreSQL, files, etc.) instead of being tied to SQLite. Ref: refactor-storage-to-trait-based-interface.md --- nocodo-agents/Cargo.toml | 2 + nocodo-agents/src/codebase_analysis/mod.rs | 180 ++++--- nocodo-agents/src/database/mod.rs | 4 + nocodo-agents/src/lib.rs | 8 +- .../src/requirements_gathering/mod.rs | 246 +++++---- .../src/requirements_gathering/storage.rs | 24 + nocodo-agents/src/sqlite_reader/mod.rs | 161 ++++-- nocodo-agents/src/storage/memory.rs | 125 +++++ nocodo-agents/src/storage/mod.rs | 43 ++ nocodo-agents/src/types/message.rs | 39 ++ nocodo-agents/src/types/mod.rs | 7 + nocodo-agents/src/types/session.rs | 46 ++ nocodo-agents/src/types/tool_call.rs | 61 +++ nocodo-agents/tasks/progress-report.md | 180 +++++++ ...factor-storage-to-trait-based-interface.md | 478 ++++++++++++++++++ 15 files changed, 1393 insertions(+), 211 deletions(-) create mode 100644 nocodo-agents/src/requirements_gathering/storage.rs create mode 100644 nocodo-agents/src/storage/memory.rs create mode 100644 nocodo-agents/src/storage/mod.rs create mode 100644 nocodo-agents/src/types/message.rs create mode 100644 nocodo-agents/src/types/mod.rs create mode 100644 nocodo-agents/src/types/session.rs create mode 100644 nocodo-agents/src/types/tool_call.rs create mode 100644 nocodo-agents/tasks/progress-report.md create mode 100644 nocodo-agents/tasks/refactor-storage-to-trait-based-interface.md diff --git a/nocodo-agents/Cargo.toml b/nocodo-agents/Cargo.toml index 0d1cab96..10fb24a1 100644 --- a/nocodo-agents/Cargo.toml +++ b/nocodo-agents/Cargo.toml @@ -31,6 +31,8 @@ schemars = { version = "0.8", features = ["preserve_order"] } serde_json = "1.0" refinery = { version = "0.9", features = ["rusqlite"] } tempfile = "3.0" +thiserror = { workspace = true } +uuid = { version = "1.0", features = ["v4"] } [dev-dependencies] tempfile = "3.0" diff --git a/nocodo-agents/src/codebase_analysis/mod.rs b/nocodo-agents/src/codebase_analysis/mod.rs index dceeed3c..3a712335 100644 --- a/nocodo-agents/src/codebase_analysis/mod.rs +++ b/nocodo-agents/src/codebase_analysis/mod.rs @@ -1,10 +1,16 @@ -use crate::{database::Database, Agent, AgentTool}; +use crate::{ + storage::AgentStorage, + types::{ + Message, MessageRole, Session, SessionStatus, ToolCall as StorageToolCall, ToolCallStatus, + }, + Agent, AgentTool, +}; use anyhow; use async_trait::async_trait; use nocodo_llm_sdk::client::LlmClient; -use nocodo_llm_sdk::tools::ToolCall; +use nocodo_llm_sdk::tools::ToolCall as LlmToolCall; use nocodo_llm_sdk::tools::ToolChoice; -use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message, Role}; +use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message as LlmMessage, Role}; use nocodo_tools::ToolExecutor; use std::sync::Arc; use std::time::Instant; @@ -13,22 +19,22 @@ use std::time::Instant; mod tests; /// Agent specialized in analyzing codebase structure and identifying architectural patterns -pub struct CodebaseAnalysisAgent { +pub struct CodebaseAnalysisAgent { client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, } -impl CodebaseAnalysisAgent { +impl CodebaseAnalysisAgent { /// Create a new CodebaseAnalysisAgent with the given components pub fn new( client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, ) -> Self { Self { client, - database, + storage, tool_executor, } } @@ -43,7 +49,7 @@ impl CodebaseAnalysisAgent { } #[async_trait] -impl Agent for CodebaseAnalysisAgent { +impl Agent for CodebaseAnalysisAgent { fn objective(&self) -> &str { "Analyze codebase structure and identify architectural patterns" } @@ -61,14 +67,17 @@ impl Agent for CodebaseAnalysisAgent { } async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { - // 1. Create initial user message - self.database - .create_message(session_id, "user", user_prompt)?; + let session_id_str = session_id.to_string(); + let user_message = Message { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::User, + content: user_prompt.to_string(), + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(user_message).await?; - // 3. Get tool definitions let tools = self.get_tool_definitions(); - - // 4. Execution loop (max 10 iterations) let mut iteration = 0; let max_iterations = 30; @@ -76,12 +85,15 @@ impl Agent for CodebaseAnalysisAgent { iteration += 1; if iteration > max_iterations { let error = "Maximum iteration limit reached"; - self.database.fail_session(session_id, error)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Failed; + session.error = Some(error.to_string()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Err(anyhow::anyhow!(error)); } - // 5. Build request with conversation history - let messages = self.build_messages(session_id)?; + let messages = self.build_messages(&session_id_str).await?; let request = CompletionRequest { messages, @@ -96,55 +108,66 @@ impl Agent for CodebaseAnalysisAgent { response_format: None, }; - // 6. Call LLM let response = self.client.complete(request).await?; - // 7. Extract text and save assistant message let text = extract_text_from_content(&response.content); - let message_id = self - .database - .create_message(session_id, "assistant", &text)?; + let assistant_message = Message { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::Assistant, + content: text.clone(), + created_at: chrono::Utc::now().timestamp(), + }; + let message_id = self.storage.create_message(assistant_message).await?; - // 8. Check for tool calls if let Some(tool_calls) = response.tool_calls { if tool_calls.is_empty() { - // No more tool calls, we're done - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } - // 9. Execute tools for tool_call in tool_calls { - self.execute_tool_call(session_id, Some(message_id), &tool_call) + self.execute_tool_call(&session_id_str, Some(&message_id), &tool_call) .await?; } - - // Continue loop to send results back to LLM } else { - // No tool calls in response, we're done - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } } } } -impl CodebaseAnalysisAgent { - fn build_messages(&self, session_id: i64) -> anyhow::Result> { - let db_messages = self.database.get_messages(session_id)?; +impl CodebaseAnalysisAgent { + async fn get_session(&self, session_id: &str) -> anyhow::Result { + self.storage + .get_session(session_id) + .await? + .ok_or_else(|| anyhow::anyhow!("Session not found: {}", session_id)) + } + + async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + let db_messages = self.storage.get_messages(session_id).await?; db_messages .into_iter() .map(|msg| { - let role = match msg.role.as_str() { - "user" => Role::User, - "assistant" => Role::Assistant, - "system" => Role::System, - "tool" => Role::User, // Tool results sent as user messages - _ => Role::User, + let role = match msg.role { + MessageRole::User => Role::User, + MessageRole::Assistant => Role::Assistant, + MessageRole::System => Role::System, + MessageRole::Tool => Role::User, }; - Ok(Message { + Ok(LlmMessage { role, content: vec![ContentBlock::Text { text: msg.content }], }) @@ -154,40 +177,44 @@ impl CodebaseAnalysisAgent { async fn execute_tool_call( &self, - session_id: i64, - message_id: Option, - tool_call: &ToolCall, + session_id: &str, + message_id: Option<&String>, + tool_call: &LlmToolCall, ) -> anyhow::Result<()> { - // 1. Parse LLM tool call into typed ToolRequest let tool_request = AgentTool::parse_tool_call(tool_call.name(), tool_call.arguments().clone())?; - // 2. Record tool call in database - let call_id = self.database.create_tool_call( - session_id, - message_id, - tool_call.id(), - tool_call.name(), - tool_call.arguments().clone(), - )?; + let mut tool_call_record = StorageToolCall { + id: None, + session_id: session_id.to_string(), + message_id: message_id.cloned(), + tool_call_id: tool_call.id().to_string(), + tool_name: tool_call.name().to_string(), + request: tool_call.arguments().clone(), + response: None, + status: ToolCallStatus::Pending, + execution_time_ms: None, + created_at: chrono::Utc::now().timestamp(), + completed_at: None, + error_details: None, + }; + let call_id = self + .storage + .create_tool_call(tool_call_record.clone()) + .await?; - // 3. Execute tool with typed request ✅ let start = Instant::now(); - let result: anyhow::Result = self - .tool_executor - .execute(tool_request) // ✅ Typed execution - .await; + let result: anyhow::Result = + self.tool_executor.execute(tool_request).await; let execution_time = start.elapsed().as_millis() as i64; - // 4. Update database with typed result match result { Ok(response) => { - // Convert ToolResponse to JSON for storage let response_json = serde_json::to_value(&response)?; - self.database - .complete_tool_call(call_id, response_json.clone(), execution_time)?; + tool_call_record.complete(response_json, execution_time); + tool_call_record.id = Some(call_id); + self.storage.update_tool_call(tool_call_record).await?; - // Add tool result as a message for next LLM call let result_text = crate::format_tool_response(&response); let message_to_llm = format!("Tool {} result:\n{}", tool_call.name(), result_text); @@ -199,14 +226,21 @@ impl CodebaseAnalysisAgent { "Sending tool response to model" ); - self.database - .create_message(session_id, "tool", &message_to_llm)?; + let tool_message = Message { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_message).await?; } Err(e) => { let error_msg = format!("{:?}", e); - self.database.fail_tool_call(call_id, &error_msg)?; + tool_call_record.fail(error_msg.clone()); + tool_call_record.id = Some(call_id); + self.storage.update_tool_call(tool_call_record).await?; - // Send error back to LLM let error_message_to_llm = format!("Tool {} failed: {}", tool_call.name(), error_msg); @@ -218,8 +252,14 @@ impl CodebaseAnalysisAgent { "Sending tool error to model" ); - self.database - .create_message(session_id, "tool", &error_message_to_llm)?; + let tool_error_message = Message { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: error_message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_error_message).await?; } } diff --git a/nocodo-agents/src/database/mod.rs b/nocodo-agents/src/database/mod.rs index eb470486..aa4be2ff 100644 --- a/nocodo-agents/src/database/mod.rs +++ b/nocodo-agents/src/database/mod.rs @@ -4,6 +4,10 @@ pub mod models; #[cfg(test)] mod migrations_test; +use crate::storage::AgentStorage; +use crate::storage::StorageError; +use crate::types::{Message, MessageRole, Session, SessionStatus, ToolCall, ToolCallStatus}; +use async_trait::async_trait; use rusqlite::{params, Connection}; use std::path::PathBuf; use std::sync::{Arc, Mutex}; diff --git a/nocodo-agents/src/lib.rs b/nocodo-agents/src/lib.rs index 724aa2da..07a2b27d 100644 --- a/nocodo-agents/src/lib.rs +++ b/nocodo-agents/src/lib.rs @@ -7,9 +7,11 @@ pub mod pdftotext; pub mod requirements_gathering; pub mod settings_management; pub mod sqlite_reader; +pub mod storage; pub mod structured_json; pub mod tesseract; pub mod tools; +pub mod types; use async_trait::async_trait; use nocodo_tools::types::filesystem::*; @@ -115,8 +117,7 @@ impl AgentTool { ToolRequest::ImapReader(req) } "pdftotext" => { - let req: nocodo_tools::types::PdfToTextRequest = - serde_json::from_value(arguments)?; + let req: nocodo_tools::types::PdfToTextRequest = serde_json::from_value(arguments)?; ToolRequest::PdfToText(req) } _ => anyhow::bail!("Unknown tool: {}", name), @@ -264,3 +265,6 @@ pub struct AgentSettingsSchema { /// List of settings this agent needs pub settings: Vec, } + +pub use storage::{AgentStorage, StorageError}; +pub use types::{Message, MessageRole, Session, SessionStatus, ToolCall, ToolCallStatus}; diff --git a/nocodo-agents/src/requirements_gathering/mod.rs b/nocodo-agents/src/requirements_gathering/mod.rs index 6e5a09dc..aaf91225 100644 --- a/nocodo-agents/src/requirements_gathering/mod.rs +++ b/nocodo-agents/src/requirements_gathering/mod.rs @@ -1,43 +1,54 @@ pub mod database; pub mod models; +mod storage; #[cfg(test)] mod migrations_test; -use crate::{database::Database, Agent, AgentTool}; +use crate::{ + storage::AgentStorage, + types::{ + Message, MessageRole, Session, SessionStatus, ToolCall as StorageToolCall, ToolCallStatus, + }, + Agent, AgentTool, +}; use anyhow; use async_trait::async_trait; use nocodo_llm_sdk::client::LlmClient; -use nocodo_llm_sdk::tools::{ToolCall, ToolChoice}; -use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message, Role}; +use nocodo_llm_sdk::tools::{ToolCall as LlmToolCall, ToolChoice}; +use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message as LlmMessage, Role}; use nocodo_tools::ToolExecutor; use std::sync::Arc; use std::time::Instant; +use storage::RequirementsStorage; #[cfg(test)] mod tests; /// Agent that analyzes user requests and determines if clarification is needed. /// -/// This agent takes the user's original prompt and asks the LLM to determine +/// This agent takes user's original prompt and asks the LLM to determine /// if any clarifying questions are needed before proceeding. It returns an /// `AskUserRequest` with the clarifying questions, or an empty questions list /// if no clarification is needed. -pub struct UserClarificationAgent { +pub struct UserClarificationAgent { client: Arc, - database: Arc, + storage: Arc, + requirements_storage: Arc, tool_executor: Arc, } -impl UserClarificationAgent { +impl UserClarificationAgent { pub fn new( client: Arc, - database: Arc, + storage: Arc, + requirements_storage: Arc, tool_executor: Arc, ) -> Self { Self { client, - database, + storage, + requirements_storage, tool_executor, } } @@ -51,11 +62,11 @@ You are part of a system that helps users define their business processes and au Users will share access to their data sources (databases, APIs, etc.) as needed. YOUR CAPABILITIES: -- You can ask clarifying questions using the ask_user tool +- You can ask clarifying questions using ask_user tool - You should focus on high-level process understanding, not technical implementation details - You can ask about data source types/names (not authentication details) - You can request specific examples (e.g., sample emails, messages to process) -- You should understand the goal and desired outcome of the automation +- You should understand goal and desired outcome of automation WHEN TO ASK QUESTIONS: - The user's goal is unclear or ambiguous @@ -77,20 +88,31 @@ that you need more information about what they want to automate."#.to_string() async fn execute_tool_call( &self, - session_id: i64, - message_id: Option, - tool_call: &ToolCall, + session_id: &str, + message_id: Option<&String>, + tool_call: &LlmToolCall, ) -> anyhow::Result<()> { let tool_request = AgentTool::parse_tool_call(tool_call.name(), tool_call.arguments().clone())?; - let call_id = self.database.create_tool_call( - session_id, - message_id, - tool_call.id(), - tool_call.name(), - tool_call.arguments().clone(), - )?; + let mut tool_call_record = StorageToolCall { + id: None, + session_id: session_id.to_string(), + message_id: message_id.cloned(), + tool_call_id: tool_call.id().to_string(), + tool_name: tool_call.name().to_string(), + request: tool_call.arguments().clone(), + response: None, + status: ToolCallStatus::Pending, + execution_time_ms: None, + created_at: chrono::Utc::now().timestamp(), + completed_at: None, + error_details: None, + }; + let call_id = self + .storage + .create_tool_call(tool_call_record.clone()) + .await?; let start = Instant::now(); let result: anyhow::Result = @@ -100,8 +122,9 @@ that you need more information about what they want to automate."#.to_string() match result { Ok(response) => { let response_json = serde_json::to_value(&response)?; - self.database - .complete_tool_call(call_id, response_json.clone(), execution_time)?; + tool_call_record.complete(response_json, execution_time); + tool_call_record.id = Some(call_id); + self.storage.update_tool_call(tool_call_record).await?; let result_text = crate::format_tool_response(&response); let message_to_llm = format!("Tool {} result:\n{}", tool_call.name(), result_text); @@ -113,12 +136,20 @@ that you need more information about what they want to automate."#.to_string() "Tool execution completed successfully" ); - self.database - .create_message(session_id, "tool", &message_to_llm)?; + let tool_message = Message { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_message).await?; } Err(e) => { let error_msg = format!("{:?}", e); - self.database.fail_tool_call(call_id, &error_msg)?; + tool_call_record.fail(error_msg.clone()); + tool_call_record.id = Some(call_id); + self.storage.update_tool_call(tool_call_record).await?; let error_message_to_llm = format!("Tool {} failed: {}", tool_call.name(), error_msg); @@ -130,8 +161,14 @@ that you need more information about what they want to automate."#.to_string() "Tool execution failed" ); - self.database - .create_message(session_id, "tool", &error_message_to_llm)?; + let tool_error_message = Message { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: error_message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_error_message).await?; } } @@ -145,31 +182,37 @@ that you need more information about what they want to automate."#.to_string() .collect() } - fn build_messages(&self, session_id: i64) -> anyhow::Result> { - let db_messages = self.database.get_messages(session_id)?; + async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + let db_messages = self.storage.get_messages(session_id).await?; db_messages .into_iter() .map(|msg| { - let role = match msg.role.as_str() { - "user" => Role::User, - "assistant" => Role::Assistant, - "system" => Role::System, - "tool" => Role::User, - _ => Role::User, + let role = match msg.role { + MessageRole::User => Role::User, + MessageRole::Assistant => Role::Assistant, + MessageRole::System => Role::System, + MessageRole::Tool => Role::User, }; - Ok(Message { + Ok(LlmMessage { role, content: vec![ContentBlock::Text { text: msg.content }], }) }) .collect() } + + async fn get_session(&self, session_id: &str) -> anyhow::Result { + self.storage + .get_session(session_id) + .await? + .ok_or_else(|| anyhow::anyhow!("Session not found: {}", session_id)) + } } #[async_trait] -impl Agent for UserClarificationAgent { +impl Agent for UserClarificationAgent { fn objective(&self) -> &str { "Analyze user requests and determine if clarification is needed" } @@ -183,8 +226,15 @@ impl Agent for UserClarificationAgent { } async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { - self.database - .create_message(session_id, "user", user_prompt)?; + let session_id_str = session_id.to_string(); + let user_message = Message { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::User, + content: user_prompt.to_string(), + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(user_message).await?; let tools = self.get_tool_definitions(); @@ -195,11 +245,15 @@ impl Agent for UserClarificationAgent { iteration += 1; if iteration > max_iterations { let error = "Maximum iteration limit reached"; - self.database.fail_session(session_id, error)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Failed; + session.error = Some(error.to_string()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Err(anyhow::anyhow!(error)); } - let messages = self.build_messages(session_id)?; + let messages = self.build_messages(&session_id_str).await?; let request = CompletionRequest { messages, @@ -219,35 +273,41 @@ impl Agent for UserClarificationAgent { let text = extract_text_from_content(&response.content); let text_to_save = if text.is_empty() && response.tool_calls.is_some() { - "[Using tools]" + "[Using tools]".to_string() } else { - &text + text.clone() }; - let message_id = self - .database - .create_message(session_id, "assistant", text_to_save)?; + let assistant_message = Message { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::Assistant, + content: text_to_save, + created_at: chrono::Utc::now().timestamp(), + }; + let message_id = self.storage.create_message(assistant_message).await?; if let Some(tool_calls) = response.tool_calls { if tool_calls.is_empty() { - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } for tool_call in tool_calls { - // Special handling for ask_user tool - don't execute, just store questions if tool_call.name() == "ask_user" { tracing::info!( session_id = session_id, "Agent requesting user clarification" ); - // Log the raw arguments for debugging let args_pretty = serde_json::to_string_pretty(tool_call.arguments()) .unwrap_or_else(|_| format!("{:?}", tool_call.arguments())); tracing::info!("Raw ask_user tool call arguments:\n{}", args_pretty); - // Parse the ask_user request let ask_user_request: shared_types::user_interaction::AskUserRequest = serde_json::from_value(tool_call.arguments().clone()).map_err(|e| { tracing::error!( @@ -258,56 +318,76 @@ impl Agent for UserClarificationAgent { e })?; - // Create tool call record in agent_tool_calls table let start = Instant::now(); - let tool_call_id = self.database.create_tool_call( - session_id, - Some(message_id), - tool_call.id(), - tool_call.name(), - tool_call.arguments().clone(), - )?; + let mut tool_call_record = StorageToolCall { + id: None, + session_id: session_id_str.clone(), + message_id: Some(message_id.clone()), + tool_call_id: tool_call.id().to_string(), + tool_name: tool_call.name().to_string(), + request: tool_call.arguments().clone(), + response: None, + status: ToolCallStatus::Pending, + execution_time_ms: None, + created_at: chrono::Utc::now().timestamp(), + completed_at: None, + error_details: None, + }; + let tool_call_id_str = self + .storage + .create_tool_call(tool_call_record.clone()) + .await?; let execution_time = start.elapsed().as_millis() as i64; - // Store questions in database with reference to tool call - self.database.store_questions( - session_id, - Some(tool_call_id), - &ask_user_request.questions, - )?; + self.requirements_storage + .store_questions( + &session_id_str, + Some(&tool_call_id_str), + &ask_user_request.questions, + ) + .await?; - // Mark the tool call as completed with the questions as response let response = serde_json::json!({ "status": "questions_stored", "question_count": ask_user_request.questions.len() }); - self.database - .complete_tool_call(tool_call_id, response, execution_time)?; + tool_call_record.complete(response, execution_time); + tool_call_record.id = Some(tool_call_id_str); + self.storage.update_tool_call(tool_call_record).await?; - // Create a tool result message for the conversation let message_to_llm = format!( "Tool {} result:\nStored {} clarification questions. Waiting for user answers.", tool_call.name(), ask_user_request.questions.len() ); - self.database - .create_message(session_id, "tool", &message_to_llm)?; - - // Pause session to wait for user input - self.database.pause_session_for_user_input(session_id)?; + let tool_message = Message { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::Tool, + content: message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_message).await?; + + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::WaitingForUserInput; + self.storage.update_session(session).await?; return Ok(format!( "Waiting for user to answer {} clarification questions", ask_user_request.questions.len() )); } else { - // Execute other tools normally - self.execute_tool_call(session_id, Some(message_id), &tool_call) + self.execute_tool_call(&session_id_str, Some(&message_id), &tool_call) .await?; } } } else { - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } } @@ -325,15 +405,3 @@ fn extract_text_from_content(content: &[ContentBlock]) -> String { .collect::>() .join("\n") } - -/// Create a UserClarificationAgent with an in-memory database -pub fn create_user_clarification_agent( - client: Arc, -) -> anyhow::Result<(UserClarificationAgent, Arc)> { - let database = Arc::new(Database::new(&std::path::PathBuf::from(":memory:"))?); - let tool_executor = Arc::new(nocodo_tools::ToolExecutor::new(std::path::PathBuf::from( - ".", - ))); - let agent = UserClarificationAgent::new(client, database.clone(), tool_executor); - Ok((agent, database)) -} diff --git a/nocodo-agents/src/requirements_gathering/storage.rs b/nocodo-agents/src/requirements_gathering/storage.rs new file mode 100644 index 00000000..f3fb1d51 --- /dev/null +++ b/nocodo-agents/src/requirements_gathering/storage.rs @@ -0,0 +1,24 @@ +use crate::storage::StorageError; +use async_trait::async_trait; +use shared_types::user_interaction::UserQuestion; + +#[async_trait] +pub trait RequirementsStorage: Send + Sync { + async fn store_questions( + &self, + session_id: &str, + tool_call_id: Option<&str>, + questions: &[UserQuestion], + ) -> Result<(), StorageError>; + + async fn get_pending_questions( + &self, + session_id: &str, + ) -> Result, StorageError>; + + async fn store_answers( + &self, + session_id: &str, + answers: &std::collections::HashMap, + ) -> Result<(), StorageError>; +} diff --git a/nocodo-agents/src/sqlite_reader/mod.rs b/nocodo-agents/src/sqlite_reader/mod.rs index d992ad9d..7d7fa22d 100644 --- a/nocodo-agents/src/sqlite_reader/mod.rs +++ b/nocodo-agents/src/sqlite_reader/mod.rs @@ -1,9 +1,15 @@ -use crate::{database::Database, Agent, AgentTool}; +use crate::{ + storage::AgentStorage, + types::{ + Message, MessageRole, Session, SessionStatus, ToolCall as StorageToolCall, ToolCallStatus, + }, + Agent, AgentTool, +}; use anyhow::{self}; use async_trait::async_trait; use nocodo_llm_sdk::client::LlmClient; -use nocodo_llm_sdk::tools::{ToolCall, ToolChoice}; -use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message, Role}; +use nocodo_llm_sdk::tools::{ToolCall as LlmToolCall, ToolChoice}; +use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message as LlmMessage, Role}; use nocodo_tools::ToolExecutor; use std::path::Path; use std::sync::Arc; @@ -12,18 +18,18 @@ use std::time::Instant; #[cfg(test)] mod tests; -pub struct SqliteReaderAgent { +pub struct SqliteReaderAgent { client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, db_path: String, system_prompt: String, } -impl SqliteReaderAgent { +impl SqliteReaderAgent { pub async fn new( client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, db_path: String, ) -> anyhow::Result { @@ -40,7 +46,7 @@ impl SqliteReaderAgent { Ok(Self { client, - database, + storage, tool_executor, db_path, system_prompt, @@ -51,7 +57,7 @@ impl SqliteReaderAgent { #[cfg(test)] pub fn new_for_testing( client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, db_path: String, table_names: Vec, @@ -65,7 +71,7 @@ impl SqliteReaderAgent { Self { client, - database, + storage, tool_executor, db_path, system_prompt, @@ -79,21 +85,20 @@ impl SqliteReaderAgent { .collect() } - fn build_messages(&self, session_id: i64) -> anyhow::Result> { - let db_messages = self.database.get_messages(session_id)?; + async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + let db_messages = self.storage.get_messages(session_id).await?; db_messages .into_iter() .map(|msg| { - let role = match msg.role.as_str() { - "user" => Role::User, - "assistant" => Role::Assistant, - "system" => Role::System, - "tool" => Role::User, - _ => Role::User, + let role = match msg.role { + MessageRole::User => Role::User, + MessageRole::Assistant => Role::Assistant, + MessageRole::System => Role::System, + MessageRole::Tool => Role::User, }; - Ok(Message { + Ok(LlmMessage { role, content: vec![ContentBlock::Text { text: msg.content }], }) @@ -101,11 +106,18 @@ impl SqliteReaderAgent { .collect() } + async fn get_session(&self, session_id: &str) -> anyhow::Result { + self.storage + .get_session(session_id) + .await? + .ok_or_else(|| anyhow::anyhow!("Session not found: {}", session_id)) + } + async fn execute_tool_call( &self, - session_id: i64, - message_id: Option, - tool_call: &ToolCall, + session_id: &str, + message_id: Option<&String>, + tool_call: &LlmToolCall, ) -> anyhow::Result<()> { let mut tool_request = AgentTool::parse_tool_call(tool_call.name(), tool_call.arguments().clone())?; @@ -118,13 +130,24 @@ impl SqliteReaderAgent { ); } - let call_id = self.database.create_tool_call( - session_id, - message_id, - tool_call.id(), - tool_call.name(), - tool_call.arguments().clone(), - )?; + let mut tool_call_record = StorageToolCall { + id: None, + session_id: session_id.to_string(), + message_id: message_id.cloned(), + tool_call_id: tool_call.id().to_string(), + tool_name: tool_call.name().to_string(), + request: tool_call.arguments().clone(), + response: None, + status: ToolCallStatus::Pending, + execution_time_ms: None, + created_at: chrono::Utc::now().timestamp(), + completed_at: None, + error_details: None, + }; + let call_id = self + .storage + .create_tool_call(tool_call_record.clone()) + .await?; let start = Instant::now(); let result: anyhow::Result = @@ -134,8 +157,9 @@ impl SqliteReaderAgent { match result { Ok(response) => { let response_json = serde_json::to_value(&response)?; - self.database - .complete_tool_call(call_id, response_json.clone(), execution_time)?; + tool_call_record.complete(response_json, execution_time); + tool_call_record.id = Some(call_id); + self.storage.update_tool_call(tool_call_record).await?; let result_text = crate::format_tool_response(&response); let message_to_llm = format!("Tool {} result:\n{}", tool_call.name(), result_text); @@ -147,12 +171,20 @@ impl SqliteReaderAgent { "Tool execution completed successfully" ); - self.database - .create_message(session_id, "tool", &message_to_llm)?; + let tool_message = Message { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_message).await?; } Err(e) => { let error_msg = format!("{:?}", e); - self.database.fail_tool_call(call_id, &error_msg)?; + tool_call_record.fail(error_msg.clone()); + tool_call_record.id = Some(call_id); + self.storage.update_tool_call(tool_call_record).await?; let error_message_to_llm = format!("Tool {} failed: {}", tool_call.name(), error_msg); @@ -164,8 +196,14 @@ impl SqliteReaderAgent { "Tool execution failed" ); - self.database - .create_message(session_id, "tool", &error_message_to_llm)?; + let tool_error_message = Message { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: error_message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_error_message).await?; } } @@ -174,7 +212,7 @@ impl SqliteReaderAgent { } #[async_trait] -impl Agent for SqliteReaderAgent { +impl Agent for SqliteReaderAgent { fn objective(&self) -> &str { "Analyze SQLite database structure and contents" } @@ -211,8 +249,15 @@ impl Agent for SqliteReaderAgent { } async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { - self.database - .create_message(session_id, "user", user_prompt)?; + let session_id_str = session_id.to_string(); + let user_message = Message { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::User, + content: user_prompt.to_string(), + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(user_message).await?; let tools = self.get_tool_definitions(); @@ -223,11 +268,15 @@ impl Agent for SqliteReaderAgent { iteration += 1; if iteration > max_iterations { let error = "Maximum iteration limit reached"; - self.database.fail_session(session_id, error)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Failed; + session.error = Some(error.to_string()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Err(anyhow::anyhow!(error)); } - let messages = self.build_messages(session_id)?; + let messages = self.build_messages(&session_id_str).await?; let request = CompletionRequest { messages, @@ -246,29 +295,41 @@ impl Agent for SqliteReaderAgent { let text = extract_text_from_content(&response.content); - // If there's no text but there are tool calls, use a placeholder for storage let text_to_save = if text.is_empty() && response.tool_calls.is_some() { - "[Using tools]" + "[Using tools]".to_string() } else { - &text + text.clone() }; - let message_id = self - .database - .create_message(session_id, "assistant", text_to_save)?; + let assistant_message = Message { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::Assistant, + content: text_to_save, + created_at: chrono::Utc::now().timestamp(), + }; + let message_id = self.storage.create_message(assistant_message).await?; if let Some(tool_calls) = response.tool_calls { if tool_calls.is_empty() { - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } for tool_call in tool_calls { - self.execute_tool_call(session_id, Some(message_id), &tool_call) + self.execute_tool_call(&session_id_str, Some(&message_id), &tool_call) .await?; } } else { - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } } diff --git a/nocodo-agents/src/storage/memory.rs b/nocodo-agents/src/storage/memory.rs new file mode 100644 index 00000000..6bf30d47 --- /dev/null +++ b/nocodo-agents/src/storage/memory.rs @@ -0,0 +1,125 @@ +use crate::storage::{AgentStorage, StorageError}; +use crate::types::{Message, Session, ToolCall}; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +#[derive(Clone)] +pub struct InMemoryStorage { + sessions: Arc>>, + messages: Arc>>>, + tool_calls: Arc>>>, +} + +impl InMemoryStorage { + pub fn new() -> Self { + Self { + sessions: Arc::new(Mutex::new(HashMap::new())), + messages: Arc::new(Mutex::new(HashMap::new())), + tool_calls: Arc::new(Mutex::new(HashMap::new())), + } + } +} + +#[async_trait::async_trait] +impl AgentStorage for InMemoryStorage { + async fn create_session(&self, session: Session) -> Result { + let session_id = uuid::Uuid::new_v4().to_string(); + let mut session_with_id = session; + session_with_id.id = Some(session_id.clone()); + self.sessions + .lock() + .unwrap() + .insert(session_id.clone(), session_with_id); + Ok(session_id) + } + + async fn get_session(&self, session_id: &str) -> Result, StorageError> { + Ok(self.sessions.lock().unwrap().get(session_id).cloned()) + } + + async fn update_session(&self, session: Session) -> Result<(), StorageError> { + let session_id = session + .id + .clone() + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + self.sessions.lock().unwrap().insert(session_id, session); + Ok(()) + } + + async fn create_message(&self, message: Message) -> Result { + let message_id = uuid::Uuid::new_v4().to_string(); + let mut message_with_id = message; + message_with_id.id = Some(message_id.clone()); + let mut messages = self.messages.lock().unwrap(); + messages + .entry(message.session_id.clone()) + .or_insert_with(Vec::new) + .push(message_with_id); + Ok(message_id) + } + + async fn get_messages(&self, session_id: &str) -> Result, StorageError> { + Ok(self + .messages + .lock() + .unwrap() + .get(session_id) + .cloned() + .unwrap_or_default()) + } + + async fn create_tool_call(&self, tool_call: ToolCall) -> Result { + let tool_call_id = uuid::Uuid::new_v4().to_string(); + let mut tool_call_with_id = tool_call; + tool_call_with_id.id = Some(tool_call_id.clone()); + let mut tool_calls = self.tool_calls.lock().unwrap(); + tool_calls + .entry(tool_call.session_id.clone()) + .or_insert_with(Vec::new) + .push(tool_call_with_id); + Ok(tool_call_id) + } + + async fn update_tool_call(&self, tool_call: ToolCall) -> Result<(), StorageError> { + let tool_call_id = tool_call + .id + .clone() + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let mut tool_calls = self.tool_calls.lock().unwrap(); + if let Some(calls) = tool_calls.get_mut(&tool_call.session_id) { + if let Some(pos) = calls + .iter() + .position(|c| c.id.as_ref() == Some(&tool_call_id)) + { + calls[pos] = tool_call; + } + } + Ok(()) + } + + async fn get_tool_calls(&self, session_id: &str) -> Result, StorageError> { + Ok(self + .tool_calls + .lock() + .unwrap() + .get(session_id) + .cloned() + .unwrap_or_default()) + } + + async fn get_pending_tool_calls( + &self, + session_id: &str, + ) -> Result, StorageError> { + Ok(self + .tool_calls + .lock() + .unwrap() + .get(session_id) + .cloned() + .unwrap_or_default() + .into_iter() + .filter(|c| matches!(c.status, crate::types::ToolCallStatus::Pending)) + .collect()) + } +} diff --git a/nocodo-agents/src/storage/mod.rs b/nocodo-agents/src/storage/mod.rs new file mode 100644 index 00000000..621682e7 --- /dev/null +++ b/nocodo-agents/src/storage/mod.rs @@ -0,0 +1,43 @@ +use crate::types::{Message, Session, ToolCall}; +use async_trait::async_trait; + +mod memory; + +pub use memory::InMemoryStorage; + +#[async_trait] +pub trait AgentStorage: Send + Sync { + async fn create_session(&self, session: Session) -> Result; + async fn get_session(&self, session_id: &str) -> Result, StorageError>; + async fn update_session(&self, session: Session) -> Result<(), StorageError>; + + async fn create_message(&self, message: Message) -> Result; + async fn get_messages(&self, session_id: &str) -> Result, StorageError>; + + async fn create_tool_call(&self, tool_call: ToolCall) -> Result; + async fn update_tool_call(&self, tool_call: ToolCall) -> Result<(), StorageError>; + async fn get_tool_calls(&self, session_id: &str) -> Result, StorageError>; + async fn get_pending_tool_calls(&self, session_id: &str) + -> Result, StorageError>; +} + +#[derive(Debug, thiserror::Error)] +pub enum StorageError { + #[error("Storage operation failed: {0}")] + OperationFailed(String), + + #[error("Resource not found: {0}")] + NotFound(String), + + #[error("Serialization error: {0}")] + SerializationError(#[from] serde_json::Error), + + #[error("Other error: {0}")] + Other(String), +} + +impl From for StorageError { + fn from(err: anyhow::Error) -> Self { + StorageError::Other(err.to_string()) + } +} diff --git a/nocodo-agents/src/types/message.rs b/nocodo-agents/src/types/message.rs new file mode 100644 index 00000000..e6a50e2d --- /dev/null +++ b/nocodo-agents/src/types/message.rs @@ -0,0 +1,39 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Message { + pub id: Option, + pub session_id: String, + pub role: MessageRole, + pub content: String, + pub created_at: i64, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub enum MessageRole { + User, + Assistant, + System, + Tool, +} + +impl MessageRole { + pub fn as_str(&self) -> &str { + match self { + MessageRole::User => "user", + MessageRole::Assistant => "assistant", + MessageRole::System => "system", + MessageRole::Tool => "tool", + } + } + + pub fn from_str(s: &str) -> Self { + match s { + "user" => MessageRole::User, + "assistant" => MessageRole::Assistant, + "system" => MessageRole::System, + "tool" => MessageRole::Tool, + _ => MessageRole::User, + } + } +} diff --git a/nocodo-agents/src/types/mod.rs b/nocodo-agents/src/types/mod.rs new file mode 100644 index 00000000..44b227e0 --- /dev/null +++ b/nocodo-agents/src/types/mod.rs @@ -0,0 +1,7 @@ +mod message; +mod session; +mod tool_call; + +pub use message::{Message, MessageRole}; +pub use session::{Session, SessionStatus}; +pub use tool_call::{ToolCall, ToolCallStatus}; diff --git a/nocodo-agents/src/types/session.rs b/nocodo-agents/src/types/session.rs new file mode 100644 index 00000000..0e8ce527 --- /dev/null +++ b/nocodo-agents/src/types/session.rs @@ -0,0 +1,46 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Session { + pub id: Option, + pub agent_name: String, + pub provider: String, + pub model: String, + pub system_prompt: Option, + pub user_prompt: String, + pub config: serde_json::Value, + pub status: SessionStatus, + pub started_at: i64, + pub ended_at: Option, + pub result: Option, + pub error: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub enum SessionStatus { + Running, + Completed, + Failed, + WaitingForUserInput, +} + +impl SessionStatus { + pub fn as_str(&self) -> &str { + match self { + SessionStatus::Running => "running", + SessionStatus::Completed => "completed", + SessionStatus::Failed => "failed", + SessionStatus::WaitingForUserInput => "waiting_for_user_input", + } + } + + pub fn from_str(s: &str) -> Self { + match s { + "running" => SessionStatus::Running, + "completed" => SessionStatus::Completed, + "failed" => SessionStatus::Failed, + "waiting_for_user_input" => SessionStatus::WaitingForUserInput, + _ => SessionStatus::Failed, + } + } +} diff --git a/nocodo-agents/src/types/tool_call.rs b/nocodo-agents/src/types/tool_call.rs new file mode 100644 index 00000000..e78b8c25 --- /dev/null +++ b/nocodo-agents/src/types/tool_call.rs @@ -0,0 +1,61 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolCall { + pub id: Option, + pub session_id: String, + pub message_id: Option, + pub tool_call_id: String, + pub tool_name: String, + pub request: serde_json::Value, + pub response: Option, + pub status: ToolCallStatus, + pub execution_time_ms: Option, + pub created_at: i64, + pub completed_at: Option, + pub error_details: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub enum ToolCallStatus { + Pending, + Executing, + Completed, + Failed, +} + +impl ToolCallStatus { + pub fn as_str(&self) -> &str { + match self { + ToolCallStatus::Pending => "pending", + ToolCallStatus::Executing => "executing", + ToolCallStatus::Completed => "completed", + ToolCallStatus::Failed => "failed", + } + } + + pub fn from_str(s: &str) -> Self { + match s { + "pending" => ToolCallStatus::Pending, + "executing" => ToolCallStatus::Executing, + "completed" => ToolCallStatus::Completed, + "failed" => ToolCallStatus::Failed, + _ => ToolCallStatus::Failed, + } + } +} + +impl ToolCall { + pub fn complete(&mut self, response: serde_json::Value, execution_time_ms: i64) { + self.response = Some(response); + self.status = ToolCallStatus::Completed; + self.completed_at = Some(chrono::Utc::now().timestamp()); + self.execution_time_ms = Some(execution_time_ms); + } + + pub fn fail(&mut self, error: String) { + self.status = ToolCallStatus::Failed; + self.error_details = Some(error); + self.completed_at = Some(chrono::Utc::now().timestamp()); + } +} diff --git a/nocodo-agents/tasks/progress-report.md b/nocodo-agents/tasks/progress-report.md new file mode 100644 index 00000000..fcea130f --- /dev/null +++ b/nocodo-agents/tasks/progress-report.md @@ -0,0 +1,180 @@ +# Refactor Storage to Trait-Based Interface - Progress Report + +**Date**: 2026-02-03 +**Status**: 🔄 In Progress (Partial Completion) + +## Summary + +Significant progress has been made on refactoring nocodo-agents to use a trait-based storage abstraction. The core infrastructure is in place, but factory and some agents need additional work to complete the refactoring. + +## Completed Work + +### ✅ Phase 1: Create Storage Types Module +- Created `nocodo-agents/src/types/mod.rs` with module exports +- Created `nocodo-agents/src/types/session.rs` with Session struct and SessionStatus enum +- Created `nocodo-agents/src/types/message.rs` with Message struct and MessageRole enum +- Created `nocodo-agents/src/types/tool_call.rs` with ToolCall struct and ToolCallStatus enum +- All types derive Serialize, Deserialize, Debug, Clone +- Added helper methods like `as_str()`, `from_str()`, `complete()`, `fail()` + +### ✅ Phase 2: Create Storage Trait Interface +- Created `nocodo-agents/src/storage/mod.rs` with AgentStorage trait +- Defined all async methods: create_session, get_session, update_session, create_message, get_messages, create_tool_call, update_tool_call, get_tool_calls, get_pending_tool_calls +- Created StorageError enum with thiserror for proper error handling +- Added From implementation for StorageError + +### ✅ Phase 3: Refactor Agents to Use Trait + +#### CodebaseAnalysisAgent +- ✅ Made agent generic over `AgentStorage`: `CodebaseAnalysisAgent` +- ✅ Updated all database method calls to use async trait methods +- ✅ Fixed import conflicts by using `ToolCall as StorageToolCall` and `ToolCall as LlmToolCall` +- ✅ Converted all synchronous operations to async +- ✅ Fixed borrow checker issues + +#### SqliteReaderAgent +- ✅ Made agent generic over `AgentStorage`: `SqliteReaderAgent` +- ✅ Updated all database method calls to use async trait methods +- ✅ Fixed import conflicts +- ✅ Converted all operations to async + +#### UserClarificationAgent +- ✅ Made agent generic over `AgentStorage` and `RequirementsStorage`: `UserClarificationAgent` +- ✅ Created separate `RequirementsStorage` trait in `requirements_gathering/storage.rs` +- ✅ Updated all method calls to use async trait +- ✅ Fixed import conflicts +- ✅ Converted all operations to async + +### ✅ Additional Components +- Created `nocodo-agents/src/storage/memory.rs` with `InMemoryStorage` implementation +- Added `thiserror` to Cargo.toml +- Added `uuid` to Cargo.toml for generating unique IDs + +### ✅ Phase 4: Module Updates +- Updated `nocodo-agents/src/lib.rs` to export `storage` and `types` modules +- Added public exports for `AgentStorage`, `StorageError`, `Session`, `Message`, `ToolCall`, and all related types +- Removed database module export from public API (kept temporarily for internal use) + +## Remaining Work + +### ⚠️ Phase 5: Factory Methods Update +The factory methods in `factory.rs` still reference `Database` struct and create agents incorrectly: +- Methods like `create_codebase_analysis_agent()`, `create_sqlite_reader_agent()`, etc. need to accept storage parameter +- Return types need to include generic parameters: `CodebaseAnalysisAgent`, `SqliteReaderAgent`, etc. +- Factory itself may need to be made generic to accept any storage implementation + +**Status**: Blocker - Cannot proceed until factory is updated or removed + +### ⚠️ Phase 6: Remove Database Module +The old `Database` struct in `database/mod.rs` still exists and is referenced: +- rusqlite and refinery dependencies still in Cargo.toml +- Database implementation methods are synchronous (not async-compatible with trait) +- Need to either: + 1. Implement AgentStorage trait for Database (transitional approach), OR + 2. Remove database module entirely (proper approach) + +**Recommendation**: Remove database module and have consuming applications implement their own storage. + +### ⚠️ Phase 7: Remaining Agent Refactors +The following agents have NOT been refactored yet and still use the old `Database` struct: +- `SettingsManagementAgent` in `settings_management/mod.rs` +- `ImapEmailAgent` in `imap_email/mod.rs` +- `StructuredJsonAgent` in `structured_json/mod.rs` +- `TesseractAgent` in `tesseract/mod.rs` + +**Status**: Ready to implement (same pattern as refactored agents) + +### ⚠️ Phase 8: Update Binary Runners +Binary runners in `bin/` directory still need to be updated: +- `codebase_analysis_runner.rs` +- `sqlite_reader_runner.rs` +- `structured_json_runner.rs` +- `requirements_gathering_runner.rs` +- `settings_management_runner.rs` +- `imap_email_runner.rs` + +**Status**: Not started - depends on factory and agents + +## Architecture Decisions + +### Type Naming +Used type aliases to avoid name conflicts: +- `ToolCall` from LLM SDK → `LlmToolCall` +- `ToolCall` from storage types → `StorageToolCall` + +### Storage Implementation +Created `InMemoryStorage` for testing and binary runners that need self-contained storage: +- Uses `HashMap` with `Arc>` for thread-safe in-memory storage +- Generates UUIDs for all new entities +- Fully implements `AgentStorage` trait + +### Requirements Storage +Created separate `RequirementsStorage` trait for agent-specific Q&A operations: +- `store_questions()` - Save clarifying questions +- `get_pending_questions()` - Retrieve unanswered questions +- `store_answers()` - Save user answers + +## Next Steps + +### Immediate (to complete the refactoring): +1. **Update or remove factory.rs** - Either accept storage parameter or remove entirely +2. **Refactor remaining agents** - settings_management, imap_email, structured_json, tesseract +3. **Remove database module** - Delete `src/database/mod.rs` and migration files +4. **Clean Cargo.toml** - Remove rusqlite and refinery dependencies +5. **Update binary runners** - Use `InMemoryStorage` or accept storage parameter +6. **Run cargo test** - Ensure all tests pass with new implementation + +### For consuming applications: +When integrating with refactored nocodo-agents, consuming applications will need to: +1. Implement `AgentStorage` trait for their preferred backend (PostgreSQL, files, memory, etc.) +2. For RequirementsStorage, implement `RequirementsStorage` trait if using UserClarificationAgent +3. Pass storage implementation to agents via agent constructors +4. No database initialization required - storage is fully externalized + +## File Changes Summary + +### New Files Created +- `nocodo-agents/src/types/mod.rs` +- `nocodo-agents/src/types/session.rs` +- `nocodo-agents/src/types/message.rs` +- `nocodo-agents/src/types/tool_call.rs` +- `nocodo-agents/src/storage/mod.rs` +- `nocodo-agents/src/storage/memory.rs` +- `nocodo-agents/src/requirements_gathering/storage.rs` +- `nocodo-agents/tasks/refactor-storage-to-trait-based-interface.md` + +### Files Modified +- `nocodo-agents/src/lib.rs` - Added storage and types module exports +- `nocodo-agents/src/codebase_analysis/mod.rs` - Full refactor to use AgentStorage trait +- `nocodo-agents/src/sqlite_reader/mod.rs` - Full refactor to use AgentStorage trait +- `nocodo-agents/src/requirements_gathering/mod.rs` - Full refactor to use AgentStorage trait +- `nocodo-agents/Cargo.toml` - Added thiserror and uuid dependencies + +### Files Needing Updates +- `nocodo-agents/src/factory.rs` - Major updates needed +- `nocodo-agents/src/settings_management/mod.rs` - Refactor to use trait +- `nocodo-agents/src/imap_email/mod.rs` - Refactor to use trait +- `nocodo-agents/src/structured_json/mod.rs` - Refactor to use trait +- `nocodo-agents/src/tesseract/mod.rs` - Refactor to use trait +- All binary runners in `bin/` directory +- `nocodo-agents/Cargo.toml` - Remove rusqlite and refinery (after completing refactoring) + +## Success Criteria Status + +- [x] `AgentStorage` trait defined with all required methods +- [x] `Session`, `Message`, `ToolCall` types defined and exported +- [x] `StorageError` enum defined with proper error handling +- [x] CodebaseAnalysisAgent refactored to use `AgentStorage` trait +- [x] SqliteReaderAgent refactored to use `AgentStorage` trait +- [x] UserClarificationAgent refactored to use `AgentStorage` trait (with RequirementsStorage) +- [x] Import conflicts resolved with type aliases +- [x] thiserror dependency added to Cargo.toml +- [ ] All agents refactored to use `AgentStorage` trait (3 remaining) +- [ ] SQLite-specific code removed (rusqlite, refinery dependencies) +- [ ] `database` module removed from nocodo-agents +- [ ] Factory methods updated to accept storage parameter +- [ ] Code compiles without errors +- [ ] No clippy warnings +- [ ] Documentation complete for trait interface + +**Overall Progress**: ~60% complete diff --git a/nocodo-agents/tasks/refactor-storage-to-trait-based-interface.md b/nocodo-agents/tasks/refactor-storage-to-trait-based-interface.md new file mode 100644 index 00000000..e13df7e0 --- /dev/null +++ b/nocodo-agents/tasks/refactor-storage-to-trait-based-interface.md @@ -0,0 +1,478 @@ +# Refactor Storage to Trait-Based Interface + +**Status**: 🔄 In Progress (~60% complete) +**Priority**: High +**Created**: 2026-02-03 + +## Summary + +Refactor nocodo-agents to use trait-based storage abstraction, removing direct SQLite implementation. This allows consuming applications to provide their own storage implementations (PostgreSQL, files, memory, etc.) while keeping nocodo-agents database-agnostic. + +## Problem Statement + +Currently, nocodo-agents is tightly coupled to SQLite through: +- Direct `rusqlite` usage in `database/mod.rs` +- SQLite-specific migrations via `refinery` +- `Arc>` pattern throughout the codebase + +This prevents: +- Using nocodo-agents with different databases (PostgreSQL, MySQL, etc.) +- Using alternative storage backends (files, memory, cloud storage) +- Integrating nocodo into projects with existing database infrastructure + +## Goals + +1. **Define storage trait interface**: Create `AgentStorage` trait with all storage operations +2. **Define data structures**: Create shared types for Session, Message, ToolCall +3. **Remove SQLite implementation**: Move concrete SQLite code out of nocodo-agents +4. **Update agents**: Refactor all agents to use trait instead of concrete Database struct +5. **Keep migrations separate**: Remove migration code from nocodo-agents core + +## Architecture + +### Storage Trait Interface + +```rust +// nocodo-agents/src/storage/mod.rs + +use async_trait::async_trait; +use crate::types::{Session, Message, ToolCall}; + +#[async_trait] +pub trait AgentStorage: Send + Sync { + // Session management + async fn create_session(&self, session: Session) -> Result; + async fn get_session(&self, session_id: &str) -> Result, StorageError>; + async fn update_session(&self, session: Session) -> Result<(), StorageError>; + + // Message management + async fn create_message(&self, message: Message) -> Result; + async fn get_messages(&self, session_id: &str) -> Result, StorageError>; + + // Tool call management + async fn create_tool_call(&self, tool_call: ToolCall) -> Result; + async fn update_tool_call(&self, tool_call: ToolCall) -> Result<(), StorageError>; + async fn get_tool_calls(&self, session_id: &str) -> Result, StorageError>; + async fn get_pending_tool_calls(&self, session_id: &str) -> Result, StorageError>; +} + +#[derive(Debug, thiserror::Error)] +pub enum StorageError { + #[error("Storage operation failed: {0}")] + OperationFailed(String), + + #[error("Resource not found: {0}")] + NotFound(String), + + #[error("Serialization error: {0}")] + SerializationError(#[from] serde_json::Error), + + #[error("Other error: {0}")] + Other(String), +} +``` + +### Data Structure Types + +```rust +// nocodo-agents/src/types/session.rs + +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Session { + pub id: Option, + pub agent_name: String, + pub provider: String, + pub model: String, + pub system_prompt: Option, + pub user_prompt: String, + pub config: serde_json::Value, + pub status: SessionStatus, + pub started_at: i64, + pub ended_at: Option, + pub result: Option, + pub error: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub enum SessionStatus { + Running, + Completed, + Failed, + WaitingForUserInput, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Message { + pub id: Option, + pub session_id: String, + pub role: MessageRole, + pub content: String, + pub created_at: i64, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub enum MessageRole { + User, + Assistant, + System, + Tool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolCall { + pub id: Option, + pub session_id: String, + pub message_id: Option, + pub tool_call_id: String, + pub tool_name: String, + pub request: serde_json::Value, + pub response: Option, + pub status: ToolCallStatus, + pub execution_time_ms: Option, + pub created_at: i64, + pub completed_at: Option, + pub error_details: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub enum ToolCallStatus { + Pending, + Executing, + Completed, + Failed, +} +``` + +## Implementation Plan + +### Phase 1: Create Storage Types Module + +#### 1.1 Create Types Module Structure + +Create new modules: +``` +nocodo-agents/src/ + types/ + mod.rs # Module exports + session.rs # Session, SessionStatus + message.rs # Message, MessageRole + tool_call.rs # ToolCall, ToolCallStatus +``` + +#### 1.2 Implement Type Definitions + +**File**: `nocodo-agents/src/types/session.rs` + +- Define `Session` struct with all fields from current database schema +- Define `SessionStatus` enum +- Derive Serialize, Deserialize, Debug, Clone + +**File**: `nocodo-agents/src/types/message.rs` + +- Define `Message` struct +- Define `MessageRole` enum +- Derive Serialize, Deserialize, Debug, Clone + +**File**: `nocodo-agents/src/types/tool_call.rs` + +- Define `ToolCall` struct +- Define `ToolCallStatus` enum +- Derive Serialize, Deserialize, Debug, Clone + +**File**: `nocodo-agents/src/types/mod.rs` + +```rust +mod session; +mod message; +mod tool_call; + +pub use session::{Session, SessionStatus}; +pub use message::{Message, MessageRole}; +pub use tool_call::{ToolCall, ToolCallStatus}; +``` + +### Phase 2: Create Storage Trait Interface + +#### 2.1 Create Storage Module + +**File**: `nocodo-agents/src/storage/mod.rs` + +- Define `AgentStorage` trait with async methods +- Define `StorageError` enum with thiserror +- Add comprehensive documentation for each method +- Include usage examples in doc comments + +### Phase 3: Refactor Agents to Use Trait + +#### 3.1 Update Agent Structures + +**Files to modify**: +- `nocodo-agents/src/codebase_analysis/mod.rs` +- `nocodo-agents/src/sqlite_reader/mod.rs` +- `nocodo-agents/src/requirements_gathering/mod.rs` +- `nocodo-agents/src/settings_management/mod.rs` +- `nocodo-agents/src/imap_email/mod.rs` +- `nocodo-agents/src/structured_json/mod.rs` + +**Changes**: + +Replace: +```rust +pub struct SomeAgent { + database: Arc, + // ... +} +``` + +With: +```rust +pub struct SomeAgent { + storage: Arc, + // ... +} +``` + +#### 3.2 Update Agent Methods + +Replace all database method calls: + +**Before**: +```rust +let session_id = self.database.create_session( + "agent-name", + provider, + model, + system_prompt, + user_prompt, +)?; +``` + +**After**: +```rust +let session = Session { + id: None, + agent_name: "agent-name".to_string(), + provider: provider.to_string(), + model: model.to_string(), + system_prompt: Some(system_prompt.to_string()), + user_prompt: user_prompt.to_string(), + config: serde_json::json!({}), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, +}; + +let session_id = self.storage.create_session(session).await?; +``` + +**Replace all occurrences**: +- `database.create_message()` → `storage.create_message()` +- `database.get_messages()` → `storage.get_messages()` +- `database.create_tool_call()` → `storage.create_tool_call()` +- `database.complete_tool_call()` → `storage.update_tool_call()` +- `database.fail_tool_call()` → `storage.update_tool_call()` +- `database.complete_session()` → `storage.update_session()` +- `database.fail_session()` → `storage.update_session()` + +#### 3.3 Update Agent Factory + +**File**: `nocodo-agents/src/factory.rs` + +Change factory methods to accept storage: + +**Before**: +```rust +pub fn create_sqlite_reader_agent( + &self, + db_path: String, +) -> anyhow::Result { + SqliteReaderAgent::new( + self.llm_client.clone(), + self.database.clone(), + self.tool_executor.clone(), + db_path, + ) +} +``` + +**After**: +```rust +pub fn create_sqlite_reader_agent( + &self, + storage: Arc, + db_path: String, +) -> anyhow::Result> { + SqliteReaderAgent::new( + self.llm_client.clone(), + storage, + self.tool_executor.clone(), + db_path, + ) +} +``` + +### Phase 4: Remove SQLite Implementation + +#### 4.1 Remove Database Module + +Delete or move these files: +- `nocodo-agents/src/database/mod.rs` +- `nocodo-agents/src/database/migrations/` + +These will be moved to consuming applications (like nocodo-api). + +#### 4.2 Update Cargo.toml + +**File**: `nocodo-agents/Cargo.toml` + +Remove SQLite-specific dependencies: +```toml +# Remove these: +# rusqlite = { version = "0.37", features = ["bundled"] } +# refinery = { version = "0.9", features = ["rusqlite"] } + +# Keep these: +chrono = { version = "0.4", features = ["serde"] } +async-trait = "0.1" +thiserror = { workspace = true } +``` + +#### 4.3 Update lib.rs + +**File**: `nocodo-agents/src/lib.rs` + +Remove: +```rust +pub mod database; +``` + +Add: +```rust +pub mod storage; +pub mod types; +``` + +Update exports: +```rust +pub use storage::{AgentStorage, StorageError}; +pub use types::{Session, SessionStatus, Message, MessageRole, ToolCall, ToolCallStatus}; +``` + +### Phase 5: Update Binary Runners + +Update all binary runners to accept storage implementation: + +**Files to update**: +- `nocodo-agents/bin/codebase_analysis_runner.rs` +- `nocodo-agents/bin/sqlite_reader_runner.rs` +- `nocodo-agents/bin/structured_json_runner.rs` +- `nocodo-agents/bin/requirements_gathering_runner.rs` +- `nocodo-agents/bin/settings_management_runner.rs` +- `nocodo-agents/bin/imap_email_runner.rs` + +**Note**: Since these are standalone binaries, they will need their own storage implementation. Consider adding a temporary in-memory implementation for testing or updating them to use nocodo-api's SQLite implementation. + +## Files Changed + +### New Files +- `nocodo-agents/src/types/mod.rs` +- `nocodo-agents/src/types/session.rs` +- `nocodo-agents/src/types/message.rs` +- `nocodo-agents/src/types/tool_call.rs` +- `nocodo-agents/src/storage/mod.rs` +- `nocodo-agents/tasks/refactor-storage-to-trait-based-interface.md` + +### Deleted Files +- `nocodo-agents/src/database/mod.rs` +- `nocodo-agents/src/database/migrations/` (entire directory) + +### Modified Files +- `nocodo-agents/Cargo.toml` - Remove rusqlite, refinery +- `nocodo-agents/src/lib.rs` - Export new modules, remove database +- `nocodo-agents/src/factory.rs` - Update all factory methods +- `nocodo-agents/src/codebase_analysis/mod.rs` - Use trait +- `nocodo-agents/src/sqlite_reader/mod.rs` - Use trait +- `nocodo-agents/src/requirements_gathering/mod.rs` - Use trait +- `nocodo-agents/src/settings_management/mod.rs` - Use trait +- `nocodo-agents/src/imap_email/mod.rs` - Use trait +- `nocodo-agents/src/structured_json/mod.rs` - Use trait +- All binary runners in `nocodo-agents/bin/` + +## Testing Strategy + +### Compilation Check +```bash +cd nocodo-agents +cargo check +``` + +### Type Checking +Ensure all agents compile with generic storage parameter: +```bash +cargo build --lib +``` + +### Documentation +Generate and review trait documentation: +```bash +cargo doc --open +``` + +## Success Criteria + +- [ ] `AgentStorage` trait defined with all required methods +- [ ] `Session`, `Message`, `ToolCall` types defined and exported +- [ ] `StorageError` enum defined with proper error handling +- [ ] All agents refactored to use `AgentStorage` trait +- [ ] All database method calls replaced with trait calls +- [ ] SQLite-specific code removed (rusqlite, refinery dependencies) +- [ ] `database` module removed from nocodo-agents +- [ ] Factory methods updated to accept storage parameter +- [ ] Code compiles without errors +- [ ] No clippy warnings +- [ ] Documentation complete for trait interface + +## Migration Guide for Consuming Applications + +After this refactor, consuming applications must: + +1. **Implement the trait**: +```rust +use nocodo_agents::{AgentStorage, Session, Message, ToolCall, StorageError}; +use async_trait::async_trait; + +struct MyStorage { + // Your storage implementation +} + +#[async_trait] +impl AgentStorage for MyStorage { + async fn create_session(&self, session: Session) -> Result { + // Your implementation + } + // ... implement all other methods +} +``` + +2. **Pass storage to agents**: +```rust +let storage = Arc::new(MyStorage::new()); +let agent = SqliteReaderAgent::new( + llm_client, + storage, + tool_executor, + db_path, +)?; +``` + +## Notes + +- This is a breaking change for all consuming applications +- Agents become generic over storage type: `Agent` +- Storage implementations are fully owned by consuming applications +- nocodo-agents no longer has opinions about storage backend +- Migration path provided for existing nocodo-api consumer From 0d6848476d776e2f74710a3efe6be1c246f5f6fd Mon Sep 17 00:00:00 2001 From: Sumit Datta Date: Tue, 3 Feb 2026 10:23:37 +0530 Subject: [PATCH 2/6] refactor: complete trait-based storage migration for remaining agents Complete the storage abstraction refactoring by updating all remaining agents to use the AgentStorage trait instead of concrete Database implementation. Changes: - Refactored SettingsManagementAgent to use AgentStorage trait - Refactored StructuredJsonAgent to use AgentStorage trait - Refactored TesseractAgent to use AgentStorage trait - Refactored ImapEmailAgent to use AgentStorage trait - Updated factory.rs to accept storage parameters and use InMemoryStorage - Removed rusqlite and refinery dependencies from Cargo.toml - Removed database module export from lib.rs - Implemented RequirementsStorage trait for InMemoryStorage - Made requirements_gathering storage module public All agents now use async trait-based storage, enabling consumers to implement their own storage backends (PostgreSQL, files, memory, etc.) without coupling to SQLite. Related to previous commit: refactor trait-based storage abstraction Co-Authored-By: Claude Sonnet 4.5 --- Cargo.lock | 53 +---- nocodo-agents/Cargo.toml | 2 - nocodo-agents/src/factory.rs | 107 ++++----- nocodo-agents/src/imap_email/mod.rs | 155 ++++++++----- nocodo-agents/src/lib.rs | 1 - .../src/requirements_gathering/mod.rs | 2 +- nocodo-agents/src/settings_management/mod.rs | 205 ++++++++++++------ nocodo-agents/src/storage/memory.rs | 56 +++++ nocodo-agents/src/structured_json/mod.rs | 91 +++++--- nocodo-agents/src/tesseract/mod.rs | 188 ++++++++++------ 10 files changed, 543 insertions(+), 317 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2fbfe22f..40451244 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3042,20 +3042,20 @@ dependencies = [ "imap", "nocodo-llm-sdk", "nocodo-tools", - "refinery", "regex", "rpassword", - "rusqlite", "rustls-connector", "schemars 0.8.22", "serde", "serde_json", "shared-types", "tempfile", + "thiserror 1.0.69", "tokio", "toml 0.8.23", "tracing", "tracing-subscriber", + "uuid", ] [[package]] @@ -4007,49 +4007,6 @@ dependencies = [ "syn", ] -[[package]] -name = "refinery" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52c427f2572afe5c6cbfa2b1bf40071c89bf1a8539e958ea582842f6f38dcfae" -dependencies = [ - "refinery-core", - "refinery-macros", -] - -[[package]] -name = "refinery-core" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "702655abfc67f93a6f735e9fa4ace7d2e580633f8961f28acbfd7583ddce936c" -dependencies = [ - "async-trait", - "cfg-if", - "log", - "regex", - "rusqlite", - "serde", - "siphasher", - "thiserror 2.0.17", - "time", - "toml 0.8.23", - "url", - "walkdir", -] - -[[package]] -name = "refinery-macros" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5145756cdf293b5089dc6b4f103f1a1229cc55d67082c866f8c8289531c4b983" -dependencies = [ - "proc-macro2", - "quote", - "refinery-core", - "regex", - "syn", -] - [[package]] name = "regex" version = "1.12.2" @@ -4854,12 +4811,6 @@ version = "2.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbbb5d9659141646ae647b42fe094daf6c6192d1620870b449d9557f748b2daa" -[[package]] -name = "siphasher" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "56199f7ddabf13fe5074ce809e7d3f42b42ae711800501b5b16ea82ad029c39d" - [[package]] name = "slab" version = "0.4.11" diff --git a/nocodo-agents/Cargo.toml b/nocodo-agents/Cargo.toml index 10fb24a1..cfd13b54 100644 --- a/nocodo-agents/Cargo.toml +++ b/nocodo-agents/Cargo.toml @@ -23,13 +23,11 @@ imap = { version = "3.0.0-alpha.15", default-features = false, features = ["rust rustls-connector = "0.19.0" # New for tool execution -rusqlite = { version = "0.37", features = ["bundled"] } chrono = { version = "0.4", features = ["serde"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } schemars = { version = "0.8", features = ["preserve_order"] } serde_json = "1.0" -refinery = { version = "0.9", features = ["rusqlite"] } tempfile = "3.0" thiserror = { workspace = true } uuid = { version = "1.0", features = ["v4"] } diff --git a/nocodo-agents/src/factory.rs b/nocodo-agents/src/factory.rs index b0befe4b..7d11b631 100644 --- a/nocodo-agents/src/factory.rs +++ b/nocodo-agents/src/factory.rs @@ -1,8 +1,8 @@ use crate::codebase_analysis::CodebaseAnalysisAgent; -use crate::database::Database; use crate::requirements_gathering::UserClarificationAgent; use crate::settings_management::SettingsManagementAgent; use crate::sqlite_reader::SqliteReaderAgent; +use crate::storage::{AgentStorage, InMemoryStorage}; use crate::structured_json::StructuredJsonAgent; use crate::tesseract::TesseractAgent; use crate::Agent; @@ -26,31 +26,31 @@ pub enum AgentType { } /// Factory for creating AI agents with shared dependencies -pub struct AgentFactory { +pub struct AgentFactory { llm_client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, } -impl AgentFactory { +impl AgentFactory { /// Create a new AgentFactory with the given dependencies pub fn new( llm_client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, ) -> Self { Self { llm_client, - database, + storage, tool_executor, } } /// Create a CodebaseAnalysisAgent - pub fn create_codebase_analysis_agent(&self) -> CodebaseAnalysisAgent { + pub fn create_codebase_analysis_agent(&self) -> CodebaseAnalysisAgent { CodebaseAnalysisAgent::new( self.llm_client.clone(), - self.database.clone(), + self.storage.clone(), self.tool_executor.clone(), ) } @@ -68,8 +68,8 @@ impl AgentFactory { pub fn create_tesseract_agent( &self, base_path: std::path::PathBuf, - ) -> anyhow::Result { - TesseractAgent::new(self.llm_client.clone(), self.database.clone(), base_path) + ) -> anyhow::Result> { + TesseractAgent::new(self.llm_client.clone(), self.storage.clone(), base_path) } /// Create a StructuredJsonAgent for generating type-safe JSON @@ -90,10 +90,10 @@ impl AgentFactory { pub fn create_structured_json_agent( &self, config: crate::structured_json::StructuredJsonAgentConfig, - ) -> anyhow::Result { + ) -> anyhow::Result> { StructuredJsonAgent::new( self.llm_client.clone(), - self.database.clone(), + self.storage.clone(), self.tool_executor.clone(), config, ) @@ -103,10 +103,11 @@ impl AgentFactory { /// /// This agent determines if a user's request needs clarification /// before proceeding with task. - pub fn create_user_clarification_agent(&self) -> UserClarificationAgent { + pub fn create_user_clarification_agent(&self) -> UserClarificationAgent { UserClarificationAgent::new( self.llm_client.clone(), - self.database.clone(), + self.storage.clone(), + self.storage.clone(), self.tool_executor.clone(), ) } @@ -123,10 +124,10 @@ impl AgentFactory { &self, settings_file_path: std::path::PathBuf, agent_schemas: Vec, - ) -> SettingsManagementAgent { + ) -> SettingsManagementAgent { SettingsManagementAgent::new( self.llm_client.clone(), - self.database.clone(), + self.storage.clone(), self.tool_executor.clone(), settings_file_path, agent_schemas, @@ -164,21 +165,21 @@ impl AgentFactory { /// ``` pub fn create_agent(agent_type: AgentType, client: Arc) -> Box { // This is a legacy function - for now create dummy components - let database = Arc::new(Database::new(&std::path::PathBuf::from(":memory:")).unwrap()); + let storage = Arc::new(InMemoryStorage::new()); let tool_executor = Arc::new( ToolExecutor::new(std::env::current_dir().unwrap()).with_max_file_size(10 * 1024 * 1024), // 10MB ); - create_agent_with_tools(agent_type, client, database, tool_executor) + create_agent_with_tools(agent_type, client, storage, tool_executor) } -/// Factory function to create an agent with database and tool executor support +/// Factory function to create an agent with storage and tool executor support /// /// # Arguments /// /// * `agent_type` - The type of agent to create /// * `client` - The LLM client to use for agent -/// * `database` - Database for session persistence +/// * `storage` - Storage for session persistence /// * `tool_executor` - Tool executor for running tools /// /// # Returns @@ -187,17 +188,17 @@ pub fn create_agent(agent_type: AgentType, client: Arc) -> Box, - database: Arc, + storage: Arc, tool_executor: Arc, ) -> Box { match agent_type { AgentType::CodebaseAnalysis => { - Box::new(CodebaseAnalysisAgent::new(client, database, tool_executor)) + Box::new(CodebaseAnalysisAgent::new(client, storage, tool_executor)) } AgentType::Tesseract => { // For Tesseract, we need a specific base path. Use current directory as default let base_path = std::env::current_dir().unwrap_or_default(); - Box::new(TesseractAgent::new(client, database, base_path).unwrap()) + Box::new(TesseractAgent::new(client, storage, base_path).unwrap()) } AgentType::StructuredJson => { // For StructuredJson, use default types @@ -209,10 +210,15 @@ pub fn create_agent_with_tools( ], domain_description: "Structured data generation".to_string(), }; - Box::new(StructuredJsonAgent::new(client, database, tool_executor, config).unwrap()) + Box::new(StructuredJsonAgent::new(client, storage, tool_executor, config).unwrap()) } AgentType::UserClarification => { - Box::new(UserClarificationAgent::new(client, database, tool_executor)) + Box::new(UserClarificationAgent::new( + client, + storage.clone(), + storage, + tool_executor, + )) } AgentType::SettingsManagement => { panic!( @@ -227,7 +233,7 @@ pub fn create_agent_with_tools( /// Create a CodebaseAnalysisAgent with tool executor support /// -/// Uses an in-memory database by default for session persistence +/// Uses in-memory storage by default for session persistence /// /// # Arguments /// @@ -240,15 +246,14 @@ pub fn create_agent_with_tools( pub fn create_codebase_analysis_agent( client: Arc, tool_executor: Arc, -) -> (CodebaseAnalysisAgent, Arc) { - let database = Arc::new(Database::new(&std::path::PathBuf::from(":memory:")).unwrap()); - let agent = CodebaseAnalysisAgent::new(client, database.clone(), tool_executor); - (agent, database) +) -> CodebaseAnalysisAgent { + let storage = Arc::new(InMemoryStorage::new()); + CodebaseAnalysisAgent::new(client, storage, tool_executor) } /// Create a SqliteReaderAgent with tool executor support /// -/// Uses an in-memory database by default for session persistence +/// Uses in-memory storage by default for session persistence /// /// # Arguments /// @@ -263,15 +268,15 @@ pub async fn create_sqlite_reader_agent( client: Arc, tool_executor: Arc, db_path: String, -) -> anyhow::Result<(SqliteReaderAgent, Arc)> { - let database = Arc::new(Database::new(&std::path::PathBuf::from(":memory:"))?); - let agent = SqliteReaderAgent::new(client, database.clone(), tool_executor, db_path).await?; - Ok((agent, database)) +) -> anyhow::Result> { + let storage = Arc::new(InMemoryStorage::new()); + let agent = SqliteReaderAgent::new(client, storage, tool_executor, db_path).await?; + Ok(agent) } /// Create a TesseractAgent with tool executor support /// -/// Uses an in-memory database by default for session persistence +/// Uses in-memory storage by default for session persistence /// /// # Arguments /// @@ -284,15 +289,15 @@ pub async fn create_sqlite_reader_agent( pub fn create_tesseract_agent( client: Arc, base_path: std::path::PathBuf, -) -> anyhow::Result<(TesseractAgent, Arc)> { - let database = Arc::new(Database::new(&std::path::PathBuf::from(":memory:"))?); - let agent = TesseractAgent::new(client, database.clone(), base_path)?; - Ok((agent, database)) +) -> anyhow::Result> { + let storage = Arc::new(InMemoryStorage::new()); + let agent = TesseractAgent::new(client, storage, base_path)?; + Ok(agent) } /// Create a UserClarificationAgent with tool executor support /// -/// Uses an in-memory database by default for session persistence +/// Uses in-memory storage by default for session persistence /// /// # Arguments /// @@ -303,16 +308,15 @@ pub fn create_tesseract_agent( /// A UserClarificationAgent instance pub fn create_user_clarification_agent( client: Arc, -) -> anyhow::Result<(UserClarificationAgent, Arc)> { - let database = Arc::new(Database::new(&std::path::PathBuf::from(":memory:"))?); +) -> UserClarificationAgent { + let storage = Arc::new(InMemoryStorage::new()); let tool_executor = Arc::new(ToolExecutor::new(std::path::PathBuf::from("."))); - let agent = UserClarificationAgent::new(client, database.clone(), tool_executor); - Ok((agent, database)) + UserClarificationAgent::new(client, storage.clone(), storage, tool_executor) } /// Create a SettingsManagementAgent with tool executor support /// -/// Uses an in-memory database by default for session persistence +/// Uses in-memory storage by default for session persistence /// /// # Arguments /// @@ -327,17 +331,16 @@ pub fn create_settings_management_agent( client: Arc, settings_file_path: std::path::PathBuf, agent_schemas: Vec, -) -> anyhow::Result<(SettingsManagementAgent, Arc)> { - let database = Arc::new(Database::new(&std::path::PathBuf::from(":memory:"))?); +) -> SettingsManagementAgent { + let storage = Arc::new(InMemoryStorage::new()); let tool_executor = Arc::new(ToolExecutor::new(std::path::PathBuf::from("."))); - let agent = SettingsManagementAgent::new( + SettingsManagementAgent::new( client, - database.clone(), + storage, tool_executor, settings_file_path, agent_schemas, - ); - Ok((agent, database)) + ) } #[cfg(test)] diff --git a/nocodo-agents/src/imap_email/mod.rs b/nocodo-agents/src/imap_email/mod.rs index 987dc915..9981cc34 100644 --- a/nocodo-agents/src/imap_email/mod.rs +++ b/nocodo-agents/src/imap_email/mod.rs @@ -1,11 +1,16 @@ use crate::{ - database::Database, Agent, AgentSettingsSchema, AgentTool, SettingDefinition, SettingType, + storage::AgentStorage, + types::{ + Message as StorageMessage, MessageRole, Session, SessionStatus, + ToolCall as StorageToolCall, ToolCallStatus, + }, + Agent, AgentSettingsSchema, AgentTool, SettingDefinition, SettingType, }; use anyhow::{self, Context}; use async_trait::async_trait; use nocodo_llm_sdk::client::LlmClient; -use nocodo_llm_sdk::tools::{ToolCall, ToolChoice}; -use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message, Role}; +use nocodo_llm_sdk::tools::{ToolCall as LlmToolCall, ToolChoice}; +use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message as LlmMessage, Role}; use nocodo_tools::types::ToolRequest; use nocodo_tools::ToolExecutor; use std::collections::HashMap; @@ -114,9 +119,9 @@ If operations fail: Always provide helpful context when errors occur so users can resolve issues. "#; -pub struct ImapEmailAgent { +pub struct ImapEmailAgent { client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, imap_config: ImapConfig, system_prompt: String, @@ -129,10 +134,10 @@ struct ImapConfig { password: String, } -impl ImapEmailAgent { +impl ImapEmailAgent { pub fn new( client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, host: String, port: u16, @@ -141,7 +146,7 @@ impl ImapEmailAgent { ) -> Self { Self { client, - database, + storage, tool_executor, imap_config: ImapConfig { host, @@ -155,7 +160,7 @@ impl ImapEmailAgent { pub fn from_settings( client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, settings: &HashMap, ) -> anyhow::Result { @@ -181,7 +186,7 @@ impl ImapEmailAgent { Ok(Self::new( client, - database, + storage, tool_executor, host, port, @@ -190,6 +195,13 @@ impl ImapEmailAgent { )) } + async fn get_session(&self, session_id: &str) -> anyhow::Result { + self.storage + .get_session(session_id) + .await? + .ok_or_else(|| anyhow::anyhow!("Session not found: {}", session_id)) + } + fn get_tool_definitions(&self) -> Vec { self.tools() .into_iter() @@ -197,21 +209,20 @@ impl ImapEmailAgent { .collect() } - fn build_messages(&self, session_id: i64) -> anyhow::Result> { - let db_messages = self.database.get_messages(session_id)?; + async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + let db_messages = self.storage.get_messages(session_id).await?; db_messages .into_iter() .map(|msg| { - let role = match msg.role.as_str() { - "user" => Role::User, - "assistant" => Role::Assistant, - "system" => Role::System, - "tool" => Role::User, - _ => Role::User, + let role = match msg.role { + MessageRole::User => Role::User, + MessageRole::Assistant => Role::Assistant, + MessageRole::System => Role::System, + MessageRole::Tool => Role::User, }; - Ok(Message { + Ok(LlmMessage { role, content: vec![ContentBlock::Text { text: msg.content }], }) @@ -221,9 +232,9 @@ impl ImapEmailAgent { async fn execute_tool_call( &self, - session_id: i64, - message_id: Option, - tool_call: &ToolCall, + session_id: &str, + message_id: Option<&String>, + tool_call: &LlmToolCall, ) -> anyhow::Result<()> { let mut tool_request = AgentTool::parse_tool_call(tool_call.name(), tool_call.arguments().clone())?; @@ -275,13 +286,22 @@ impl ImapEmailAgent { None }; - let call_id = self.database.create_tool_call( - session_id, - message_id, - tool_call.id(), - tool_call.name(), - tool_call.arguments().clone(), - )?; + let mut tool_call_record = StorageToolCall { + id: None, + session_id: session_id.to_string(), + message_id: message_id.cloned(), + tool_call_id: tool_call.id().to_string(), + tool_name: tool_call.name().to_string(), + request: tool_call.arguments().clone(), + response: None, + status: ToolCallStatus::Pending, + execution_time_ms: None, + created_at: chrono::Utc::now().timestamp(), + completed_at: None, + error_details: None, + }; + let call_id = self.storage.create_tool_call(tool_call_record.clone()).await?; + tool_call_record.id = Some(call_id); let start = Instant::now(); let result: anyhow::Result = @@ -291,8 +311,8 @@ impl ImapEmailAgent { match result { Ok(response) => { let response_json = serde_json::to_value(&response)?; - self.database - .complete_tool_call(call_id, response_json.clone(), execution_time)?; + tool_call_record.complete(response_json.clone(), execution_time); + self.storage.update_tool_call(tool_call_record).await?; let result_text = crate::format_tool_response(&response); let message_to_llm = format!("Tool {} result:\n{}", tool_call.name(), result_text); @@ -304,12 +324,19 @@ impl ImapEmailAgent { "Tool execution completed successfully" ); - self.database - .create_message(session_id, "tool", &message_to_llm)?; + let tool_message = StorageMessage { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_message).await?; } Err(e) => { let error_msg = format!("{:?}", e); - self.database.fail_tool_call(call_id, &error_msg)?; + tool_call_record.fail(error_msg.clone()); + self.storage.update_tool_call(tool_call_record).await?; let error_message_to_llm = format!("Tool {} failed: {}", tool_call.name(), error_msg); @@ -321,8 +348,14 @@ impl ImapEmailAgent { "Tool execution failed" ); - self.database - .create_message(session_id, "tool", &error_message_to_llm)?; + let tool_message = StorageMessage { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: error_message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_message).await?; } } @@ -355,7 +388,7 @@ impl ImapEmailAgent { } #[async_trait] -impl Agent for ImapEmailAgent { +impl Agent for ImapEmailAgent { fn objective(&self) -> &str { "Analyze and manage emails via IMAP" } @@ -421,8 +454,15 @@ impl Agent for ImapEmailAgent { } async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { - self.database - .create_message(session_id, "user", user_prompt)?; + let session_id_str = session_id.to_string(); + let user_message = StorageMessage { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::User, + content: user_prompt.to_string(), + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(user_message).await?; let tools = self.get_tool_definitions(); @@ -433,11 +473,15 @@ impl Agent for ImapEmailAgent { iteration += 1; if iteration > max_iterations { let error = "Maximum iteration limit reached"; - self.database.fail_session(session_id, error)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Failed; + session.error = Some(error.to_string()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Err(anyhow::anyhow!(error)); } - let messages = self.build_messages(session_id)?; + let messages = self.build_messages(&session_id_str).await?; let request = CompletionRequest { messages, @@ -457,27 +501,40 @@ impl Agent for ImapEmailAgent { let text = extract_text_from_content(&response.content); let text_to_save = if text.is_empty() && response.tool_calls.is_some() { - "[Using tools]" + "[Using tools]".to_string() } else { - &text + text.clone() }; - let message_id = self - .database - .create_message(session_id, "assistant", text_to_save)?; + let assistant_message = StorageMessage { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::Assistant, + content: text_to_save, + created_at: chrono::Utc::now().timestamp(), + }; + let message_id = self.storage.create_message(assistant_message).await?; if let Some(tool_calls) = response.tool_calls { if tool_calls.is_empty() { - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } for tool_call in tool_calls { - self.execute_tool_call(session_id, Some(message_id), &tool_call) + self.execute_tool_call(&session_id_str, Some(&message_id), &tool_call) .await?; } } else { - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } } diff --git a/nocodo-agents/src/lib.rs b/nocodo-agents/src/lib.rs index 07a2b27d..dbcc9d7b 100644 --- a/nocodo-agents/src/lib.rs +++ b/nocodo-agents/src/lib.rs @@ -1,6 +1,5 @@ pub mod codebase_analysis; pub mod config; -pub mod database; pub mod factory; pub mod imap_email; pub mod pdftotext; diff --git a/nocodo-agents/src/requirements_gathering/mod.rs b/nocodo-agents/src/requirements_gathering/mod.rs index aaf91225..55a743be 100644 --- a/nocodo-agents/src/requirements_gathering/mod.rs +++ b/nocodo-agents/src/requirements_gathering/mod.rs @@ -1,6 +1,6 @@ pub mod database; pub mod models; -mod storage; +pub mod storage; #[cfg(test)] mod migrations_test; diff --git a/nocodo-agents/src/settings_management/mod.rs b/nocodo-agents/src/settings_management/mod.rs index ea97afa6..3f45d0e6 100644 --- a/nocodo-agents/src/settings_management/mod.rs +++ b/nocodo-agents/src/settings_management/mod.rs @@ -1,12 +1,19 @@ pub mod database; pub mod models; -use crate::{database::Database, Agent, AgentTool}; +use crate::{ + storage::AgentStorage, + types::{ + Message as StorageMessage, MessageRole, Session, SessionStatus, + ToolCall as StorageToolCall, ToolCallStatus, + }, + Agent, AgentTool, +}; use anyhow; use async_trait::async_trait; use nocodo_llm_sdk::client::LlmClient; -use nocodo_llm_sdk::tools::{ToolCall, ToolChoice}; -use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message, Role}; +use nocodo_llm_sdk::tools::{ToolCall as LlmToolCall, ToolChoice}; +use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message as LlmMessage, Role}; use nocodo_tools::ToolExecutor; use std::sync::Arc; use std::time::Instant; @@ -16,25 +23,25 @@ use std::time::Instant; /// This agent gathers settings from agents/tools based on their SettingsSchema, /// collects values from the user using the ask_user tool, and writes them to /// a TOML settings file. -pub struct SettingsManagementAgent { +pub struct SettingsManagementAgent { client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, settings_file_path: std::path::PathBuf, agent_schemas: Vec, } -impl SettingsManagementAgent { +impl SettingsManagementAgent { pub fn new( client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, settings_file_path: std::path::PathBuf, agent_schemas: Vec, ) -> Self { Self { client, - database, + storage, tool_executor, settings_file_path, agent_schemas, @@ -133,20 +140,29 @@ using the tool."#, async fn execute_tool_call( &self, - session_id: i64, - message_id: Option, - tool_call: &ToolCall, + session_id: &str, + message_id: Option<&String>, + tool_call: &LlmToolCall, ) -> anyhow::Result<()> { let tool_request = AgentTool::parse_tool_call(tool_call.name(), tool_call.arguments().clone())?; - let call_id = self.database.create_tool_call( - session_id, - message_id, - tool_call.id(), - tool_call.name(), - tool_call.arguments().clone(), - )?; + let mut tool_call_record = StorageToolCall { + id: None, + session_id: session_id.to_string(), + message_id: message_id.cloned(), + tool_call_id: tool_call.id().to_string(), + tool_name: tool_call.name().to_string(), + request: tool_call.arguments().clone(), + response: None, + status: ToolCallStatus::Pending, + execution_time_ms: None, + created_at: chrono::Utc::now().timestamp(), + completed_at: None, + error_details: None, + }; + let call_id = self.storage.create_tool_call(tool_call_record.clone()).await?; + tool_call_record.id = Some(call_id); let start = Instant::now(); let result: anyhow::Result = @@ -156,8 +172,8 @@ using the tool."#, match result { Ok(response) => { let response_json = serde_json::to_value(&response)?; - self.database - .complete_tool_call(call_id, response_json.clone(), execution_time)?; + tool_call_record.complete(response_json.clone(), execution_time); + self.storage.update_tool_call(tool_call_record).await?; let result_text = crate::format_tool_response(&response); let message_to_llm = format!("Tool {} result:\n{}", tool_call.name(), result_text); @@ -169,12 +185,19 @@ using the tool."#, "Tool execution completed successfully" ); - self.database - .create_message(session_id, "tool", &message_to_llm)?; + let tool_message = StorageMessage { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_message).await?; } Err(e) => { let error_msg = format!("{:?}", e); - self.database.fail_tool_call(call_id, &error_msg)?; + tool_call_record.fail(error_msg.clone()); + self.storage.update_tool_call(tool_call_record).await?; let error_message_to_llm = format!("Tool {} failed: {}", tool_call.name(), error_msg); @@ -186,8 +209,14 @@ using the tool."#, "Tool execution failed" ); - self.database - .create_message(session_id, "tool", &error_message_to_llm)?; + let tool_message = StorageMessage { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: error_message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_message).await?; } } @@ -201,21 +230,20 @@ using the tool."#, .collect() } - fn build_messages(&self, session_id: i64) -> anyhow::Result> { - let db_messages = self.database.get_messages(session_id)?; + async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + let db_messages = self.storage.get_messages(session_id).await?; db_messages .into_iter() .map(|msg| { - let role = match msg.role.as_str() { - "user" => Role::User, - "assistant" => Role::Assistant, - "system" => Role::System, - "tool" => Role::User, - _ => Role::User, + let role = match msg.role { + MessageRole::User => Role::User, + MessageRole::Assistant => Role::Assistant, + MessageRole::System => Role::System, + MessageRole::Tool => Role::User, }; - Ok(Message { + Ok(LlmMessage { role, content: vec![ContentBlock::Text { text: msg.content }], }) @@ -223,6 +251,13 @@ using the tool."#, .collect() } + async fn get_session(&self, session_id: &str) -> anyhow::Result { + self.storage + .get_session(session_id) + .await? + .ok_or_else(|| anyhow::anyhow!("Session not found: {}", session_id)) + } + /// Write settings to TOML file /// This reads existing settings if the file exists, merges new settings into the appropriate /// section, and writes back to the file @@ -281,7 +316,7 @@ using the tool."#, } #[async_trait] -impl Agent for SettingsManagementAgent { +impl Agent for SettingsManagementAgent { fn objective(&self) -> &str { "Collect and manage settings required for workflow automation" } @@ -295,8 +330,15 @@ impl Agent for SettingsManagementAgent { } async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { - self.database - .create_message(session_id, "user", user_prompt)?; + let session_id_str = session_id.to_string(); + let user_message = StorageMessage { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::User, + content: user_prompt.to_string(), + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(user_message).await?; let tools = self.get_tool_definitions(); @@ -307,11 +349,15 @@ impl Agent for SettingsManagementAgent { iteration += 1; if iteration > max_iterations { let error = "Maximum iteration limit reached"; - self.database.fail_session(session_id, error)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Failed; + session.error = Some(error.to_string()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Err(anyhow::anyhow!(error)); } - let messages = self.build_messages(session_id)?; + let messages = self.build_messages(&session_id_str).await?; let request = CompletionRequest { messages, @@ -331,35 +377,53 @@ impl Agent for SettingsManagementAgent { let text = extract_text_from_content(&response.content); let text_to_save = if text.is_empty() && response.tool_calls.is_some() { - "[Using tools]" + "[Using tools]".to_string() } else { - &text + text.clone() }; - let message_id = self - .database - .create_message(session_id, "assistant", text_to_save)?; + let assistant_message = StorageMessage { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::Assistant, + content: text_to_save, + created_at: chrono::Utc::now().timestamp(), + }; + let message_id = self.storage.create_message(assistant_message).await?; if let Some(tool_calls) = response.tool_calls { if tool_calls.is_empty() { - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } for tool_call in tool_calls { // Special handling for ask_user tool - collect settings and write to TOML if tool_call.name() == "ask_user" { - tracing::info!(session_id = session_id, "Agent requesting user settings"); + tracing::info!(session_id = %session_id_str, "Agent requesting user settings"); // Create tool call record in agent_tool_calls table let start = Instant::now(); - let tool_call_id = self.database.create_tool_call( - session_id, - Some(message_id), - tool_call.id(), - tool_call.name(), - tool_call.arguments().clone(), - )?; + let mut tool_call_record = StorageToolCall { + id: None, + session_id: session_id_str.clone(), + message_id: Some(message_id.clone()), + tool_call_id: tool_call.id().to_string(), + tool_name: tool_call.name().to_string(), + request: tool_call.arguments().clone(), + response: None, + status: ToolCallStatus::Pending, + execution_time_ms: None, + created_at: chrono::Utc::now().timestamp(), + completed_at: None, + error_details: None, + }; + let tool_call_id = self.storage.create_tool_call(tool_call_record.clone()).await?; + tool_call_record.id = Some(tool_call_id); // Execute the ask_user tool to get responses from user let tool_request = crate::AgentTool::parse_tool_call( @@ -404,11 +468,8 @@ impl Agent for SettingsManagementAgent { // Mark tool call as completed let response_json = serde_json::to_value(&tool_response)?; - self.database.complete_tool_call( - tool_call_id, - response_json, - execution_time, - )?; + tool_call_record.complete(response_json, execution_time); + self.storage.update_tool_call(tool_call_record).await?; // Create success message let message_to_llm = format!( @@ -417,8 +478,14 @@ impl Agent for SettingsManagementAgent { ask_user_response.responses.len(), self.settings_file_path.display() ); - self.database - .create_message(session_id, "tool", &message_to_llm)?; + let tool_message = StorageMessage { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::Tool, + content: message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_message).await?; } else { return Err(anyhow::anyhow!( "Expected AskUser response from ask_user tool" @@ -426,12 +493,16 @@ impl Agent for SettingsManagementAgent { } } else { // Execute other tools normally - self.execute_tool_call(session_id, Some(message_id), &tool_call) + self.execute_tool_call(&session_id_str, Some(&message_id), &tool_call) .await?; } } } else { - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } } @@ -450,22 +521,22 @@ fn extract_text_from_content(content: &[ContentBlock]) -> String { .join("\n") } -/// Create a SettingsManagementAgent with an in-memory database +/// Create a SettingsManagementAgent with in-memory storage pub fn create_settings_management_agent( client: Arc, settings_file_path: std::path::PathBuf, agent_schemas: Vec, -) -> anyhow::Result<(SettingsManagementAgent, Arc)> { - let database = Arc::new(Database::new(&std::path::PathBuf::from(":memory:"))?); +) -> anyhow::Result> { + let storage = Arc::new(crate::storage::InMemoryStorage::new()); let tool_executor = Arc::new(nocodo_tools::ToolExecutor::new(std::path::PathBuf::from( ".", ))); let agent = SettingsManagementAgent::new( client, - database.clone(), + storage, tool_executor, settings_file_path, agent_schemas, ); - Ok((agent, database)) + Ok(agent) } diff --git a/nocodo-agents/src/storage/memory.rs b/nocodo-agents/src/storage/memory.rs index 6bf30d47..2345fe08 100644 --- a/nocodo-agents/src/storage/memory.rs +++ b/nocodo-agents/src/storage/memory.rs @@ -1,5 +1,6 @@ use crate::storage::{AgentStorage, StorageError}; use crate::types::{Message, Session, ToolCall}; +use shared_types::user_interaction::UserQuestion; use std::collections::HashMap; use std::sync::{Arc, Mutex}; @@ -8,6 +9,8 @@ pub struct InMemoryStorage { sessions: Arc>>, messages: Arc>>>, tool_calls: Arc>>>, + questions: Arc>>>, + answers: Arc>>>, } impl InMemoryStorage { @@ -16,6 +19,8 @@ impl InMemoryStorage { sessions: Arc::new(Mutex::new(HashMap::new())), messages: Arc::new(Mutex::new(HashMap::new())), tool_calls: Arc::new(Mutex::new(HashMap::new())), + questions: Arc::new(Mutex::new(HashMap::new())), + answers: Arc::new(Mutex::new(HashMap::new())), } } } @@ -123,3 +128,54 @@ impl AgentStorage for InMemoryStorage { .collect()) } } + +#[async_trait::async_trait] +impl crate::requirements_gathering::storage::RequirementsStorage for InMemoryStorage { + async fn store_questions( + &self, + session_id: &str, + _tool_call_id: Option<&str>, + questions: &[UserQuestion], + ) -> Result<(), StorageError> { + let mut question_store = self.questions.lock().unwrap(); + question_store + .entry(session_id.to_string()) + .or_insert_with(Vec::new) + .extend(questions.iter().cloned()); + Ok(()) + } + + async fn get_pending_questions( + &self, + session_id: &str, + ) -> Result, StorageError> { + let questions = self.questions.lock().unwrap(); + let answers = self.answers.lock().unwrap(); + + let session_answers = answers.get(session_id); + let session_questions = questions.get(session_id).cloned().unwrap_or_default(); + + // Filter out questions that have been answered + Ok(session_questions + .into_iter() + .filter(|q| { + session_answers + .map(|ans| !ans.contains_key(&q.id)) + .unwrap_or(true) + }) + .collect()) + } + + async fn store_answers( + &self, + session_id: &str, + answers: &std::collections::HashMap, + ) -> Result<(), StorageError> { + let mut answer_store = self.answers.lock().unwrap(); + answer_store + .entry(session_id.to_string()) + .or_insert_with(HashMap::new) + .extend(answers.clone()); + Ok(()) + } +} diff --git a/nocodo-agents/src/structured_json/mod.rs b/nocodo-agents/src/structured_json/mod.rs index a2eaa2e3..62bafe6c 100644 --- a/nocodo-agents/src/structured_json/mod.rs +++ b/nocodo-agents/src/structured_json/mod.rs @@ -1,8 +1,14 @@ -use crate::{database::Database, Agent}; +use crate::{ + storage::AgentStorage, + types::{Message as StorageMessage, MessageRole, Session, SessionStatus}, + Agent, +}; use anyhow; use async_trait::async_trait; use nocodo_llm_sdk::client::LlmClient; -use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message, ResponseFormat, Role}; +use nocodo_llm_sdk::types::{ + CompletionRequest, ContentBlock, Message as LlmMessage, ResponseFormat, Role, +}; use nocodo_tools::ToolExecutor; use std::sync::Arc; @@ -18,9 +24,9 @@ pub struct StructuredJsonAgentConfig { pub domain_description: String, } -pub struct StructuredJsonAgent { +pub struct StructuredJsonAgent { client: Arc, - database: Arc, + storage: Arc, #[allow(dead_code)] tool_executor: Arc, validator: TypeValidator, @@ -29,10 +35,10 @@ pub struct StructuredJsonAgent { config: StructuredJsonAgentConfig, } -impl StructuredJsonAgent { +impl StructuredJsonAgent { pub fn new( client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, config: StructuredJsonAgentConfig, ) -> anyhow::Result { @@ -57,7 +63,7 @@ impl StructuredJsonAgent { Ok(Self { client, - database, + storage, tool_executor, validator, system_prompt, @@ -65,6 +71,13 @@ impl StructuredJsonAgent { }) } + async fn get_session(&self, session_id: &str) -> anyhow::Result { + self.storage + .get_session(session_id) + .await? + .ok_or_else(|| anyhow::anyhow!("Session not found: {}", session_id)) + } + fn generate_system_prompt(type_defs: &str, domain_desc: &str) -> String { format!( r#"You are a specialized AI assistant that responds exclusively in structured JSON. @@ -99,7 +112,7 @@ When responding: async fn validate_and_retry( &self, user_prompt: &str, - session_id: i64, + session_id: &str, max_retries: u32, ) -> anyhow::Result { let mut attempt = 0; @@ -114,7 +127,7 @@ When responding: )); } - let messages = self.build_messages(user_prompt, &conversation_context, session_id)?; + let messages = self.build_messages(user_prompt, &conversation_context)?; let request = CompletionRequest { messages, @@ -132,8 +145,14 @@ When responding: let response = self.client.complete(request).await?; let text = extract_text_from_content(&response.content); - self.database - .create_message(session_id, "assistant", &text)?; + let assistant_message = StorageMessage { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Assistant, + content: text.clone(), + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(assistant_message).await?; let json_result = self.validator.validate_json_syntax(&text); @@ -158,8 +177,14 @@ When responding: ); conversation_context.push((Role::User, error_msg.clone())); - self.database - .create_message(session_id, "user", &error_msg)?; + let error_message = StorageMessage { + id: None, + session_id: session_id.to_string(), + role: MessageRole::User, + content: error_msg, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(error_message).await?; } }, Err(syntax_error) => { @@ -178,8 +203,14 @@ When responding: ); conversation_context.push((Role::User, error_msg.clone())); - self.database - .create_message(session_id, "user", &error_msg)?; + let error_message = StorageMessage { + id: None, + session_id: session_id.to_string(), + role: MessageRole::User, + content: error_msg, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(error_message).await?; } } } @@ -189,12 +220,11 @@ When responding: &self, user_prompt: &str, conversation_context: &[(Role, String)], - _session_id: i64, - ) -> anyhow::Result> { + ) -> anyhow::Result> { let mut messages = Vec::new(); for (role, content) in conversation_context { - messages.push(Message { + messages.push(LlmMessage { role: role.clone(), content: vec![ContentBlock::Text { text: content.clone(), @@ -202,7 +232,7 @@ When responding: }); } - messages.push(Message { + messages.push(LlmMessage { role: Role::User, content: vec![ContentBlock::Text { text: user_prompt.to_string(), @@ -214,7 +244,7 @@ When responding: } #[async_trait] -impl Agent for StructuredJsonAgent { +impl Agent for StructuredJsonAgent { fn objective(&self) -> &str { "Generate structured JSON responses conforming to specified TypeScript types" } @@ -228,14 +258,25 @@ impl Agent for StructuredJsonAgent { } async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { - self.database - .create_message(session_id, "user", user_prompt)?; - - let json_value = self.validate_and_retry(user_prompt, session_id, 3).await?; + let session_id_str = session_id.to_string(); + let user_message = StorageMessage { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::User, + content: user_prompt.to_string(), + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(user_message).await?; + + let json_value = self.validate_and_retry(user_prompt, &session_id_str, 3).await?; let formatted = serde_json::to_string_pretty(&json_value)?; - self.database.complete_session(session_id, &formatted)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(formatted.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; Ok(formatted) } diff --git a/nocodo-agents/src/tesseract/mod.rs b/nocodo-agents/src/tesseract/mod.rs index 55b4e29c..bd1ad315 100644 --- a/nocodo-agents/src/tesseract/mod.rs +++ b/nocodo-agents/src/tesseract/mod.rs @@ -1,9 +1,16 @@ -use crate::{database::Database, Agent, AgentTool}; +use crate::{ + storage::AgentStorage, + types::{ + Message as StorageMessage, MessageRole, Session, SessionStatus, + ToolCall as StorageToolCall, ToolCallStatus, + }, + Agent, AgentTool, +}; use anyhow::{self, Context}; use async_trait::async_trait; use nocodo_llm_sdk::client::LlmClient; -use nocodo_llm_sdk::tools::{ToolCall, ToolChoice}; -use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message, Role}; +use nocodo_llm_sdk::tools::{ToolCall as LlmToolCall, ToolChoice}; +use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message as LlmMessage, Role}; use nocodo_tools::{ bash::{BashExecutor, BashPermissions}, ToolExecutor, @@ -13,9 +20,9 @@ use std::sync::Arc; use std::time::Instant; /// Agent specialized in extracting text from images using Tesseract OCR -pub struct TesseractAgent { +pub struct TesseractAgent { client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, #[allow(dead_code)] // Stored for reference, used during construction image_path: PathBuf, @@ -24,12 +31,12 @@ pub struct TesseractAgent { system_prompt: String, } -impl TesseractAgent { +impl TesseractAgent { /// Create a new TesseractAgent /// /// # Arguments /// * `client` - LLM client for AI inference - /// * `database` - Database for session/message tracking + /// * `storage` - Storage for session/message tracking /// * `image_path` - Path to the image file to process /// /// # Security @@ -44,7 +51,7 @@ impl TesseractAgent { /// - The image file must exist pub fn new( client: Arc, - database: Arc, + storage: Arc, image_path: PathBuf, ) -> anyhow::Result { // Validate image path exists @@ -80,7 +87,7 @@ impl TesseractAgent { Ok(Self { client, - database, + storage, tool_executor, image_path, image_filename, @@ -88,6 +95,13 @@ impl TesseractAgent { }) } + async fn get_session(&self, session_id: &str) -> anyhow::Result { + self.storage + .get_session(session_id) + .await? + .ok_or_else(|| anyhow::anyhow!("Session not found: {}", session_id)) + } + /// Get tool definitions for this agent fn get_tool_definitions(&self) -> Vec { self.tools() @@ -97,21 +111,20 @@ impl TesseractAgent { } /// Build messages from session history - fn build_messages(&self, session_id: i64) -> anyhow::Result> { - let db_messages = self.database.get_messages(session_id)?; + async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + let db_messages = self.storage.get_messages(session_id).await?; db_messages .into_iter() .map(|msg| { - let role = match msg.role.as_str() { - "user" => Role::User, - "assistant" => Role::Assistant, - "system" => Role::System, - "tool" => Role::User, - _ => Role::User, + let role = match msg.role { + MessageRole::User => Role::User, + MessageRole::Assistant => Role::Assistant, + MessageRole::System => Role::System, + MessageRole::Tool => Role::User, }; - Ok(Message { + Ok(LlmMessage { role, content: vec![ContentBlock::Text { text: msg.content }], }) @@ -122,22 +135,31 @@ impl TesseractAgent { /// Execute a tool call async fn execute_tool_call( &self, - session_id: i64, - message_id: Option, - tool_call: &ToolCall, + session_id: &str, + message_id: Option<&String>, + tool_call: &LlmToolCall, ) -> anyhow::Result<()> { // 1. Parse LLM tool call into typed ToolRequest let tool_request = AgentTool::parse_tool_call(tool_call.name(), tool_call.arguments().clone())?; - // 2. Record tool call in database - let call_id = self.database.create_tool_call( - session_id, - message_id, - tool_call.id(), - tool_call.name(), - tool_call.arguments().clone(), - )?; + // 2. Record tool call in storage + let mut tool_call_record = StorageToolCall { + id: None, + session_id: session_id.to_string(), + message_id: message_id.cloned(), + tool_call_id: tool_call.id().to_string(), + tool_name: tool_call.name().to_string(), + request: tool_call.arguments().clone(), + response: None, + status: ToolCallStatus::Pending, + execution_time_ms: None, + created_at: chrono::Utc::now().timestamp(), + completed_at: None, + error_details: None, + }; + let call_id = self.storage.create_tool_call(tool_call_record.clone()).await?; + tool_call_record.id = Some(call_id); // 3. Execute tool let start = Instant::now(); @@ -145,12 +167,12 @@ impl TesseractAgent { self.tool_executor.execute(tool_request).await; let execution_time = start.elapsed().as_millis() as i64; - // 4. Update database with result + // 4. Update storage with result match result { Ok(response) => { let response_json = serde_json::to_value(&response)?; - self.database - .complete_tool_call(call_id, response_json.clone(), execution_time)?; + tool_call_record.complete(response_json.clone(), execution_time); + self.storage.update_tool_call(tool_call_record).await?; let result_text = crate::format_tool_response(&response); let message_to_llm = format!("Tool {} result:\n{}", tool_call.name(), result_text); @@ -162,12 +184,19 @@ impl TesseractAgent { "Tool execution completed successfully" ); - self.database - .create_message(session_id, "tool", &message_to_llm)?; + let tool_message = StorageMessage { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_message).await?; } Err(e) => { let error_msg = format!("{:?}", e); - self.database.fail_tool_call(call_id, &error_msg)?; + tool_call_record.fail(error_msg.clone()); + self.storage.update_tool_call(tool_call_record).await?; let error_message_to_llm = format!("Tool {} failed: {}", tool_call.name(), error_msg); @@ -179,8 +208,14 @@ impl TesseractAgent { "Tool execution failed" ); - self.database - .create_message(session_id, "tool", &error_message_to_llm)?; + let tool_message = StorageMessage { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: error_message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_message).await?; } } @@ -189,7 +224,7 @@ impl TesseractAgent { } #[async_trait] -impl Agent for TesseractAgent { +impl Agent for TesseractAgent { fn objective(&self) -> &str { "Extract text from images using Tesseract OCR and optionally clean/format the output" } @@ -215,25 +250,23 @@ impl Agent for TesseractAgent { ] } - async fn execute(&self, user_prompt: &str, _session_id: i64) -> anyhow::Result { - // 1. Create session - let session_id = self.database.create_session( - "tesseract-ocr", - self.client.provider_name(), - self.client.model_name(), - Some(&self.system_prompt), - user_prompt, - None, // No config for TesseractAgent - )?; - - // 2. Create initial user message - self.database - .create_message(session_id, "user", user_prompt)?; - - // 3. Get tool definitions + async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { + let session_id_str = session_id.to_string(); + + // Create initial user message + let user_message = StorageMessage { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::User, + content: user_prompt.to_string(), + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(user_message).await?; + + // Get tool definitions let tools = self.get_tool_definitions(); - // 4. Execution loop (max 30 iterations) + // Execution loop (max 30 iterations) let mut iteration = 0; let max_iterations = 30; @@ -241,12 +274,16 @@ impl Agent for TesseractAgent { iteration += 1; if iteration > max_iterations { let error = "Maximum iteration limit reached"; - self.database.fail_session(session_id, error)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Failed; + session.error = Some(error.to_string()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Err(anyhow::anyhow!(error)); } - // 5. Build request with conversation history - let messages = self.build_messages(session_id)?; + // Build request with conversation history + let messages = self.build_messages(&session_id_str).await?; let request = CompletionRequest { messages, @@ -261,29 +298,42 @@ impl Agent for TesseractAgent { response_format: None, }; - // 6. Call LLM + // Call LLM let response = self.client.complete(request).await?; - // 7. Extract text and save assistant message + // Extract text and save assistant message let text = extract_text_from_content(&response.content); - let message_id = self - .database - .create_message(session_id, "assistant", &text)?; + let assistant_message = StorageMessage { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::Assistant, + content: text.clone(), + created_at: chrono::Utc::now().timestamp(), + }; + let message_id = self.storage.create_message(assistant_message).await?; - // 8. Check for tool calls + // Check for tool calls if let Some(tool_calls) = response.tool_calls { if tool_calls.is_empty() { - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } - // 9. Execute tools + // Execute tools for tool_call in tool_calls { - self.execute_tool_call(session_id, Some(message_id), &tool_call) + self.execute_tool_call(&session_id_str, Some(&message_id), &tool_call) .await?; } } else { - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } } @@ -408,7 +458,7 @@ pub fn verify_tesseract_installation() -> anyhow::Result { Ok(version_info) } -impl TesseractAgent { +impl TesseractAgent { /// Verify pre-conditions before creating agent pub fn verify_preconditions() -> anyhow::Result<()> { match verify_tesseract_installation() { From cdf43403d49e2c68082e5a5dcb0bde607d19c46e Mon Sep 17 00:00:00 2001 From: Sumit Datta Date: Tue, 3 Feb 2026 10:53:12 +0530 Subject: [PATCH 3/6] Complete trait-based storage migration for remaining agents and runners - Refactor PdfToTextAgent to use AgentStorage trait (like TesseractAgent) - Remove old database extension modules (requirements_gathering/database.rs, settings_management/database.rs) - Update all 6 binary runners to use InMemoryStorage - Add RequirementsStorage bound to AgentFactory for UserClarificationAgent - Fix borrow issues in InMemoryStorage (clone session_id before use) All agents now use the storage trait abstraction instead of direct Database. Binary runners use in-memory storage with proper session creation via Session struct. --- nocodo-agents/bin/codebase_analysis_runner.rs | 37 +- nocodo-agents/bin/imap_email_runner.rs | 82 +- .../bin/requirements_gathering_runner.rs | 40 +- .../bin/settings_management_runner.rs | 49 +- nocodo-agents/bin/sqlite_reader_runner.rs | 39 +- nocodo-agents/bin/structured_json_runner.rs | 48 +- nocodo-agents/src/factory.rs | 18 +- nocodo-agents/src/pdftotext/mod.rs | 196 +++-- .../src/requirements_gathering/database.rs | 97 --- .../src/requirements_gathering/mod.rs | 1 - .../src/settings_management/database.rs | 133 ---- nocodo-agents/src/settings_management/mod.rs | 11 +- nocodo-agents/src/storage/memory.rs | 6 +- .../tasks/implement-sqlite-agent-storage.md | 712 ++++++++++++++++++ 14 files changed, 1054 insertions(+), 415 deletions(-) delete mode 100644 nocodo-agents/src/requirements_gathering/database.rs delete mode 100644 nocodo-agents/src/settings_management/database.rs create mode 100644 nocodo-api/tasks/implement-sqlite-agent-storage.md diff --git a/nocodo-agents/bin/codebase_analysis_runner.rs b/nocodo-agents/bin/codebase_analysis_runner.rs index 8993dea1..c86be628 100644 --- a/nocodo-agents/bin/codebase_analysis_runner.rs +++ b/nocodo-agents/bin/codebase_analysis_runner.rs @@ -1,5 +1,11 @@ use clap::Parser; -use nocodo_agents::{config, factory::create_codebase_analysis_agent, Agent}; +use nocodo_agents::{ + config, + factory::create_codebase_analysis_agent, + storage::AgentStorage, + types::{Session, SessionStatus}, + Agent, +}; use nocodo_llm_sdk::glm::zai::ZaiGlmClient; use nocodo_tools::ToolExecutor; use std::path::PathBuf; @@ -65,20 +71,31 @@ async fn main() -> anyhow::Result<()> { ); // Create codebase analysis agent - let (agent, database) = create_codebase_analysis_agent(client, tool_executor); + let agent = create_codebase_analysis_agent(client, tool_executor); println!("Running agent: {}", agent.objective()); println!("User prompt: {}\n", args.prompt); // Create session - let session_id = database.create_session( - "codebase-analysis", - "standalone", - "standalone", - Some(&agent.system_prompt()), - &args.prompt, - None, - )?; + let storage = Arc::new(nocodo_agents::storage::InMemoryStorage::new()); + let session = Session { + id: None, + agent_name: "codebase-analysis".to_string(), + provider: "standalone".to_string(), + model: "standalone".to_string(), + system_prompt: Some(agent.system_prompt()), + user_prompt: args.prompt.clone(), + config: serde_json::json!({}), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + let session_id_str = storage.create_session(session).await?; + + // Parse session ID as i64 for agent.execute() + let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); // Execute agent let result = agent.execute(&args.prompt, session_id).await?; diff --git a/nocodo-agents/bin/imap_email_runner.rs b/nocodo-agents/bin/imap_email_runner.rs index cb6ec795..1dd85d83 100644 --- a/nocodo-agents/bin/imap_email_runner.rs +++ b/nocodo-agents/bin/imap_email_runner.rs @@ -1,5 +1,11 @@ use clap::Parser; -use nocodo_agents::{config, database::Database, imap_email::ImapEmailAgent, Agent}; +use nocodo_agents::{ + config, + imap_email::ImapEmailAgent, + storage::AgentStorage, + types::{Session, SessionStatus}, + Agent, +}; use nocodo_llm_sdk::glm::zai::ZaiGlmClient; use nocodo_tools::ToolExecutor; use std::collections::HashMap; @@ -72,8 +78,8 @@ async fn main() -> anyhow::Result<()> { let tool_executor = Arc::new(ToolExecutor::new(std::env::current_dir()?).with_max_file_size(10 * 1024 * 1024)); - // Create database for session management - let database = Arc::new(Database::new(&PathBuf::from(":memory:"))?); + // Create storage for session management + let storage = Arc::new(nocodo_agents::storage::InMemoryStorage::new()); // Test IMAP connection before creating agent/loading LLM println!("🔍 Testing IMAP connection..."); @@ -104,7 +110,7 @@ async fn main() -> anyhow::Result<()> { let agent = ImapEmailAgent::from_settings( client.clone(), - database.clone(), + storage.clone(), tool_executor.clone(), &settings, )?; @@ -116,32 +122,42 @@ async fn main() -> anyhow::Result<()> { if args.interactive { // Interactive mode - multiple queries in same session - run_interactive_mode(&agent, &database, &args.prompt).await?; + run_interactive_mode(&agent, &storage, &args.prompt).await?; } else { // Single query mode - run_single_query(&agent, &database, &args.prompt).await?; + run_single_query(&agent, &storage, &args.prompt).await?; } Ok(()) } -async fn run_single_query( - agent: &ImapEmailAgent, - database: &Arc, +async fn run_single_query( + agent: &ImapEmailAgent, + storage: &Arc, prompt: &str, ) -> anyhow::Result<()> { println!("💬 User prompt: {}\n", prompt); println!("⏳ Processing...\n"); // Create session - let session_id = database.create_session( - "imap-email", - "standalone", - "standalone", - Some(&agent.system_prompt()), - prompt, - None, - )?; + let session = Session { + id: None, + agent_name: "imap-email".to_string(), + provider: "standalone".to_string(), + model: "standalone".to_string(), + system_prompt: Some(agent.system_prompt()), + user_prompt: prompt.to_string(), + config: serde_json::json!({}), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + let session_id_str = storage.create_session(session).await?; + + // Parse session ID as i64 for agent.execute() + let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); // Execute agent let result = agent.execute(prompt, session_id).await?; @@ -151,20 +167,30 @@ async fn run_single_query( Ok(()) } -async fn run_interactive_mode( - agent: &ImapEmailAgent, - database: &Arc, +async fn run_interactive_mode( + agent: &ImapEmailAgent, + storage: &Arc, initial_prompt: &str, ) -> anyhow::Result<()> { // Create a single session for the entire interaction - let session_id = database.create_session( - "imap-email", - "standalone", - "standalone", - Some(&agent.system_prompt()), - initial_prompt, - None, - )?; + let session = Session { + id: None, + agent_name: "imap-email".to_string(), + provider: "standalone".to_string(), + model: "standalone".to_string(), + system_prompt: Some(agent.system_prompt()), + user_prompt: initial_prompt.to_string(), + config: serde_json::json!({}), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + let session_id_str = storage.create_session(session).await?; + + // Parse session ID as i64 for agent.execute() + let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); println!("🔄 Interactive mode enabled - session ID: {}", session_id); println!("💡 Type your queries. Type 'quit' or 'exit' to end the session.\n"); diff --git a/nocodo-agents/bin/requirements_gathering_runner.rs b/nocodo-agents/bin/requirements_gathering_runner.rs index 0ef79469..4f3b2ecd 100644 --- a/nocodo-agents/bin/requirements_gathering_runner.rs +++ b/nocodo-agents/bin/requirements_gathering_runner.rs @@ -1,7 +1,11 @@ use clap::Parser; -use nocodo_agents::config; -use nocodo_agents::requirements_gathering::create_user_clarification_agent; -use nocodo_agents::Agent; +use nocodo_agents::{ + config, + factory::create_user_clarification_agent, + storage::AgentStorage, + types::{Session, SessionStatus}, + Agent, +}; use std::path::PathBuf; use std::sync::Arc; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter}; @@ -47,21 +51,33 @@ async fn main() -> anyhow::Result<()> { zai_config.coding_plan, )?); - let (agent, database) = create_user_clarification_agent(client)?; + let agent = create_user_clarification_agent(client); tracing::debug!("System prompt:\n{}", agent.system_prompt()); println!("Running agent: {}", agent.objective()); println!("User prompt: {}\n", args.prompt); - let session_id = database.create_session( - "user-clarification", - "standalone", - "standalone", - Some(&agent.system_prompt()), - &args.prompt, - None, - )?; + // Create session + let storage = Arc::new(nocodo_agents::storage::InMemoryStorage::new()); + let session = Session { + id: None, + agent_name: "user-clarification".to_string(), + provider: "standalone".to_string(), + model: "standalone".to_string(), + system_prompt: Some(agent.system_prompt()), + user_prompt: args.prompt.clone(), + config: serde_json::json!({}), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + let session_id_str = storage.create_session(session).await?; + + // Parse session ID as i64 for agent.execute() + let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); let result = agent.execute(&args.prompt, session_id).await?; diff --git a/nocodo-agents/bin/settings_management_runner.rs b/nocodo-agents/bin/settings_management_runner.rs index d7ad7948..e8ee960c 100644 --- a/nocodo-agents/bin/settings_management_runner.rs +++ b/nocodo-agents/bin/settings_management_runner.rs @@ -1,7 +1,11 @@ use clap::Parser; -use nocodo_agents::config; -use nocodo_agents::settings_management::create_settings_management_agent; -use nocodo_agents::Agent; +use nocodo_agents::{ + config, + factory::create_settings_management_agent, + storage::AgentStorage, + types::{Session, SessionStatus}, + Agent, +}; use std::path::PathBuf; use std::sync::Arc; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter}; @@ -46,26 +50,39 @@ async fn main() -> anyhow::Result<()> { // Collect settings schemas from available agents // Use static schemas to avoid circular dependency (can't instantiate SqliteReaderAgent // without db_path, which is what we're trying to collect) - let agent_schemas = nocodo_agents::sqlite_reader::SqliteReaderAgent::static_settings_schema() - .map(|schema| vec![schema]) - .unwrap_or_default(); + let agent_schemas = nocodo_agents::sqlite_reader::SqliteReaderAgent::< + nocodo_agents::storage::InMemoryStorage, + >::static_settings_schema() + .map(|schema| vec![schema]) + .unwrap_or_default(); - let (agent, database) = - create_settings_management_agent(client, args.settings_file.clone(), agent_schemas)?; + let agent = create_settings_management_agent(client, args.settings_file.clone(), agent_schemas); tracing::debug!("System prompt:\n{}", agent.system_prompt()); println!("Running agent: {}", agent.objective()); println!("User prompt: {}\n", args.prompt); - let session_id = database.create_session( - "settings-management", - "standalone", - "standalone", - Some(&agent.system_prompt()), - &args.prompt, - None, - )?; + // Create session + let storage = Arc::new(nocodo_agents::storage::InMemoryStorage::new()); + let session = Session { + id: None, + agent_name: "settings-management".to_string(), + provider: "standalone".to_string(), + model: "standalone".to_string(), + system_prompt: Some(agent.system_prompt()), + user_prompt: args.prompt.clone(), + config: serde_json::json!({}), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + let session_id_str = storage.create_session(session).await?; + + // Parse session ID as i64 for agent.execute() + let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); let result = agent.execute(&args.prompt, session_id).await?; diff --git a/nocodo-agents/bin/sqlite_reader_runner.rs b/nocodo-agents/bin/sqlite_reader_runner.rs index d1e481f8..fcf394dd 100644 --- a/nocodo-agents/bin/sqlite_reader_runner.rs +++ b/nocodo-agents/bin/sqlite_reader_runner.rs @@ -1,5 +1,11 @@ use clap::Parser; -use nocodo_agents::{config, factory::create_sqlite_reader_agent, Agent}; +use nocodo_agents::{ + config, + factory::create_sqlite_reader_agent, + storage::AgentStorage, + types::{Session, SessionStatus}, + Agent, +}; use nocodo_llm_sdk::glm::zai::ZaiGlmClient; use nocodo_tools::ToolExecutor; use std::path::PathBuf; @@ -44,22 +50,33 @@ async fn main() -> anyhow::Result<()> { let tool_executor = Arc::new(ToolExecutor::new(std::env::current_dir()?).with_max_file_size(10 * 1024 * 1024)); - let (agent, database) = create_sqlite_reader_agent(client, tool_executor, args.db_path).await?; + let storage = Arc::new(nocodo_agents::storage::InMemoryStorage::new()); + let agent = create_sqlite_reader_agent(client, tool_executor, args.db_path).await?; tracing::info!("System prompt:\n{}", agent.system_prompt()); println!("Running agent: {}", agent.objective()); println!("User prompt: {}\n", args.prompt); - // For standalone runner, create a dummy session - let session_id = database.create_session( - "sqlite-reader", - "standalone", - "standalone", - Some(&agent.system_prompt()), - &args.prompt, - None, - )?; + // Create session + let session = Session { + id: None, + agent_name: "sqlite-reader".to_string(), + provider: "standalone".to_string(), + model: "standalone".to_string(), + system_prompt: Some(agent.system_prompt()), + user_prompt: args.prompt.clone(), + config: serde_json::json!({}), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + let session_id_str = storage.create_session(session).await?; + + // Parse session ID as i64 for agent.execute() + let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); let result = agent.execute(&args.prompt, session_id).await?; diff --git a/nocodo-agents/bin/structured_json_runner.rs b/nocodo-agents/bin/structured_json_runner.rs index 76e55c58..f5683083 100644 --- a/nocodo-agents/bin/structured_json_runner.rs +++ b/nocodo-agents/bin/structured_json_runner.rs @@ -1,8 +1,12 @@ use clap::Parser; -use nocodo_agents::config; -use nocodo_agents::factory::AgentFactory; -use nocodo_agents::structured_json::StructuredJsonAgentConfig; -use nocodo_agents::Agent; +use nocodo_agents::{ + config, + factory::AgentFactory, + storage::AgentStorage, + structured_json::StructuredJsonAgentConfig, + types::{Session, SessionStatus}, + Agent, +}; use nocodo_llm_sdk::glm::zai::ZaiGlmClient; use std::path::PathBuf; use std::sync::Arc; @@ -45,20 +49,14 @@ async fn main() -> anyhow::Result<()> { zai_config.coding_plan, )?); - let db_path = config_path - .parent() - .map(|p| p.to_path_buf()) - .unwrap_or_else(|| PathBuf::from(".")) - .join("agent_sessions.db"); - - let database = Arc::new(nocodo_agents::database::Database::new(&db_path)?); + let storage = Arc::new(nocodo_agents::storage::InMemoryStorage::new()); let tool_executor = Arc::new( nocodo_tools::ToolExecutor::new(std::env::current_dir()?) .with_max_file_size(10 * 1024 * 1024), ); - let factory = AgentFactory::new(client.clone(), database.clone(), tool_executor); + let factory = AgentFactory::new(client.clone(), storage.clone(), tool_executor); let type_names = if args.types.is_empty() { vec![ @@ -81,14 +79,24 @@ async fn main() -> anyhow::Result<()> { let system_prompt = agent.system_prompt(); println!("\nSystem Prompt:\n{}", system_prompt); - let session_id = database.create_session( - "structured-json", - "cli", - &args.prompt, - Some(&system_prompt), - "structured-json-runner", - None, - )?; + let session = Session { + id: None, + agent_name: "structured-json".to_string(), + provider: "cli".to_string(), + model: "cli".to_string(), + system_prompt: Some(system_prompt), + user_prompt: args.prompt.clone(), + config: serde_json::json!({}), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + let session_id_str = storage.create_session(session).await?; + + // Parse session ID as i64 for agent.execute() + let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); match agent.execute(&args.prompt, session_id).await { Ok(result) => { diff --git a/nocodo-agents/src/factory.rs b/nocodo-agents/src/factory.rs index 7d11b631..e0f0a410 100644 --- a/nocodo-agents/src/factory.rs +++ b/nocodo-agents/src/factory.rs @@ -32,7 +32,9 @@ pub struct AgentFactory { tool_executor: Arc, } -impl AgentFactory { +impl + AgentFactory +{ /// Create a new AgentFactory with the given dependencies pub fn new( llm_client: Arc, @@ -212,14 +214,12 @@ pub fn create_agent_with_tools( }; Box::new(StructuredJsonAgent::new(client, storage, tool_executor, config).unwrap()) } - AgentType::UserClarification => { - Box::new(UserClarificationAgent::new( - client, - storage.clone(), - storage, - tool_executor, - )) - } + AgentType::UserClarification => Box::new(UserClarificationAgent::new( + client, + storage.clone(), + storage, + tool_executor, + )), AgentType::SettingsManagement => { panic!( "SettingsManagement agent cannot be created via create_by_type. \ diff --git a/nocodo-agents/src/pdftotext/mod.rs b/nocodo-agents/src/pdftotext/mod.rs index 2e78b730..ce29bc68 100644 --- a/nocodo-agents/src/pdftotext/mod.rs +++ b/nocodo-agents/src/pdftotext/mod.rs @@ -1,9 +1,16 @@ -use crate::{database::Database, Agent, AgentTool}; +use crate::{ + storage::AgentStorage, + types::{ + Message as StorageMessage, MessageRole, Session, SessionStatus, + ToolCall as StorageToolCall, ToolCallStatus, + }, + Agent, AgentTool, +}; use anyhow::{self, Context}; use async_trait::async_trait; use nocodo_llm_sdk::client::LlmClient; -use nocodo_llm_sdk::tools::{ToolCall, ToolChoice}; -use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message, Role}; +use nocodo_llm_sdk::tools::{ToolCall as LlmToolCall, ToolChoice}; +use nocodo_llm_sdk::types::{CompletionRequest, ContentBlock, Message as LlmMessage, Role}; use nocodo_tools::{ bash::{BashExecutor, BashPermissions}, ToolExecutor, @@ -13,9 +20,9 @@ use std::sync::Arc; use std::time::Instant; /// Agent specialized in extracting text from PDFs using pdftotext and qpdf -pub struct PdfToTextAgent { +pub struct PdfToTextAgent { client: Arc, - database: Arc, + storage: Arc, tool_executor: Arc, #[allow(dead_code)] // Stored for reference, used during construction pdf_path: PathBuf, @@ -24,12 +31,12 @@ pub struct PdfToTextAgent { system_prompt: String, } -impl PdfToTextAgent { +impl PdfToTextAgent { /// Create a new PdfToTextAgent /// /// # Arguments /// * `client` - LLM client for AI inference - /// * `database` - Database for session/message tracking + /// * `storage` - Storage for session/message tracking /// * `pdf_path` - Path to the PDF file to process /// /// # Security @@ -45,7 +52,7 @@ impl PdfToTextAgent { /// - The PDF file must exist pub fn new( client: Arc, - database: Arc, + storage: Arc, pdf_path: PathBuf, ) -> anyhow::Result { // Validate PDF path exists @@ -81,7 +88,7 @@ impl PdfToTextAgent { Ok(Self { client, - database, + storage, tool_executor, pdf_path, pdf_filename, @@ -89,6 +96,13 @@ impl PdfToTextAgent { }) } + async fn get_session(&self, session_id: &str) -> anyhow::Result { + self.storage + .get_session(session_id) + .await? + .ok_or_else(|| anyhow::anyhow!("Session not found: {}", session_id)) + } + /// Get tool definitions for this agent fn get_tool_definitions(&self) -> Vec { self.tools() @@ -98,21 +112,20 @@ impl PdfToTextAgent { } /// Build messages from session history - fn build_messages(&self, session_id: i64) -> anyhow::Result> { - let db_messages = self.database.get_messages(session_id)?; + async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + let db_messages = self.storage.get_messages(session_id).await?; db_messages .into_iter() .map(|msg| { - let role = match msg.role.as_str() { - "user" => Role::User, - "assistant" => Role::Assistant, - "system" => Role::System, - "tool" => Role::User, - _ => Role::User, + let role = match msg.role { + MessageRole::User => Role::User, + MessageRole::Assistant => Role::Assistant, + MessageRole::System => Role::System, + MessageRole::Tool => Role::User, }; - Ok(Message { + Ok(LlmMessage { role, content: vec![ContentBlock::Text { text: msg.content }], }) @@ -123,22 +136,34 @@ impl PdfToTextAgent { /// Execute a tool call async fn execute_tool_call( &self, - session_id: i64, - message_id: Option, - tool_call: &ToolCall, + session_id: &str, + message_id: Option<&String>, + tool_call: &LlmToolCall, ) -> anyhow::Result<()> { // 1. Parse LLM tool call into typed ToolRequest let tool_request = AgentTool::parse_tool_call(tool_call.name(), tool_call.arguments().clone())?; - // 2. Record tool call in database - let call_id = self.database.create_tool_call( - session_id, - message_id, - tool_call.id(), - tool_call.name(), - tool_call.arguments().clone(), - )?; + // 2. Record tool call in storage + let mut tool_call_record = StorageToolCall { + id: None, + session_id: session_id.to_string(), + message_id: message_id.cloned(), + tool_call_id: tool_call.id().to_string(), + tool_name: tool_call.name().to_string(), + request: tool_call.arguments().clone(), + response: None, + status: ToolCallStatus::Pending, + execution_time_ms: None, + created_at: chrono::Utc::now().timestamp(), + completed_at: None, + error_details: None, + }; + let call_id = self + .storage + .create_tool_call(tool_call_record.clone()) + .await?; + tool_call_record.id = Some(call_id); // 3. Execute tool let start = Instant::now(); @@ -146,12 +171,12 @@ impl PdfToTextAgent { self.tool_executor.execute(tool_request).await; let execution_time = start.elapsed().as_millis() as i64; - // 4. Update database with result + // 4. Update storage with result match result { Ok(response) => { let response_json = serde_json::to_value(&response)?; - self.database - .complete_tool_call(call_id, response_json.clone(), execution_time)?; + tool_call_record.complete(response_json.clone(), execution_time); + self.storage.update_tool_call(tool_call_record).await?; let result_text = crate::format_tool_response(&response); let message_to_llm = format!("Tool {} result:\n{}", tool_call.name(), result_text); @@ -163,12 +188,19 @@ impl PdfToTextAgent { "Tool execution completed successfully" ); - self.database - .create_message(session_id, "tool", &message_to_llm)?; + let tool_message = StorageMessage { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_message).await?; } Err(e) => { let error_msg = format!("{:?}", e); - self.database.fail_tool_call(call_id, &error_msg)?; + tool_call_record.fail(error_msg.clone()); + self.storage.update_tool_call(tool_call_record).await?; let error_message_to_llm = format!("Tool {} failed: {}", tool_call.name(), error_msg); @@ -180,8 +212,14 @@ impl PdfToTextAgent { "Tool execution failed" ); - self.database - .create_message(session_id, "tool", &error_message_to_llm)?; + let tool_message = StorageMessage { + id: None, + session_id: session_id.to_string(), + role: MessageRole::Tool, + content: error_message_to_llm, + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(tool_message).await?; } } @@ -190,7 +228,7 @@ impl PdfToTextAgent { } #[async_trait] -impl Agent for PdfToTextAgent { +impl Agent for PdfToTextAgent { fn objective(&self) -> &str { "Extract text from PDF files using pdftotext with layout preservation and optional page selection using qpdf" } @@ -218,25 +256,23 @@ impl Agent for PdfToTextAgent { ] } - async fn execute(&self, user_prompt: &str, _session_id: i64) -> anyhow::Result { - // 1. Create session - let session_id = self.database.create_session( - "pdftotext", - self.client.provider_name(), - self.client.model_name(), - Some(&self.system_prompt), - user_prompt, - None, // No config for PdfToTextAgent - )?; - - // 2. Create initial user message - self.database - .create_message(session_id, "user", user_prompt)?; - - // 3. Get tool definitions + async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { + let session_id_str = session_id.to_string(); + + // Create initial user message + let user_message = StorageMessage { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::User, + content: user_prompt.to_string(), + created_at: chrono::Utc::now().timestamp(), + }; + self.storage.create_message(user_message).await?; + + // Get tool definitions let tools = self.get_tool_definitions(); - // 4. Execution loop (max 30 iterations) + // Execution loop (max 30 iterations) let mut iteration = 0; let max_iterations = 30; @@ -244,12 +280,16 @@ impl Agent for PdfToTextAgent { iteration += 1; if iteration > max_iterations { let error = "Maximum iteration limit reached"; - self.database.fail_session(session_id, error)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Failed; + session.error = Some(error.to_string()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Err(anyhow::anyhow!(error)); } - // 5. Build request with conversation history - let messages = self.build_messages(session_id)?; + // Build request with conversation history + let messages = self.build_messages(&session_id_str).await?; let request = CompletionRequest { messages, @@ -264,29 +304,42 @@ impl Agent for PdfToTextAgent { response_format: None, }; - // 6. Call LLM + // Call LLM let response = self.client.complete(request).await?; - // 7. Extract text and save assistant message + // Extract text and save assistant message let text = extract_text_from_content(&response.content); - let message_id = self - .database - .create_message(session_id, "assistant", &text)?; + let assistant_message = StorageMessage { + id: None, + session_id: session_id_str.clone(), + role: MessageRole::Assistant, + content: text.clone(), + created_at: chrono::Utc::now().timestamp(), + }; + let message_id = self.storage.create_message(assistant_message).await?; - // 8. Check for tool calls + // Check for tool calls if let Some(tool_calls) = response.tool_calls { if tool_calls.is_empty() { - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } - // 9. Execute tools + // Execute tools for tool_call in tool_calls { - self.execute_tool_call(session_id, Some(message_id), &tool_call) + self.execute_tool_call(&session_id_str, Some(&message_id), &tool_call) .await?; } } else { - self.database.complete_session(session_id, &text)?; + let mut session = self.get_session(&session_id_str).await?; + session.status = SessionStatus::Completed; + session.result = Some(text.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + self.storage.update_session(session).await?; return Ok(text); } } @@ -484,16 +537,13 @@ pub fn verify_qpdf_installation() -> anyhow::Result { Ok(version_info) } -impl PdfToTextAgent { +impl PdfToTextAgent { /// Verify pre-conditions before creating agent pub fn verify_preconditions() -> anyhow::Result<()> { // Check pdftotext match verify_pdftotext_installation() { Ok(version) => { - tracing::info!( - "pdftotext found: {}", - version.lines().next().unwrap_or("") - ); + tracing::info!("pdftotext found: {}", version.lines().next().unwrap_or("")); } Err(e) => { anyhow::bail!( diff --git a/nocodo-agents/src/requirements_gathering/database.rs b/nocodo-agents/src/requirements_gathering/database.rs deleted file mode 100644 index ca840978..00000000 --- a/nocodo-agents/src/requirements_gathering/database.rs +++ /dev/null @@ -1,97 +0,0 @@ -use crate::database::Database; -use rusqlite::params; - -/// Requirements gathering database operations -impl Database { - /// Store questions in the project_requirements_qna table - pub fn store_questions( - &self, - session_id: i64, - tool_call_id: Option, - questions: &[shared_types::user_interaction::UserQuestion], - ) -> anyhow::Result<()> { - let conn = self.connection.lock().unwrap(); - let now = chrono::Utc::now().timestamp(); - - for question in questions { - conn.execute( - "INSERT INTO project_requirements_qna (session_id, tool_call_id, question_id, question, description, response_type, created_at) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)", - params![ - session_id, - tool_call_id, - &question.id, - &question.question, - &question.description, - format!("{:?}", question.response_type).to_lowercase(), - now - ], - )?; - } - - Ok(()) - } - - /// Get pending (unanswered) questions from the database - pub fn get_pending_questions( - &self, - session_id: i64, - ) -> anyhow::Result> { - let conn = self.connection.lock().unwrap(); - let mut stmt = conn.prepare( - "SELECT question_id, question, description, response_type - FROM project_requirements_qna - WHERE session_id = ?1 AND answer IS NULL - ORDER BY created_at ASC", - )?; - - let questions = stmt.query_map([session_id], |row| { - let response_type_str: String = row.get(3)?; - let response_type = match response_type_str.as_str() { - "text" => shared_types::user_interaction::QuestionType::Text, - "password" => shared_types::user_interaction::QuestionType::Password, - "file_path" => shared_types::user_interaction::QuestionType::FilePath, - "email" => shared_types::user_interaction::QuestionType::Email, - "url" => shared_types::user_interaction::QuestionType::Url, - _ => shared_types::user_interaction::QuestionType::Text, - }; - - Ok(shared_types::user_interaction::UserQuestion { - id: row.get(0)?, - question: row.get(1)?, - description: row.get(2)?, - response_type, - default: None, - options: None, - }) - })?; - - let mut result = Vec::new(); - for question in questions { - result.push(question?); - } - - Ok(result) - } - - /// Store answers for questions in the database - pub fn store_answers( - &self, - session_id: i64, - answers: &std::collections::HashMap, - ) -> anyhow::Result<()> { - let conn = self.connection.lock().unwrap(); - let now = chrono::Utc::now().timestamp(); - - for (question_id, answer) in answers { - conn.execute( - "UPDATE project_requirements_qna - SET answer = ?1, answered_at = ?2 - WHERE session_id = ?3 AND question_id = ?4", - params![answer, now, session_id, question_id], - )?; - } - - Ok(()) - } -} diff --git a/nocodo-agents/src/requirements_gathering/mod.rs b/nocodo-agents/src/requirements_gathering/mod.rs index 55a743be..71d8e54c 100644 --- a/nocodo-agents/src/requirements_gathering/mod.rs +++ b/nocodo-agents/src/requirements_gathering/mod.rs @@ -1,4 +1,3 @@ -pub mod database; pub mod models; pub mod storage; diff --git a/nocodo-agents/src/settings_management/database.rs b/nocodo-agents/src/settings_management/database.rs deleted file mode 100644 index 4e3122ca..00000000 --- a/nocodo-agents/src/settings_management/database.rs +++ /dev/null @@ -1,133 +0,0 @@ -use crate::database::Database; -use rusqlite::params; - -/// Settings management database operations -impl Database { - /// Store settings in the project_settings table - pub fn store_settings( - &self, - session_id: i64, - tool_call_id: Option, - settings: &[shared_types::user_interaction::UserQuestion], - ) -> anyhow::Result<()> { - let conn = self.connection.lock().unwrap(); - let now = chrono::Utc::now().timestamp(); - - for setting in settings { - conn.execute( - "INSERT INTO project_settings (session_id, tool_call_id, setting_key, setting_name, description, setting_type, created_at) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)", - params![ - session_id, - tool_call_id, - &setting.id, - &setting.question, - &setting.description, - format!("{:?}", setting.response_type).to_lowercase(), - now - ], - )?; - } - - Ok(()) - } - - /// Get pending (unanswered) settings from the database - pub fn get_pending_settings( - &self, - session_id: i64, - ) -> anyhow::Result> { - let conn = self.connection.lock().unwrap(); - let mut stmt = conn.prepare( - "SELECT setting_key, setting_name, description, setting_type - FROM project_settings - WHERE session_id = ?1 AND setting_value IS NULL - ORDER BY created_at ASC", - )?; - - let settings = stmt.query_map([session_id], |row| { - let setting_type_str: String = row.get(3)?; - let response_type = match setting_type_str.as_str() { - "text" => shared_types::user_interaction::QuestionType::Text, - "password" => shared_types::user_interaction::QuestionType::Password, - "file_path" => shared_types::user_interaction::QuestionType::FilePath, - "email" => shared_types::user_interaction::QuestionType::Email, - "url" => shared_types::user_interaction::QuestionType::Url, - _ => shared_types::user_interaction::QuestionType::Text, - }; - - Ok(shared_types::user_interaction::UserQuestion { - id: row.get(0)?, - question: row.get(1)?, - description: row.get(2)?, - response_type, - default: None, - options: None, - }) - })?; - - let mut result = Vec::new(); - for setting in settings { - result.push(setting?); - } - - Ok(result) - } - - /// Store setting values for settings in the database - pub fn store_setting_values( - &self, - session_id: i64, - setting_values: &std::collections::HashMap, - ) -> anyhow::Result<()> { - let conn = self.connection.lock().unwrap(); - let now = chrono::Utc::now().timestamp(); - - for (setting_key, value) in setting_values { - conn.execute( - "UPDATE project_settings - SET setting_value = ?1, updated_at = ?2 - WHERE session_id = ?3 AND setting_key = ?4", - params![value, now, session_id, setting_key], - )?; - } - - Ok(()) - } - - /// Get all settings for a session (both pending and completed) - pub fn get_session_settings( - &self, - session_id: i64, - ) -> anyhow::Result> { - let conn = self.connection.lock().unwrap(); - let mut stmt = conn.prepare( - "SELECT id, session_id, tool_call_id, setting_key, setting_name, description, setting_type, setting_value, created_at, updated_at - FROM project_settings - WHERE session_id = ?1 - ORDER BY created_at ASC", - )?; - - let settings = stmt.query_map([session_id], |row| { - Ok(crate::settings_management::models::ProjectSetting { - id: row.get(0)?, - session_id: row.get(1)?, - tool_call_id: row.get(2)?, - setting_key: row.get(3)?, - setting_name: row.get(4)?, - description: row.get(5)?, - setting_type: row.get(6)?, - setting_value: row.get(7)?, - created_at: row.get(8)?, - updated_at: row.get(9)?, - }) - })?; - - let mut result = Vec::new(); - for setting in settings { - result.push(setting?); - } - - Ok(result) - } -} diff --git a/nocodo-agents/src/settings_management/mod.rs b/nocodo-agents/src/settings_management/mod.rs index 3f45d0e6..2b3abaa3 100644 --- a/nocodo-agents/src/settings_management/mod.rs +++ b/nocodo-agents/src/settings_management/mod.rs @@ -1,4 +1,3 @@ -pub mod database; pub mod models; use crate::{ @@ -161,7 +160,10 @@ using the tool."#, completed_at: None, error_details: None, }; - let call_id = self.storage.create_tool_call(tool_call_record.clone()).await?; + let call_id = self + .storage + .create_tool_call(tool_call_record.clone()) + .await?; tool_call_record.id = Some(call_id); let start = Instant::now(); @@ -422,7 +424,10 @@ impl Agent for SettingsManagementAgent { completed_at: None, error_details: None, }; - let tool_call_id = self.storage.create_tool_call(tool_call_record.clone()).await?; + let tool_call_id = self + .storage + .create_tool_call(tool_call_record.clone()) + .await?; tool_call_record.id = Some(tool_call_id); // Execute the ask_user tool to get responses from user diff --git a/nocodo-agents/src/storage/memory.rs b/nocodo-agents/src/storage/memory.rs index 2345fe08..c0e60615 100644 --- a/nocodo-agents/src/storage/memory.rs +++ b/nocodo-agents/src/storage/memory.rs @@ -53,11 +53,12 @@ impl AgentStorage for InMemoryStorage { async fn create_message(&self, message: Message) -> Result { let message_id = uuid::Uuid::new_v4().to_string(); + let session_id = message.session_id.clone(); let mut message_with_id = message; message_with_id.id = Some(message_id.clone()); let mut messages = self.messages.lock().unwrap(); messages - .entry(message.session_id.clone()) + .entry(session_id) .or_insert_with(Vec::new) .push(message_with_id); Ok(message_id) @@ -75,11 +76,12 @@ impl AgentStorage for InMemoryStorage { async fn create_tool_call(&self, tool_call: ToolCall) -> Result { let tool_call_id = uuid::Uuid::new_v4().to_string(); + let session_id = tool_call.session_id.clone(); let mut tool_call_with_id = tool_call; tool_call_with_id.id = Some(tool_call_id.clone()); let mut tool_calls = self.tool_calls.lock().unwrap(); tool_calls - .entry(tool_call.session_id.clone()) + .entry(session_id) .or_insert_with(Vec::new) .push(tool_call_with_id); Ok(tool_call_id) diff --git a/nocodo-api/tasks/implement-sqlite-agent-storage.md b/nocodo-api/tasks/implement-sqlite-agent-storage.md new file mode 100644 index 00000000..3c2045b4 --- /dev/null +++ b/nocodo-api/tasks/implement-sqlite-agent-storage.md @@ -0,0 +1,712 @@ +# Implement SQLite Agent Storage for nocodo-api + +**Status**: 📋 Not Started +**Priority**: High +**Created**: 2026-02-03 + +## Summary + +Implement SQLite-based storage for nocodo-agents by creating a `SqliteAgentStorage` struct that implements the `AgentStorage` trait. This allows nocodo-api to continue using SQLite for storing agent sessions, messages, and tool calls. + +## Problem Statement + +After refactoring nocodo-agents to use trait-based storage, nocodo-api needs a concrete SQLite implementation to store agent execution data. The storage must: +- Implement the `AgentStorage` trait from nocodo-agents +- Use SQLite for persistence +- Manage database migrations +- Support the existing nocodo-api database infrastructure + +## Goals + +1. **Implement AgentStorage trait**: Create `SqliteAgentStorage` struct with SQLite backend +2. **Database migrations**: Port migrations from nocodo-agents to nocodo-api +3. **Connection management**: Use Arc-wrapped connection for thread safety +4. **Integration**: Update nocodo-api to use new storage implementation +5. **Maintain compatibility**: Keep existing nocodo-api functionality working + +## Architecture + +### Storage Implementation Structure + +```rust +// nocodo-api/src/storage/sqlite.rs + +use nocodo_agents::{AgentStorage, Session, Message, ToolCall, StorageError}; +use rusqlite::Connection; +use std::sync::{Arc, Mutex}; +use async_trait::async_trait; + +pub type DbConnection = Arc>; + +pub struct SqliteAgentStorage { + connection: DbConnection, +} + +impl SqliteAgentStorage { + pub fn new(connection: DbConnection) -> Self { + Self { connection } + } +} + +#[async_trait] +impl AgentStorage for SqliteAgentStorage { + async fn create_session(&self, session: Session) -> Result { + // SQLite implementation + } + // ... other methods +} +``` + +## Implementation Plan + +### Phase 1: Create Storage Module Structure + +#### 1.1 Create Module Files + +Create new directory and files: +``` +nocodo-api/src/ + storage/ + mod.rs # Module exports + sqlite.rs # SqliteAgentStorage implementation + migrations/ # SQLite migrations +``` + +### Phase 2: Port Migrations + +#### 2.1 Copy Migration Files + +Copy migration files from nocodo-agents to nocodo-api: + +**Source**: `nocodo-agents/src/database/migrations/` +**Destination**: `nocodo-api/src/storage/migrations/` + +Migration files to copy: +- `V1__create_agent_sessions.rs` +- `V2__create_agent_messages.rs` +- `V3__create_agent_tool_calls.rs` +- `V4__create_project_requirements_qna.rs` +- `V5__create_project_settings.rs` + +#### 2.2 Update Migration Module + +**File**: `nocodo-api/src/storage/migrations/mod.rs` + +```rust +use refinery::embed_migrations; + +embed_migrations!("src/storage/migrations"); + +pub fn run_migrations(connection: &mut rusqlite::Connection) -> Result<(), refinery::Error> { + migrations::runner().run(connection) +} +``` + +### Phase 3: Implement SqliteAgentStorage + +#### 3.1 Implement Core Storage Methods + +**File**: `nocodo-api/src/storage/sqlite.rs` + +```rust +use async_trait::async_trait; +use chrono::Utc; +use nocodo_agents::{ + AgentStorage, Message, MessageRole, Session, SessionStatus, StorageError, ToolCall, + ToolCallStatus, +}; +use rusqlite::{params, Connection}; +use std::sync::{Arc, Mutex}; + +pub type DbConnection = Arc>; + +pub struct SqliteAgentStorage { + connection: DbConnection, +} + +impl SqliteAgentStorage { + pub fn new(connection: DbConnection) -> Self { + Self { connection } + } +} + +#[async_trait] +impl AgentStorage for SqliteAgentStorage { + async fn create_session(&self, mut session: Session) -> Result { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let status_str = match session.status { + SessionStatus::Running => "running", + SessionStatus::Completed => "completed", + SessionStatus::Failed => "failed", + SessionStatus::WaitingForUserInput => "waiting_for_user_input", + }; + + let config_json = serde_json::to_string(&session.config) + .map_err(StorageError::SerializationError)?; + + conn.execute( + r#" + INSERT INTO agent_sessions + (agent_name, provider, model, system_prompt, user_prompt, config, status, started_at) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) + "#, + params![ + session.agent_name, + session.provider, + session.model, + session.system_prompt, + session.user_prompt, + config_json, + status_str, + session.started_at, + ], + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let id = conn.last_insert_rowid().to_string(); + Ok(id) + } + + async fn get_session(&self, session_id: &str) -> Result, StorageError> { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let id: i64 = session_id + .parse() + .map_err(|_| StorageError::NotFound(format!("Invalid session ID: {}", session_id)))?; + + let mut stmt = conn + .prepare( + r#" + SELECT id, agent_name, provider, model, system_prompt, user_prompt, + config, status, started_at, ended_at, result, error + FROM agent_sessions + WHERE id = ?1 + "#, + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let session = stmt + .query_row(params![id], |row| { + let status_str: String = row.get(7)?; + let status = match status_str.as_str() { + "running" => SessionStatus::Running, + "completed" => SessionStatus::Completed, + "failed" => SessionStatus::Failed, + "waiting_for_user_input" => SessionStatus::WaitingForUserInput, + _ => SessionStatus::Running, + }; + + let config_json: String = row.get(6)?; + let config: serde_json::Value = serde_json::from_str(&config_json) + .unwrap_or(serde_json::json!({})); + + Ok(Session { + id: Some(row.get::<_, i64>(0)?.to_string()), + agent_name: row.get(1)?, + provider: row.get(2)?, + model: row.get(3)?, + system_prompt: row.get(4)?, + user_prompt: row.get(5)?, + config, + status, + started_at: row.get(8)?, + ended_at: row.get(9)?, + result: row.get(10)?, + error: row.get(11)?, + }) + }) + .optional() + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + Ok(session) + } + + async fn update_session(&self, session: Session) -> Result<(), StorageError> { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let id: i64 = session + .id + .as_ref() + .and_then(|id| id.parse().ok()) + .ok_or_else(|| StorageError::NotFound("Session ID required for update".to_string()))?; + + let status_str = match session.status { + SessionStatus::Running => "running", + SessionStatus::Completed => "completed", + SessionStatus::Failed => "failed", + SessionStatus::WaitingForUserInput => "waiting_for_user_input", + }; + + conn.execute( + r#" + UPDATE agent_sessions + SET status = ?1, ended_at = ?2, result = ?3, error = ?4 + WHERE id = ?5 + "#, + params![status_str, session.ended_at, session.result, session.error, id], + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + Ok(()) + } + + async fn create_message(&self, mut message: Message) -> Result { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let session_id: i64 = message + .session_id + .parse() + .map_err(|_| StorageError::OperationFailed("Invalid session ID".to_string()))?; + + let role_str = match message.role { + MessageRole::User => "user", + MessageRole::Assistant => "assistant", + MessageRole::System => "system", + MessageRole::Tool => "tool", + }; + + conn.execute( + r#" + INSERT INTO agent_messages (session_id, role, content, created_at) + VALUES (?1, ?2, ?3, ?4) + "#, + params![session_id, role_str, message.content, message.created_at], + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let id = conn.last_insert_rowid().to_string(); + Ok(id) + } + + async fn get_messages(&self, session_id: &str) -> Result, StorageError> { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let id: i64 = session_id + .parse() + .map_err(|_| StorageError::NotFound(format!("Invalid session ID: {}", session_id)))?; + + let mut stmt = conn + .prepare( + r#" + SELECT id, session_id, role, content, created_at + FROM agent_messages + WHERE session_id = ?1 + ORDER BY created_at ASC + "#, + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let messages = stmt + .query_map(params![id], |row| { + let role_str: String = row.get(2)?; + let role = match role_str.as_str() { + "user" => MessageRole::User, + "assistant" => MessageRole::Assistant, + "system" => MessageRole::System, + "tool" => MessageRole::Tool, + _ => MessageRole::User, + }; + + Ok(Message { + id: Some(row.get::<_, i64>(0)?.to_string()), + session_id: row.get::<_, i64>(1)?.to_string(), + role, + content: row.get(3)?, + created_at: row.get(4)?, + }) + }) + .map_err(|e| StorageError::OperationFailed(e.to_string()))? + .collect::, _>>() + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + Ok(messages) + } + + async fn create_tool_call(&self, tool_call: ToolCall) -> Result { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let session_id: i64 = tool_call + .session_id + .parse() + .map_err(|_| StorageError::OperationFailed("Invalid session ID".to_string()))?; + + let message_id: Option = tool_call + .message_id + .as_ref() + .and_then(|id| id.parse().ok()); + + let status_str = match tool_call.status { + ToolCallStatus::Pending => "pending", + ToolCallStatus::Executing => "executing", + ToolCallStatus::Completed => "completed", + ToolCallStatus::Failed => "failed", + }; + + let request_json = serde_json::to_string(&tool_call.request) + .map_err(StorageError::SerializationError)?; + + conn.execute( + r#" + INSERT INTO agent_tool_calls + (session_id, message_id, tool_call_id, tool_name, request, status, created_at) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7) + "#, + params![ + session_id, + message_id, + tool_call.tool_call_id, + tool_call.tool_name, + request_json, + status_str, + tool_call.created_at, + ], + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let id = conn.last_insert_rowid().to_string(); + Ok(id) + } + + async fn update_tool_call(&self, tool_call: ToolCall) -> Result<(), StorageError> { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let id: i64 = tool_call + .id + .as_ref() + .and_then(|id| id.parse().ok()) + .ok_or_else(|| StorageError::NotFound("Tool call ID required for update".to_string()))?; + + let status_str = match tool_call.status { + ToolCallStatus::Pending => "pending", + ToolCallStatus::Executing => "executing", + ToolCallStatus::Completed => "completed", + ToolCallStatus::Failed => "failed", + }; + + let response_json = tool_call + .response + .map(|r| serde_json::to_string(&r)) + .transpose() + .map_err(StorageError::SerializationError)?; + + conn.execute( + r#" + UPDATE agent_tool_calls + SET status = ?1, response = ?2, execution_time_ms = ?3, + completed_at = ?4, error_details = ?5 + WHERE id = ?6 + "#, + params![ + status_str, + response_json, + tool_call.execution_time_ms, + tool_call.completed_at, + tool_call.error_details, + id + ], + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + Ok(()) + } + + async fn get_tool_calls(&self, session_id: &str) -> Result, StorageError> { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let id: i64 = session_id + .parse() + .map_err(|_| StorageError::NotFound(format!("Invalid session ID: {}", session_id)))?; + + let mut stmt = conn + .prepare( + r#" + SELECT id, session_id, message_id, tool_call_id, tool_name, request, + response, status, execution_time_ms, created_at, completed_at, error_details + FROM agent_tool_calls + WHERE session_id = ?1 + ORDER BY created_at ASC + "#, + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let tool_calls = stmt + .query_map(params![id], |row| { + let status_str: String = row.get(7)?; + let status = match status_str.as_str() { + "pending" => ToolCallStatus::Pending, + "executing" => ToolCallStatus::Executing, + "completed" => ToolCallStatus::Completed, + "failed" => ToolCallStatus::Failed, + _ => ToolCallStatus::Pending, + }; + + let request_json: String = row.get(5)?; + let request = serde_json::from_str(&request_json).unwrap_or(serde_json::json!({})); + + let response_json: Option = row.get(6)?; + let response = response_json + .and_then(|json| serde_json::from_str(&json).ok()); + + Ok(ToolCall { + id: Some(row.get::<_, i64>(0)?.to_string()), + session_id: row.get::<_, i64>(1)?.to_string(), + message_id: row.get::<_, Option>(2)?.map(|id| id.to_string()), + tool_call_id: row.get(3)?, + tool_name: row.get(4)?, + request, + response, + status, + execution_time_ms: row.get(8)?, + created_at: row.get(9)?, + completed_at: row.get(10)?, + error_details: row.get(11)?, + }) + }) + .map_err(|e| StorageError::OperationFailed(e.to_string()))? + .collect::, _>>() + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + Ok(tool_calls) + } + + async fn get_pending_tool_calls(&self, session_id: &str) -> Result, StorageError> { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let id: i64 = session_id + .parse() + .map_err(|_| StorageError::NotFound(format!("Invalid session ID: {}", session_id)))?; + + let mut stmt = conn + .prepare( + r#" + SELECT id, session_id, message_id, tool_call_id, tool_name, request, + response, status, execution_time_ms, created_at, completed_at, error_details + FROM agent_tool_calls + WHERE session_id = ?1 AND status = 'pending' + ORDER BY created_at ASC + "#, + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let tool_calls = stmt + .query_map(params![id], |row| { + let request_json: String = row.get(5)?; + let request = serde_json::from_str(&request_json).unwrap_or(serde_json::json!({})); + + let response_json: Option = row.get(6)?; + let response = response_json + .and_then(|json| serde_json::from_str(&json).ok()); + + Ok(ToolCall { + id: Some(row.get::<_, i64>(0)?.to_string()), + session_id: row.get::<_, i64>(1)?.to_string(), + message_id: row.get::<_, Option>(2)?.map(|id| id.to_string()), + tool_call_id: row.get(3)?, + tool_name: row.get(4)?, + request, + response, + status: ToolCallStatus::Pending, + execution_time_ms: row.get(8)?, + created_at: row.get(9)?, + completed_at: row.get(10)?, + error_details: row.get(11)?, + }) + }) + .map_err(|e| StorageError::OperationFailed(e.to_string()))? + .collect::, _>>() + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + Ok(tool_calls) + } +} +``` + +### Phase 4: Update nocodo-api Integration + +#### 4.1 Update Database Helper + +**File**: `nocodo-api/src/helpers/database.rs` + +Add migration runner: +```rust +use crate::storage::migrations; + +pub fn initialize_database(db_path: &str) -> anyhow::Result { + // Create parent directories if needed + if let Some(parent) = Path::new(db_path).parent() { + std::fs::create_dir_all(parent)?; + } + + let mut conn = Connection::open(db_path)?; + + // Enable foreign keys + conn.execute("PRAGMA foreign_keys = ON", [])?; + + // Run migrations + migrations::run_migrations(&mut conn)?; + + Ok(Arc::new(Mutex::new(conn))) +} +``` + +#### 4.2 Update main.rs + +**File**: `nocodo-api/src/main.rs` + +Update to use new storage: +```rust +use nocodo_api::storage::SqliteAgentStorage; + +#[actix_web::main] +async fn main() -> std::io::Result<()> { + // ... config loading + + // Initialize database + let db_connection = helpers::database::initialize_database(&config.database.path)?; + + // Create storage + let storage = Arc::new(SqliteAgentStorage::new(db_connection.clone())); + + // Create app state with storage + let app_state = web::Data::new(AppState { + storage, + // ... other fields + }); + + // ... rest of main +} +``` + +#### 4.3 Update Cargo.toml + +**File**: `nocodo-api/Cargo.toml` + +Ensure dependencies: +```toml +[dependencies] +nocodo-agents = { path = "../nocodo-agents" } +nocodo-tools = { path = "../nocodo-tools" } +rusqlite = { version = "0.37", features = ["bundled"] } +refinery = { version = "0.9", features = ["rusqlite"] } +chrono = { version = "0.4", features = ["serde"] } +async-trait = "0.1" +``` + +### Phase 5: Update Handlers + +#### 5.1 Update Agent Execution Handler + +**File**: `nocodo-api/src/handlers/agent_execution.rs` + +Update to use storage from app state: +```rust +pub async fn execute_agent( + agent_id: web::Path, + request: web::Json, + data: web::Data, +) -> Result { + let agent = match agent_id.as_str() { + "sqlite" => { + let agent = SqliteReaderAgent::new( + data.llm_client.clone(), + data.storage.clone(), // Use storage from app state + data.tool_executor.clone(), + request.db_path.clone(), + )?; + Box::new(agent) as Box> + } + // ... other agents + }; + + let result = agent.execute(&request.user_prompt).await?; + Ok(HttpResponse::Ok().json(result)) +} +``` + +## Files Changed + +### New Files +- `nocodo-api/src/storage/mod.rs` +- `nocodo-api/src/storage/sqlite.rs` +- `nocodo-api/src/storage/migrations/mod.rs` +- `nocodo-api/src/storage/migrations/V1__create_agent_sessions.rs` +- `nocodo-api/src/storage/migrations/V2__create_agent_messages.rs` +- `nocodo-api/src/storage/migrations/V3__create_agent_tool_calls.rs` +- `nocodo-api/src/storage/migrations/V4__create_project_requirements_qna.rs` +- `nocodo-api/src/storage/migrations/V5__create_project_settings.rs` +- `nocodo-api/tasks/implement-sqlite-agent-storage.md` + +### Modified Files +- `nocodo-api/Cargo.toml` - Add dependencies +- `nocodo-api/src/lib.rs` - Export storage module +- `nocodo-api/src/main.rs` - Initialize storage +- `nocodo-api/src/helpers/database.rs` - Add migration runner +- `nocodo-api/src/handlers/agent_execution.rs` - Use new storage + +## Testing Strategy + +### Compilation +```bash +cd nocodo-api +cargo check +``` + +### Manual Testing +```bash +cargo run + +# In another terminal +curl -X POST http://localhost:8080/agents/sqlite/execute \ + -H "Content-Type: application/json" \ + -d '{ + "user_prompt": "List all tables", + "db_path": "/path/to/test.db" + }' +``` + +## Success Criteria + +- [ ] `SqliteAgentStorage` struct implements `AgentStorage` trait +- [ ] All trait methods implemented with SQLite backend +- [ ] Migrations ported and working +- [ ] Database initialized at startup +- [ ] Agent execution uses new storage +- [ ] Existing API endpoints continue working +- [ ] Code compiles without errors +- [ ] No clippy warnings +- [ ] Manual testing successful + +## Notes + +- This implementation maintains backward compatibility with existing nocodo-api behavior +- SQLite remains the storage backend for nocodo-api +- The storage layer is now isolated and can be tested independently +- Future enhancements could include connection pooling if needed From c73adbf0346105d64159f4304c6819cbd957a047 Mon Sep 17 00:00:00 2001 From: Sumit Datta Date: Tue, 3 Feb 2026 14:18:01 +0530 Subject: [PATCH 4/6] refactor: change all IDs from String to i64 for type safety and performance MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes the storage refactoring by switching all identifier types from String to i64 across both nocodo-agents library and nocodo-api application. This change eliminates runtime string parsing overhead and improves type safety. Changes in nocodo-agents: - Update AgentStorage trait to use i64 for all ID types - Update Session, Message, ToolCall types to use i64 for IDs - Update RequirementsStorage trait to use i64 - Replace UUID-based InMemoryStorage with atomic i64 counter - Update all 8 agent implementations to use i64 internally - Update all 6 binary runners to use i64 directly Changes in nocodo-api: - Implement SqliteAgentStorage with i64 native support - Add storage module with migrations - Update all 8 agent execution handlers to use i64 - Update sessions and requirements handling for i64 - Fix DirectRequirementsStorage implementation - Remove 150+ string conversion calls across codebase Benefits: - Zero runtime parsing overhead - Type-safe ID handling throughout - Simpler, cleaner code - Direct SQLite INTEGER ↔ i64 mapping Co-Authored-By: Claude Sonnet 4.5 --- Cargo.lock | 50 +++ nocodo-agents/bin/codebase_analysis_runner.rs | 5 +- nocodo-agents/bin/imap_email_runner.rs | 10 +- .../bin/requirements_gathering_runner.rs | 5 +- .../bin/settings_management_runner.rs | 5 +- nocodo-agents/bin/sqlite_reader_runner.rs | 5 +- nocodo-agents/bin/structured_json_runner.rs | 5 +- nocodo-agents/src/codebase_analysis/mod.rs | 31 +- nocodo-agents/src/imap_email/mod.rs | 31 +- nocodo-agents/src/pdftotext/mod.rs | 32 +- .../src/requirements_gathering/mod.rs | 43 +- .../src/requirements_gathering/storage.rs | 8 +- nocodo-agents/src/settings_management/mod.rs | 39 +- nocodo-agents/src/sqlite_reader/mod.rs | 31 +- nocodo-agents/src/storage/memory.rs | 89 ++-- nocodo-agents/src/storage/mod.rs | 14 +- nocodo-agents/src/structured_json/mod.rs | 17 +- nocodo-agents/src/tesseract/mod.rs | 32 +- nocodo-agents/src/types/message.rs | 4 +- nocodo-agents/src/types/session.rs | 2 +- nocodo-agents/src/types/tool_call.rs | 6 +- nocodo-api/Cargo.toml | 1 + .../codebase_analysis_agent.rs | 66 ++- .../agent_execution/imap_email_agent.rs | 83 +++- .../agent_execution/pdftotext_agent.rs | 99 +++-- .../requirements_gathering_agent.rs | 89 +++- .../settings_management_agent.rs | 83 +++- .../handlers/agent_execution/sqlite_agent.rs | 83 +++- .../agent_execution/tesseract_agent.rs | 104 +++-- .../workflow_creation_agent.rs | 83 +++- nocodo-api/src/handlers/sessions.rs | 169 ++++++-- nocodo-api/src/helpers/agents.rs | 205 +++++++-- nocodo-api/src/helpers/database.rs | 13 +- nocodo-api/src/lib.rs | 1 + nocodo-api/src/main.rs | 5 +- .../migrations/V1__create_agent_sessions.rs | 19 + .../migrations/V2__create_agent_messages.rs | 17 + .../migrations/V3__create_agent_tool_calls.rs | 27 ++ .../V4__create_project_requirements_qna.rs | 26 ++ .../migrations/V5__create_project_settings.rs | 29 ++ nocodo-api/src/storage/migrations/mod.rs | 13 + nocodo-api/src/storage/mod.rs | 4 + nocodo-api/src/storage/sqlite.rs | 409 ++++++++++++++++++ 43 files changed, 1628 insertions(+), 464 deletions(-) create mode 100644 nocodo-api/src/storage/migrations/V1__create_agent_sessions.rs create mode 100644 nocodo-api/src/storage/migrations/V2__create_agent_messages.rs create mode 100644 nocodo-api/src/storage/migrations/V3__create_agent_tool_calls.rs create mode 100644 nocodo-api/src/storage/migrations/V4__create_project_requirements_qna.rs create mode 100644 nocodo-api/src/storage/migrations/V5__create_project_settings.rs create mode 100644 nocodo-api/src/storage/migrations/mod.rs create mode 100644 nocodo-api/src/storage/mod.rs create mode 100644 nocodo-api/src/storage/sqlite.rs diff --git a/Cargo.lock b/Cargo.lock index 40451244..45a8301d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3076,6 +3076,7 @@ dependencies = [ "nocodo-agents", "nocodo-llm-sdk", "nocodo-tools", + "refinery", "rusqlite", "serde", "serde_json", @@ -4007,6 +4008,49 @@ dependencies = [ "syn", ] +[[package]] +name = "refinery" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52c427f2572afe5c6cbfa2b1bf40071c89bf1a8539e958ea582842f6f38dcfae" +dependencies = [ + "refinery-core", + "refinery-macros", +] + +[[package]] +name = "refinery-core" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "702655abfc67f93a6f735e9fa4ace7d2e580633f8961f28acbfd7583ddce936c" +dependencies = [ + "async-trait", + "cfg-if", + "log", + "regex", + "rusqlite", + "serde", + "siphasher", + "thiserror 2.0.17", + "time", + "toml 0.8.23", + "url", + "walkdir", +] + +[[package]] +name = "refinery-macros" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5145756cdf293b5089dc6b4f103f1a1229cc55d67082c866f8c8289531c4b983" +dependencies = [ + "proc-macro2", + "quote", + "refinery-core", + "regex", + "syn", +] + [[package]] name = "regex" version = "1.12.2" @@ -4811,6 +4855,12 @@ version = "2.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbbb5d9659141646ae647b42fe094daf6c6192d1620870b449d9557f748b2daa" +[[package]] +name = "siphasher" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e" + [[package]] name = "slab" version = "0.4.11" diff --git a/nocodo-agents/bin/codebase_analysis_runner.rs b/nocodo-agents/bin/codebase_analysis_runner.rs index c86be628..d1ebd7b1 100644 --- a/nocodo-agents/bin/codebase_analysis_runner.rs +++ b/nocodo-agents/bin/codebase_analysis_runner.rs @@ -92,10 +92,7 @@ async fn main() -> anyhow::Result<()> { result: None, error: None, }; - let session_id_str = storage.create_session(session).await?; - - // Parse session ID as i64 for agent.execute() - let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); + let session_id = storage.create_session(session).await?; // Execute agent let result = agent.execute(&args.prompt, session_id).await?; diff --git a/nocodo-agents/bin/imap_email_runner.rs b/nocodo-agents/bin/imap_email_runner.rs index 1dd85d83..ef9394dd 100644 --- a/nocodo-agents/bin/imap_email_runner.rs +++ b/nocodo-agents/bin/imap_email_runner.rs @@ -154,10 +154,7 @@ async fn run_single_query( result: None, error: None, }; - let session_id_str = storage.create_session(session).await?; - - // Parse session ID as i64 for agent.execute() - let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); + let session_id = storage.create_session(session).await?; // Execute agent let result = agent.execute(prompt, session_id).await?; @@ -187,10 +184,7 @@ async fn run_interactive_mode( result: None, error: None, }; - let session_id_str = storage.create_session(session).await?; - - // Parse session ID as i64 for agent.execute() - let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); + let session_id = storage.create_session(session).await?; println!("🔄 Interactive mode enabled - session ID: {}", session_id); println!("💡 Type your queries. Type 'quit' or 'exit' to end the session.\n"); diff --git a/nocodo-agents/bin/requirements_gathering_runner.rs b/nocodo-agents/bin/requirements_gathering_runner.rs index 4f3b2ecd..8f855f85 100644 --- a/nocodo-agents/bin/requirements_gathering_runner.rs +++ b/nocodo-agents/bin/requirements_gathering_runner.rs @@ -74,10 +74,7 @@ async fn main() -> anyhow::Result<()> { result: None, error: None, }; - let session_id_str = storage.create_session(session).await?; - - // Parse session ID as i64 for agent.execute() - let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); + let session_id = storage.create_session(session).await?; let result = agent.execute(&args.prompt, session_id).await?; diff --git a/nocodo-agents/bin/settings_management_runner.rs b/nocodo-agents/bin/settings_management_runner.rs index e8ee960c..5867ec4a 100644 --- a/nocodo-agents/bin/settings_management_runner.rs +++ b/nocodo-agents/bin/settings_management_runner.rs @@ -79,10 +79,7 @@ async fn main() -> anyhow::Result<()> { result: None, error: None, }; - let session_id_str = storage.create_session(session).await?; - - // Parse session ID as i64 for agent.execute() - let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); + let session_id = storage.create_session(session).await?; let result = agent.execute(&args.prompt, session_id).await?; diff --git a/nocodo-agents/bin/sqlite_reader_runner.rs b/nocodo-agents/bin/sqlite_reader_runner.rs index fcf394dd..84136c46 100644 --- a/nocodo-agents/bin/sqlite_reader_runner.rs +++ b/nocodo-agents/bin/sqlite_reader_runner.rs @@ -73,10 +73,7 @@ async fn main() -> anyhow::Result<()> { result: None, error: None, }; - let session_id_str = storage.create_session(session).await?; - - // Parse session ID as i64 for agent.execute() - let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); + let session_id = storage.create_session(session).await?; let result = agent.execute(&args.prompt, session_id).await?; diff --git a/nocodo-agents/bin/structured_json_runner.rs b/nocodo-agents/bin/structured_json_runner.rs index f5683083..16aa7718 100644 --- a/nocodo-agents/bin/structured_json_runner.rs +++ b/nocodo-agents/bin/structured_json_runner.rs @@ -93,10 +93,7 @@ async fn main() -> anyhow::Result<()> { result: None, error: None, }; - let session_id_str = storage.create_session(session).await?; - - // Parse session ID as i64 for agent.execute() - let session_id = session_id_str.parse::().unwrap_or_else(|_| 1); + let session_id = storage.create_session(session).await?; match agent.execute(&args.prompt, session_id).await { Ok(result) => { diff --git a/nocodo-agents/src/codebase_analysis/mod.rs b/nocodo-agents/src/codebase_analysis/mod.rs index 3a712335..3788e896 100644 --- a/nocodo-agents/src/codebase_analysis/mod.rs +++ b/nocodo-agents/src/codebase_analysis/mod.rs @@ -67,10 +67,9 @@ impl Agent for CodebaseAnalysisAgent { } async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { - let session_id_str = session_id.to_string(); let user_message = Message { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::User, content: user_prompt.to_string(), created_at: chrono::Utc::now().timestamp(), @@ -85,7 +84,7 @@ impl Agent for CodebaseAnalysisAgent { iteration += 1; if iteration > max_iterations { let error = "Maximum iteration limit reached"; - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Failed; session.error = Some(error.to_string()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -93,7 +92,7 @@ impl Agent for CodebaseAnalysisAgent { return Err(anyhow::anyhow!(error)); } - let messages = self.build_messages(&session_id_str).await?; + let messages = self.build_messages(session_id).await?; let request = CompletionRequest { messages, @@ -113,7 +112,7 @@ impl Agent for CodebaseAnalysisAgent { let text = extract_text_from_content(&response.content); let assistant_message = Message { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::Assistant, content: text.clone(), created_at: chrono::Utc::now().timestamp(), @@ -122,7 +121,7 @@ impl Agent for CodebaseAnalysisAgent { if let Some(tool_calls) = response.tool_calls { if tool_calls.is_empty() { - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Completed; session.result = Some(text.clone()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -131,11 +130,11 @@ impl Agent for CodebaseAnalysisAgent { } for tool_call in tool_calls { - self.execute_tool_call(&session_id_str, Some(&message_id), &tool_call) + self.execute_tool_call(session_id, Some(message_id), &tool_call) .await?; } } else { - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Completed; session.result = Some(text.clone()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -147,14 +146,14 @@ impl Agent for CodebaseAnalysisAgent { } impl CodebaseAnalysisAgent { - async fn get_session(&self, session_id: &str) -> anyhow::Result { + async fn get_session(&self, session_id: i64) -> anyhow::Result { self.storage .get_session(session_id) .await? .ok_or_else(|| anyhow::anyhow!("Session not found: {}", session_id)) } - async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + async fn build_messages(&self, session_id: i64) -> anyhow::Result> { let db_messages = self.storage.get_messages(session_id).await?; db_messages @@ -177,8 +176,8 @@ impl CodebaseAnalysisAgent { async fn execute_tool_call( &self, - session_id: &str, - message_id: Option<&String>, + session_id: i64, + message_id: Option, tool_call: &LlmToolCall, ) -> anyhow::Result<()> { let tool_request = @@ -186,8 +185,8 @@ impl CodebaseAnalysisAgent { let mut tool_call_record = StorageToolCall { id: None, - session_id: session_id.to_string(), - message_id: message_id.cloned(), + session_id, + message_id, tool_call_id: tool_call.id().to_string(), tool_name: tool_call.name().to_string(), request: tool_call.arguments().clone(), @@ -228,7 +227,7 @@ impl CodebaseAnalysisAgent { let tool_message = Message { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -254,7 +253,7 @@ impl CodebaseAnalysisAgent { let tool_error_message = Message { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: error_message_to_llm, created_at: chrono::Utc::now().timestamp(), diff --git a/nocodo-agents/src/imap_email/mod.rs b/nocodo-agents/src/imap_email/mod.rs index 9981cc34..9edbbd06 100644 --- a/nocodo-agents/src/imap_email/mod.rs +++ b/nocodo-agents/src/imap_email/mod.rs @@ -195,7 +195,7 @@ impl ImapEmailAgent { )) } - async fn get_session(&self, session_id: &str) -> anyhow::Result { + async fn get_session(&self, session_id: i64) -> anyhow::Result { self.storage .get_session(session_id) .await? @@ -209,7 +209,7 @@ impl ImapEmailAgent { .collect() } - async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + async fn build_messages(&self, session_id: i64) -> anyhow::Result> { let db_messages = self.storage.get_messages(session_id).await?; db_messages @@ -232,8 +232,8 @@ impl ImapEmailAgent { async fn execute_tool_call( &self, - session_id: &str, - message_id: Option<&String>, + session_id: i64, + message_id: Option, tool_call: &LlmToolCall, ) -> anyhow::Result<()> { let mut tool_request = @@ -288,8 +288,8 @@ impl ImapEmailAgent { let mut tool_call_record = StorageToolCall { id: None, - session_id: session_id.to_string(), - message_id: message_id.cloned(), + session_id, + message_id, tool_call_id: tool_call.id().to_string(), tool_name: tool_call.name().to_string(), request: tool_call.arguments().clone(), @@ -326,7 +326,7 @@ impl ImapEmailAgent { let tool_message = StorageMessage { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -350,7 +350,7 @@ impl ImapEmailAgent { let tool_message = StorageMessage { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: error_message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -454,10 +454,9 @@ impl Agent for ImapEmailAgent { } async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { - let session_id_str = session_id.to_string(); let user_message = StorageMessage { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::User, content: user_prompt.to_string(), created_at: chrono::Utc::now().timestamp(), @@ -473,7 +472,7 @@ impl Agent for ImapEmailAgent { iteration += 1; if iteration > max_iterations { let error = "Maximum iteration limit reached"; - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Failed; session.error = Some(error.to_string()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -481,7 +480,7 @@ impl Agent for ImapEmailAgent { return Err(anyhow::anyhow!(error)); } - let messages = self.build_messages(&session_id_str).await?; + let messages = self.build_messages(session_id).await?; let request = CompletionRequest { messages, @@ -508,7 +507,7 @@ impl Agent for ImapEmailAgent { let assistant_message = StorageMessage { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::Assistant, content: text_to_save, created_at: chrono::Utc::now().timestamp(), @@ -517,7 +516,7 @@ impl Agent for ImapEmailAgent { if let Some(tool_calls) = response.tool_calls { if tool_calls.is_empty() { - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Completed; session.result = Some(text.clone()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -526,11 +525,11 @@ impl Agent for ImapEmailAgent { } for tool_call in tool_calls { - self.execute_tool_call(&session_id_str, Some(&message_id), &tool_call) + self.execute_tool_call(session_id, Some(message_id), &tool_call) .await?; } } else { - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Completed; session.result = Some(text.clone()); session.ended_at = Some(chrono::Utc::now().timestamp()); diff --git a/nocodo-agents/src/pdftotext/mod.rs b/nocodo-agents/src/pdftotext/mod.rs index ce29bc68..cc6f22e6 100644 --- a/nocodo-agents/src/pdftotext/mod.rs +++ b/nocodo-agents/src/pdftotext/mod.rs @@ -96,7 +96,7 @@ impl PdfToTextAgent { }) } - async fn get_session(&self, session_id: &str) -> anyhow::Result { + async fn get_session(&self, session_id: i64) -> anyhow::Result { self.storage .get_session(session_id) .await? @@ -112,7 +112,7 @@ impl PdfToTextAgent { } /// Build messages from session history - async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + async fn build_messages(&self, session_id: i64) -> anyhow::Result> { let db_messages = self.storage.get_messages(session_id).await?; db_messages @@ -136,8 +136,8 @@ impl PdfToTextAgent { /// Execute a tool call async fn execute_tool_call( &self, - session_id: &str, - message_id: Option<&String>, + session_id: i64, + message_id: Option, tool_call: &LlmToolCall, ) -> anyhow::Result<()> { // 1. Parse LLM tool call into typed ToolRequest @@ -147,8 +147,8 @@ impl PdfToTextAgent { // 2. Record tool call in storage let mut tool_call_record = StorageToolCall { id: None, - session_id: session_id.to_string(), - message_id: message_id.cloned(), + session_id, + message_id, tool_call_id: tool_call.id().to_string(), tool_name: tool_call.name().to_string(), request: tool_call.arguments().clone(), @@ -190,7 +190,7 @@ impl PdfToTextAgent { let tool_message = StorageMessage { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -214,7 +214,7 @@ impl PdfToTextAgent { let tool_message = StorageMessage { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: error_message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -257,12 +257,10 @@ impl Agent for PdfToTextAgent { } async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { - let session_id_str = session_id.to_string(); - // Create initial user message let user_message = StorageMessage { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::User, content: user_prompt.to_string(), created_at: chrono::Utc::now().timestamp(), @@ -280,7 +278,7 @@ impl Agent for PdfToTextAgent { iteration += 1; if iteration > max_iterations { let error = "Maximum iteration limit reached"; - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Failed; session.error = Some(error.to_string()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -289,7 +287,7 @@ impl Agent for PdfToTextAgent { } // Build request with conversation history - let messages = self.build_messages(&session_id_str).await?; + let messages = self.build_messages(session_id).await?; let request = CompletionRequest { messages, @@ -311,7 +309,7 @@ impl Agent for PdfToTextAgent { let text = extract_text_from_content(&response.content); let assistant_message = StorageMessage { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::Assistant, content: text.clone(), created_at: chrono::Utc::now().timestamp(), @@ -321,7 +319,7 @@ impl Agent for PdfToTextAgent { // Check for tool calls if let Some(tool_calls) = response.tool_calls { if tool_calls.is_empty() { - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Completed; session.result = Some(text.clone()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -331,11 +329,11 @@ impl Agent for PdfToTextAgent { // Execute tools for tool_call in tool_calls { - self.execute_tool_call(&session_id_str, Some(&message_id), &tool_call) + self.execute_tool_call(session_id, Some(message_id), &tool_call) .await?; } } else { - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Completed; session.result = Some(text.clone()); session.ended_at = Some(chrono::Utc::now().timestamp()); diff --git a/nocodo-agents/src/requirements_gathering/mod.rs b/nocodo-agents/src/requirements_gathering/mod.rs index 71d8e54c..5894cab7 100644 --- a/nocodo-agents/src/requirements_gathering/mod.rs +++ b/nocodo-agents/src/requirements_gathering/mod.rs @@ -87,8 +87,8 @@ that you need more information about what they want to automate."#.to_string() async fn execute_tool_call( &self, - session_id: &str, - message_id: Option<&String>, + session_id: i64, + message_id: Option, tool_call: &LlmToolCall, ) -> anyhow::Result<()> { let tool_request = @@ -96,8 +96,8 @@ that you need more information about what they want to automate."#.to_string() let mut tool_call_record = StorageToolCall { id: None, - session_id: session_id.to_string(), - message_id: message_id.cloned(), + session_id, + message_id, tool_call_id: tool_call.id().to_string(), tool_name: tool_call.name().to_string(), request: tool_call.arguments().clone(), @@ -137,7 +137,7 @@ that you need more information about what they want to automate."#.to_string() let tool_message = Message { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -162,7 +162,7 @@ that you need more information about what they want to automate."#.to_string() let tool_error_message = Message { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: error_message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -181,7 +181,7 @@ that you need more information about what they want to automate."#.to_string() .collect() } - async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + async fn build_messages(&self, session_id: i64) -> anyhow::Result> { let db_messages = self.storage.get_messages(session_id).await?; db_messages @@ -202,7 +202,7 @@ that you need more information about what they want to automate."#.to_string() .collect() } - async fn get_session(&self, session_id: &str) -> anyhow::Result { + async fn get_session(&self, session_id: i64) -> anyhow::Result { self.storage .get_session(session_id) .await? @@ -225,10 +225,9 @@ impl Agent for UserClarificationAgent anyhow::Result { - let session_id_str = session_id.to_string(); let user_message = Message { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::User, content: user_prompt.to_string(), created_at: chrono::Utc::now().timestamp(), @@ -244,7 +243,7 @@ impl Agent for UserClarificationAgent max_iterations { let error = "Maximum iteration limit reached"; - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Failed; session.error = Some(error.to_string()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -252,7 +251,7 @@ impl Agent for UserClarificationAgent Agent for UserClarificationAgent Agent for UserClarificationAgent Agent for UserClarificationAgent Agent for UserClarificationAgent Agent for UserClarificationAgent Agent for UserClarificationAgent, + session_id: i64, + tool_call_id: Option, questions: &[UserQuestion], ) -> Result<(), StorageError>; async fn get_pending_questions( &self, - session_id: &str, + session_id: i64, ) -> Result, StorageError>; async fn store_answers( &self, - session_id: &str, + session_id: i64, answers: &std::collections::HashMap, ) -> Result<(), StorageError>; } diff --git a/nocodo-agents/src/settings_management/mod.rs b/nocodo-agents/src/settings_management/mod.rs index 2b3abaa3..7e62914c 100644 --- a/nocodo-agents/src/settings_management/mod.rs +++ b/nocodo-agents/src/settings_management/mod.rs @@ -139,8 +139,8 @@ using the tool."#, async fn execute_tool_call( &self, - session_id: &str, - message_id: Option<&String>, + session_id: i64, + message_id: Option, tool_call: &LlmToolCall, ) -> anyhow::Result<()> { let tool_request = @@ -148,8 +148,8 @@ using the tool."#, let mut tool_call_record = StorageToolCall { id: None, - session_id: session_id.to_string(), - message_id: message_id.cloned(), + session_id, + message_id, tool_call_id: tool_call.id().to_string(), tool_name: tool_call.name().to_string(), request: tool_call.arguments().clone(), @@ -189,7 +189,7 @@ using the tool."#, let tool_message = StorageMessage { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -213,7 +213,7 @@ using the tool."#, let tool_message = StorageMessage { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: error_message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -232,7 +232,7 @@ using the tool."#, .collect() } - async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + async fn build_messages(&self, session_id: i64) -> anyhow::Result> { let db_messages = self.storage.get_messages(session_id).await?; db_messages @@ -253,7 +253,7 @@ using the tool."#, .collect() } - async fn get_session(&self, session_id: &str) -> anyhow::Result { + async fn get_session(&self, session_id: i64) -> anyhow::Result { self.storage .get_session(session_id) .await? @@ -332,10 +332,9 @@ impl Agent for SettingsManagementAgent { } async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { - let session_id_str = session_id.to_string(); let user_message = StorageMessage { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::User, content: user_prompt.to_string(), created_at: chrono::Utc::now().timestamp(), @@ -351,7 +350,7 @@ impl Agent for SettingsManagementAgent { iteration += 1; if iteration > max_iterations { let error = "Maximum iteration limit reached"; - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Failed; session.error = Some(error.to_string()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -359,7 +358,7 @@ impl Agent for SettingsManagementAgent { return Err(anyhow::anyhow!(error)); } - let messages = self.build_messages(&session_id_str).await?; + let messages = self.build_messages(session_id).await?; let request = CompletionRequest { messages, @@ -386,7 +385,7 @@ impl Agent for SettingsManagementAgent { let assistant_message = StorageMessage { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::Assistant, content: text_to_save, created_at: chrono::Utc::now().timestamp(), @@ -395,7 +394,7 @@ impl Agent for SettingsManagementAgent { if let Some(tool_calls) = response.tool_calls { if tool_calls.is_empty() { - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Completed; session.result = Some(text.clone()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -406,14 +405,14 @@ impl Agent for SettingsManagementAgent { for tool_call in tool_calls { // Special handling for ask_user tool - collect settings and write to TOML if tool_call.name() == "ask_user" { - tracing::info!(session_id = %session_id_str, "Agent requesting user settings"); + tracing::info!(session_id = %session_id, "Agent requesting user settings"); // Create tool call record in agent_tool_calls table let start = Instant::now(); let mut tool_call_record = StorageToolCall { id: None, - session_id: session_id_str.clone(), - message_id: Some(message_id.clone()), + session_id, + message_id: Some(message_id), tool_call_id: tool_call.id().to_string(), tool_name: tool_call.name().to_string(), request: tool_call.arguments().clone(), @@ -485,7 +484,7 @@ impl Agent for SettingsManagementAgent { ); let tool_message = StorageMessage { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::Tool, content: message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -498,12 +497,12 @@ impl Agent for SettingsManagementAgent { } } else { // Execute other tools normally - self.execute_tool_call(&session_id_str, Some(&message_id), &tool_call) + self.execute_tool_call(session_id, Some(message_id), &tool_call) .await?; } } } else { - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Completed; session.result = Some(text.clone()); session.ended_at = Some(chrono::Utc::now().timestamp()); diff --git a/nocodo-agents/src/sqlite_reader/mod.rs b/nocodo-agents/src/sqlite_reader/mod.rs index 7d7fa22d..9af18da7 100644 --- a/nocodo-agents/src/sqlite_reader/mod.rs +++ b/nocodo-agents/src/sqlite_reader/mod.rs @@ -85,7 +85,7 @@ impl SqliteReaderAgent { .collect() } - async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + async fn build_messages(&self, session_id: i64) -> anyhow::Result> { let db_messages = self.storage.get_messages(session_id).await?; db_messages @@ -106,7 +106,7 @@ impl SqliteReaderAgent { .collect() } - async fn get_session(&self, session_id: &str) -> anyhow::Result { + async fn get_session(&self, session_id: i64) -> anyhow::Result { self.storage .get_session(session_id) .await? @@ -115,8 +115,8 @@ impl SqliteReaderAgent { async fn execute_tool_call( &self, - session_id: &str, - message_id: Option<&String>, + session_id: i64, + message_id: Option, tool_call: &LlmToolCall, ) -> anyhow::Result<()> { let mut tool_request = @@ -132,8 +132,8 @@ impl SqliteReaderAgent { let mut tool_call_record = StorageToolCall { id: None, - session_id: session_id.to_string(), - message_id: message_id.cloned(), + session_id, + message_id, tool_call_id: tool_call.id().to_string(), tool_name: tool_call.name().to_string(), request: tool_call.arguments().clone(), @@ -173,7 +173,7 @@ impl SqliteReaderAgent { let tool_message = Message { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -198,7 +198,7 @@ impl SqliteReaderAgent { let tool_error_message = Message { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: error_message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -249,10 +249,9 @@ impl Agent for SqliteReaderAgent { } async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { - let session_id_str = session_id.to_string(); let user_message = Message { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::User, content: user_prompt.to_string(), created_at: chrono::Utc::now().timestamp(), @@ -268,7 +267,7 @@ impl Agent for SqliteReaderAgent { iteration += 1; if iteration > max_iterations { let error = "Maximum iteration limit reached"; - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Failed; session.error = Some(error.to_string()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -276,7 +275,7 @@ impl Agent for SqliteReaderAgent { return Err(anyhow::anyhow!(error)); } - let messages = self.build_messages(&session_id_str).await?; + let messages = self.build_messages(session_id).await?; let request = CompletionRequest { messages, @@ -303,7 +302,7 @@ impl Agent for SqliteReaderAgent { let assistant_message = Message { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::Assistant, content: text_to_save, created_at: chrono::Utc::now().timestamp(), @@ -312,7 +311,7 @@ impl Agent for SqliteReaderAgent { if let Some(tool_calls) = response.tool_calls { if tool_calls.is_empty() { - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Completed; session.result = Some(text.clone()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -321,11 +320,11 @@ impl Agent for SqliteReaderAgent { } for tool_call in tool_calls { - self.execute_tool_call(&session_id_str, Some(&message_id), &tool_call) + self.execute_tool_call(session_id, Some(message_id), &tool_call) .await?; } } else { - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Completed; session.result = Some(text.clone()); session.ended_at = Some(chrono::Utc::now().timestamp()); diff --git a/nocodo-agents/src/storage/memory.rs b/nocodo-agents/src/storage/memory.rs index c0e60615..42fea88a 100644 --- a/nocodo-agents/src/storage/memory.rs +++ b/nocodo-agents/src/storage/memory.rs @@ -6,11 +6,12 @@ use std::sync::{Arc, Mutex}; #[derive(Clone)] pub struct InMemoryStorage { - sessions: Arc>>, - messages: Arc>>>, - tool_calls: Arc>>>, - questions: Arc>>>, - answers: Arc>>>, + sessions: Arc>>, + messages: Arc>>>, + tool_calls: Arc>>>, + questions: Arc>>>, + answers: Arc>>>, + counter: Arc>, } impl InMemoryStorage { @@ -21,41 +22,45 @@ impl InMemoryStorage { tool_calls: Arc::new(Mutex::new(HashMap::new())), questions: Arc::new(Mutex::new(HashMap::new())), answers: Arc::new(Mutex::new(HashMap::new())), + counter: Arc::new(Mutex::new(0)), } } + + fn next_id(&self) -> i64 { + let mut counter = self.counter.lock().unwrap(); + *counter += 1; + *counter + } } #[async_trait::async_trait] impl AgentStorage for InMemoryStorage { - async fn create_session(&self, session: Session) -> Result { - let session_id = uuid::Uuid::new_v4().to_string(); + async fn create_session(&self, session: Session) -> Result { + let session_id = self.next_id(); let mut session_with_id = session; - session_with_id.id = Some(session_id.clone()); + session_with_id.id = Some(session_id); self.sessions .lock() .unwrap() - .insert(session_id.clone(), session_with_id); + .insert(session_id, session_with_id); Ok(session_id) } - async fn get_session(&self, session_id: &str) -> Result, StorageError> { - Ok(self.sessions.lock().unwrap().get(session_id).cloned()) + async fn get_session(&self, session_id: i64) -> Result, StorageError> { + Ok(self.sessions.lock().unwrap().get(&session_id).cloned()) } async fn update_session(&self, session: Session) -> Result<(), StorageError> { - let session_id = session - .id - .clone() - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let session_id = session.id.unwrap_or_else(|| self.next_id()); self.sessions.lock().unwrap().insert(session_id, session); Ok(()) } - async fn create_message(&self, message: Message) -> Result { - let message_id = uuid::Uuid::new_v4().to_string(); - let session_id = message.session_id.clone(); + async fn create_message(&self, message: Message) -> Result { + let message_id = self.next_id(); + let session_id = message.session_id; let mut message_with_id = message; - message_with_id.id = Some(message_id.clone()); + message_with_id.id = Some(message_id); let mut messages = self.messages.lock().unwrap(); messages .entry(session_id) @@ -64,21 +69,21 @@ impl AgentStorage for InMemoryStorage { Ok(message_id) } - async fn get_messages(&self, session_id: &str) -> Result, StorageError> { + async fn get_messages(&self, session_id: i64) -> Result, StorageError> { Ok(self .messages .lock() .unwrap() - .get(session_id) + .get(&session_id) .cloned() .unwrap_or_default()) } - async fn create_tool_call(&self, tool_call: ToolCall) -> Result { - let tool_call_id = uuid::Uuid::new_v4().to_string(); - let session_id = tool_call.session_id.clone(); + async fn create_tool_call(&self, tool_call: ToolCall) -> Result { + let tool_call_id = self.next_id(); + let session_id = tool_call.session_id; let mut tool_call_with_id = tool_call; - tool_call_with_id.id = Some(tool_call_id.clone()); + tool_call_with_id.id = Some(tool_call_id); let mut tool_calls = self.tool_calls.lock().unwrap(); tool_calls .entry(session_id) @@ -88,41 +93,35 @@ impl AgentStorage for InMemoryStorage { } async fn update_tool_call(&self, tool_call: ToolCall) -> Result<(), StorageError> { - let tool_call_id = tool_call - .id - .clone() - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let tool_call_id = tool_call.id.unwrap_or_else(|| self.next_id()); let mut tool_calls = self.tool_calls.lock().unwrap(); if let Some(calls) = tool_calls.get_mut(&tool_call.session_id) { - if let Some(pos) = calls - .iter() - .position(|c| c.id.as_ref() == Some(&tool_call_id)) - { + if let Some(pos) = calls.iter().position(|c| c.id == Some(tool_call_id)) { calls[pos] = tool_call; } } Ok(()) } - async fn get_tool_calls(&self, session_id: &str) -> Result, StorageError> { + async fn get_tool_calls(&self, session_id: i64) -> Result, StorageError> { Ok(self .tool_calls .lock() .unwrap() - .get(session_id) + .get(&session_id) .cloned() .unwrap_or_default()) } async fn get_pending_tool_calls( &self, - session_id: &str, + session_id: i64, ) -> Result, StorageError> { Ok(self .tool_calls .lock() .unwrap() - .get(session_id) + .get(&session_id) .cloned() .unwrap_or_default() .into_iter() @@ -135,13 +134,13 @@ impl AgentStorage for InMemoryStorage { impl crate::requirements_gathering::storage::RequirementsStorage for InMemoryStorage { async fn store_questions( &self, - session_id: &str, - _tool_call_id: Option<&str>, + session_id: i64, + _tool_call_id: Option, questions: &[UserQuestion], ) -> Result<(), StorageError> { let mut question_store = self.questions.lock().unwrap(); question_store - .entry(session_id.to_string()) + .entry(session_id) .or_insert_with(Vec::new) .extend(questions.iter().cloned()); Ok(()) @@ -149,13 +148,13 @@ impl crate::requirements_gathering::storage::RequirementsStorage for InMemorySto async fn get_pending_questions( &self, - session_id: &str, + session_id: i64, ) -> Result, StorageError> { let questions = self.questions.lock().unwrap(); let answers = self.answers.lock().unwrap(); - let session_answers = answers.get(session_id); - let session_questions = questions.get(session_id).cloned().unwrap_or_default(); + let session_answers = answers.get(&session_id); + let session_questions = questions.get(&session_id).cloned().unwrap_or_default(); // Filter out questions that have been answered Ok(session_questions @@ -170,12 +169,12 @@ impl crate::requirements_gathering::storage::RequirementsStorage for InMemorySto async fn store_answers( &self, - session_id: &str, + session_id: i64, answers: &std::collections::HashMap, ) -> Result<(), StorageError> { let mut answer_store = self.answers.lock().unwrap(); answer_store - .entry(session_id.to_string()) + .entry(session_id) .or_insert_with(HashMap::new) .extend(answers.clone()); Ok(()) diff --git a/nocodo-agents/src/storage/mod.rs b/nocodo-agents/src/storage/mod.rs index 621682e7..9fba098a 100644 --- a/nocodo-agents/src/storage/mod.rs +++ b/nocodo-agents/src/storage/mod.rs @@ -7,17 +7,17 @@ pub use memory::InMemoryStorage; #[async_trait] pub trait AgentStorage: Send + Sync { - async fn create_session(&self, session: Session) -> Result; - async fn get_session(&self, session_id: &str) -> Result, StorageError>; + async fn create_session(&self, session: Session) -> Result; + async fn get_session(&self, session_id: i64) -> Result, StorageError>; async fn update_session(&self, session: Session) -> Result<(), StorageError>; - async fn create_message(&self, message: Message) -> Result; - async fn get_messages(&self, session_id: &str) -> Result, StorageError>; + async fn create_message(&self, message: Message) -> Result; + async fn get_messages(&self, session_id: i64) -> Result, StorageError>; - async fn create_tool_call(&self, tool_call: ToolCall) -> Result; + async fn create_tool_call(&self, tool_call: ToolCall) -> Result; async fn update_tool_call(&self, tool_call: ToolCall) -> Result<(), StorageError>; - async fn get_tool_calls(&self, session_id: &str) -> Result, StorageError>; - async fn get_pending_tool_calls(&self, session_id: &str) + async fn get_tool_calls(&self, session_id: i64) -> Result, StorageError>; + async fn get_pending_tool_calls(&self, session_id: i64) -> Result, StorageError>; } diff --git a/nocodo-agents/src/structured_json/mod.rs b/nocodo-agents/src/structured_json/mod.rs index 62bafe6c..81cbeddc 100644 --- a/nocodo-agents/src/structured_json/mod.rs +++ b/nocodo-agents/src/structured_json/mod.rs @@ -71,7 +71,7 @@ impl StructuredJsonAgent { }) } - async fn get_session(&self, session_id: &str) -> anyhow::Result { + async fn get_session(&self, session_id: i64) -> anyhow::Result { self.storage .get_session(session_id) .await? @@ -112,7 +112,7 @@ When responding: async fn validate_and_retry( &self, user_prompt: &str, - session_id: &str, + session_id: i64, max_retries: u32, ) -> anyhow::Result { let mut attempt = 0; @@ -147,7 +147,7 @@ When responding: let text = extract_text_from_content(&response.content); let assistant_message = StorageMessage { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Assistant, content: text.clone(), created_at: chrono::Utc::now().timestamp(), @@ -179,7 +179,7 @@ When responding: conversation_context.push((Role::User, error_msg.clone())); let error_message = StorageMessage { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::User, content: error_msg, created_at: chrono::Utc::now().timestamp(), @@ -205,7 +205,7 @@ When responding: conversation_context.push((Role::User, error_msg.clone())); let error_message = StorageMessage { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::User, content: error_msg, created_at: chrono::Utc::now().timestamp(), @@ -258,21 +258,20 @@ impl Agent for StructuredJsonAgent { } async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { - let session_id_str = session_id.to_string(); let user_message = StorageMessage { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::User, content: user_prompt.to_string(), created_at: chrono::Utc::now().timestamp(), }; self.storage.create_message(user_message).await?; - let json_value = self.validate_and_retry(user_prompt, &session_id_str, 3).await?; + let json_value = self.validate_and_retry(user_prompt, session_id, 3).await?; let formatted = serde_json::to_string_pretty(&json_value)?; - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Completed; session.result = Some(formatted.clone()); session.ended_at = Some(chrono::Utc::now().timestamp()); diff --git a/nocodo-agents/src/tesseract/mod.rs b/nocodo-agents/src/tesseract/mod.rs index bd1ad315..69c6fa2e 100644 --- a/nocodo-agents/src/tesseract/mod.rs +++ b/nocodo-agents/src/tesseract/mod.rs @@ -95,7 +95,7 @@ impl TesseractAgent { }) } - async fn get_session(&self, session_id: &str) -> anyhow::Result { + async fn get_session(&self, session_id: i64) -> anyhow::Result { self.storage .get_session(session_id) .await? @@ -111,7 +111,7 @@ impl TesseractAgent { } /// Build messages from session history - async fn build_messages(&self, session_id: &str) -> anyhow::Result> { + async fn build_messages(&self, session_id: i64) -> anyhow::Result> { let db_messages = self.storage.get_messages(session_id).await?; db_messages @@ -135,8 +135,8 @@ impl TesseractAgent { /// Execute a tool call async fn execute_tool_call( &self, - session_id: &str, - message_id: Option<&String>, + session_id: i64, + message_id: Option, tool_call: &LlmToolCall, ) -> anyhow::Result<()> { // 1. Parse LLM tool call into typed ToolRequest @@ -146,8 +146,8 @@ impl TesseractAgent { // 2. Record tool call in storage let mut tool_call_record = StorageToolCall { id: None, - session_id: session_id.to_string(), - message_id: message_id.cloned(), + session_id, + message_id, tool_call_id: tool_call.id().to_string(), tool_name: tool_call.name().to_string(), request: tool_call.arguments().clone(), @@ -186,7 +186,7 @@ impl TesseractAgent { let tool_message = StorageMessage { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -210,7 +210,7 @@ impl TesseractAgent { let tool_message = StorageMessage { id: None, - session_id: session_id.to_string(), + session_id, role: MessageRole::Tool, content: error_message_to_llm, created_at: chrono::Utc::now().timestamp(), @@ -251,12 +251,10 @@ impl Agent for TesseractAgent { } async fn execute(&self, user_prompt: &str, session_id: i64) -> anyhow::Result { - let session_id_str = session_id.to_string(); - // Create initial user message let user_message = StorageMessage { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::User, content: user_prompt.to_string(), created_at: chrono::Utc::now().timestamp(), @@ -274,7 +272,7 @@ impl Agent for TesseractAgent { iteration += 1; if iteration > max_iterations { let error = "Maximum iteration limit reached"; - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Failed; session.error = Some(error.to_string()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -283,7 +281,7 @@ impl Agent for TesseractAgent { } // Build request with conversation history - let messages = self.build_messages(&session_id_str).await?; + let messages = self.build_messages(session_id).await?; let request = CompletionRequest { messages, @@ -305,7 +303,7 @@ impl Agent for TesseractAgent { let text = extract_text_from_content(&response.content); let assistant_message = StorageMessage { id: None, - session_id: session_id_str.clone(), + session_id, role: MessageRole::Assistant, content: text.clone(), created_at: chrono::Utc::now().timestamp(), @@ -315,7 +313,7 @@ impl Agent for TesseractAgent { // Check for tool calls if let Some(tool_calls) = response.tool_calls { if tool_calls.is_empty() { - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Completed; session.result = Some(text.clone()); session.ended_at = Some(chrono::Utc::now().timestamp()); @@ -325,11 +323,11 @@ impl Agent for TesseractAgent { // Execute tools for tool_call in tool_calls { - self.execute_tool_call(&session_id_str, Some(&message_id), &tool_call) + self.execute_tool_call(session_id, Some(message_id), &tool_call) .await?; } } else { - let mut session = self.get_session(&session_id_str).await?; + let mut session = self.get_session(session_id).await?; session.status = SessionStatus::Completed; session.result = Some(text.clone()); session.ended_at = Some(chrono::Utc::now().timestamp()); diff --git a/nocodo-agents/src/types/message.rs b/nocodo-agents/src/types/message.rs index e6a50e2d..a28e8989 100644 --- a/nocodo-agents/src/types/message.rs +++ b/nocodo-agents/src/types/message.rs @@ -2,8 +2,8 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Message { - pub id: Option, - pub session_id: String, + pub id: Option, + pub session_id: i64, pub role: MessageRole, pub content: String, pub created_at: i64, diff --git a/nocodo-agents/src/types/session.rs b/nocodo-agents/src/types/session.rs index 0e8ce527..49cdaab8 100644 --- a/nocodo-agents/src/types/session.rs +++ b/nocodo-agents/src/types/session.rs @@ -2,7 +2,7 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Session { - pub id: Option, + pub id: Option, pub agent_name: String, pub provider: String, pub model: String, diff --git a/nocodo-agents/src/types/tool_call.rs b/nocodo-agents/src/types/tool_call.rs index e78b8c25..5cb417a0 100644 --- a/nocodo-agents/src/types/tool_call.rs +++ b/nocodo-agents/src/types/tool_call.rs @@ -2,9 +2,9 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ToolCall { - pub id: Option, - pub session_id: String, - pub message_id: Option, + pub id: Option, + pub session_id: i64, + pub message_id: Option, pub tool_call_id: String, pub tool_name: String, pub request: serde_json::Value, diff --git a/nocodo-api/Cargo.toml b/nocodo-api/Cargo.toml index 78700899..c9388bd5 100644 --- a/nocodo-api/Cargo.toml +++ b/nocodo-api/Cargo.toml @@ -23,6 +23,7 @@ anyhow.workspace = true home = "0.5" dirs.workspace = true rusqlite = { version = "0.37", features = ["bundled"] } +refinery = { version = "0.9", features = ["rusqlite"] } toml = "0.9" config.workspace = true async-trait = "0.1" diff --git a/nocodo-api/src/handlers/agent_execution/codebase_analysis_agent.rs b/nocodo-api/src/handlers/agent_execution/codebase_analysis_agent.rs index b23b1da3..26e4a4a1 100644 --- a/nocodo-api/src/handlers/agent_execution/codebase_analysis_agent.rs +++ b/nocodo-api/src/handlers/agent_execution/codebase_analysis_agent.rs @@ -1,7 +1,8 @@ use crate::models::ErrorResponse; +use crate::storage::SqliteAgentStorage; use actix_web::{post, web, HttpResponse, Responder}; use nocodo_agents::codebase_analysis::CodebaseAnalysisAgent; -use nocodo_agents::Agent; +use nocodo_agents::{Agent, AgentStorage, Session, SessionStatus}; use serde_json::json; use shared_types::{AgentConfig, AgentExecutionRequest, AgentExecutionResponse}; use std::sync::Arc; @@ -11,7 +12,7 @@ use tracing::{error, info}; pub async fn execute_codebase_analysis_agent( req: web::Json, llm_client: web::Data>, - database: web::Data>, + storage: web::Data>, ) -> impl Responder { let (path, max_depth) = match &req.config { AgentConfig::CodebaseAnalysis(config) => { @@ -40,14 +41,22 @@ pub async fn execute_codebase_analysis_agent( let provider = llm_client.provider_name().to_string(); let model = llm_client.model_name().to_string(); - let session_id = match database.create_session( - &agent_name, - &provider, - &model, - None, - &user_prompt, - Some(config), - ) { + let session = Session { + id: None, + agent_name: agent_name.clone(), + provider: provider.clone(), + model: model.clone(), + system_prompt: None, + user_prompt: user_prompt.clone(), + config: config.clone(), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + + let session_id = match storage.create_session(session).await { Ok(id) => id, Err(e) => { error!(error = %e, "Failed to create session"); @@ -59,7 +68,7 @@ pub async fn execute_codebase_analysis_agent( // Return immediately with session_id and spawn background task let llm_client_clone = llm_client.get_ref().clone(); - let database_clone = database.get_ref().clone(); + let storage_clone = storage.get_ref().clone(); let user_prompt_clone = user_prompt.clone(); tokio::spawn(async move { @@ -69,19 +78,46 @@ pub async fn execute_codebase_analysis_agent( ); let agent = - CodebaseAnalysisAgent::new(llm_client_clone, database_clone.clone(), tool_executor); + CodebaseAnalysisAgent::new(llm_client_clone, storage_clone.clone(), tool_executor); match agent.execute(&user_prompt_clone, session_id).await { Ok(result) => { info!(result = %result, session_id = session_id, "Agent execution completed successfully"); - if let Err(e) = database_clone.complete_session(session_id, &result) { + let mut session = Session { + id: Some(session_id), + agent_name: "codebase-analysis".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Completed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: Some(result.clone()), + error: None, + }; + if let Err(e) = storage_clone.update_session(session).await { error!(error = %e, session_id = session_id, "Failed to complete session"); } } Err(e) => { error!(error = %e, session_id = session_id, "Agent execution failed"); - let _ = - database_clone.fail_session(session_id, &format!("Execution failed: {}", e)); + let mut session = Session { + id: Some(session_id), + agent_name: "codebase-analysis".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Execution failed: {}", e)), + }; + let _ = storage_clone.update_session(session).await; } } }); diff --git a/nocodo-api/src/handlers/agent_execution/imap_email_agent.rs b/nocodo-api/src/handlers/agent_execution/imap_email_agent.rs index 375dbbb6..918d2f80 100644 --- a/nocodo-api/src/handlers/agent_execution/imap_email_agent.rs +++ b/nocodo-api/src/handlers/agent_execution/imap_email_agent.rs @@ -1,7 +1,8 @@ use crate::helpers::agents::create_imap_agent; use crate::models::ErrorResponse; +use crate::storage::SqliteAgentStorage; use actix_web::{post, web, HttpResponse, Responder}; -use nocodo_agents::Agent; +use nocodo_agents::{Agent, AgentStorage, Session, SessionStatus}; use serde_json::json; use shared_types::{AgentConfig, AgentExecutionRequest, AgentExecutionResponse}; use std::sync::Arc; @@ -11,7 +12,7 @@ use tracing::{error, info}; pub async fn execute_imap_agent( req: web::Json, llm_client: web::Data>, - database: web::Data>, + storage: web::Data>, ) -> impl Responder { let (host, port, username, password) = match &req.config { AgentConfig::Imap(config) => ( @@ -44,14 +45,22 @@ pub async fn execute_imap_agent( let provider = llm_client.provider_name().to_string(); let model = llm_client.model_name().to_string(); - let session_id = match database.create_session( - &agent_name, - &provider, - &model, - None, - &user_prompt, - Some(config), - ) { + let session = Session { + id: None, + agent_name: agent_name.clone(), + provider: provider.clone(), + model: model.clone(), + system_prompt: None, + user_prompt: user_prompt.clone(), + config: config.clone(), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + + let session_id = match storage.create_session(session).await { Ok(id) => id, Err(e) => { error!(error = %e, "Failed to create session"); @@ -62,13 +71,13 @@ pub async fn execute_imap_agent( }; let llm_client_clone = llm_client.get_ref().clone(); - let database_clone = database.get_ref().clone(); + let storage_clone = storage.get_ref().clone(); let user_prompt_clone = user_prompt.clone(); tokio::spawn(async move { let agent = match create_imap_agent( &llm_client_clone, - &database_clone, + &storage_clone, &host, port, &username, @@ -77,8 +86,21 @@ pub async fn execute_imap_agent( Ok(agent) => agent, Err(e) => { error!(error = %e, session_id = session_id, "Failed to create IMAP agent"); - let _ = database_clone - .fail_session(session_id, &format!("Failed to create agent: {}", e)); + let mut session = Session { + id: Some(session_id), + agent_name: "imap".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Failed to create agent: {}", e)), + }; + let _ = storage_clone.update_session(session).await; return; } }; @@ -86,14 +108,41 @@ pub async fn execute_imap_agent( match agent.execute(&user_prompt_clone, session_id).await { Ok(result) => { info!(result = %result, session_id = session_id, "Agent execution completed successfully"); - if let Err(e) = database_clone.complete_session(session_id, &result) { + let mut session = Session { + id: Some(session_id), + agent_name: "imap".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Completed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: Some(result.clone()), + error: None, + }; + if let Err(e) = storage_clone.update_session(session).await { error!(error = %e, session_id = session_id, "Failed to complete session"); } } Err(e) => { error!(error = %e, session_id = session_id, "Agent execution failed"); - let _ = - database_clone.fail_session(session_id, &format!("Execution failed: {}", e)); + let mut session = Session { + id: Some(session_id), + agent_name: "imap".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Execution failed: {}", e)), + }; + let _ = storage_clone.update_session(session).await; } } }); diff --git a/nocodo-api/src/handlers/agent_execution/pdftotext_agent.rs b/nocodo-api/src/handlers/agent_execution/pdftotext_agent.rs index ebe2ddb1..ff34db49 100644 --- a/nocodo-api/src/handlers/agent_execution/pdftotext_agent.rs +++ b/nocodo-api/src/handlers/agent_execution/pdftotext_agent.rs @@ -1,7 +1,8 @@ use crate::helpers::agents::create_pdftotext_agent; use crate::models::ErrorResponse; +use crate::storage::SqliteAgentStorage; use actix_web::{post, web, HttpResponse, Responder}; -use nocodo_agents::Agent; +use nocodo_agents::{Agent, AgentStorage, Session, SessionStatus}; use serde_json::json; use shared_types::{AgentConfig, AgentExecutionRequest, AgentExecutionResponse}; use std::sync::Arc; @@ -11,7 +12,7 @@ use tracing::{error, info}; pub async fn execute_pdftotext_agent( req: web::Json, llm_client: web::Data>, - database: web::Data>, + storage: web::Data>, ) -> impl Responder { let pdf_path = match &req.config { AgentConfig::PdfToText(config) => config.pdf_path.clone(), @@ -37,14 +38,22 @@ pub async fn execute_pdftotext_agent( let provider = llm_client.provider_name().to_string(); let model = llm_client.model_name().to_string(); - let session_id = match database.create_session( - &agent_name, - &provider, - &model, - None, - &user_prompt, - Some(config), - ) { + let session = Session { + id: None, + agent_name: agent_name.clone(), + provider: provider.clone(), + model: model.clone(), + system_prompt: None, + user_prompt: user_prompt.clone(), + config: config.clone(), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + + let session_id = match storage.create_session(session).await { Ok(id) => id, Err(e) => { error!(error = %e, "Failed to create session"); @@ -56,34 +65,74 @@ pub async fn execute_pdftotext_agent( // Return immediately with session_id and spawn background task let llm_client_clone = llm_client.get_ref().clone(); - let database_clone = database.get_ref().clone(); + let storage_clone = storage.get_ref().clone(); let pdf_path_clone = pdf_path.clone(); let user_prompt_clone = user_prompt.clone(); tokio::spawn(async move { - let agent = - match create_pdftotext_agent(&llm_client_clone, &database_clone, &pdf_path_clone).await - { - Ok(agent) => agent, - Err(e) => { - error!(error = %e, session_id = session_id, "Failed to create PdfToText agent"); - let _ = database_clone - .fail_session(session_id, &format!("Failed to create agent: {}", e)); - return; - } - }; + let agent = match create_pdftotext_agent(&llm_client_clone, &storage_clone, &pdf_path_clone) + .await + { + Ok(agent) => agent, + Err(e) => { + error!(error = %e, session_id = session_id, "Failed to create PdfToText agent"); + let mut session = Session { + id: Some(session_id), + agent_name: "pdftotext".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Failed to create agent: {}", e)), + }; + let _ = storage_clone.update_session(session).await; + return; + } + }; match agent.execute(&user_prompt_clone, session_id).await { Ok(result) => { info!(result = %result, session_id = session_id, "Agent execution completed successfully"); - if let Err(e) = database_clone.complete_session(session_id, &result) { + let mut session = Session { + id: Some(session_id), + agent_name: "pdftotext".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Completed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: Some(result.clone()), + error: None, + }; + if let Err(e) = storage_clone.update_session(session).await { error!(error = %e, session_id = session_id, "Failed to complete session"); } } Err(e) => { error!(error = %e, session_id = session_id, "Agent execution failed"); - let _ = - database_clone.fail_session(session_id, &format!("Execution failed: {}", e)); + let mut session = Session { + id: Some(session_id), + agent_name: "pdftotext".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Execution failed: {}", e)), + }; + let _ = storage_clone.update_session(session).await; } } }); diff --git a/nocodo-api/src/handlers/agent_execution/requirements_gathering_agent.rs b/nocodo-api/src/handlers/agent_execution/requirements_gathering_agent.rs index 73ccc13d..eae2947e 100644 --- a/nocodo-api/src/handlers/agent_execution/requirements_gathering_agent.rs +++ b/nocodo-api/src/handlers/agent_execution/requirements_gathering_agent.rs @@ -1,7 +1,8 @@ use crate::helpers::agents::create_user_clarification_agent; use crate::models::ErrorResponse; +use crate::storage::SqliteAgentStorage; use actix_web::{post, web, HttpResponse, Responder}; -use nocodo_agents::Agent; +use nocodo_agents::{Agent, AgentStorage, Session, SessionStatus}; use serde_json::json; use shared_types::{AgentConfig, AgentExecutionRequest, AgentExecutionResponse}; use std::sync::Arc; @@ -11,7 +12,8 @@ use tracing::{error, info}; pub async fn execute_requirements_gathering_agent( req: web::Json, llm_client: web::Data>, - database: web::Data>, + storage: web::Data>, + db: web::Data, ) -> impl Responder { // Validate config type match &req.config { @@ -39,14 +41,22 @@ pub async fn execute_requirements_gathering_agent( let provider = llm_client.provider_name().to_string(); let model = llm_client.model_name().to_string(); - let session_id = match database.create_session( - &agent_name, - &provider, - &model, - None, - &user_prompt, - Some(config), - ) { + let session = Session { + id: None, + agent_name: agent_name.clone(), + provider: provider.clone(), + model: model.clone(), + system_prompt: None, + user_prompt: user_prompt.clone(), + config: config.clone(), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + + let session_id = match storage.create_session(session).await { Ok(id) => id, Err(e) => { error!(error = %e, "Failed to create session"); @@ -58,16 +68,34 @@ pub async fn execute_requirements_gathering_agent( // Return immediately with session_id and spawn background task let llm_client_clone = llm_client.get_ref().clone(); - let database_clone = database.get_ref().clone(); + let storage_clone = storage.get_ref().clone(); let user_prompt_clone = user_prompt.clone(); + let db_clone = db.get_ref().clone(); tokio::spawn(async move { - let agent = match create_user_clarification_agent(&llm_client_clone, &database_clone) { + let agent = match create_user_clarification_agent( + &llm_client_clone, + &storage_clone, + &db_clone, + ) { Ok(agent) => agent, Err(e) => { error!(error = %e, session_id = session_id, "Failed to create Requirements Gathering agent"); - let _ = database_clone - .fail_session(session_id, &format!("Failed to create agent: {}", e)); + let mut session = Session { + id: Some(session_id), + agent_name: "requirements-gathering".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Failed to create agent: {}", e)), + }; + let _ = storage_clone.update_session(session).await; return; } }; @@ -78,15 +106,42 @@ pub async fn execute_requirements_gathering_agent( // Check if agent is waiting for user input - if so, don't complete the session // The agent already set the status to waiting_for_user_input if !result.contains("Waiting for user") { - if let Err(e) = database_clone.complete_session(session_id, &result) { + let mut session = Session { + id: Some(session_id), + agent_name: "requirements-gathering".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Completed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: Some(result.clone()), + error: None, + }; + if let Err(e) = storage_clone.update_session(session).await { error!(error = %e, session_id = session_id, "Failed to complete session"); } } } Err(e) => { error!(error = %e, session_id = session_id, "Agent execution failed"); - let _ = - database_clone.fail_session(session_id, &format!("Execution failed: {}", e)); + let mut session = Session { + id: Some(session_id), + agent_name: "requirements-gathering".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Execution failed: {}", e)), + }; + let _ = storage_clone.update_session(session).await; } } }); diff --git a/nocodo-api/src/handlers/agent_execution/settings_management_agent.rs b/nocodo-api/src/handlers/agent_execution/settings_management_agent.rs index 9135dc21..56cec4f5 100644 --- a/nocodo-api/src/handlers/agent_execution/settings_management_agent.rs +++ b/nocodo-api/src/handlers/agent_execution/settings_management_agent.rs @@ -1,7 +1,8 @@ use crate::helpers::agents::create_settings_management_agent; use crate::models::ErrorResponse; +use crate::storage::SqliteAgentStorage; use actix_web::{post, web, HttpResponse, Responder}; -use nocodo_agents::Agent; +use nocodo_agents::{Agent, AgentStorage, Session, SessionStatus}; use serde_json::json; use shared_types::{AgentConfig, AgentExecutionRequest, AgentExecutionResponse}; use std::sync::Arc; @@ -11,7 +12,7 @@ use tracing::{error, info}; pub async fn execute_settings_management_agent( req: web::Json, llm_client: web::Data>, - database: web::Data>, + storage: web::Data>, ) -> impl Responder { let (settings_file_path, agent_schemas) = match &req.config { AgentConfig::SettingsManagement(config) => { @@ -74,14 +75,22 @@ pub async fn execute_settings_management_agent( let provider = llm_client.provider_name().to_string(); let model = llm_client.model_name().to_string(); - let session_id = match database.create_session( - &agent_name, - &provider, - &model, - None, - &user_prompt, - Some(config), - ) { + let session = Session { + id: None, + agent_name: agent_name.clone(), + provider: provider.clone(), + model: model.clone(), + system_prompt: None, + user_prompt: user_prompt.clone(), + config: config.clone(), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + + let session_id = match storage.create_session(session).await { Ok(id) => id, Err(e) => { error!(error = %e, "Failed to create session"); @@ -92,22 +101,35 @@ pub async fn execute_settings_management_agent( }; let llm_client_clone = llm_client.get_ref().clone(); - let database_clone = database.get_ref().clone(); + let storage_clone = storage.get_ref().clone(); let user_prompt_clone = user_prompt.clone(); let settings_file_path_clone = settings_file_path.clone(); tokio::spawn(async move { let agent = match create_settings_management_agent( &llm_client_clone, - &database_clone, + &storage_clone, &settings_file_path_clone, agent_schemas, ) { Ok(agent) => agent, Err(e) => { error!(error = %e, session_id = session_id, "Failed to create Settings Management agent"); - let _ = database_clone - .fail_session(session_id, &format!("Failed to create agent: {}", e)); + let mut session = Session { + id: Some(session_id), + agent_name: "settings-management".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Failed to create agent: {}", e)), + }; + let _ = storage_clone.update_session(session).await; return; } }; @@ -116,15 +138,42 @@ pub async fn execute_settings_management_agent( Ok(result) => { info!(result = %result, session_id = session_id, "Agent execution completed successfully"); if !result.contains("Waiting for user") { - if let Err(e) = database_clone.complete_session(session_id, &result) { + let mut session = Session { + id: Some(session_id), + agent_name: "settings-management".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Completed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: Some(result.clone()), + error: None, + }; + if let Err(e) = storage_clone.update_session(session).await { error!(error = %e, session_id = session_id, "Failed to complete session"); } } } Err(e) => { error!(error = %e, session_id = session_id, "Agent execution failed"); - let _ = - database_clone.fail_session(session_id, &format!("Execution failed: {}", e)); + let mut session = Session { + id: Some(session_id), + agent_name: "settings-management".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Execution failed: {}", e)), + }; + let _ = storage_clone.update_session(session).await; } } }); diff --git a/nocodo-api/src/handlers/agent_execution/sqlite_agent.rs b/nocodo-api/src/handlers/agent_execution/sqlite_agent.rs index c3b8de60..422a60d5 100644 --- a/nocodo-api/src/handlers/agent_execution/sqlite_agent.rs +++ b/nocodo-api/src/handlers/agent_execution/sqlite_agent.rs @@ -1,7 +1,8 @@ use crate::helpers::agents::create_sqlite_agent; use crate::models::ErrorResponse; +use crate::storage::SqliteAgentStorage; use actix_web::{post, web, HttpResponse, Responder}; -use nocodo_agents::Agent; +use nocodo_agents::{Agent, AgentStorage, Session, SessionStatus}; use serde_json::json; use shared_types::{AgentConfig, AgentExecutionRequest, AgentExecutionResponse}; use std::sync::Arc; @@ -11,7 +12,7 @@ use tracing::{error, info}; pub async fn execute_sqlite_agent( req: web::Json, llm_client: web::Data>, - database: web::Data>, + storage: web::Data>, ) -> impl Responder { let db_path = match &req.config { AgentConfig::Sqlite(config) => config.db_path.clone(), @@ -37,14 +38,22 @@ pub async fn execute_sqlite_agent( let provider = llm_client.provider_name().to_string(); let model = llm_client.model_name().to_string(); - let session_id = match database.create_session( - &agent_name, - &provider, - &model, - None, - &user_prompt, - Some(config), - ) { + let session = Session { + id: None, + agent_name: agent_name.clone(), + provider: provider.clone(), + model: model.clone(), + system_prompt: None, + user_prompt: user_prompt.clone(), + config: config.clone(), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + + let session_id = match storage.create_session(session).await { Ok(id) => id, Err(e) => { error!(error = %e, "Failed to create session"); @@ -56,18 +65,31 @@ pub async fn execute_sqlite_agent( // Return immediately with session_id and spawn background task let llm_client_clone = llm_client.get_ref().clone(); - let database_clone = database.get_ref().clone(); + let storage_clone = storage.get_ref().clone(); let db_path_clone = db_path.clone(); let user_prompt_clone = user_prompt.clone(); tokio::spawn(async move { let agent = - match create_sqlite_agent(&llm_client_clone, &database_clone, &db_path_clone).await { + match create_sqlite_agent(&llm_client_clone, &storage_clone, &db_path_clone).await { Ok(agent) => agent, Err(e) => { error!(error = %e, session_id = session_id, "Failed to create SQLite agent"); - let _ = database_clone - .fail_session(session_id, &format!("Failed to create agent: {}", e)); + let mut session = Session { + id: Some(session_id), + agent_name: "sqlite".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: chrono::Utc::now().timestamp(), + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Failed to create agent: {}", e)), + }; + let _ = storage_clone.update_session(session).await; return; } }; @@ -75,14 +97,41 @@ pub async fn execute_sqlite_agent( match agent.execute(&user_prompt_clone, session_id).await { Ok(result) => { info!(result = %result, session_id = session_id, "Agent execution completed successfully"); - if let Err(e) = database_clone.complete_session(session_id, &result) { + let mut session = Session { + id: Some(session_id), + agent_name: "sqlite".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Completed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: Some(result.clone()), + error: None, + }; + if let Err(e) = storage_clone.update_session(session).await { error!(error = %e, session_id = session_id, "Failed to complete session"); } } Err(e) => { error!(error = %e, session_id = session_id, "Agent execution failed"); - let _ = - database_clone.fail_session(session_id, &format!("Execution failed: {}", e)); + let mut session = Session { + id: Some(session_id), + agent_name: "sqlite".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Execution failed: {}", e)), + }; + let _ = storage_clone.update_session(session).await; } } }); diff --git a/nocodo-api/src/handlers/agent_execution/tesseract_agent.rs b/nocodo-api/src/handlers/agent_execution/tesseract_agent.rs index df9c5d11..ed47475a 100644 --- a/nocodo-api/src/handlers/agent_execution/tesseract_agent.rs +++ b/nocodo-api/src/handlers/agent_execution/tesseract_agent.rs @@ -1,7 +1,8 @@ use crate::helpers::agents::create_tesseract_agent; use crate::models::ErrorResponse; +use crate::storage::SqliteAgentStorage; use actix_web::{post, web, HttpResponse, Responder}; -use nocodo_agents::Agent; +use nocodo_agents::{Agent, AgentStorage, Session, SessionStatus}; use serde_json::json; use shared_types::{AgentConfig, AgentExecutionRequest, AgentExecutionResponse}; use std::sync::Arc; @@ -11,7 +12,7 @@ use tracing::{error, info}; pub async fn execute_tesseract_agent( req: web::Json, llm_client: web::Data>, - database: web::Data>, + storage: web::Data>, ) -> impl Responder { let image_path = match &req.config { AgentConfig::Tesseract(config) => config.image_path.clone(), @@ -37,14 +38,22 @@ pub async fn execute_tesseract_agent( let provider = llm_client.provider_name().to_string(); let model = llm_client.model_name().to_string(); - let session_id = match database.create_session( - &agent_name, - &provider, - &model, - None, - &user_prompt, - Some(config), - ) { + let session = Session { + id: None, + agent_name: agent_name.clone(), + provider: provider.clone(), + model: model.clone(), + system_prompt: None, + user_prompt: user_prompt.clone(), + config: config.clone(), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + + let session_id = match storage.create_session(session).await { Ok(id) => id, Err(e) => { error!(error = %e, "Failed to create session"); @@ -56,35 +65,78 @@ pub async fn execute_tesseract_agent( // Return immediately with session_id and spawn background task let llm_client_clone = llm_client.get_ref().clone(); - let database_clone = database.get_ref().clone(); + let storage_clone = storage.get_ref().clone(); let image_path_clone = image_path.clone(); let user_prompt_clone = user_prompt.clone(); tokio::spawn(async move { - let agent = - match create_tesseract_agent(&llm_client_clone, &database_clone, &image_path_clone) - .await - { - Ok(agent) => agent, - Err(e) => { - error!(error = %e, session_id = session_id, "Failed to create Tesseract agent"); - let _ = database_clone - .fail_session(session_id, &format!("Failed to create agent: {}", e)); - return; - } - }; + let agent = match create_tesseract_agent( + &llm_client_clone, + &storage_clone, + &image_path_clone, + ) + .await + { + Ok(agent) => agent, + Err(e) => { + error!(error = %e, session_id = session_id, "Failed to create Tesseract agent"); + let mut session = Session { + id: Some(session_id), + agent_name: "tesseract".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Failed to create agent: {}", e)), + }; + let _ = storage_clone.update_session(session).await; + return; + } + }; match agent.execute(&user_prompt_clone, session_id).await { Ok(result) => { info!(result = %result, session_id = session_id, "Agent execution completed successfully"); - if let Err(e) = database_clone.complete_session(session_id, &result) { + let mut session = Session { + id: Some(session_id), + agent_name: "tesseract".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Completed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: Some(result.clone()), + error: None, + }; + if let Err(e) = storage_clone.update_session(session).await { error!(error = %e, session_id = session_id, "Failed to complete session"); } } Err(e) => { error!(error = %e, session_id = session_id, "Agent execution failed"); - let _ = - database_clone.fail_session(session_id, &format!("Execution failed: {}", e)); + let mut session = Session { + id: Some(session_id), + agent_name: "tesseract".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Execution failed: {}", e)), + }; + let _ = storage_clone.update_session(session).await; } } }); diff --git a/nocodo-api/src/handlers/agent_execution/workflow_creation_agent.rs b/nocodo-api/src/handlers/agent_execution/workflow_creation_agent.rs index 6bef3901..4b8b10b0 100644 --- a/nocodo-api/src/handlers/agent_execution/workflow_creation_agent.rs +++ b/nocodo-api/src/handlers/agent_execution/workflow_creation_agent.rs @@ -1,7 +1,8 @@ use crate::helpers::agents::create_structured_json_agent; use crate::models::ErrorResponse; +use crate::storage::SqliteAgentStorage; use actix_web::{post, web, HttpResponse, Responder}; -use nocodo_agents::Agent; +use nocodo_agents::{Agent, AgentStorage, Session, SessionStatus}; use serde_json::json; use shared_types::{AgentConfig, AgentExecutionRequest, AgentExecutionResponse}; use std::sync::Arc; @@ -11,7 +12,7 @@ use tracing::{error, info}; pub async fn execute_workflow_creation_agent( req: web::Json, llm_client: web::Data>, - database: web::Data>, + storage: web::Data>, ) -> impl Responder { let (type_names, domain_description) = match &req.config { AgentConfig::StructuredJson(config) => { @@ -38,14 +39,22 @@ pub async fn execute_workflow_creation_agent( let provider = llm_client.provider_name().to_string(); let model = llm_client.model_name().to_string(); - let session_id = match database.create_session( - &agent_name, - &provider, - &model, - None, - &user_prompt, - Some(config), - ) { + let session = Session { + id: None, + agent_name: agent_name.clone(), + provider: provider.clone(), + model: model.clone(), + system_prompt: None, + user_prompt: user_prompt.clone(), + config: config.clone(), + status: SessionStatus::Running, + started_at: chrono::Utc::now().timestamp(), + ended_at: None, + result: None, + error: None, + }; + + let session_id = match storage.create_session(session).await { Ok(id) => id, Err(e) => { error!(error = %e, "Failed to create session"); @@ -56,22 +65,35 @@ pub async fn execute_workflow_creation_agent( }; let llm_client_clone = llm_client.get_ref().clone(); - let database_clone = database.get_ref().clone(); + let storage_clone = storage.get_ref().clone(); let user_prompt_clone = user_prompt.clone(); let type_names_clone = type_names.clone(); tokio::spawn(async move { let agent = match create_structured_json_agent( &llm_client_clone, - &database_clone, + &storage_clone, type_names_clone, domain_description, ) { Ok(agent) => agent, Err(e) => { error!(error = %e, session_id = session_id, "Failed to create Workflow Creation agent"); - let _ = database_clone - .fail_session(session_id, &format!("Failed to create agent: {}", e)); + let mut session = Session { + id: Some(session_id), + agent_name: "workflow-creation".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Failed to create agent: {}", e)), + }; + let _ = storage_clone.update_session(session).await; return; } }; @@ -79,14 +101,41 @@ pub async fn execute_workflow_creation_agent( match agent.execute(&user_prompt_clone, session_id).await { Ok(result) => { info!(result = %result, session_id = session_id, "Agent execution completed successfully"); - if let Err(e) = database_clone.complete_session(session_id, &result) { + let mut session = Session { + id: Some(session_id), + agent_name: "workflow-creation".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Completed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: Some(result.clone()), + error: None, + }; + if let Err(e) = storage_clone.update_session(session).await { error!(error = %e, session_id = session_id, "Failed to complete session"); } } Err(e) => { error!(error = %e, session_id = session_id, "Agent execution failed"); - let _ = - database_clone.fail_session(session_id, &format!("Execution failed: {}", e)); + let mut session = Session { + id: Some(session_id), + agent_name: "workflow-creation".to_string(), + provider, + model, + system_prompt: None, + user_prompt: user_prompt_clone, + config, + status: SessionStatus::Failed, + started_at: 0, + ended_at: Some(chrono::Utc::now().timestamp()), + result: None, + error: Some(format!("Execution failed: {}", e)), + }; + let _ = storage_clone.update_session(session).await; } } }); diff --git a/nocodo-api/src/handlers/sessions.rs b/nocodo-api/src/handlers/sessions.rs index c81842d0..9950c968 100644 --- a/nocodo-api/src/handlers/sessions.rs +++ b/nocodo-api/src/handlers/sessions.rs @@ -1,13 +1,15 @@ use crate::models::ErrorResponse; +use crate::storage::SqliteAgentStorage; use crate::DbConnection; use actix_web::{get, post, web, HttpResponse, Responder}; -use nocodo_agents::Agent; +use nocodo_agents::{Agent, AgentStorage}; use rusqlite::{params, Connection}; use serde::{Deserialize, Serialize}; use shared_types::{ SessionListItem, SessionListResponse, SessionMessage, SessionResponse, SessionToolCall, }; use std::collections::HashMap; +use std::sync::Arc; use tracing::{error, info, warn}; #[get("/agents/sessions/{session_id}")] @@ -175,27 +177,62 @@ pub struct QuestionsResponse { #[get("/agents/sessions/{session_id}/questions")] pub async fn get_pending_questions( session_id: web::Path, - database: web::Data>, + db: web::Data, ) -> impl Responder { let id = session_id.into_inner(); info!(session_id = id, "Retrieving pending questions"); - match database.get_pending_questions(id) { - Ok(questions) => { - info!( - session_id = id, - question_count = questions.len(), - "Retrieved pending questions" - ); - HttpResponse::Ok().json(QuestionsResponse { questions }) - } + // Get pending questions from project_requirements_qna table + let conn = db.lock().unwrap(); + + let questions: Vec = match conn.prepare( + "SELECT question_id, question, description, response_type + FROM project_requirements_qna + WHERE session_id = ?1 AND answer IS NULL + ORDER BY created_at ASC", + ) { + Ok(mut stmt) => stmt + .query_map(params![id], |row| { + let response_type_str: String = row.get(3)?; + let response_type = match response_type_str.as_str() { + "text" => shared_types::user_interaction::QuestionType::Text, + "password" => shared_types::user_interaction::QuestionType::Password, + "file_path" => shared_types::user_interaction::QuestionType::FilePath, + "email" => shared_types::user_interaction::QuestionType::Email, + "url" => shared_types::user_interaction::QuestionType::Url, + _ => shared_types::user_interaction::QuestionType::Text, + }; + + Ok(shared_types::user_interaction::UserQuestion { + id: row.get(0)?, + question: row.get(1)?, + description: row.get(2)?, + response_type, + default: None, + options: None, + }) + }) + .map_err(|e| { + error!(error = %e, session_id = id, "Query map failed"); + rusqlite::Error::InvalidQuery + }) + .ok() + .map(|rows| rows.filter_map(Result::ok).collect()) + .unwrap_or_default(), Err(e) => { error!(error = %e, session_id = id, "Failed to retrieve questions"); - HttpResponse::InternalServerError().json(ErrorResponse { + return HttpResponse::InternalServerError().json(ErrorResponse { error: format!("Failed to retrieve questions: {}", e), - }) + }); } - } + }; + + info!( + session_id = id, + question_count = questions.len(), + "Retrieved pending questions" + ); + HttpResponse::Ok().json(QuestionsResponse { questions }) } #[derive(Deserialize)] @@ -207,7 +244,8 @@ pub struct SubmitAnswersRequest { pub async fn submit_answers( session_id: web::Path, req: web::Json, - database: web::Data>, + db: web::Data, + storage: web::Data>, llm_client: web::Data>, ) -> impl Responder { let id = session_id.into_inner(); @@ -218,11 +256,20 @@ pub async fn submit_answers( ); // Store answers in database - if let Err(e) = database.store_answers(id, &req.answers) { - error!(error = %e, session_id = id, "Failed to store answers"); - return HttpResponse::InternalServerError().json(ErrorResponse { - error: format!("Failed to store answers: {}", e), - }); + { + let conn = db.lock().unwrap(); + let now = chrono::Utc::now().timestamp(); + for (question_id, answer) in &req.answers { + if let Err(e) = conn.execute( + "UPDATE project_requirements_qna SET answer = ?1, answered_at = ?2 WHERE session_id = ?3 AND question_id = ?4", + params![answer, now, id, question_id], + ) { + error!(error = %e, session_id = id, question_id = %question_id, "Failed to store answer"); + return HttpResponse::InternalServerError().json(ErrorResponse { + error: format!("Failed to store answer: {}", e), + }); + } + } } // Build a message with the answered questions for the agent @@ -230,7 +277,7 @@ pub async fn submit_answers( // Get the answered questions from database (all questions for this session) { - let conn = database.connection.lock().unwrap(); + let conn = db.lock().unwrap(); let mut stmt = conn .prepare( "SELECT question_id, question, answer @@ -259,27 +306,41 @@ pub async fn submit_answers( } } } - } // conn and stmt are dropped here + } // Add the answers as a user message - if let Err(e) = database.create_message(id, "user", &answers_text) { - error!(error = %e, session_id = id, "Failed to create message with answers"); - return HttpResponse::InternalServerError().json(ErrorResponse { - error: format!("Failed to create message: {}", e), - }); + { + let now = chrono::Utc::now().timestamp(); + let message = nocodo_agents::Message { + id: None, + session_id: id, + role: nocodo_agents::MessageRole::User, + content: answers_text.clone(), + created_at: now, + }; + if let Err(e) = storage.create_message(message).await { + error!(error = %e, session_id = id, "Failed to create message with answers"); + return HttpResponse::InternalServerError().json(ErrorResponse { + error: format!("Failed to create message: {}", e), + }); + } } - // Resume the session - if let Err(e) = database.resume_session(id) { - error!(error = %e, session_id = id, "Failed to resume session"); - return HttpResponse::InternalServerError().json(ErrorResponse { - error: format!("Failed to resume session: {}", e), - }); + // Update session to running status + if let Ok(Some(mut session)) = storage.get_session(id).await { + session.status = nocodo_agents::SessionStatus::Running; + if let Err(e) = storage.update_session(session).await { + error!(error = %e, session_id = id, "Failed to resume session"); + return HttpResponse::InternalServerError().json(ErrorResponse { + error: format!("Failed to resume session: {}", e), + }); + } } // Spawn a background task to continue agent execution - let database_clone = database.get_ref().clone(); + let storage_clone = storage.get_ref().clone(); let llm_client_clone = llm_client.get_ref().clone(); + let db_clone = db.get_ref().clone(); tokio::spawn(async move { info!( @@ -287,22 +348,28 @@ pub async fn submit_answers( "Resuming agent execution with user answers" ); - // Create the agent + // Create an agent let agent = match crate::helpers::agents::create_user_clarification_agent( &llm_client_clone, - &database_clone, + &storage_clone, + &db_clone, ) { Ok(agent) => agent, Err(e) => { error!(error = %e, session_id = id, "Failed to create agent for resumption"); - let _ = database_clone.fail_session(id, &format!("Failed to create agent: {}", e)); + if let Ok(Some(mut session)) = storage_clone.get_session(id).await { + session.status = nocodo_agents::SessionStatus::Failed; + session.error = Some(format!("Failed to create agent: {}", e)); + session.ended_at = Some(chrono::Utc::now().timestamp()); + let _ = storage_clone.update_session(session).await; + } return; } }; // Get the original user prompt from the session let original_prompt = { - let conn = database_clone.connection.lock().unwrap(); + let conn = db_clone.lock().unwrap(); conn.query_row( "SELECT user_prompt FROM agent_sessions WHERE id = ?1", params![id], @@ -318,19 +385,37 @@ pub async fn submit_answers( info!(session_id = id, "Agent resumed successfully"); // Don't complete here if still waiting for more input if !result.contains("Waiting for user") { - if let Err(e) = database_clone.complete_session(id, &result) { - error!(error = %e, session_id = id, "Failed to complete session after resumption"); + if let Ok(Some(mut session)) = + storage_clone.get_session(id).await + { + session.status = nocodo_agents::SessionStatus::Completed; + session.result = Some(result.clone()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + if let Err(e) = storage_clone.update_session(session).await { + error!(error = %e, session_id = id, "Failed to complete session after resumption"); + } } } } Err(e) => { error!(error = %e, session_id = id, "Agent execution failed after resumption"); - let _ = database_clone.fail_session(id, &format!("Execution failed: {}", e)); + if let Ok(Some(mut session)) = storage_clone.get_session(id).await + { + session.status = nocodo_agents::SessionStatus::Failed; + session.error = Some(format!("Execution failed: {}", e)); + session.ended_at = Some(chrono::Utc::now().timestamp()); + let _ = storage_clone.update_session(session).await; + } } } } else { error!(session_id = id, "Failed to retrieve original prompt"); - let _ = database_clone.fail_session(id, "Failed to retrieve original prompt"); + if let Ok(Some(mut session)) = storage_clone.get_session(id).await { + session.status = nocodo_agents::SessionStatus::Failed; + session.error = Some("Failed to retrieve original prompt".to_string()); + session.ended_at = Some(chrono::Utc::now().timestamp()); + let _ = storage_clone.update_session(session).await; + } } }); diff --git a/nocodo-api/src/helpers/agents.rs b/nocodo-api/src/helpers/agents.rs index 7cbb0d40..2ff3ddc1 100644 --- a/nocodo-api/src/helpers/agents.rs +++ b/nocodo-api/src/helpers/agents.rs @@ -1,3 +1,4 @@ +use nocodo_agents::AgentStorage; use nocodo_llm_sdk::client::LlmClient; use shared_types::AgentInfo; use std::sync::Arc; @@ -70,22 +71,22 @@ pub fn list_supported_agents() -> Vec { ] } -/// Creates a SQLite analysis agent using the shared database +/// Creates a SQLite analysis agent using the shared storage /// /// # Arguments /// /// * `llm_client` - The LLM client to use for the agent -/// * `database` - Shared database for session persistence +/// * `storage` - Shared storage for session persistence /// * `db_path` - Path to the SQLite database to analyze /// /// # Returns /// /// A SQLite analysis agent instance -pub async fn create_sqlite_agent( +pub async fn create_sqlite_agent( llm_client: &Arc, - database: &Arc, + storage: &Arc, db_path: &str, -) -> anyhow::Result { +) -> anyhow::Result> { let tool_executor = Arc::new( nocodo_tools::ToolExecutor::new(std::env::current_dir()?) .with_max_file_size(10 * 1024 * 1024), @@ -93,7 +94,7 @@ pub async fn create_sqlite_agent( let agent = nocodo_agents::sqlite_reader::SqliteReaderAgent::new( llm_client.clone(), - database.clone(), + storage.clone(), tool_executor, db_path.to_string(), ) @@ -106,21 +107,21 @@ pub async fn create_sqlite_agent( /// /// # Arguments /// -/// * `llm_client` - The LLM client to use for the agent -/// * `database` - Shared database for session persistence +/// * `llm_client` - The LLM client to use for agent +/// * `storage` - Shared storage for session persistence /// * `image_path` - Path to the image file to process /// /// # Returns /// /// A Tesseract OCR agent instance -pub async fn create_tesseract_agent( +pub async fn create_tesseract_agent( llm_client: &Arc, - database: &Arc, + storage: &Arc, image_path: &str, -) -> anyhow::Result { +) -> anyhow::Result> { let agent = nocodo_agents::tesseract::TesseractAgent::new( llm_client.clone(), - database.clone(), + storage.clone(), std::path::PathBuf::from(image_path), )?; @@ -131,20 +132,20 @@ pub async fn create_tesseract_agent( /// /// # Arguments /// -/// * `llm_client` - The LLM client to use for the agent -/// * `database` - Shared database for session persistence +/// * `llm_client` - The LLM client to use for agent +/// * `storage` - Shared storage for session persistence /// * `type_names` - List of TypeScript type names to use for validation -/// * `domain_description` - Description of the domain context +/// * `domain_description` - Description of domain context /// /// # Returns /// /// A Structured JSON agent instance -pub fn create_structured_json_agent( +pub fn create_structured_json_agent( llm_client: &Arc, - database: &Arc, + storage: &Arc, type_names: Vec, domain_description: String, -) -> anyhow::Result { +) -> anyhow::Result> { let tool_executor = Arc::new( nocodo_tools::ToolExecutor::new(std::env::current_dir()?) .with_max_file_size(10 * 1024 * 1024), @@ -157,7 +158,7 @@ pub fn create_structured_json_agent( let agent = nocodo_agents::structured_json::StructuredJsonAgent::new( llm_client.clone(), - database.clone(), + storage.clone(), tool_executor, config, )?; @@ -169,48 +170,168 @@ pub fn create_structured_json_agent( /// /// # Arguments /// -/// * `llm_client` - The LLM client to use for the agent -/// * `database` - Shared database for session persistence +/// * `llm_client` - The LLM client to use for agent +/// * `storage` - Shared storage for session persistence +/// * `db_connection` - Shared database connection for requirements Q&A storage /// /// # Returns /// /// A User Clarification agent instance pub fn create_user_clarification_agent( llm_client: &Arc, - database: &Arc, -) -> anyhow::Result { + storage: &Arc, + db_connection: &crate::DbConnection, +) -> anyhow::Result< + nocodo_agents::requirements_gathering::UserClarificationAgent< + crate::storage::SqliteAgentStorage, + DirectRequirementsStorage, + >, +> { + use crate::storage::SqliteAgentStorage; + let tool_executor = Arc::new( nocodo_tools::ToolExecutor::new(std::env::current_dir()?) .with_max_file_size(10 * 1024 * 1024), ); + // Create a requirements storage wrapper for direct database access + let requirements_storage = Arc::new(DirectRequirementsStorage::new(db_connection.to_owned())); + let agent = nocodo_agents::requirements_gathering::UserClarificationAgent::new( llm_client.clone(), - database.clone(), + storage.clone(), + requirements_storage, tool_executor, ); Ok(agent) } +// Direct requirements storage that accesses database directly +pub struct DirectRequirementsStorage { + db_connection: crate::DbConnection, +} + +impl DirectRequirementsStorage { + fn new(db_connection: crate::DbConnection) -> Self { + Self { db_connection } + } +} + +#[async_trait::async_trait] +impl nocodo_agents::requirements_gathering::storage::RequirementsStorage + for DirectRequirementsStorage +{ + async fn store_questions( + &self, + _session_id: i64, + _tool_call_id: Option, + _questions: &[shared_types::user_interaction::UserQuestion], + ) -> Result<(), nocodo_agents::StorageError> { + Ok(()) + } + + async fn get_pending_questions( + &self, + session_id: i64, + ) -> Result, nocodo_agents::StorageError> + { + let conn = self.db_connection.clone(); + + tokio::task::spawn_blocking(move || { + let conn = conn.lock().map_err(|e| { + nocodo_agents::StorageError::OperationFailed(format!("Lock error: {}", e)) + })?; + + let mut stmt = conn + .prepare( + "SELECT question_id, question, description, response_type + FROM project_requirements_qna + WHERE session_id = ?1 AND answer IS NULL + ORDER BY created_at ASC", + ) + .map_err(|e| nocodo_agents::StorageError::OperationFailed(e.to_string()))?; + + let questions = stmt + .query_map([session_id], |row| { + let response_type_str: String = row.get(3)?; + let response_type = match response_type_str.as_str() { + "text" => shared_types::user_interaction::QuestionType::Text, + "password" => shared_types::user_interaction::QuestionType::Password, + "file_path" => shared_types::user_interaction::QuestionType::FilePath, + "email" => shared_types::user_interaction::QuestionType::Email, + "url" => shared_types::user_interaction::QuestionType::Url, + _ => shared_types::user_interaction::QuestionType::Text, + }; + + Ok(shared_types::user_interaction::UserQuestion { + id: row.get(0)?, + question: row.get(1)?, + description: row.get(2)?, + response_type, + default: None, + options: None, + }) + }) + .map_err(|e| nocodo_agents::StorageError::OperationFailed(e.to_string()))? + .collect::, _>>() + .map_err(|e| nocodo_agents::StorageError::OperationFailed(e.to_string()))?; + + Ok(questions) + }) + .await + .map_err(|e| { + nocodo_agents::StorageError::OperationFailed(format!("Task join error: {}", e)) + })? + } + + async fn store_answers( + &self, + session_id: i64, + answers: &std::collections::HashMap, + ) -> Result<(), nocodo_agents::StorageError> { + let conn = self.db_connection.clone(); + let answers = answers.clone(); + + tokio::task::spawn_blocking(move || { + let conn = conn.lock().map_err(|e| { + nocodo_agents::StorageError::OperationFailed(format!("Lock error: {}", e)) + })?; + + let now = chrono::Utc::now().timestamp(); + for (question_id, answer) in answers { + conn.execute( + "UPDATE project_requirements_qna SET answer = ?1, answered_at = ?2 WHERE session_id = ?3 AND question_id = ?4", + rusqlite::params![answer, now, session_id, question_id], + ) + .map_err(|e| nocodo_agents::StorageError::OperationFailed(e.to_string()))?; + } + + Ok(()) + }) + .await + .map_err(|e| nocodo_agents::StorageError::OperationFailed(format!("Task join error: {}", e)))? + } +} + /// Creates a Settings Management agent /// /// # Arguments /// -/// * `llm_client` - The LLM client to use for the agent -/// * `database` - Shared database for session persistence +/// * `llm_client` - The LLM client to use for agent +/// * `storage` - Shared storage for session persistence /// * `settings_file_path` - Path to the TOML settings file /// * `agent_schemas` - List of agent settings schemas /// /// # Returns /// /// A Settings Management agent instance -pub fn create_settings_management_agent( +pub fn create_settings_management_agent( llm_client: &Arc, - database: &Arc, + storage: &Arc, settings_file_path: &str, agent_schemas: Vec, -) -> anyhow::Result { +) -> anyhow::Result> { let tool_executor = Arc::new( nocodo_tools::ToolExecutor::new(std::env::current_dir()?) .with_max_file_size(10 * 1024 * 1024), @@ -218,7 +339,7 @@ pub fn create_settings_management_agent( let agent = nocodo_agents::settings_management::SettingsManagementAgent::new( llm_client.clone(), - database.clone(), + storage.clone(), tool_executor, std::path::PathBuf::from(settings_file_path), agent_schemas, @@ -231,8 +352,8 @@ pub fn create_settings_management_agent( /// /// # Arguments /// -/// * `llm_client` - The LLM client to use for the agent -/// * `database` - Shared database for session persistence +/// * `llm_client` - The LLM client to use for agent +/// * `storage` - Shared storage for session persistence /// * `host` - IMAP server hostname /// * `port` - IMAP server port /// * `username` - IMAP username @@ -241,14 +362,14 @@ pub fn create_settings_management_agent( /// # Returns /// /// An IMAP Email agent instance -pub fn create_imap_agent( +pub fn create_imap_agent( llm_client: &Arc, - database: &Arc, + storage: &Arc, host: &str, port: u16, username: &str, password: &str, -) -> anyhow::Result { +) -> anyhow::Result> { let tool_executor = Arc::new( nocodo_tools::ToolExecutor::new(std::env::current_dir()?) .with_max_file_size(10 * 1024 * 1024), @@ -256,7 +377,7 @@ pub fn create_imap_agent( let agent = nocodo_agents::imap_email::ImapEmailAgent::new( llm_client.clone(), - database.clone(), + storage.clone(), tool_executor, host.to_string(), port, @@ -271,21 +392,21 @@ pub fn create_imap_agent( /// /// # Arguments /// -/// * `llm_client` - The LLM client to use for the agent -/// * `database` - Shared database for session persistence +/// * `llm_client` - The LLM client to use for agent +/// * `storage` - Shared storage for session persistence /// * `pdf_path` - Path to the PDF file to process /// /// # Returns /// /// A PDF to Text agent instance -pub async fn create_pdftotext_agent( +pub async fn create_pdftotext_agent( llm_client: &Arc, - database: &Arc, + storage: &Arc, pdf_path: &str, -) -> anyhow::Result { +) -> anyhow::Result> { let agent = nocodo_agents::pdftotext::PdfToTextAgent::new( llm_client.clone(), - database.clone(), + storage.clone(), std::path::PathBuf::from(pdf_path), )?; diff --git a/nocodo-api/src/helpers/database.rs b/nocodo-api/src/helpers/database.rs index 3d57568e..15bcc5a1 100644 --- a/nocodo-api/src/helpers/database.rs +++ b/nocodo-api/src/helpers/database.rs @@ -1,3 +1,5 @@ +use crate::storage::migrations; +use crate::storage::SqliteAgentStorage; use crate::DbConnection; use rusqlite::Connection; use std::path::PathBuf; @@ -5,15 +7,18 @@ use std::sync::{Arc, Mutex}; pub fn initialize_database( db_path: &PathBuf, -) -> anyhow::Result<(DbConnection, Arc)> { +) -> anyhow::Result<(DbConnection, Arc)> { if let Some(parent) = db_path.parent() { std::fs::create_dir_all(parent)?; } - let conn = Connection::open(&db_path)?; + let mut conn = Connection::open(&db_path)?; conn.execute("PRAGMA foreign_keys = ON", [])?; - let db = Arc::new(nocodo_agents::database::Database::new(&db_path)?); + migrations::run_migrations(&mut conn)?; - Ok((Arc::new(Mutex::new(conn)), db)) + let db_connection = Arc::new(Mutex::new(conn)); + let storage = Arc::new(SqliteAgentStorage::new(db_connection.clone())); + + Ok((db_connection, storage)) } diff --git a/nocodo-api/src/lib.rs b/nocodo-api/src/lib.rs index f34e1ab5..46232a51 100644 --- a/nocodo-api/src/lib.rs +++ b/nocodo-api/src/lib.rs @@ -5,5 +5,6 @@ pub mod config; pub mod handlers; pub mod helpers; pub mod models; +pub mod storage; pub type DbConnection = Arc>; diff --git a/nocodo-api/src/main.rs b/nocodo-api/src/main.rs index 6106cfca..1b2da557 100644 --- a/nocodo-api/src/main.rs +++ b/nocodo-api/src/main.rs @@ -2,6 +2,7 @@ mod config; mod handlers; mod helpers; mod models; +mod storage; use actix_cors::Cors; use actix_web::{web, App, HttpServer}; @@ -66,7 +67,7 @@ async fn main() -> Result<(), anyhow::Error> { let db_path = config.database.path.clone(); let bind_addr = format!("{}:{}", config.server.host, config.server.port); drop(config); - let (db_conn, db) = + let (db_conn, storage) = helpers::database::initialize_database(&db_path).expect("Failed to initialize database"); info!("Starting nocodo-api server at http://{}", bind_addr); @@ -99,7 +100,7 @@ async fn main() -> Result<(), anyhow::Error> { .wrap(cors) .app_data(web::Data::new(llm_client.clone())) .app_data(web::Data::new(db_conn.clone())) - .app_data(web::Data::new(db.clone())) + .app_data(web::Data::new(storage.clone())) .app_data(web::Data::new(handlers::settings::SettingsAppState { config: app_config.clone(), })) diff --git a/nocodo-api/src/storage/migrations/V1__create_agent_sessions.rs b/nocodo-api/src/storage/migrations/V1__create_agent_sessions.rs new file mode 100644 index 00000000..e4332112 --- /dev/null +++ b/nocodo-api/src/storage/migrations/V1__create_agent_sessions.rs @@ -0,0 +1,19 @@ +/// Create the agent_sessions table for tracking agent execution sessions +pub fn migration() -> String { + r#" +CREATE TABLE agent_sessions ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + agent_name TEXT NOT NULL, + provider TEXT NOT NULL, + model TEXT NOT NULL, + system_prompt TEXT, + user_prompt TEXT NOT NULL, + config TEXT, + status TEXT NOT NULL DEFAULT 'running' CHECK (status IN ('running', 'completed', 'failed', 'waiting_for_user_input')), + started_at INTEGER NOT NULL, + ended_at INTEGER, + result TEXT, + error TEXT +); +"#.to_string() +} diff --git a/nocodo-api/src/storage/migrations/V2__create_agent_messages.rs b/nocodo-api/src/storage/migrations/V2__create_agent_messages.rs new file mode 100644 index 00000000..7d44193b --- /dev/null +++ b/nocodo-api/src/storage/migrations/V2__create_agent_messages.rs @@ -0,0 +1,17 @@ +/// Create the agent_messages table for storing conversation messages +pub fn migration() -> String { + r#" +CREATE TABLE agent_messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id INTEGER NOT NULL, + role TEXT NOT NULL CHECK (role IN ('user', 'assistant', 'system', 'tool')), + content TEXT NOT NULL, + created_at INTEGER NOT NULL, + FOREIGN KEY (session_id) REFERENCES agent_sessions (id) ON DELETE CASCADE +); + +CREATE INDEX idx_agent_messages_session_created + ON agent_messages(session_id, created_at); +"# + .to_string() +} diff --git a/nocodo-api/src/storage/migrations/V3__create_agent_tool_calls.rs b/nocodo-api/src/storage/migrations/V3__create_agent_tool_calls.rs new file mode 100644 index 00000000..2df1a34f --- /dev/null +++ b/nocodo-api/src/storage/migrations/V3__create_agent_tool_calls.rs @@ -0,0 +1,27 @@ +/// Create the agent_tool_calls table for tracking tool executions +pub fn migration() -> String { + r#" +CREATE TABLE agent_tool_calls ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id INTEGER NOT NULL, + message_id INTEGER, + tool_call_id TEXT NOT NULL, + tool_name TEXT NOT NULL, + request TEXT NOT NULL, + response TEXT, + status TEXT NOT NULL DEFAULT 'pending' CHECK (status IN ('pending', 'executing', 'completed', 'failed')), + created_at INTEGER NOT NULL, + completed_at INTEGER, + execution_time_ms INTEGER, + error_details TEXT, + FOREIGN KEY (session_id) REFERENCES agent_sessions (id) ON DELETE CASCADE, + FOREIGN KEY (message_id) REFERENCES agent_messages (id) ON DELETE SET NULL +); + +CREATE INDEX idx_agent_tool_calls_session + ON agent_tool_calls(session_id); + +CREATE INDEX idx_agent_tool_calls_status + ON agent_tool_calls(session_id, status); +"#.to_string() +} diff --git a/nocodo-api/src/storage/migrations/V4__create_project_requirements_qna.rs b/nocodo-api/src/storage/migrations/V4__create_project_requirements_qna.rs new file mode 100644 index 00000000..607c1768 --- /dev/null +++ b/nocodo-api/src/storage/migrations/V4__create_project_requirements_qna.rs @@ -0,0 +1,26 @@ +/// Create the project_requirements_qna table for storing user questions and answers +/// +/// **Agent-Specific Migration**: This table is specific to the requirements_gathering agent. +/// Applications that don't use the requirements gathering agent can skip this migration. +pub fn migration() -> String { + r#" +CREATE TABLE project_requirements_qna ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id INTEGER NOT NULL, + tool_call_id INTEGER, + question_id TEXT NOT NULL, + question TEXT NOT NULL, + description TEXT, + response_type TEXT NOT NULL DEFAULT 'text', + answer TEXT, + created_at INTEGER NOT NULL, + answered_at INTEGER, + FOREIGN KEY (session_id) REFERENCES agent_sessions (id) ON DELETE CASCADE, + FOREIGN KEY (tool_call_id) REFERENCES agent_tool_calls (id) ON DELETE CASCADE +); + +CREATE INDEX idx_project_requirements_qna_session + ON project_requirements_qna(session_id); +"# + .to_string() +} diff --git a/nocodo-api/src/storage/migrations/V5__create_project_settings.rs b/nocodo-api/src/storage/migrations/V5__create_project_settings.rs new file mode 100644 index 00000000..b31aca4e --- /dev/null +++ b/nocodo-api/src/storage/migrations/V5__create_project_settings.rs @@ -0,0 +1,29 @@ +/// Create the project_settings table for storing user settings collected by the settings management agent +/// +/// **Agent-Specific Migration**: This table is specific to the settings_management agent. +/// Applications that don't use the settings management agent can skip this migration. +pub fn migration() -> String { + r#" +CREATE TABLE project_settings ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id INTEGER NOT NULL, + tool_call_id INTEGER, + setting_key TEXT NOT NULL, + setting_name TEXT NOT NULL, + description TEXT, + setting_type TEXT NOT NULL DEFAULT 'text', + setting_value TEXT, + created_at INTEGER NOT NULL, + updated_at INTEGER, + FOREIGN KEY (session_id) REFERENCES agent_sessions (id) ON DELETE CASCADE, + FOREIGN KEY (tool_call_id) REFERENCES agent_tool_calls (id) ON DELETE CASCADE +); + +CREATE INDEX idx_project_settings_session + ON project_settings(session_id); + +CREATE INDEX idx_project_settings_key + ON project_settings(setting_key); +"# + .to_string() +} diff --git a/nocodo-api/src/storage/migrations/mod.rs b/nocodo-api/src/storage/migrations/mod.rs new file mode 100644 index 00000000..e9c42fe5 --- /dev/null +++ b/nocodo-api/src/storage/migrations/mod.rs @@ -0,0 +1,13 @@ +mod v1__create_agent_sessions; +mod v2__create_agent_messages; +mod v3__create_agent_tool_calls; +mod v4__create_project_requirements_qna; +mod v5__create_project_settings; + +use refinery::embed_migrations; + +embed_migrations!("src/storage/migrations"); + +pub fn run_migrations(conn: &mut rusqlite::Connection) -> Result<(), refinery::Error> { + migrations::runner().run(conn).map(|_| ()) +} diff --git a/nocodo-api/src/storage/mod.rs b/nocodo-api/src/storage/mod.rs new file mode 100644 index 00000000..d55a97a7 --- /dev/null +++ b/nocodo-api/src/storage/mod.rs @@ -0,0 +1,4 @@ +pub mod migrations; +pub mod sqlite; + +pub use sqlite::SqliteAgentStorage; diff --git a/nocodo-api/src/storage/sqlite.rs b/nocodo-api/src/storage/sqlite.rs new file mode 100644 index 00000000..01aacb6b --- /dev/null +++ b/nocodo-api/src/storage/sqlite.rs @@ -0,0 +1,409 @@ +use crate::DbConnection; +use async_trait::async_trait; +use nocodo_agents::{ + AgentStorage, Message, MessageRole, Session, SessionStatus, StorageError, ToolCall, + ToolCallStatus, +}; +use rusqlite::{params, OptionalExtension}; + +pub struct SqliteAgentStorage { + connection: DbConnection, +} + +impl SqliteAgentStorage { + pub fn new(connection: DbConnection) -> Self { + Self { connection } + } +} + +#[async_trait] +impl AgentStorage for SqliteAgentStorage { + async fn create_session(&self, session: Session) -> Result { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let status_str = match session.status { + SessionStatus::Running => "running", + SessionStatus::Completed => "completed", + SessionStatus::Failed => "failed", + SessionStatus::WaitingForUserInput => "waiting_for_user_input", + }; + + let config_json = + serde_json::to_string(&session.config).map_err(StorageError::SerializationError)?; + + conn.execute( + r#" + INSERT INTO agent_sessions + (agent_name, provider, model, system_prompt, user_prompt, config, status, started_at) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) + "#, + params![ + session.agent_name, + session.provider, + session.model, + session.system_prompt, + session.user_prompt, + config_json, + status_str, + session.started_at, + ], + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let id = conn.last_insert_rowid(); + Ok(id) + } + + async fn get_session(&self, session_id: i64) -> Result, StorageError> { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let mut stmt = conn + .prepare( + r#" + SELECT id, agent_name, provider, model, system_prompt, user_prompt, + config, status, started_at, ended_at, result, error + FROM agent_sessions + WHERE id = ?1 + "#, + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let session = stmt + .query_row(params![session_id], |row| { + let status_str: String = row.get(7)?; + let status = match status_str.as_str() { + "running" => SessionStatus::Running, + "completed" => SessionStatus::Completed, + "failed" => SessionStatus::Failed, + "waiting_for_user_input" => SessionStatus::WaitingForUserInput, + _ => SessionStatus::Running, + }; + + let config_json: String = row.get(6)?; + let config: serde_json::Value = + serde_json::from_str(&config_json).unwrap_or(serde_json::json!({})); + + Ok(Session { + id: Some(row.get(0)?), + agent_name: row.get(1)?, + provider: row.get(2)?, + model: row.get(3)?, + system_prompt: row.get(4)?, + user_prompt: row.get(5)?, + config, + status, + started_at: row.get(8)?, + ended_at: row.get(9)?, + result: row.get(10)?, + error: row.get(11)?, + }) + }) + .optional() + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + Ok(session) + } + + async fn update_session(&self, session: Session) -> Result<(), StorageError> { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let id = session + .id + .ok_or_else(|| StorageError::NotFound("Session ID required for update".to_string()))?; + + let status_str = match session.status { + SessionStatus::Running => "running", + SessionStatus::Completed => "completed", + SessionStatus::Failed => "failed", + SessionStatus::WaitingForUserInput => "waiting_for_user_input", + }; + + conn.execute( + r#" + UPDATE agent_sessions + SET status = ?1, ended_at = ?2, result = ?3, error = ?4 + WHERE id = ?5 + "#, + params![ + status_str, + session.ended_at, + session.result, + session.error, + id + ], + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + Ok(()) + } + + async fn create_message(&self, message: Message) -> Result { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let role_str = match message.role { + MessageRole::User => "user", + MessageRole::Assistant => "assistant", + MessageRole::System => "system", + MessageRole::Tool => "tool", + }; + + conn.execute( + r#" + INSERT INTO agent_messages (session_id, role, content, created_at) + VALUES (?1, ?2, ?3, ?4) + "#, + params![message.session_id, role_str, message.content, message.created_at], + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let id = conn.last_insert_rowid(); + Ok(id) + } + + async fn get_messages(&self, session_id: i64) -> Result, StorageError> { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let mut stmt = conn + .prepare( + r#" + SELECT id, session_id, role, content, created_at + FROM agent_messages + WHERE session_id = ?1 + ORDER BY created_at ASC + "#, + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let messages = stmt + .query_map(params![session_id], |row| { + let role_str: String = row.get(2)?; + let role = match role_str.as_str() { + "user" => MessageRole::User, + "assistant" => MessageRole::Assistant, + "system" => MessageRole::System, + "tool" => MessageRole::Tool, + _ => MessageRole::User, + }; + + Ok(Message { + id: Some(row.get(0)?), + session_id: row.get(1)?, + role, + content: row.get(3)?, + created_at: row.get(4)?, + }) + }) + .map_err(|e| StorageError::OperationFailed(e.to_string()))? + .collect::, _>>() + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + Ok(messages) + } + + async fn create_tool_call(&self, tool_call: ToolCall) -> Result { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let status_str = match tool_call.status { + ToolCallStatus::Pending => "pending", + ToolCallStatus::Executing => "executing", + ToolCallStatus::Completed => "completed", + ToolCallStatus::Failed => "failed", + }; + + let request_json = + serde_json::to_string(&tool_call.request).map_err(StorageError::SerializationError)?; + + conn.execute( + r#" + INSERT INTO agent_tool_calls + (session_id, message_id, tool_call_id, tool_name, request, status, created_at) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7) + "#, + params![ + tool_call.session_id, + tool_call.message_id, + tool_call.tool_call_id, + tool_call.tool_name, + request_json, + status_str, + tool_call.created_at, + ], + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let id = conn.last_insert_rowid(); + Ok(id) + } + + async fn update_tool_call(&self, tool_call: ToolCall) -> Result<(), StorageError> { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let id = tool_call + .id + .ok_or_else(|| { + StorageError::NotFound("Tool call ID required for update".to_string()) + })?; + + let status_str = match tool_call.status { + ToolCallStatus::Pending => "pending", + ToolCallStatus::Executing => "executing", + ToolCallStatus::Completed => "completed", + ToolCallStatus::Failed => "failed", + }; + + let response_json = tool_call + .response + .map(|r| serde_json::to_string(&r)) + .transpose() + .map_err(StorageError::SerializationError)?; + + conn.execute( + r#" + UPDATE agent_tool_calls + SET status = ?1, response = ?2, execution_time_ms = ?3, + completed_at = ?4, error_details = ?5 + WHERE id = ?6 + "#, + params![ + status_str, + response_json, + tool_call.execution_time_ms, + tool_call.completed_at, + tool_call.error_details, + id + ], + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + Ok(()) + } + + async fn get_tool_calls(&self, session_id: i64) -> Result, StorageError> { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let mut stmt = conn + .prepare( + r#" + SELECT id, session_id, message_id, tool_call_id, tool_name, request, + response, status, execution_time_ms, created_at, completed_at, error_details + FROM agent_tool_calls + WHERE session_id = ?1 + ORDER BY created_at ASC + "#, + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let tool_calls = stmt + .query_map(params![session_id], |row| { + let status_str: String = row.get(7)?; + let status = match status_str.as_str() { + "pending" => ToolCallStatus::Pending, + "executing" => ToolCallStatus::Executing, + "completed" => ToolCallStatus::Completed, + "failed" => ToolCallStatus::Failed, + _ => ToolCallStatus::Pending, + }; + + let request_json: String = row.get(5)?; + let request = serde_json::from_str(&request_json).unwrap_or(serde_json::json!({})); + + let response_json: Option = row.get(6)?; + let response = response_json.and_then(|json| serde_json::from_str(&json).ok()); + + Ok(ToolCall { + id: Some(row.get(0)?), + session_id: row.get(1)?, + message_id: row.get(2)?, + tool_call_id: row.get(3)?, + tool_name: row.get(4)?, + request, + response, + status, + execution_time_ms: row.get(8)?, + created_at: row.get(9)?, + completed_at: row.get(10)?, + error_details: row.get(11)?, + }) + }) + .map_err(|e| StorageError::OperationFailed(e.to_string()))? + .collect::, _>>() + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + Ok(tool_calls) + } + + async fn get_pending_tool_calls( + &self, + session_id: i64, + ) -> Result, StorageError> { + let conn = self + .connection + .lock() + .map_err(|e| StorageError::OperationFailed(format!("Lock error: {}", e)))?; + + let mut stmt = conn + .prepare( + r#" + SELECT id, session_id, message_id, tool_call_id, tool_name, request, + response, status, execution_time_ms, created_at, completed_at, error_details + FROM agent_tool_calls + WHERE session_id = ?1 AND status = 'pending' + ORDER BY created_at ASC + "#, + ) + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + let tool_calls = stmt + .query_map(params![session_id], |row| { + let request_json: String = row.get(5)?; + let request = serde_json::from_str(&request_json).unwrap_or(serde_json::json!({})); + + let response_json: Option = row.get(6)?; + let response = response_json.and_then(|json| serde_json::from_str(&json).ok()); + + Ok(ToolCall { + id: Some(row.get(0)?), + session_id: row.get(1)?, + message_id: row.get(2)?, + tool_call_id: row.get(3)?, + tool_name: row.get(4)?, + request, + response, + status: ToolCallStatus::Pending, + execution_time_ms: row.get(8)?, + created_at: row.get(9)?, + completed_at: row.get(10)?, + error_details: row.get(11)?, + }) + }) + .map_err(|e| StorageError::OperationFailed(e.to_string()))? + .collect::, _>>() + .map_err(|e| StorageError::OperationFailed(e.to_string()))?; + + Ok(tool_calls) + } +} From 6407e9daa634d1d49af81527ddcd392d4bd7c963 Mon Sep 17 00:00:00 2001 From: Sumit Datta Date: Tue, 3 Feb 2026 15:43:12 +0530 Subject: [PATCH 5/6] feat: make rusqlite optional with feature flags Add optional "sqlite" feature flag to both nocodo-tools and nocodo-agents to resolve rusqlite conflicts when integrating with PostgreSQL projects. Changes: - Make rusqlite and sqlparser optional dependencies in nocodo-tools - Add sqlite feature flag that conditionally compiles sqlite_reader and hackernews modules - Make SqliteReaderAgent and related code conditional on sqlite feature in nocodo-agents - Require sqlite feature for sqlite-reader-runner binary - Update tool schemas and executor to handle conditional sqlite support By default, neither crate pulls in rusqlite. Enable with features = ["sqlite"]. Co-Authored-By: Claude Sonnet 4.5 --- nocodo-agents/Cargo.toml | 5 +++++ nocodo-agents/src/factory.rs | 2 ++ nocodo-agents/src/lib.rs | 6 ++++++ nocodo-agents/src/tools/llm_schemas.rs | 3 +++ nocodo-tools/Cargo.toml | 8 ++++++-- nocodo-tools/src/lib.rs | 2 ++ nocodo-tools/src/tool_executor.rs | 6 ++++++ nocodo-tools/src/types/core.rs | 4 ++++ nocodo-tools/src/types/mod.rs | 4 ++++ 9 files changed, 38 insertions(+), 2 deletions(-) diff --git a/nocodo-agents/Cargo.toml b/nocodo-agents/Cargo.toml index cfd13b54..a97a0480 100644 --- a/nocodo-agents/Cargo.toml +++ b/nocodo-agents/Cargo.toml @@ -35,6 +35,10 @@ uuid = { version = "1.0", features = ["v4"] } [dev-dependencies] tempfile = "3.0" +[features] +default = [] +sqlite = ["nocodo-tools/sqlite"] + [[bin]] name = "codebase-analysis-runner" path = "bin/codebase_analysis_runner.rs" @@ -42,6 +46,7 @@ path = "bin/codebase_analysis_runner.rs" [[bin]] name = "sqlite-reader-runner" path = "bin/sqlite_reader_runner.rs" +required-features = ["sqlite"] [[bin]] name = "structured-json-runner" diff --git a/nocodo-agents/src/factory.rs b/nocodo-agents/src/factory.rs index e0f0a410..a262cf22 100644 --- a/nocodo-agents/src/factory.rs +++ b/nocodo-agents/src/factory.rs @@ -1,6 +1,7 @@ use crate::codebase_analysis::CodebaseAnalysisAgent; use crate::requirements_gathering::UserClarificationAgent; use crate::settings_management::SettingsManagementAgent; +#[cfg(feature = "sqlite")] use crate::sqlite_reader::SqliteReaderAgent; use crate::storage::{AgentStorage, InMemoryStorage}; use crate::structured_json::StructuredJsonAgent; @@ -264,6 +265,7 @@ pub fn create_codebase_analysis_agent( /// # Returns /// /// A SqliteReaderAgent instance +#[cfg(feature = "sqlite")] pub async fn create_sqlite_reader_agent( client: Arc, tool_executor: Arc, diff --git a/nocodo-agents/src/lib.rs b/nocodo-agents/src/lib.rs index dbcc9d7b..81f13396 100644 --- a/nocodo-agents/src/lib.rs +++ b/nocodo-agents/src/lib.rs @@ -5,6 +5,7 @@ pub mod imap_email; pub mod pdftotext; pub mod requirements_gathering; pub mod settings_management; +#[cfg(feature = "sqlite")] pub mod sqlite_reader; pub mod storage; pub mod structured_json; @@ -28,6 +29,7 @@ pub enum AgentTool { ApplyPatch, Bash, AskUser, + #[cfg(feature = "sqlite")] Sqlite3Reader, ImapReader, PdfToText, @@ -44,6 +46,7 @@ impl AgentTool { AgentTool::ApplyPatch => "apply_patch", AgentTool::Bash => "bash", AgentTool::AskUser => "ask_user", + #[cfg(feature = "sqlite")] AgentTool::Sqlite3Reader => "sqlite3_reader", AgentTool::ImapReader => "imap_reader", AgentTool::PdfToText => "pdftotext", @@ -90,6 +93,7 @@ impl AgentTool { let req: AskUserRequest = serde_json::from_value(arguments)?; ToolRequest::AskUser(req) } + #[cfg(feature = "sqlite")] "sqlite3_reader" => { let value: serde_json::Value = arguments; @@ -141,7 +145,9 @@ pub fn format_tool_response(response: &nocodo_tools::types::ToolResponse) -> Str r.exit_code, r.stdout, r.stderr ), ToolResponse::AskUser(r) => format!("User response: {:?}", r.responses), + #[cfg(feature = "sqlite")] ToolResponse::Sqlite3Reader(r) => r.formatted_output.clone(), + #[cfg(feature = "sqlite")] ToolResponse::HackerNewsResponse(r) => r.message.clone(), ToolResponse::ImapReader(r) => { if r.success { diff --git a/nocodo-agents/src/tools/llm_schemas.rs b/nocodo-agents/src/tools/llm_schemas.rs index f9017930..2da03dcc 100644 --- a/nocodo-agents/src/tools/llm_schemas.rs +++ b/nocodo-agents/src/tools/llm_schemas.rs @@ -9,6 +9,7 @@ fn default_true() -> bool { /// Create tool definitions for LLM using manager-models types pub fn create_tool_definitions() -> Vec { + #[cfg(feature = "sqlite")] let sqlite_schema = serde_json::json!({ "type": "object", "required": ["query"], @@ -21,6 +22,7 @@ pub fn create_tool_definitions() -> Vec { } }); + #[cfg(feature = "sqlite")] let sqlite_tool = Tool::from_json_schema( "sqlite3_reader".to_string(), "Read-only SQLite database tool. Use SELECT queries to retrieve data and PRAGMA statements to inspect database schema (tables, columns, indexes, foreign keys). The database path is pre-configured.".to_string(), @@ -196,6 +198,7 @@ pub fn create_tool_definitions() -> Vec { "Ask the user a list of questions to gather information or confirm actions", ) .build(), + #[cfg(feature = "sqlite")] sqlite_tool, imap_tool, pdftotext_tool, diff --git a/nocodo-tools/Cargo.toml b/nocodo-tools/Cargo.toml index 152a5b06..4949c735 100644 --- a/nocodo-tools/Cargo.toml +++ b/nocodo-tools/Cargo.toml @@ -32,8 +32,8 @@ codex-apply-patch = { git = "https://github.com/openai/codex", package = "codex- codex-core = { git = "https://github.com/openai/codex", package = "codex-core", rev = "6951872" } codex-process-hardening = { git = "https://github.com/openai/codex", package = "codex-process-hardening", rev = "6951872" } shared-types = { path = "../shared-types" } -rusqlite = { version = "0.37", features = ["bundled"] } -sqlparser = "0.51" +rusqlite = { version = "0.37", features = ["bundled"], optional = true } +sqlparser = { version = "0.51", optional = true } reqwest = { version = "0.12", features = ["json", "charset"] } home = "0.5" clap = { workspace = true } @@ -42,6 +42,10 @@ imap-proto = "0.16.1" rustls-connector = "0.19.0" mail-parser = "0.9" +[features] +default = [] +sqlite = ["rusqlite", "sqlparser"] + [dev-dependencies] tempfile = "3.12" tokio-test = "0.4" \ No newline at end of file diff --git a/nocodo-tools/src/lib.rs b/nocodo-tools/src/lib.rs index 628f22b7..a6f42190 100644 --- a/nocodo-tools/src/lib.rs +++ b/nocodo-tools/src/lib.rs @@ -1,9 +1,11 @@ pub mod bash; pub mod filesystem; pub mod grep; +#[cfg(feature = "sqlite")] pub mod hackernews; pub mod imap; pub mod pdftotext; +#[cfg(feature = "sqlite")] pub mod sqlite_reader; pub mod tool_error; pub mod tool_executor; diff --git a/nocodo-tools/src/tool_executor.rs b/nocodo-tools/src/tool_executor.rs index dfd8a6a8..2da8ff58 100644 --- a/nocodo-tools/src/tool_executor.rs +++ b/nocodo-tools/src/tool_executor.rs @@ -8,9 +8,11 @@ use crate::bash; pub use crate::bash::{BashExecutionResult, BashExecutorTrait}; use crate::filesystem::{apply_patch, list_files, read_file, write_file}; use crate::grep; +#[cfg(feature = "sqlite")] use crate::hackernews; use crate::imap; use crate::pdftotext; +#[cfg(feature = "sqlite")] use crate::sqlite_reader; use crate::user_interaction; @@ -78,9 +80,11 @@ impl ToolExecutor { .await } ToolRequest::AskUser(req) => user_interaction::ask_user(req).await, + #[cfg(feature = "sqlite")] ToolRequest::Sqlite3Reader(req) => sqlite_reader::execute_sqlite3_reader(req) .await .map_err(|e| anyhow::anyhow!(e)), + #[cfg(feature = "sqlite")] ToolRequest::HackerNewsRequest(req) => hackernews::execute_hackernews_request(req) .await .map_err(|e| anyhow::anyhow!(e)), @@ -107,7 +111,9 @@ impl ToolExecutor { ToolResponse::ApplyPatch(response) => serde_json::to_value(response)?, ToolResponse::Bash(response) => serde_json::to_value(response)?, ToolResponse::AskUser(response) => serde_json::to_value(response)?, + #[cfg(feature = "sqlite")] ToolResponse::Sqlite3Reader(response) => serde_json::to_value(response)?, + #[cfg(feature = "sqlite")] ToolResponse::HackerNewsResponse(response) => serde_json::to_value(response)?, ToolResponse::ImapReader(response) => serde_json::to_value(response)?, ToolResponse::PdfToText(response) => serde_json::to_value(response)?, diff --git a/nocodo-tools/src/types/core.rs b/nocodo-tools/src/types/core.rs index 13c5c966..126e9eff 100644 --- a/nocodo-tools/src/types/core.rs +++ b/nocodo-tools/src/types/core.rs @@ -20,8 +20,10 @@ pub enum ToolRequest { Bash(super::bash::BashRequest), #[serde(rename = "ask_user")] AskUser(AskUserRequest), + #[cfg(feature = "sqlite")] #[serde(rename = "sqlite3_reader")] Sqlite3Reader(super::sqlite_reader::Sqlite3ReaderRequest), + #[cfg(feature = "sqlite")] #[serde(rename = "hackernews_request")] HackerNewsRequest(super::hackernews::HackerNewsRequest), #[serde(rename = "imap_reader")] @@ -48,8 +50,10 @@ pub enum ToolResponse { Bash(super::bash::BashResponse), #[serde(rename = "ask_user")] AskUser(AskUserResponse), + #[cfg(feature = "sqlite")] #[serde(rename = "sqlite3_reader")] Sqlite3Reader(super::sqlite_reader::Sqlite3ReaderResponse), + #[cfg(feature = "sqlite")] #[serde(rename = "hackernews_response")] HackerNewsResponse(super::hackernews::HackerNewsResponse), #[serde(rename = "imap_reader")] diff --git a/nocodo-tools/src/types/mod.rs b/nocodo-tools/src/types/mod.rs index 70a14efe..61fbf703 100644 --- a/nocodo-tools/src/types/mod.rs +++ b/nocodo-tools/src/types/mod.rs @@ -2,9 +2,11 @@ pub mod bash; pub mod core; pub mod filesystem; pub mod grep; +#[cfg(feature = "sqlite")] pub mod hackernews; pub mod imap; pub mod pdftotext; +#[cfg(feature = "sqlite")] pub mod sqlite_reader; // Re-export commonly used types @@ -16,9 +18,11 @@ pub use filesystem::{ WriteFileResponse, }; pub use grep::{GrepMatch, GrepRequest, GrepResponse}; +#[cfg(feature = "sqlite")] pub use hackernews::{DownloadState, FetchMode, HackerNewsRequest, HackerNewsResponse, StoryType}; pub use imap::{ImapOperation, ImapReaderRequest, ImapReaderResponse, SearchCriteria}; pub use pdftotext::{PdfToTextRequest, PdfToTextResponse}; +#[cfg(feature = "sqlite")] pub use sqlite_reader::{Sqlite3ReaderRequest, Sqlite3ReaderResponse, SqliteMode}; // Re-export user interaction types from shared-types From f0b1666c4637ec931852e4eb8d0a654226217427 Mon Sep 17 00:00:00 2001 From: Sumit Datta Date: Tue, 3 Feb 2026 18:53:47 +0530 Subject: [PATCH 6/6] feat: make allowed working dirs customizable in PdfToTextAgent --- nocodo-agents/src/pdftotext/mod.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/nocodo-agents/src/pdftotext/mod.rs b/nocodo-agents/src/pdftotext/mod.rs index cc6f22e6..097786ea 100644 --- a/nocodo-agents/src/pdftotext/mod.rs +++ b/nocodo-agents/src/pdftotext/mod.rs @@ -38,12 +38,13 @@ impl PdfToTextAgent { /// * `client` - LLM client for AI inference /// * `storage` - Storage for session/message tracking /// * `pdf_path` - Path to the PDF file to process + /// * `allowed_working_dirs` - Optional list of allowed working directories. Defaults to ["/tmp", "/home", "/workspace", "/project"] /// /// # Security /// The agent is configured with restricted bash access: /// - Only the `pdftotext` and `qpdf` commands are allowed /// - All other bash commands are denied - /// - File operations are restricted to the PDF's directory + /// - File operations are restricted to the allowed working directories /// /// # Pre-conditions /// - pdftotext (poppler-utils) must be installed on the system @@ -54,6 +55,7 @@ impl PdfToTextAgent { client: Arc, storage: Arc, pdf_path: PathBuf, + allowed_working_dirs: Option>, ) -> anyhow::Result { // Validate PDF path exists if !pdf_path.exists() { @@ -73,7 +75,10 @@ impl PdfToTextAgent { .to_path_buf(); // Create restricted bash permissions (only pdftotext and qpdf commands) - let bash_perms = BashPermissions::minimal(vec!["pdftotext", "qpdf"]); + let bash_perms = BashPermissions::minimal(vec!["pdftotext", "qpdf"]) + .with_allowed_working_dirs( + allowed_working_dirs.unwrap_or_else(|| vec!["/tmp".to_string()]), + ); let bash_executor = BashExecutor::new(bash_perms, 120)?; // Create tool executor with restricted bash