Merge pull request #8 from Ammar-Alnagar/stream

added streaming

Author: Ammar-Alnagar

rust_tui_coder/Cargo.lock

@@ -9,7 +9,8 @@ edition = "2021"
 ratatui = { version = "0.26.0", features = ["all-widgets"] }
 crossterm = "0.27.0"
 tokio = { version = "1.35.1", features = ["full"] }
-reqwest = { version = "0.11.23", features = ["json"] }
+reqwest = { version = "0.11.23", features = ["json", "stream"] }
 serde = { version = "1.0.195", features = ["derive"] }
 serde_json = "1.0.111"
 toml = "0.8.8"
+futures-util = "0.3.30"

rust_tui_coder/Cargo.toml

@@ -0,0 +1,12 @@
+# Task Plan: Create test directory with 10 files of various extensions
+
+## Checklist:
+
+- [ ] Step 1: Create the 'test' directory
+- [ ] Step 2: Create 10 test files with different extensions (standard and dummy)
+- [ ] Step 3: Verify directory and file creation
+
+## Progress:
+- Total Steps: 0
+- Completed: 0
+- Remaining: 0

rust_tui_coder/plan.md

@@ -3,7 +3,10 @@ use std::fs;
 use std::io;
 use std::process::Command;
 use std::path::Path;
-use crate::llm::{self, Message};
+use std::sync::Arc;
+use futures_util::StreamExt;
+use tokio::sync::Mutex;
+use crate::llm::Message;
 use crate::config::LlmConfig;
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -683,6 +686,7 @@ edition = "2021"
     }
 }
 
+#[derive(Clone)]
 pub struct Agent {
     messages: Vec<Message>,
 }
@@ -1066,7 +1070,11 @@ TOOL: {"name": "RUN_COMMAND", "parameters": {"command": "cargo build --release"}
         None
     }
 
