diff --git a/.sqlx/query-0eff137ea88269d3b5b4cd1e973f8e5a2ee3d996101fa6541f3fc9ddbe874ea8.json b/.sqlx/query-0eff137ea88269d3b5b4cd1e973f8e5a2ee3d996101fa6541f3fc9ddbe874ea8.json new file mode 100644 index 0000000..4ad8d71 --- /dev/null +++ b/.sqlx/query-0eff137ea88269d3b5b4cd1e973f8e5a2ee3d996101fa6541f3fc9ddbe874ea8.json @@ -0,0 +1,50 @@ +{ + "db_name": "SQLite", + "query": "SELECT * FROM flowcharts WHERE session_id = ? ORDER BY created_at DESC", + "describe": { + "columns": [ + { + "name": "id", + "ordinal": 0, + "type_info": "Text" + }, + { + "name": "session_id", + "ordinal": 1, + "type_info": "Text" + }, + { + "name": "nodes", + "ordinal": 2, + "type_info": "Text" + }, + { + "name": "edges", + "ordinal": 3, + "type_info": "Text" + }, + { + "name": "created_at", + "ordinal": 4, + "type_info": "Text" + }, + { + "name": "token_usage", + "ordinal": 5, + "type_info": "Int64" + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + true, + false, + false, + false, + false, + true + ] + }, + "hash": "0eff137ea88269d3b5b4cd1e973f8e5a2ee3d996101fa6541f3fc9ddbe874ea8" +} diff --git a/.sqlx/query-1302f7c038525bce6da5d8aa8dc38debb5c4887ec8aeaa5fd44c04bec4d9dfe6.json b/.sqlx/query-1302f7c038525bce6da5d8aa8dc38debb5c4887ec8aeaa5fd44c04bec4d9dfe6.json new file mode 100644 index 0000000..8e6f3ea --- /dev/null +++ b/.sqlx/query-1302f7c038525bce6da5d8aa8dc38debb5c4887ec8aeaa5fd44c04bec4d9dfe6.json @@ -0,0 +1,12 @@ +{ + "db_name": "SQLite", + "query": "DELETE FROM flowcharts WHERE id = ?", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + }, + "hash": "1302f7c038525bce6da5d8aa8dc38debb5c4887ec8aeaa5fd44c04bec4d9dfe6" +} diff --git a/.sqlx/query-17371301b37431d07bc19575cc78ebe7ce777507f61182589eada5e7c181c147.json b/.sqlx/query-17371301b37431d07bc19575cc78ebe7ce777507f61182589eada5e7c181c147.json new file mode 100644 index 0000000..f73c5b2 --- /dev/null +++ b/.sqlx/query-17371301b37431d07bc19575cc78ebe7ce777507f61182589eada5e7c181c147.json @@ -0,0 +1,12 @@ +{ + "db_name": "SQLite", + "query": "\n INSERT INTO flowcharts (\n id, session_id, nodes, edges, created_at, token_usage\n ) VALUES (?, ?, ?, ?, ?, ?)\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 6 + }, + "nullable": [] + }, + "hash": "17371301b37431d07bc19575cc78ebe7ce777507f61182589eada5e7c181c147" +} diff --git a/.sqlx/query-56ef8e0d3dc92983df605fc3a53b9f05c13677a1209ddf7a4e4990c15b0a4edd.json b/.sqlx/query-56ef8e0d3dc92983df605fc3a53b9f05c13677a1209ddf7a4e4990c15b0a4edd.json new file mode 100644 index 0000000..e341acf --- /dev/null +++ b/.sqlx/query-56ef8e0d3dc92983df605fc3a53b9f05c13677a1209ddf7a4e4990c15b0a4edd.json @@ -0,0 +1,50 @@ +{ + "db_name": "SQLite", + "query": "SELECT * FROM flowcharts WHERE id = ?", + "describe": { + "columns": [ + { + "name": "id", + "ordinal": 0, + "type_info": "Text" + }, + { + "name": "session_id", + "ordinal": 1, + "type_info": "Text" + }, + { + "name": "nodes", + "ordinal": 2, + "type_info": "Text" + }, + { + "name": "edges", + "ordinal": 3, + "type_info": "Text" + }, + { + "name": "created_at", + "ordinal": 4, + "type_info": "Text" + }, + { + "name": "token_usage", + "ordinal": 5, + "type_info": "Int64" + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + true, + false, + false, + false, + false, + true + ] + }, + "hash": "56ef8e0d3dc92983df605fc3a53b9f05c13677a1209ddf7a4e4990c15b0a4edd" +} diff --git a/.sqlx/query-88a5c1b4029e9839a3988eb09182c4a04b9f2c846027ba5b114d2528206891e1.json b/.sqlx/query-88a5c1b4029e9839a3988eb09182c4a04b9f2c846027ba5b114d2528206891e1.json new file mode 100644 index 0000000..d1fc2e1 --- /dev/null +++ b/.sqlx/query-88a5c1b4029e9839a3988eb09182c4a04b9f2c846027ba5b114d2528206891e1.json @@ -0,0 +1,12 @@ +{ + "db_name": "SQLite", + "query": "DELETE FROM flowcharts WHERE session_id = ?", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + }, + "hash": "88a5c1b4029e9839a3988eb09182c4a04b9f2c846027ba5b114d2528206891e1" +} diff --git a/Cargo.lock b/Cargo.lock index 7e48ad9..ac5146a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2126,6 +2126,7 @@ dependencies = [ "criterion", "crossterm 0.27.0", "dirs", + "dotenvy", "flate2", "futures", "hex", diff --git a/Cargo.toml b/Cargo.toml index 9a2328d..e83a838 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -41,6 +41,7 @@ inquire = "0.7" console = "0.15" regex = "1.10" lazy_static = "1.4" +dotenvy = "0.15" [features] default = ["reqwest"] diff --git a/migrations/012_add_flowcharts.sql b/migrations/012_add_flowcharts.sql new file mode 100644 index 0000000..2a923f7 --- /dev/null +++ b/migrations/012_add_flowcharts.sql @@ -0,0 +1,14 @@ +-- Add flowcharts table for storing session context flow analysis +CREATE TABLE flowcharts ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + nodes TEXT NOT NULL, -- JSON array of FlowchartNode + edges TEXT NOT NULL, -- JSON array of FlowchartEdge + created_at TEXT NOT NULL DEFAULT (datetime('now', 'utc')), + token_usage INTEGER, + FOREIGN KEY (session_id) REFERENCES chat_sessions(id) ON DELETE CASCADE +); + +-- Indexes for performance +CREATE INDEX idx_flowcharts_session_id ON flowcharts(session_id); +CREATE INDEX idx_flowcharts_created_at ON flowcharts(created_at); diff --git a/src/cli/flowchart.rs b/src/cli/flowchart.rs new file mode 100644 index 0000000..18a4e33 --- /dev/null +++ b/src/cli/flowchart.rs @@ -0,0 +1,139 @@ +use anyhow::{Context, Result}; +use clap::Subcommand; +use std::sync::Arc; + +use crate::database::DatabaseManager; +use crate::services::{FlowchartService, GoogleAiClient, GoogleAiConfig}; + +#[derive(Subcommand)] +pub enum FlowchartCommands { + /// Generate flowchart for a session + Generate { + /// Session ID to generate flowchart for + session_id: String, + + /// Force regenerate even if flowchart exists + #[arg(short, long)] + force: bool, + }, + /// Show flowchart for a session + Show { + /// Session ID to show flowchart for + session_id: String, + }, + /// Delete flowchart for a session + Delete { + /// Session ID to delete flowchart for + session_id: String, + }, +} + +pub async fn handle_flowchart_command(command: FlowchartCommands) -> Result<()> { + match command { + FlowchartCommands::Generate { session_id, force } => { + generate_flowchart(&session_id, force).await + } + FlowchartCommands::Show { session_id } => show_flowchart(&session_id).await, + FlowchartCommands::Delete { session_id } => delete_flowchart(&session_id).await, + } +} + +async fn generate_flowchart(session_id: &str, force: bool) -> Result<()> { + let db_path = crate::database::config::get_default_db_path()?; + let db_manager = Arc::new(DatabaseManager::new(&db_path).await?); + + // Get Google AI API key + let api_key = std::env::var("GOOGLE_AI_API_KEY") + .context("GOOGLE_AI_API_KEY not set. Please set this environment variable.")?; + + let config = GoogleAiConfig::new(api_key); + let client = GoogleAiClient::new(config)?; + let service = FlowchartService::new(db_manager, client); + + println!("Generating flowchart for session: {session_id}"); + + if force { + // Delete existing flowchart first + service + .delete_flowchart(session_id) + .await + .map_err(|e| anyhow::anyhow!("Failed to delete existing flowchart: {e}"))?; + } + + let flowchart = service + .get_or_generate_flowchart(session_id) + .await + .map_err(|e| anyhow::anyhow!("Failed to generate flowchart: {e}"))?; + + println!("✓ Flowchart generated successfully!"); + println!(" - {} nodes", flowchart.nodes.len()); + println!(" - {} edges", flowchart.edges.len()); + if let Some(tokens) = flowchart.token_usage { + println!(" - {tokens} tokens used"); + } + println!("\nUse 'retrochat flowchart show {session_id}' to view the flowchart"); + Ok(()) +} + +async fn show_flowchart(session_id: &str) -> Result<()> { + let db_path = crate::database::config::get_default_db_path()?; + let db_manager = Arc::new(DatabaseManager::new(&db_path).await?); + let flowchart_repo = crate::database::FlowchartRepository::new(db_manager); + + let flowcharts = flowchart_repo + .get_by_session_id(session_id) + .await + .map_err(|e| anyhow::anyhow!("Error loading flowchart: {e}"))?; + + if let Some(flowchart) = flowcharts.first() { + println!("Flowchart for session: {session_id}\n"); + + // Render using the same renderer as TUI + use crate::tui::flowchart_renderer::FlowchartRenderer; + let renderer = FlowchartRenderer::new(80); + let lines = renderer.render(flowchart); + + for line in lines { + // Extract text from Line (ratatui type) + let text: String = line + .spans + .iter() + .map(|span| span.content.as_ref()) + .collect(); + println!("{text}"); + } + + println!("\nMetadata:"); + println!( + " Created: {}", + flowchart.created_at.format("%Y-%m-%d %H:%M:%S") + ); + if let Some(tokens) = flowchart.token_usage { + println!(" Tokens: {tokens}"); + } + println!(" Nodes: {}", flowchart.nodes.len()); + println!(" Edges: {}", flowchart.edges.len()); + + Ok(()) + } else { + println!("No flowchart found for session: {session_id}"); + println!("Generate one with: retrochat flowchart generate {session_id}"); + Ok(()) + } +} + +async fn delete_flowchart(session_id: &str) -> Result<()> { + let db_path = crate::database::config::get_default_db_path()?; + let db_manager = Arc::new(DatabaseManager::new(&db_path).await?); + let flowchart_repo = crate::database::FlowchartRepository::new(db_manager); + + println!("Deleting flowchart for session: {session_id}"); + + flowchart_repo + .delete_by_session_id(session_id) + .await + .map_err(|e| anyhow::anyhow!("Failed to delete flowchart: {e}"))?; + + println!("✓ Flowchart deleted successfully"); + Ok(()) +} diff --git a/src/cli/mod.rs b/src/cli/mod.rs index f5489f7..3c5625f 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -1,4 +1,5 @@ pub mod analytics; +pub mod flowchart; pub mod help; pub mod import; pub mod init; @@ -14,6 +15,7 @@ use tokio::runtime::Runtime; use crate::env::apis as env_vars; use crate::models::Provider; +use flowchart::FlowchartCommands; use retrospect::RetrospectCommands; #[derive(Parser)] @@ -94,6 +96,11 @@ pub enum Commands { #[command(subcommand)] command: RetrospectCommands, }, + /// Generate context flowchart for chat sessions + Flowchart { + #[command(subcommand)] + command: FlowchartCommands, + }, /// Interactive setup wizard for first-time users Setup, /// [Alias for 'import'] Add chat files interactively or from providers @@ -326,6 +333,9 @@ impl Cli { retrospect::handle_cancel_command(request_id, all).await } }, + Commands::Flowchart { command } => { + flowchart::handle_flowchart_command(command).await + } // New commands Commands::Setup => setup::run_setup_wizard().await, Commands::Add { diff --git a/src/database/flowchart_repo.rs b/src/database/flowchart_repo.rs new file mode 100644 index 0000000..d2cc6c1 --- /dev/null +++ b/src/database/flowchart_repo.rs @@ -0,0 +1,215 @@ +use chrono::{DateTime, Utc}; +use std::sync::Arc; + +use crate::database::DatabaseManager; +use crate::models::Flowchart; + +#[derive(Clone)] +pub struct FlowchartRepository { + db_manager: Arc, +} + +impl FlowchartRepository { + pub fn new(db_manager: Arc) -> Self { + Self { db_manager } + } + + pub async fn create( + &self, + flowchart: &Flowchart, + ) -> Result<(), Box> { + let pool = self.db_manager.pool(); + + let created_at_str = flowchart.created_at.to_rfc3339(); + let nodes_json = serde_json::to_string(&flowchart.nodes)?; + let edges_json = serde_json::to_string(&flowchart.edges)?; + + sqlx::query!( + r#" + INSERT INTO flowcharts ( + id, session_id, nodes, edges, created_at, token_usage + ) VALUES (?, ?, ?, ?, ?, ?) + "#, + flowchart.id, + flowchart.session_id, + nodes_json, + edges_json, + created_at_str, + flowchart.token_usage + ) + .execute(pool) + .await?; + + Ok(()) + } + + pub async fn find_by_id( + &self, + id: &str, + ) -> Result, Box> { + let pool = self.db_manager.pool(); + + let row = sqlx::query!("SELECT * FROM flowcharts WHERE id = ?", id) + .fetch_optional(pool) + .await?; + + if let Some(row) = row { + let created_at = DateTime::parse_from_rfc3339(&row.created_at)?.with_timezone(&Utc); + let nodes = serde_json::from_str(&row.nodes)?; + let edges = serde_json::from_str(&row.edges)?; + + Ok(Some(Flowchart { + id: row.id.expect("flowchart id should not be null"), + session_id: row.session_id, + nodes, + edges, + created_at, + token_usage: row.token_usage.map(|t| t as u32), + })) + } else { + Ok(None) + } + } + + pub async fn get_by_session_id( + &self, + session_id: &str, + ) -> Result, Box> { + let pool = self.db_manager.pool(); + + let rows = sqlx::query!( + "SELECT * FROM flowcharts WHERE session_id = ? ORDER BY created_at DESC", + session_id + ) + .fetch_all(pool) + .await?; + + let mut flowcharts = Vec::new(); + for row in rows { + let created_at = DateTime::parse_from_rfc3339(&row.created_at)?.with_timezone(&Utc); + let nodes = serde_json::from_str(&row.nodes)?; + let edges = serde_json::from_str(&row.edges)?; + + flowcharts.push(Flowchart { + id: row.id.expect("flowchart id should not be null"), + session_id: row.session_id, + nodes, + edges, + created_at, + token_usage: row.token_usage.map(|t| t as u32), + }); + } + + Ok(flowcharts) + } + + pub async fn delete(&self, id: &str) -> Result<(), Box> { + let pool = self.db_manager.pool(); + + sqlx::query!("DELETE FROM flowcharts WHERE id = ?", id) + .execute(pool) + .await?; + + Ok(()) + } + + pub async fn delete_by_session_id( + &self, + session_id: &str, + ) -> Result<(), Box> { + let pool = self.db_manager.pool(); + + sqlx::query!("DELETE FROM flowcharts WHERE session_id = ?", session_id) + .execute(pool) + .await?; + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::database::Database; + use crate::models::{EdgeType, FlowchartEdge, FlowchartNode, MessageRef, NodeType}; + + #[tokio::test] + async fn test_create_and_find_flowchart() { + let db = Database::new_in_memory().await.unwrap(); + db.initialize().await.unwrap(); + + // Create a test session first + let session_repo = db.chat_session_repo(); + let session = crate::models::ChatSession::new( + crate::models::Provider::ClaudeCode, + "test-provider".to_string(), + "test-hash".to_string(), + chrono::Utc::now(), + ); + session_repo.create(&session).await.unwrap(); + + let repo = FlowchartRepository::new(Arc::new(db.manager)); + + let nodes = vec![FlowchartNode { + id: "1".to_string(), + label: "Test node".to_string(), + message_refs: vec![MessageRef { + message_id: "msg-1".to_string(), + sequence_number: 1, + portion: None, + }], + node_type: NodeType::Action, + description: None, + }]; + + let edges = vec![FlowchartEdge { + from_node: "1".to_string(), + to_node: "2".to_string(), + edge_type: EdgeType::Sequential, + label: None, + }]; + + let flowchart = Flowchart::new(session.id.to_string(), nodes, edges); + let flowchart_id = flowchart.id.clone(); + + repo.create(&flowchart).await.unwrap(); + + let found = repo.find_by_id(&flowchart_id).await.unwrap(); + assert!(found.is_some()); + + let found_flowchart = found.unwrap(); + assert_eq!(found_flowchart.session_id, session.id.to_string()); + assert_eq!(found_flowchart.nodes.len(), 1); + assert_eq!(found_flowchart.edges.len(), 1); + } + + #[tokio::test] + async fn test_get_by_session_id() { + let db = Database::new_in_memory().await.unwrap(); + db.initialize().await.unwrap(); + + // Create a test session first + let session_repo = db.chat_session_repo(); + let session = crate::models::ChatSession::new( + crate::models::Provider::ClaudeCode, + "test-provider".to_string(), + "test-hash".to_string(), + chrono::Utc::now(), + ); + session_repo.create(&session).await.unwrap(); + + let repo = FlowchartRepository::new(Arc::new(db.manager)); + + let flowchart1 = Flowchart::new(session.id.to_string(), vec![], vec![]); + let flowchart2 = Flowchart::new(session.id.to_string(), vec![], vec![]); + + repo.create(&flowchart1).await.unwrap(); + repo.create(&flowchart2).await.unwrap(); + + let flowcharts = repo + .get_by_session_id(&session.id.to_string()) + .await + .unwrap(); + assert_eq!(flowcharts.len(), 2); + } +} diff --git a/src/database/mod.rs b/src/database/mod.rs index ace983f..d6fb2aa 100644 --- a/src/database/mod.rs +++ b/src/database/mod.rs @@ -2,6 +2,7 @@ pub mod analytics_repo; pub mod chat_session_repo; pub mod config; pub mod connection; +pub mod flowchart_repo; pub mod message_repo; pub mod migrations; pub mod project_repo; @@ -17,6 +18,7 @@ pub use analytics_repo::{ }; pub use chat_session_repo::ChatSessionRepository; pub use connection::DatabaseManager; +pub use flowchart_repo::FlowchartRepository; pub use message_repo::MessageRepository; pub use migrations::{MigrationManager, MigrationStatus}; pub use project_repo::ProjectRepository; @@ -76,6 +78,10 @@ impl Database { RetrospectionRepository::new(std::sync::Arc::new(self.manager.clone())) } + pub fn flowchart_repo(&self) -> FlowchartRepository { + FlowchartRepository::new(std::sync::Arc::new(self.manager.clone())) + } + pub fn tool_operation_repo(&self) -> ToolOperationRepository { ToolOperationRepository::new(&self.manager) } diff --git a/src/main.rs b/src/main.rs index f3512cf..94d71c1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,6 +4,14 @@ use retrochat::logging::LoggingConfig; use std::path::PathBuf; fn main() -> anyhow::Result<()> { + // Load environment variables from .env file if it exists + if let Err(e) = dotenvy::dotenv() { + // Only warn if .env file exists but couldn't be loaded + if std::path::Path::new(".env").exists() { + eprintln!("Warning: Could not load .env file: {e}"); + } + } + let cli = Cli::parse(); // Configure logging based on command diff --git a/src/models/flowchart.rs b/src/models/flowchart.rs new file mode 100644 index 0000000..edee0a6 --- /dev/null +++ b/src/models/flowchart.rs @@ -0,0 +1,307 @@ +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +/// A reference to a message or part of a message that belongs to this node +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct MessageRef { + /// Message UUID + pub message_id: String, + /// Sequence number in the session + pub sequence_number: u32, + /// If this node only covers part of the message content, describe which portion + pub portion: Option, +} + +/// Type of flowchart node +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub enum NodeType { + /// A context or action step + #[serde(rename = "context")] + Context, + /// A decision point + #[serde(rename = "decision")] + Decision, + /// An action performed + #[serde(rename = "action")] + Action, + /// Start of the flow + #[serde(rename = "start")] + Start, + /// End of the flow + #[serde(rename = "end")] + End, + /// Tool usage + #[serde(rename = "tool_use")] + ToolUse, + /// An event that occurred + #[serde(rename = "event")] + Event, +} + +impl std::fmt::Display for NodeType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + NodeType::Context => write!(f, "context"), + NodeType::Decision => write!(f, "decision"), + NodeType::Action => write!(f, "action"), + NodeType::Start => write!(f, "start"), + NodeType::End => write!(f, "end"), + NodeType::ToolUse => write!(f, "tool_use"), + NodeType::Event => write!(f, "event"), + } + } +} + +/// A node in the flowchart representing a single clear step or context +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct FlowchartNode { + /// Unique node ID (within this flowchart) + pub id: String, + /// Clear, concise label for this step (e.g., "Create todo-list") + pub label: String, + /// Messages that contribute to this node + pub message_refs: Vec, + /// Type of node + pub node_type: NodeType, + /// Optional detailed description + pub description: Option, +} + +/// Type of edge connecting nodes +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "lowercase")] +pub enum EdgeType { + /// Normal sequential flow + Sequential, + /// Multiple flows merging into one (2+ inputs → 1 output) + Merge, + /// Flow branching into multiple paths (1 input → 2+ outputs) + Branch, + /// Flow returning to a previous node (loop) + Loop, +} + +impl std::fmt::Display for EdgeType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + EdgeType::Sequential => write!(f, "sequential"), + EdgeType::Merge => write!(f, "merge"), + EdgeType::Branch => write!(f, "branch"), + EdgeType::Loop => write!(f, "loop"), + } + } +} + +/// An edge connecting two nodes in the flowchart +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct FlowchartEdge { + /// Source node ID + pub from_node: String, + /// Destination node ID + pub to_node: String, + /// Type of edge + pub edge_type: EdgeType, + /// Optional label for the edge + pub label: Option, +} + +/// A complete flowchart for a session +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Flowchart { + /// Unique flowchart ID + pub id: String, + /// Session this flowchart belongs to + pub session_id: String, + /// All nodes in the flowchart + pub nodes: Vec, + /// All edges in the flowchart + pub edges: Vec, + /// When this flowchart was created + pub created_at: DateTime, + /// Optional token usage for generation + pub token_usage: Option, +} + +impl Flowchart { + pub fn new(session_id: String, nodes: Vec, edges: Vec) -> Self { + Self { + id: Uuid::new_v4().to_string(), + session_id, + nodes, + edges, + created_at: Utc::now(), + token_usage: None, + } + } + + pub fn with_token_usage(mut self, token_usage: u32) -> Self { + self.token_usage = Some(token_usage); + self + } + + /// Get a node by ID + pub fn get_node(&self, node_id: &str) -> Option<&FlowchartNode> { + self.nodes.iter().find(|n| n.id == node_id) + } + + /// Get all edges originating from a node + pub fn get_outgoing_edges(&self, node_id: &str) -> Vec<&FlowchartEdge> { + self.edges + .iter() + .filter(|e| e.from_node == node_id) + .collect() + } + + /// Get all edges pointing to a node + pub fn get_incoming_edges(&self, node_id: &str) -> Vec<&FlowchartEdge> { + self.edges.iter().filter(|e| e.to_node == node_id).collect() + } + + /// Check if this is a valid DAG (no cycles) + pub fn is_valid_dag(&self) -> bool { + let mut visited = std::collections::HashSet::new(); + let mut rec_stack = std::collections::HashSet::new(); + + for node in &self.nodes { + if !visited.contains(&node.id) && self.has_cycle(&node.id, &mut visited, &mut rec_stack) + { + return false; + } + } + + true + } + + fn has_cycle( + &self, + node_id: &str, + visited: &mut std::collections::HashSet, + rec_stack: &mut std::collections::HashSet, + ) -> bool { + visited.insert(node_id.to_string()); + rec_stack.insert(node_id.to_string()); + + for edge in self.get_outgoing_edges(node_id) { + if !visited.contains(&edge.to_node) { + if self.has_cycle(&edge.to_node, visited, rec_stack) { + return true; + } + } else if rec_stack.contains(&edge.to_node) { + return true; + } + } + + rec_stack.remove(node_id); + false + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_flowchart_creation() { + let nodes = vec![ + FlowchartNode { + id: "1".to_string(), + label: "Start".to_string(), + message_refs: vec![], + node_type: NodeType::Start, + description: None, + }, + FlowchartNode { + id: "2".to_string(), + label: "Create todo-list".to_string(), + message_refs: vec![MessageRef { + message_id: Uuid::new_v4().to_string(), + sequence_number: 1, + portion: None, + }], + node_type: NodeType::Action, + description: None, + }, + ]; + + let edges = vec![FlowchartEdge { + from_node: "1".to_string(), + to_node: "2".to_string(), + edge_type: EdgeType::Sequential, + label: None, + }]; + + let flowchart = Flowchart::new("session-123".to_string(), nodes, edges); + + assert_eq!(flowchart.session_id, "session-123"); + assert_eq!(flowchart.nodes.len(), 2); + assert_eq!(flowchart.edges.len(), 1); + assert!(flowchart.is_valid_dag()); + } + + #[test] + fn test_dag_validation() { + let nodes = vec![ + FlowchartNode { + id: "1".to_string(), + label: "A".to_string(), + message_refs: vec![], + node_type: NodeType::Context, + description: None, + }, + FlowchartNode { + id: "2".to_string(), + label: "B".to_string(), + message_refs: vec![], + node_type: NodeType::Context, + description: None, + }, + ]; + + let edges = vec![FlowchartEdge { + from_node: "1".to_string(), + to_node: "2".to_string(), + edge_type: EdgeType::Sequential, + label: None, + }]; + + let flowchart = Flowchart::new("session-123".to_string(), nodes, edges); + assert!(flowchart.is_valid_dag()); + + let edges_with_cycle = vec![ + FlowchartEdge { + from_node: "1".to_string(), + to_node: "2".to_string(), + edge_type: EdgeType::Sequential, + label: None, + }, + FlowchartEdge { + from_node: "2".to_string(), + to_node: "1".to_string(), + edge_type: EdgeType::Sequential, + label: None, + }, + ]; + + let nodes_clone = vec![ + FlowchartNode { + id: "1".to_string(), + label: "A".to_string(), + message_refs: vec![], + node_type: NodeType::Context, + description: None, + }, + FlowchartNode { + id: "2".to_string(), + label: "B".to_string(), + message_refs: vec![], + node_type: NodeType::Context, + description: None, + }, + ]; + + let flowchart_with_cycle = + Flowchart::new("session-123".to_string(), nodes_clone, edges_with_cycle); + assert!(!flowchart_with_cycle.is_valid_dag()); + } +} diff --git a/src/models/mod.rs b/src/models/mod.rs index 9c7d3c4..e0715c6 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,5 +1,6 @@ pub mod bash_metadata; pub mod chat_session; +pub mod flowchart; pub mod message; pub mod project; pub mod provider; @@ -9,6 +10,7 @@ pub mod tool_operation; pub use bash_metadata::BashMetadata; pub use chat_session::{ChatSession, SessionState}; +pub use flowchart::{EdgeType, Flowchart, FlowchartEdge, FlowchartNode, MessageRef, NodeType}; pub use message::{Message, MessageRole, ToolCall, ToolResult, ToolUse}; pub use project::Project; pub use provider::{ParserType, Provider, ProviderConfig, ProviderRegistry}; diff --git a/src/services/flowchart_service.rs b/src/services/flowchart_service.rs new file mode 100644 index 0000000..f6f679a --- /dev/null +++ b/src/services/flowchart_service.rs @@ -0,0 +1,240 @@ +use std::sync::Arc; + +use crate::database::{DatabaseManager, FlowchartRepository, MessageRepository}; +use crate::models::{Flowchart, Message}; +use crate::services::google_ai::{GenerateContentRequest, GoogleAiClient}; + +pub struct FlowchartService { + google_ai_client: GoogleAiClient, + flowchart_repo: FlowchartRepository, + message_repo: MessageRepository, +} + +impl FlowchartService { + pub fn new(db_manager: Arc, google_ai_client: GoogleAiClient) -> Self { + let flowchart_repo = FlowchartRepository::new(db_manager.clone()); + let message_repo = MessageRepository::new(&db_manager); + + Self { + google_ai_client, + flowchart_repo, + message_repo, + } + } + + /// Generate or retrieve a flowchart for a session + /// Returns cached flowchart if exists, otherwise generates a new one + pub async fn get_or_generate_flowchart( + &self, + session_id: &str, + ) -> Result> { + // Check if flowchart already exists + let existing = self.flowchart_repo.get_by_session_id(session_id).await?; + if let Some(flowchart) = existing.first() { + return Ok(flowchart.clone()); + } + + // Generate new flowchart + self.generate_flowchart(session_id).await + } + + /// Force generate a new flowchart for a session + pub async fn generate_flowchart( + &self, + session_id: &str, + ) -> Result> { + // Get all messages for this session + let session_uuid = uuid::Uuid::parse_str(session_id)?; + let messages = self.message_repo.get_by_session_id(&session_uuid).await?; + + if messages.is_empty() { + return Err("No messages found for this session".into()); + } + + // Build prompt for Google AI + let prompt = self.build_flowchart_prompt(&messages); + + // Call Google AI + let request = GenerateContentRequest::new(prompt); + let response = self.google_ai_client.generate_content(request).await?; + + // Parse response + let response_text = response + .extract_text() + .ok_or("Failed to extract text from AI response")?; + + // Parse JSON response + let flowchart_data: FlowchartResponse = serde_json::from_str(&response_text) + .or_else(|_| self.extract_json_from_markdown(&response_text))?; + + // Create Flowchart model + let mut flowchart = Flowchart::new( + session_id.to_string(), + flowchart_data.nodes, + flowchart_data.edges, + ); + + if let Some(token_usage) = response.get_token_usage() { + flowchart = flowchart.with_token_usage(token_usage); + } + + // Validate DAG + if !flowchart.is_valid_dag() { + return Err("Generated flowchart contains cycles (not a valid DAG)".into()); + } + + // Save to database + self.flowchart_repo.create(&flowchart).await?; + + Ok(flowchart) + } + + fn build_flowchart_prompt(&self, messages: &[Message]) -> String { + let mut conversation = String::new(); + + for (idx, msg) in messages.iter().enumerate() { + // Truncate long messages to save tokens + let content = if msg.content.len() > 200 { + format!("{}...", &msg.content[..200]) + } else { + msg.content.clone() + }; + conversation.push_str(&format!( + "\n[Message {}] Role: {:?}\nContent: {}\n", + idx + 1, + msg.role, + content + )); + } + + format!( + r#"Create a simple flowchart with MAXIMUM 8 NODES. Each node represents a major work phase. + +CONVERSATION: +{conversation} + +Return ONLY this JSON (no markdown, no extra text): +{{ + "nodes": [ + {{"id": "1", "label": "Phase Name", "message_refs": [{{"message_id": "1", "sequence_number": 1, "portion": null}}], "node_type": "action", "description": null}} + ], + "edges": [ + {{"from_node": "1", "to_node": "2", "edge_type": "sequential", "label": null}} + ] +}} + +RULES: +- MAX 8 nodes only +- Each message_refs array: MAX 3 items only +- Always "portion": null +- node_type: "action", "context", "decision", "start", "end" +- edge_type: "sequential", "merge", "branch" +"# + ) + } + + /// Extract JSON from markdown code blocks if the AI wrapped it + fn extract_json_from_markdown( + &self, + text: &str, + ) -> Result> { + // Try to find JSON in markdown code blocks + if let Some(start) = text.find("```json") { + let search_start = start + 7; // Start after "```json" + if let Some(end_offset) = text[search_start..].find("```") { + let end = search_start + end_offset; + let json_str = &text[search_start..end].trim(); + return Ok(serde_json::from_str(json_str)?); + } + } + + // Try to find JSON without markdown wrapper + if let Some(start) = text.find('{') { + if let Some(end) = text.rfind('}') { + let json_str = &text[start..=end].trim(); + return Ok(serde_json::from_str(json_str)?); + } + } + + Err("Could not extract valid JSON from AI response".into()) + } + + /// Delete flowchart for a session (to force regeneration) + pub async fn delete_flowchart( + &self, + session_id: &str, + ) -> Result<(), Box> { + self.flowchart_repo.delete_by_session_id(session_id).await?; + Ok(()) + } +} + +/// Temporary structure for parsing AI response +#[derive(serde::Deserialize)] +struct FlowchartResponse { + nodes: Vec, + edges: Vec, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::database::Database; + use crate::models::{MessageRole, Project}; + use crate::services::google_ai::GoogleAiConfig; + + #[tokio::test] + #[ignore] // Requires GOOGLE_AI_API_KEY + async fn test_generate_flowchart() { + let db = Database::new_in_memory().await.unwrap(); + db.initialize().await.unwrap(); + + // Create test project and session + let project_repo = db.project_repo(); + let project = Project::new("test-project".to_string()); + project_repo.create(&project).await.unwrap(); + + let session_repo = db.chat_session_repo(); + let session = crate::models::ChatSession::new( + crate::models::Provider::ClaudeCode, + "test-provider".to_string(), + "test-hash".to_string(), + chrono::Utc::now(), + ); + session_repo.create(&session).await.unwrap(); + + // Create test messages + let message_repo = db.message_repo(); + let msg1 = Message::new( + session.id, + MessageRole::User, + "Create a todo list application".to_string(), + chrono::Utc::now(), + 1, + ); + let msg2 = Message::new( + session.id, + MessageRole::Assistant, + "I'll create a todo list app for you.".to_string(), + chrono::Utc::now(), + 2, + ); + message_repo.create(&msg1).await.unwrap(); + message_repo.create(&msg2).await.unwrap(); + + // Create service + let api_key = std::env::var("GOOGLE_AI_API_KEY").unwrap(); + let config = GoogleAiConfig::new(api_key); + let client = GoogleAiClient::new(config).unwrap(); + let service = FlowchartService::new(Arc::new(db.manager), client); + + // Generate flowchart + let flowchart = service + .generate_flowchart(&session.id.to_string()) + .await + .unwrap(); + + assert!(!flowchart.nodes.is_empty()); + assert!(flowchart.is_valid_dag()); + } +} diff --git a/src/services/mod.rs b/src/services/mod.rs index ba75018..824c891 100644 --- a/src/services/mod.rs +++ b/src/services/mod.rs @@ -1,6 +1,7 @@ pub mod analytics; pub mod analytics_service; pub mod auto_detect; +pub mod flowchart_service; pub mod google_ai; pub mod import_service; pub mod parser_service; @@ -19,6 +20,7 @@ pub use analytics::{ }; pub use analytics_service::AnalyticsService; pub use auto_detect::{AutoDetectService, DetectedProvider}; +pub use flowchart_service::FlowchartService; pub use google_ai::{ GenerateContentRequest, GenerateContentResponse, GoogleAiClient, GoogleAiConfig, GoogleAiError, }; diff --git a/src/tui/app.rs b/src/tui/app.rs index 5d12a8f..f67df05 100644 --- a/src/tui/app.rs +++ b/src/tui/app.rs @@ -16,7 +16,7 @@ use tokio::time::timeout; use crate::database::DatabaseManager; use crate::env::apis as env_vars; use crate::services::google_ai::{GoogleAiClient, GoogleAiConfig}; -use crate::services::{AnalyticsService, QueryService, RetrospectionService}; +use crate::services::{AnalyticsService, FlowchartService, QueryService, RetrospectionService}; use super::{ analytics::AnalyticsWidget, @@ -146,6 +146,7 @@ pub struct App { pub query_service: QueryService, pub analytics_service: AnalyticsService, pub retrospection_service: Option>, + pub flowchart_service: Option>, pub event_handler: EventHandler, } @@ -154,19 +155,23 @@ impl App { let query_service = QueryService::with_database(db_manager.clone()); let analytics_service = AnalyticsService::new((*db_manager).clone()); - // Try to create retrospection service if Google AI API key is available - let retrospection_service = if std::env::var(env_vars::GOOGLE_AI_API_KEY).is_ok() { - let config = GoogleAiConfig::default(); - match GoogleAiClient::new(config) { - Ok(client) => Some(Arc::new(RetrospectionService::new( - db_manager.clone(), - client, - ))), - Err(_) => None, - } - } else { - None - }; + // Try to create retrospection and flowchart services if Google AI API key is available + let (retrospection_service, flowchart_service) = + if std::env::var(env_vars::GOOGLE_AI_API_KEY).is_ok() { + let config = GoogleAiConfig::default(); + match GoogleAiClient::new(config) { + Ok(client) => ( + Some(Arc::new(RetrospectionService::new( + db_manager.clone(), + client.clone(), + ))), + Some(Arc::new(FlowchartService::new(db_manager.clone(), client))), + ), + Err(_) => (None, None), + } + } else { + (None, None) + }; Ok(Self { state: AppState::new(), @@ -176,6 +181,7 @@ impl App { query_service, analytics_service, retrospection_service, + flowchart_service, event_handler: EventHandler::new(), }) } @@ -379,6 +385,14 @@ impl App { SessionDetailToggleRetrospection => { self.session_detail.state.toggle_retrospection(); } + SessionDetailToggleFlowchart => { + self.session_detail.state.toggle_flowchart(); + if self.session_detail.state.show_flowchart { + if let Some(session_id) = self.session_detail.state.session_id.clone() { + self.session_detail.load_flowchart_cached(&session_id).await; + } + } + } // Analytics actions AnalyticsNavigate(_direction) => { @@ -398,9 +412,9 @@ impl App { } async fn handle_start_analysis(&mut self, session_id: String) -> Result<()> { - if let Some(ref service) = self.retrospection_service { - // Start actual analysis - match service + if let Some(ref retro_service) = self.retrospection_service { + // Start retrospection analysis + match retro_service .create_analysis_request(session_id.clone(), None, None) .await { @@ -412,20 +426,34 @@ impl App { tracing::error!(error = %e, "Failed to refresh session list after analysis start"); } - // Execute the analysis in background task - let service_clone = service.clone(); + // Execute retrospection analysis in background task + let service_clone = retro_service.clone(); let request_id = request.id.clone(); task::spawn(async move { if let Err(e) = service_clone.execute_analysis(request_id).await { - tracing::error!(error = %e, "Background analysis failed"); + tracing::error!(error = %e, "Background retrospection analysis failed"); } }); } Err(e) => { self.state - .show_error(format!("Failed to start analysis: {e}")); + .show_error(format!("Failed to start retrospection: {e}")); } } + + // Also start flowchart generation if service is available + if let Some(ref flowchart_service) = self.flowchart_service { + let flowchart_service_clone = flowchart_service.clone(); + let session_id_clone = session_id.clone(); + task::spawn(async move { + if let Err(e) = flowchart_service_clone + .generate_flowchart(&session_id_clone) + .await + { + tracing::error!(error = %e, "Background flowchart generation failed"); + } + }); + } } else { // Show message that Google AI API key is required self.state.show_error(format!( diff --git a/src/tui/events/event.rs b/src/tui/events/event.rs index 0581916..0462634 100644 --- a/src/tui/events/event.rs +++ b/src/tui/events/event.rs @@ -43,6 +43,7 @@ pub enum UserAction { SessionDetailEnd, SessionDetailToggleWrap, SessionDetailToggleRetrospection, + SessionDetailToggleFlowchart, // Analytics actions AnalyticsNavigate(NavigationDirection), diff --git a/src/tui/events/handler.rs b/src/tui/events/handler.rs index 109156f..f673baf 100644 --- a/src/tui/events/handler.rs +++ b/src/tui/events/handler.rs @@ -118,6 +118,7 @@ impl EventHandler { KeyCode::Home => vec![UserAction::SessionDetailHome], KeyCode::End => vec![UserAction::SessionDetailEnd], KeyCode::Char('w') => vec![UserAction::SessionDetailToggleWrap], + KeyCode::Char('f') => vec![UserAction::SessionDetailToggleFlowchart], KeyCode::Char('t') | KeyCode::Char('r') if key.modifiers.contains(KeyModifiers::CONTROL) => { diff --git a/src/tui/flowchart_renderer.rs b/src/tui/flowchart_renderer.rs new file mode 100644 index 0000000..651a144 --- /dev/null +++ b/src/tui/flowchart_renderer.rs @@ -0,0 +1,348 @@ +use ratatui::{ + style::{Color, Modifier, Style}, + text::{Line, Span}, +}; +use std::collections::{HashMap, HashSet}; + +use crate::models::{EdgeType, Flowchart, FlowchartNode}; + +/// Renders a flowchart as ASCII art with Unicode box characters +pub struct FlowchartRenderer { + max_width: usize, +} + +impl FlowchartRenderer { + pub fn new(max_width: usize) -> Self { + Self { max_width } + } + + /// Render flowchart to lines for TUI display + pub fn render(&self, flowchart: &Flowchart) -> Vec> { + let mut lines = Vec::new(); + + if flowchart.nodes.is_empty() { + lines.push(Line::from(vec![Span::styled( + "Empty flowchart", + Style::default().fg(Color::Gray), + )])); + return lines; + } + + // Build adjacency lists + let mut outgoing: HashMap<&str, Vec<&str>> = HashMap::new(); + let mut incoming: HashMap<&str, Vec<&str>> = HashMap::new(); + let mut edge_types: HashMap<(&str, &str), EdgeType> = HashMap::new(); + + for edge in &flowchart.edges { + outgoing + .entry(&edge.from_node) + .or_default() + .push(&edge.to_node); + incoming + .entry(&edge.to_node) + .or_default() + .push(&edge.from_node); + edge_types.insert((&edge.from_node, &edge.to_node), edge.edge_type.clone()); + } + + // Topological sort to get rendering order + let order = self.topological_sort(&flowchart.nodes, &outgoing); + + // Render each node in order + for (idx, node_id) in order.iter().enumerate() { + if let Some(node) = flowchart.get_node(node_id) { + // Render the node + self.render_node(&mut lines, node, idx + 1); + + // Render connections to next nodes based on edge types + if let Some(targets) = outgoing.get(node_id.as_str()) { + if targets.len() == 1 { + // Check if this is a loop + let target = targets[0]; + if let Some(EdgeType::Loop) = edge_types.get(&(node_id.as_str(), target)) { + self.render_loop(&mut lines); + } else { + self.render_sequential_arrow(&mut lines); + } + } else if targets.len() > 1 { + // Branch + self.render_branch(&mut lines, targets.len()); + } + } + + // Check if this is a merge point + if let Some(sources) = incoming.get(node_id.as_str()) { + if sources.len() > 1 && idx > 0 { + // This was a merge, add visual indicator before the node + // (We'll handle this by detecting it before rendering the node) + } + } + } + } + + lines + } + + fn render_node(&self, lines: &mut Vec>, node: &FlowchartNode, number: usize) { + let label = self.truncate_label(&node.label, self.max_width.saturating_sub(8)); + + // Calculate proper box width based on content + let content = format!("{number}. {label}"); + let box_width = content.chars().count() + 4; // 2 spaces on each side + + // Top border + lines.push(Line::from(vec![Span::styled( + format!("┌{}┐", "─".repeat(box_width.saturating_sub(2))), + Style::default().fg(Color::Cyan), + )])); + + // Content with proper padding + lines.push(Line::from(vec![Span::styled( + format!("│ {content} │"), + Style::default() + .fg(Color::White) + .add_modifier(Modifier::BOLD), + )])); + + // Bottom border + lines.push(Line::from(vec![Span::styled( + format!("└{}┘", "─".repeat(box_width.saturating_sub(2))), + Style::default().fg(Color::Cyan), + )])); + } + + fn render_sequential_arrow(&self, lines: &mut Vec>) { + lines.push(Line::from(vec![Span::styled( + " │", + Style::default().fg(Color::Gray), + )])); + lines.push(Line::from(vec![Span::styled( + " ▼", + Style::default().fg(Color::Yellow), + )])); + } + + fn render_branch(&self, lines: &mut Vec>, branch_count: usize) { + lines.push(Line::from(vec![Span::styled( + " │", + Style::default().fg(Color::Gray), + )])); + + // Branch point + if branch_count == 2 { + lines.push(Line::from(vec![Span::styled( + " ┌───┴───┐", + Style::default().fg(Color::Magenta), + )])); + lines.push(Line::from(vec![Span::styled( + " │ │", + Style::default().fg(Color::Magenta), + )])); + lines.push(Line::from(vec![Span::styled( + " ▼ ▼", + Style::default().fg(Color::Yellow), + )])); + } else { + // More than 2 branches - simplified representation + let branch_line = format!(" ├{}┤ {} branches", "─".repeat(5), branch_count); + lines.push(Line::from(vec![Span::styled( + branch_line, + Style::default().fg(Color::Magenta), + )])); + } + } + + fn render_loop(&self, lines: &mut Vec>) { + lines.push(Line::from(vec![Span::styled( + " │", + Style::default().fg(Color::Gray), + )])); + lines.push(Line::from(vec![Span::styled( + " ▼", + Style::default().fg(Color::Yellow), + )])); + lines.push(Line::from(vec![Span::styled( + " ┌─────┐", + Style::default().fg(Color::Red), + )])); + lines.push(Line::from(vec![Span::styled( + " │ LOOP│", + Style::default().fg(Color::Red), + )])); + lines.push(Line::from(vec![Span::styled( + " └─────┘", + Style::default().fg(Color::Red), + )])); + } + + fn truncate_label(&self, label: &str, max_len: usize) -> String { + if label.len() <= max_len { + label.to_string() + } else { + format!("{}...", &label[..max_len.saturating_sub(3)]) + } + } + + /// Simple topological sort for rendering order + fn topological_sort( + &self, + nodes: &[FlowchartNode], + outgoing: &HashMap<&str, Vec<&str>>, + ) -> Vec { + let mut visited = HashSet::new(); + let mut result = Vec::new(); + let mut in_degree: HashMap<&str, usize> = HashMap::new(); + + // Calculate in-degrees + for node in nodes { + in_degree.entry(&node.id).or_insert(0); + } + + for targets in outgoing.values() { + for target in targets { + *in_degree.entry(target).or_insert(0) += 1; + } + } + + // Find nodes with in-degree 0 (start nodes) + let mut queue: Vec<&str> = in_degree + .iter() + .filter(|(_, °ree)| degree == 0) + .map(|(&node, _)| node) + .collect(); + + // BFS + while let Some(node_id) = queue.pop() { + if visited.contains(node_id) { + continue; + } + + visited.insert(node_id); + result.push(node_id.to_string()); + + if let Some(targets) = outgoing.get(node_id) { + for &target in targets { + if let Some(degree) = in_degree.get_mut(target) { + *degree = degree.saturating_sub(1); + if *degree == 0 { + queue.push(target); + } + } + } + } + } + + // Add any remaining nodes (disconnected or cyclic) + for node in nodes { + if !visited.contains(node.id.as_str()) { + result.push(node.id.clone()); + } + } + + result + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::models::{EdgeType, Flowchart, FlowchartEdge, FlowchartNode, NodeType}; + + #[test] + fn test_render_simple_flowchart() { + let nodes = vec![ + FlowchartNode { + id: "1".to_string(), + label: "Start".to_string(), + message_refs: vec![], + node_type: NodeType::Start, + description: None, + }, + FlowchartNode { + id: "2".to_string(), + label: "Process data".to_string(), + message_refs: vec![], + node_type: NodeType::Action, + description: None, + }, + FlowchartNode { + id: "3".to_string(), + label: "End".to_string(), + message_refs: vec![], + node_type: NodeType::End, + description: None, + }, + ]; + + let edges = vec![ + FlowchartEdge { + from_node: "1".to_string(), + to_node: "2".to_string(), + edge_type: EdgeType::Sequential, + label: None, + }, + FlowchartEdge { + from_node: "2".to_string(), + to_node: "3".to_string(), + edge_type: EdgeType::Sequential, + label: None, + }, + ]; + + let flowchart = Flowchart::new("test-session".to_string(), nodes, edges); + let renderer = FlowchartRenderer::new(40); + let lines = renderer.render(&flowchart); + + assert!(!lines.is_empty()); + // Should have boxes for 3 nodes + arrows + assert!(lines.len() > 9); // At least 3 nodes * 3 lines each + } + + #[test] + fn test_render_branch() { + let nodes = vec![ + FlowchartNode { + id: "1".to_string(), + label: "Start".to_string(), + message_refs: vec![], + node_type: NodeType::Start, + description: None, + }, + FlowchartNode { + id: "2".to_string(), + label: "Branch A".to_string(), + message_refs: vec![], + node_type: NodeType::Action, + description: None, + }, + FlowchartNode { + id: "3".to_string(), + label: "Branch B".to_string(), + message_refs: vec![], + node_type: NodeType::Action, + description: None, + }, + ]; + + let edges = vec![ + FlowchartEdge { + from_node: "1".to_string(), + to_node: "2".to_string(), + edge_type: EdgeType::Branch, + label: None, + }, + FlowchartEdge { + from_node: "1".to_string(), + to_node: "3".to_string(), + edge_type: EdgeType::Branch, + label: None, + }, + ]; + + let flowchart = Flowchart::new("test-session".to_string(), nodes, edges); + let renderer = FlowchartRenderer::new(40); + let lines = renderer.render(&flowchart); + + assert!(!lines.is_empty()); + } +} diff --git a/src/tui/mod.rs b/src/tui/mod.rs index 5a21402..cd16838 100644 --- a/src/tui/mod.rs +++ b/src/tui/mod.rs @@ -2,6 +2,7 @@ pub mod analytics; pub mod app; pub mod components; pub mod events; +pub mod flowchart_renderer; pub mod retrospection; pub mod session_detail; pub mod session_list; diff --git a/src/tui/session_detail.rs b/src/tui/session_detail.rs index 7066e2c..cb07a8d 100644 --- a/src/tui/session_detail.rs +++ b/src/tui/session_detail.rs @@ -9,10 +9,14 @@ use ratatui::{ }; use std::sync::Arc; -use crate::database::{DatabaseManager, RetrospectionRepository}; +use crate::database::{DatabaseManager, FlowchartRepository, RetrospectionRepository}; use crate::models::{Message, MessageRole}; -use crate::services::{MessageGroup, QueryService, SessionDetailRequest}; +use crate::services::{ + FlowchartService, GoogleAiClient, GoogleAiConfig, MessageGroup, QueryService, + SessionDetailRequest, +}; +use super::flowchart_renderer::FlowchartRenderer; use super::state::SessionDetailState; use super::tool_display::{ToolDisplayConfig, ToolDisplayFormatter}; use super::utils::text::wrap_text; @@ -21,15 +25,29 @@ pub struct SessionDetailWidget { pub state: SessionDetailState, query_service: QueryService, retrospection_repo: RetrospectionRepository, + flowchart_repo: FlowchartRepository, + flowchart_service: Option, tool_formatter: ToolDisplayFormatter, } impl SessionDetailWidget { pub fn new(db_manager: Arc) -> Self { + // Try to initialize flowchart service with Google AI + let flowchart_service = if let Ok(api_key) = std::env::var("GOOGLE_AI_API_KEY") { + let config = GoogleAiConfig::new(api_key); + GoogleAiClient::new(config) + .ok() + .map(|client| FlowchartService::new(db_manager.clone(), client)) + } else { + None + }; + Self { state: SessionDetailState::new(), query_service: QueryService::with_database(db_manager.clone()), - retrospection_repo: RetrospectionRepository::new(db_manager), + retrospection_repo: RetrospectionRepository::new(db_manager.clone()), + flowchart_repo: FlowchartRepository::new(db_manager.clone()), + flowchart_service, tool_formatter: ToolDisplayFormatter::new(), } } @@ -61,6 +79,9 @@ impl SessionDetailWidget { // Load retrospection results for this session self.load_retrospections().await; + + // Load flowchart if exists (cached only, no generation) + self.load_flowchart_cached(session_id).await; } Err(e) => { tracing::error!(error = %e, "Failed to load session details"); @@ -117,6 +138,17 @@ impl SessionDetailWidget { // T: Toggle retrospection view self.state.toggle_retrospection(); } + KeyCode::Char('f') => { + // F: Toggle flowchart view + self.state.toggle_flowchart(); + + // Load flowchart data if showing flowchart and we have a session + if self.state.show_flowchart { + if let Some(session_id) = self.state.session_id.clone() { + self.load_flowchart_cached(&session_id).await; + } + } + } KeyCode::Char('d') => { // D: Toggle tool details (expand/collapse) self.state.toggle_tool_details(); @@ -145,18 +177,28 @@ impl SessionDetailWidget { self.render_session_header(f, chunks[0]); // Render main content area - if self.state.show_retrospection && !self.state.retrospections.is_empty() { - // Split horizontally for messages and retrospection + let show_right_panel = (self.state.show_retrospection + && !self.state.retrospections.is_empty()) + || self.state.show_flowchart; + + if show_right_panel { + // Split horizontally for messages and right panel let main_chunks = Layout::default() .direction(Direction::Horizontal) .constraints([ Constraint::Percentage(60), // Messages - Constraint::Percentage(40), // Retrospection + Constraint::Percentage(40), // Right panel ]) .split(chunks[1]); self.render_messages(f, main_chunks[0]); - self.render_retrospections(f, main_chunks[1]); + + // Render appropriate right panel + if self.state.show_flowchart { + self.render_flowchart(f, main_chunks[1]); + } else if self.state.show_retrospection { + self.render_retrospections(f, main_chunks[1]); + } } else { // Full width for messages self.render_messages(f, chunks[1]); @@ -182,7 +224,7 @@ impl SessionDetailWidget { }; format!( - "Provider: {} | Project: {} | Messages: {} | Tokens: {} | Started: {} | Status: {} | {} | Keys: 'w'=wrap, 't'=retro, 'd'=tool-details", + "Provider: {} | Project: {} | Messages: {} | Tokens: {} | Started: {} | Status: {} | {} | Keys: 'w'=wrap, 't'=retro, 'f'=flowchart, 'd'=tool-details", session.provider, project_str, session.message_count, @@ -493,6 +535,20 @@ impl SessionDetailWidget { } } + pub async fn load_flowchart_cached(&mut self, session_id: &str) { + // Only load cached flowchart from DB, don't generate + match self.flowchart_repo.get_by_session_id(session_id).await { + Ok(flowcharts) => { + if let Some(flowchart) = flowcharts.first() { + self.state.update_flowchart(Some(flowchart.clone())); + } + } + Err(e) => { + tracing::error!(error = %e, "Failed to load flowchart from DB"); + } + } + } + fn render_retrospections(&mut self, f: &mut Frame, area: Rect) { if self.state.retrospections.is_empty() { let empty_msg = Paragraph::new("No retrospection analysis available\n\nUse 'retrochat retrospect execute' to analyze this session") @@ -603,4 +659,61 @@ impl SessionDetailWidget { } } } + + fn render_flowchart(&mut self, f: &mut Frame, area: Rect) { + if self.state.flowchart_loading { + let loading_msg = Paragraph::new("Generating flowchart...\n\nThis may take a moment.") + .block( + Block::default() + .borders(Borders::ALL) + .title("Context Flowchart"), + ) + .style(Style::default().fg(Color::Yellow)) + .wrap(Wrap { trim: true }); + + f.render_widget(loading_msg, area); + return; + } + + if let Some(flowchart) = &self.state.flowchart { + let renderer = FlowchartRenderer::new(area.width.saturating_sub(4) as usize); + let flowchart_lines = renderer.render(flowchart); + + // Handle scrolling for flowchart panel + let available_height = area.height.saturating_sub(2) as usize; + let visible_lines: Vec = flowchart_lines + .into_iter() + .skip(self.state.flowchart_scroll) + .take(available_height) + .collect(); + + let flowchart_block = Paragraph::new(visible_lines) + .block( + Block::default() + .borders(Borders::ALL) + .title("Context Flowchart"), + ) + .wrap(Wrap { trim: true }) + .scroll((0, 0)); + + f.render_widget(flowchart_block, area); + } else { + let empty_msg = if self.flowchart_service.is_some() { + "No flowchart generated yet\n\nGenerate with: retrochat flowchart generate \nPress 'f' to toggle this panel" + } else { + "Flowchart unavailable\n\nSet GOOGLE_AI_API_KEY environment variable to enable flowchart generation" + }; + + let paragraph = Paragraph::new(empty_msg) + .block( + Block::default() + .borders(Borders::ALL) + .title("Context Flowchart"), + ) + .style(Style::default().fg(Color::Gray)) + .wrap(Wrap { trim: true }); + + f.render_widget(paragraph, area); + } + } } diff --git a/src/tui/state/session_detail_state.rs b/src/tui/state/session_detail_state.rs index 7a0a2ec..2d90c90 100644 --- a/src/tui/state/session_detail_state.rs +++ b/src/tui/state/session_detail_state.rs @@ -1,6 +1,6 @@ use ratatui::widgets::ScrollbarState; -use crate::models::{ChatSession, Message, Retrospection}; +use crate::models::{ChatSession, Flowchart, Message, Retrospection}; /// State for the session detail view #[derive(Debug)] @@ -11,6 +11,8 @@ pub struct SessionDetailState { pub messages: Vec, /// Retrospection analyses for this session pub retrospections: Vec, + /// Flowchart for this session + pub flowchart: Option, /// Currently selected session ID pub session_id: Option, /// Scrollbar state for messages @@ -19,12 +21,18 @@ pub struct SessionDetailState { pub current_scroll: usize, /// Scroll position for retrospection panel pub retrospection_scroll: usize, + /// Scroll position for flowchart panel + pub flowchart_scroll: usize, /// Loading indicator pub loading: bool, + /// Whether flowchart is currently being generated + pub flowchart_loading: bool, /// Whether to wrap message text pub message_wrap: bool, /// Whether to show the retrospection panel pub show_retrospection: bool, + /// Whether to show the flowchart panel + pub show_flowchart: bool, /// Whether to show detailed tool output (expanded view) pub show_tool_details: bool, } @@ -36,13 +44,17 @@ impl SessionDetailState { session: None, messages: Vec::new(), retrospections: Vec::new(), + flowchart: None, session_id: None, scroll_state: ScrollbarState::default(), current_scroll: 0, retrospection_scroll: 0, + flowchart_scroll: 0, loading: false, + flowchart_loading: false, message_wrap: true, show_retrospection: false, + show_flowchart: false, show_tool_details: false, } } @@ -55,8 +67,10 @@ impl SessionDetailState { self.session = None; self.messages.clear(); self.retrospections.clear(); + self.flowchart = None; self.current_scroll = 0; self.retrospection_scroll = 0; + self.flowchart_scroll = 0; } } @@ -132,6 +146,22 @@ impl SessionDetailState { self.show_tool_details = !self.show_tool_details; } + /// Toggle flowchart panel visibility + pub fn toggle_flowchart(&mut self) { + self.show_flowchart = !self.show_flowchart; + } + + /// Update flowchart data + pub fn update_flowchart(&mut self, flowchart: Option) { + self.flowchart = flowchart; + self.flowchart_loading = false; + } + + /// Mark flowchart as loading + pub fn set_flowchart_loading(&mut self, loading: bool) { + self.flowchart_loading = loading; + } + /// Update the scrollbar state pub fn update_scroll_state(&mut self, total_lines: usize) { self.scroll_state = self.scroll_state.content_length(total_lines);