-    pub async fn run(&mut self, config: &LlmConfig, user_prompt: String, app: &mut crate::app::App) -> Result<(String, Vec<String>), Box<dyn std::error::Error>> {
+    pub async fn run(&mut self, config: &LlmConfig, user_prompt: String, app: Arc<Mutex<crate::app::App>>) -> Result<(String, Vec<String>), Box<dyn std::error::Error + Send + Sync>> {
+        self.run_with_streaming(config, user_prompt, app).await
+    }
+
+    pub async fn run_with_streaming(&mut self, config: &LlmConfig, user_prompt: String, app: Arc<Mutex<crate::app::App>>) -> Result<(String, Vec<String>), Box<dyn std::error::Error + Send + Sync>> {
         // Add system message if this is the first interaction
         if self.messages.is_empty() {
             self.messages.push(Message {
@@ -1110,20 +1118,62 @@ TOOL: {"name": "RUN_COMMAND", "parameters": {"command": "cargo build --release"}
 
         let mut all_tool_logs = Vec::new();
         let mut attempts = 0;
-        const MAX_ATTEMPTS: usize = 8; // Increased for better error recovery
+        const MAX_ATTEMPTS: usize = 8;
 
         loop {
             attempts += 1;
-            
-            // Get response from LLM
-            let (response, tokens_used) = llm::ask_llm_with_messages(config, &self.messages).await?;
-            app.increment_tokens(tokens_used);
-            app.increment_requests();
+
+            // Start streaming for this response
+            {
+                let mut app_guard = app.lock().await;
+                app_guard.start_streaming();
+            }
+
+            // Create a string to collect the full response
+            let mut full_response = String::new();
+
+            // Get streaming response from LLM
+            let mut stream = match crate::llm::stream_llm_response(config, &self.messages).await {
+                Ok(stream) => stream,
+                Err(e) => {
+                    let mut app_guard = app.lock().await;
+                    app_guard.finish_streaming("Error: Failed to start streaming response".to_string());
+                    return Err(Box::new(e));
+                }
+            };
+
+            // Collect tokens from the stream
+            let mut token_stream = String::new();
+            while let Some(chunk_result) = stream.next().await {
+                match chunk_result {
+                    Ok(chunk) => {
+                        if !chunk.is_empty() {
+                            token_stream.push_str(&chunk);
+                            full_response.push_str(&chunk);
+                            // Update streaming message with brief lock
+                            {
+                                let mut app_guard = app.lock().await;
+                                app_guard.update_streaming_message(&chunk);
+                            }
+                        }
+                    }
+                    Err(e) => {
+                        let mut app_guard = app.lock().await;
+                        app_guard.finish_streaming(format!("Error in streaming: {}", e));
+                        return Err(Box::new(e));
+                    }
+                }
+            }
+
+            {
+                let mut app_guard = app.lock().await;
+                app_guard.increment_requests();
+            }
 
             // Check if response contains a tool call
-            if let Some(tool) = self.parse_tool_call(&response) {
+            if let Some(tool) = self.parse_tool_call(&full_response) {
                 let mut tool_logs = Vec::new();
-                
+
                 // Log the tool execution
                 let tool_name = match &tool {
                     Tool::ReadFile { path } => format!("READ_FILE {}", path),
@@ -1150,40 +1200,45 @@ TOOL: {"name": "RUN_COMMAND", "parameters": {"command": "cargo build --release"}
                     Tool::ClearPlan => "CLEAR_PLAN".to_string(),
                 };
                 tool_logs.push(format!("🔧 Attempt {}: Executing {}", attempts, tool_name));
-                
+
                 // Execute the tool
                 let tool_result = match tool.execute() {
                     Ok(result) => {
-                        app.increment_tools_executed();
+                        {
+                            let mut app_guard = app.lock().await;
+                            app_guard.increment_tools_executed();
+                        }
                         tool_logs.push(format!("✅ Success: {}", result));
                         result
                     }
                     Err(e) => {
                         let error_msg = format!("❌ Error: {}", e);
                         tool_logs.push(error_msg.clone());
-                        // Add error context to help the model understand what went wrong
                         format!("Tool failed: {}. Please try a different approach or check if the path/command is correct.", e)
                     }
                 };
                 all_tool_logs.extend(tool_logs);
-                
+
                 // Add assistant message and tool result to conversation
                 self.messages.push(Message {
                     role: "assistant".to_string(),
-                    content: response,
+                    content: full_response.clone(),
                 });
-                
+
                 self.messages.push(Message {
                     role: "user".to_string(),
                     content: format!("Tool result: {}", tool_result),
                 });
 
                 // Check if we should continue or if the task is complete
                 if attempts >= MAX_ATTEMPTS {
-                    // Get final response after max attempts
-                    let (final_response, final_tokens) = llm::ask_llm_with_messages(config, &self.messages).await?;
-                    app.increment_tokens(final_tokens);
-                    app.increment_requests();
+                    // Get final response after max attempts (non-streaming for final response)
+                    let (final_response, final_tokens) = crate::llm::ask_llm_with_messages(config, &self.messages).await?;
+                    {
+                        let mut app_guard = app.lock().await;
+                        app_guard.increment_tokens(final_tokens);
+                        app_guard.finish_streaming(final_response.clone());
+                    }
 
                     self.messages.push(Message {
                         role: "assistant".to_string(),
@@ -1200,14 +1255,18 @@ TOOL: {"name": "RUN_COMMAND", "parameters": {"command": "cargo build --release"}
                 // No tool call, task appears to be complete
                 self.messages.push(Message {
                     role: "assistant".to_string(),
-                    content: response.clone(),
+                    content: full_response.clone(),
                 });
-                
+
                 if attempts > 1 {
                     all_tool_logs.push(format!("✅ Task completed after {} attempts", attempts));
                 }
-                
-                return Ok((response, all_tool_logs));
+
+                {
+                    let mut app_guard = app.lock().await;
+                    app_guard.finish_streaming(full_response.clone());
+                }
+                return Ok((full_response, all_tool_logs));
             }
         }
     }

rust_tui_coder/src/agent.rs

@@ -10,14 +10,19 @@ pub struct App {
     pub tokens_used: u64,
     pub total_requests: u64,
     pub total_tools_executed: u64,
+    // Scrolling and streaming state
+    pub conversation_scroll_position: usize,
+    pub tool_logs_scroll_position: usize,
+    pub is_streaming: bool,
+    pub current_streaming_message: String,
 }
 
 impl App {
     pub fn new() -> Self {
         Self {
             user_input: String::new(),
             conversation: Vec::new(),
-            status_message: "Commands: /quit (exit), /stats (usage stats)".to_string(),
+            status_message: "Commands: /quit (exit), /stats (usage stats), ↑↓ (scroll conversation)".to_string(),
             tool_logs: Vec::new(),
             is_executing_tool: false,
             current_tool: String::new(),
@@ -26,6 +31,11 @@ impl App {
             tokens_used: 0,
             total_requests: 0,
             total_tools_executed: 0,
+            // Initialize scrolling and streaming state
+            conversation_scroll_position: 0,
+            tool_logs_scroll_position: 0,
+            is_streaming: false,
+            current_streaming_message: String::new(),
         }
     }
 
@@ -76,4 +86,49 @@ impl App {
             }
         )
     }
+
+    // Scroll management methods
+    pub fn scroll_conversation_up(&mut self) {
+        if self.conversation_scroll_position > 0 {
+            self.conversation_scroll_position -= 1;
+        }
+    }
+
+    pub fn scroll_conversation_down(&mut self) {
+        // Don't increment if we're already at or near max (prevent overflow)
+        if self.conversation_scroll_position < usize::MAX - 1 {
+            self.conversation_scroll_position += 1;
+        }
+    }
+
+    pub fn scroll_conversation_to_top(&mut self) {
+        self.conversation_scroll_position = 0;
+    }
+
+    pub fn scroll_conversation_to_bottom(&mut self) {
+        self.conversation_scroll_position = usize::MAX; // Set to max, will be clamped in UI
+    }
+
+
+    // Streaming state management
+    pub fn start_streaming(&mut self) {
+        self.is_streaming = true;
+        self.current_streaming_message = String::new();
+        self.status_message = "🤔 Thinking... (streaming response)".to_string();
+    }
+
+    pub fn update_streaming_message(&mut self, new_content: &str) {
+        self.current_streaming_message.push_str(new_content);
+    }
+
+    pub fn finish_streaming(&mut self, final_message: String) {
+        if !self.current_streaming_message.is_empty() {
+            self.conversation.push(format!("Agent: {}", final_message));
+        }
+        self.is_streaming = false;
+        self.current_streaming_message = String::new();
+        self.status_message = "✅ Done.".to_string();
+        self.scroll_conversation_to_bottom();
+    }
+
 }