rustGPT/src/cli.rs

768 lines
30 KiB
Rust

use anyhow::Result;
use crate::config::Config;
use crate::core::{
create_client, get_provider_for_model,
provider::{get_display_name_for_model, get_model_id_from_display_name, get_model_info_list},
ChatClient, Session,
};
use crate::utils::{Display, InputHandler, SessionAction};
use std::future::Future;
use std::pin::Pin;
pub struct ChatCLI {
session: Session,
client: Option<ChatClient>,
current_model: Option<String>,
display: Display,
input: InputHandler,
config: Config,
}
impl ChatCLI {
pub fn new(session: Session, config: Config) -> Result<Self> {
Ok(Self {
session,
client: None,
current_model: None,
display: Display::new(),
input: InputHandler::new()?,
config,
})
}
fn get_client(&mut self) -> Result<&ChatClient> {
if self.client.is_none() || self.current_model.as_ref() != Some(&self.session.model) {
let client = create_client(&self.session.model, &self.config)?;
self.current_model = Some(self.session.model.clone());
self.client = Some(client);
}
Ok(self.client.as_ref().unwrap())
}
pub async fn run(&mut self) -> Result<()> {
self.display.print_header();
self.display
.print_info("Type your message and press Enter. Commands start with '/'.");
self.display.print_info("Type /help for help.");
let provider = get_provider_for_model(&self.session.model);
let display_name = get_display_name_for_model(&self.session.model);
self.display
.print_model_info(&display_name, provider.as_str());
self.display.print_session_info(&self.session.name);
println!();
loop {
match self.input.read_line("👤> ")? {
Some(line) => {
let line = line.trim();
if line.is_empty() {
continue;
}
if line.starts_with('/') {
if !self.handle_command(line).await? {
break;
}
} else {
if let Err(e) = self.handle_user_message(line).await {
self.display.print_error(&format!("Error: {}", e));
}
println!(); // Add padding before next prompt
}
}
None => {
self.display.print_info("Goodbye!");
break;
}
}
}
self.save_and_cleanup()
}
pub fn save_and_cleanup(&mut self) -> Result<()> {
self.session.save()?;
self.input.cleanup()?; // Use cleanup instead of just save_history
Ok(())
}
async fn handle_user_message(&mut self, message: &str) -> Result<()> {
self.session.add_user_message(message.to_string());
self.session.save()?;
// Clone data needed for the API call before getting mutable client reference
let model = self.session.model.clone();
let messages = self.session.messages.clone();
let enable_web_search = self.session.enable_web_search;
let enable_reasoning_summary = self.session.enable_reasoning_summary;
let reasoning_effort = self.session.reasoning_effort.clone();
let enable_extended_thinking = self.session.enable_extended_thinking;
let thinking_budget_tokens = self.session.thinking_budget_tokens;
// Check if we should use streaming before getting client
let should_use_streaming = {
let client = self.get_client()?;
client.supports_streaming()
};
if should_use_streaming {
println!(); // Add padding before AI response
print!("{}> ", console::style("🤖").magenta());
use std::io::{self, Write};
io::stdout().flush().ok();
let stream_callback = {
use crate::core::StreamCallback;
Box::new(move |chunk: &str| {
print!("{}", chunk);
use std::io::{self, Write};
io::stdout().flush().ok();
Box::pin(async move {}) as Pin<Box<dyn Future<Output = ()> + Send>>
}) as StreamCallback
};
let client = self.get_client()?;
match client
.chat_completion_stream(
&model,
&messages,
enable_web_search,
enable_reasoning_summary,
&reasoning_effort,
enable_extended_thinking,
thinking_budget_tokens,
stream_callback,
)
.await
{
Ok(response) => {
println!(); // Add newline after streaming
self.session.add_assistant_message(response);
self.session.save()?;
}
Err(e) => {
println!(); // Add newline after failed streaming
self.display
.print_error(&format!("Streaming failed: {}", e));
return Err(e);
}
}
} else {
// Fallback to non-streaming
let spinner = self.display.show_spinner("Thinking");
let client = self.get_client()?;
match client
.chat_completion(
&model,
&messages,
enable_web_search,
enable_reasoning_summary,
&reasoning_effort,
enable_extended_thinking,
thinking_budget_tokens,
)
.await
{
Ok(response) => {
spinner.finish("Done");
self.display.print_assistant_response(&response);
self.session.add_assistant_message(response);
self.session.save()?;
}
Err(e) => {
spinner.finish_with_error("Failed");
return Err(e);
}
}
}
Ok(())
}
async fn handle_command(&mut self, command: &str) -> Result<bool> {
let parts: Vec<&str> = command.split_whitespace().collect();
if parts.is_empty() {
return Ok(true);
}
match parts[0].to_lowercase().as_str() {
"/help" => {
self.display.print_help();
}
"/exit" => {
self.session.save()?;
self.input.cleanup()?; // Clean up terminal state
self.display.print_info("Session saved. Goodbye!");
return Ok(false);
}
"/model" => {
self.model_switcher().await?;
}
"/new" => {
self.handle_new_session(&parts)?;
}
"/switch" => {
self.session_manager().await?;
}
"/clear" => {
self.session.clear_messages();
self.session.save()?;
self.display.print_command_result("Conversation cleared");
}
"/tools" => {
self.tools_manager().await?;
}
"/history" => {
self.handle_history_command(&parts)?;
}
"/export" => {
self.handle_export_command(&parts)?;
}
"/save" => {
self.handle_save_command(&parts)?;
}
_ => {
self.display
.print_error(&format!("Unknown command: {} (see /help)", parts[0]));
}
}
Ok(true)
}
async fn model_switcher(&mut self) -> Result<()> {
let model_info_list = get_model_info_list();
let display_names: Vec<String> = model_info_list
.iter()
.map(|info| info.display_name.to_string())
.collect();
let current_display_name = get_display_name_for_model(&self.session.model);
let selection = self.input.select_from_list(
"Select a model:",
&display_names,
Some(&current_display_name),
)?;
match selection {
Some(display_name) => {
if let Some(model_id) = get_model_id_from_display_name(&display_name) {
if model_id == self.session.model {
self.display.print_info("Already using that model");
} else {
self.session.model = model_id.clone();
let provider = get_provider_for_model(&self.session.model);
self.display.print_command_result(&format!(
"Model switched to {} ({})",
display_name,
provider.as_str()
));
self.client = None; // Force client recreation
self.session.save()?; // Save the model change
}
} else {
self.display.print_error("Invalid model selection");
}
}
None => {
self.display.print_info("Model selection cancelled");
}
}
Ok(())
}
fn handle_new_session(&mut self, parts: &[&str]) -> Result<()> {
if parts.len() != 2 {
self.display.print_error("Usage: /new <session_name>");
return Ok(());
}
self.session.save()?;
let new_session = Session::new(parts[1].to_string(), self.session.model.clone());
self.session = new_session;
self.display
.print_command_result(&format!("New session '{}' started", self.session.name));
Ok(())
}
async fn session_manager(&mut self) -> Result<()> {
loop {
let sessions = Session::list_sessions()?;
let session_names: Vec<String> = sessions.into_iter().map(|(name, _)| name).collect();
if session_names.is_empty() {
self.display.print_info("No sessions available");
return Ok(());
}
let action = self.input.session_manager(
"Session Manager:",
&session_names,
Some(&self.session.name),
)?;
match action {
SessionAction::Switch(session_name) => {
if session_name == self.session.name {
self.display.print_info("Already in that session");
return Ok(());
}
self.session.save()?;
match Session::load(&session_name) {
Ok(session) => {
self.session = session;
let display_name = get_display_name_for_model(&self.session.model);
self.display.print_command_result(&format!(
"Switched to session '{}' (model={})",
self.session.name, display_name
));
self.client = None; // Force client recreation
return Ok(());
}
Err(e) => {
self.display
.print_error(&format!("Failed to load session: {}", e));
// Don't return, allow user to try again or cancel
}
}
}
SessionAction::Delete(session_name) => {
match Session::delete_session(&session_name) {
Ok(()) => {
self.display.print_command_result(&format!(
"Session '{}' deleted",
session_name
));
// If we deleted the current session, we need to handle this specially
if session_name == self.session.name {
// Try to switch to another session or create a default one
let remaining_sessions = Session::list_sessions()?;
let remaining_names: Vec<String> = remaining_sessions
.into_iter()
.map(|(name, _)| name)
.collect();
if remaining_names.is_empty() {
// No sessions left, create a default one
self.session = Session::new(
"default".to_string(),
self.session.model.clone(),
);
self.display
.print_command_result("Created new default session");
return Ok(());
} else {
// Switch to the first available session
match Session::load(&remaining_names[0]) {
Ok(session) => {
self.session = session;
let display_name =
get_display_name_for_model(&self.session.model);
self.display.print_command_result(&format!(
"Switched to session '{}' (model={})",
self.session.name, display_name
));
self.client = None;
return Ok(());
}
Err(e) => {
self.display.print_error(&format!(
"Failed to load fallback session: {}",
e
));
// Create a new default session as fallback
self.session = Session::new(
"default".to_string(),
self.session.model.clone(),
);
self.display.print_command_result(
"Created new default session",
);
return Ok(());
}
}
}
}
// Continue to show updated session list if we didn't delete current session
}
Err(e) => {
self.display
.print_error(&format!("Failed to delete session: {}", e));
// Continue to allow retry
}
}
}
SessionAction::SetAsDefault(session_name) => {
let mut config = crate::config::Config::load().unwrap_or_default();
match config.set_default_session(session_name.clone()) {
Ok(()) => {
self.display.print_command_result(&format!(
"Session '{}' is now the default session",
session_name
));
}
Err(e) => {
self.display
.print_error(&format!("Failed to set default session: {}", e));
}
}
// Continue to show session list
}
SessionAction::Cancel => {
return Ok(());
}
}
}
}
async fn tools_manager(&mut self) -> Result<()> {
loop {
// Show current tool status
self.display.print_info("Tool Management:");
let web_status = if self.session.enable_web_search {
"✓ enabled"
} else {
"✗ disabled"
};
let reasoning_status = if self.session.enable_reasoning_summary {
"✓ enabled"
} else {
"✗ disabled"
};
let extended_thinking_status = if self.session.enable_extended_thinking {
"✓ enabled"
} else {
"✗ disabled"
};
println!(" Web Search: {}", web_status);
println!(" Reasoning Summaries: {}", reasoning_status);
println!(" Reasoning Effort: {}", self.session.reasoning_effort);
println!(" Extended Thinking: {}", extended_thinking_status);
println!(
" Thinking Budget: {} tokens",
self.session.thinking_budget_tokens
);
// Check model compatibility
let model = self.session.model.clone();
let provider = get_provider_for_model(&model);
let reasoning_enabled = self.session.enable_reasoning_summary;
// Show compatibility warnings based on provider
match provider {
crate::core::provider::Provider::Anthropic => {
if reasoning_enabled {
self.display.print_warning(
"Reasoning summaries are not supported by Anthropic models",
);
}
if self.session.enable_extended_thinking {
// Extended thinking is supported by Anthropic models
}
// Web search is now supported by Anthropic models
}
crate::core::provider::Provider::OpenAI => {
if self.session.enable_extended_thinking {
self.display
.print_warning("Extended thinking is not supported by OpenAI models");
}
// OpenAI models generally support other features
}
}
// Tool management options
let options = vec![
"Toggle Web Search",
"Toggle Reasoning Summaries",
"Set Reasoning Effort",
"Toggle Extended Thinking",
"Set Thinking Budget",
"Done",
];
let selection = self
.input
.select_from_list("Select an option:", &options, None)?;
match selection.as_deref() {
Some("Toggle Web Search") => {
self.session.enable_web_search = !self.session.enable_web_search;
let state = if self.session.enable_web_search {
"enabled"
} else {
"disabled"
};
self.display
.print_command_result(&format!("Web search {}", state));
}
Some("Toggle Reasoning Summaries") => {
self.session.enable_reasoning_summary = !self.session.enable_reasoning_summary;
let state = if self.session.enable_reasoning_summary {
"enabled"
} else {
"disabled"
};
self.display
.print_command_result(&format!("Reasoning summaries {}", state));
}
Some("Set Reasoning Effort") => {
let effort_options = vec!["low", "medium", "high"];
if let Some(effort) = self.input.select_from_list(
"Select reasoning effort:",
&effort_options,
Some(&self.session.reasoning_effort),
)? {
self.session.reasoning_effort = effort.to_string();
self.display
.print_command_result(&format!("Reasoning effort set to {}", effort));
if !self.session.model.starts_with("gpt-5") {
self.display.print_warning(
"Reasoning effort is only supported by GPT-5 models",
);
}
}
}
Some("Toggle Extended Thinking") => {
self.session.enable_extended_thinking = !self.session.enable_extended_thinking;
let state = if self.session.enable_extended_thinking {
"enabled"
} else {
"disabled"
};
self.display
.print_command_result(&format!("Extended thinking {}", state));
let provider = get_provider_for_model(&self.session.model);
match provider {
crate::core::provider::Provider::OpenAI => {
self.display.print_warning(
"Extended thinking is not supported by OpenAI models",
);
}
crate::core::provider::Provider::Anthropic => {
// Supported
}
}
}
Some("Set Thinking Budget") => {
let budget_options = vec!["1024", "2500", "5000", "10000", "16000"];
let current_budget = self.session.thinking_budget_tokens.to_string();
if let Some(budget_str) = self.input.select_from_list(
"Select thinking budget (tokens):",
&budget_options,
Some(&current_budget),
)? {
if let Ok(budget) = budget_str.parse::<u32>() {
self.session.thinking_budget_tokens = budget;
self.display.print_command_result(&format!(
"Thinking budget set to {} tokens",
budget
));
let provider = get_provider_for_model(&self.session.model);
match provider {
crate::core::provider::Provider::OpenAI => {
self.display.print_warning(
"Extended thinking is not supported by OpenAI models",
);
}
crate::core::provider::Provider::Anthropic => {
if budget < 1024 {
self.display.print_warning("Minimum thinking budget is 1024 tokens for Anthropic models");
}
}
}
}
}
}
Some("Done") | None => {
break;
}
_ => {}
}
self.session.save()?; // Save changes after each modification
println!(); // Add spacing
}
Ok(())
}
fn handle_history_command(&mut self, parts: &[&str]) -> Result<()> {
let mut filter_role: Option<&str> = None;
let mut limit: Option<usize> = None;
// Parse parameters
for &part in parts.iter().skip(1) {
match part {
"user" | "assistant" => filter_role = Some(part),
_ => {
if let Ok(num) = part.parse::<usize>() {
limit = Some(num);
} else {
self.display
.print_error(&format!("Invalid parameter: {}", part));
self.display
.print_info("Usage: /history [user|assistant] [number]");
return Ok(());
}
}
}
}
// Filter messages (skip system prompt at index 0)
let mut messages: Vec<(usize, &crate::core::Message)> = self
.session
.messages
.iter()
.enumerate()
.skip(1) // Skip system prompt
.collect();
// Apply role filter
if let Some(role) = filter_role {
messages.retain(|(_, msg)| msg.role == role);
}
// Apply limit
if let Some(limit_count) = limit {
let start_index = messages.len().saturating_sub(limit_count);
messages = messages[start_index..].to_vec();
}
if messages.is_empty() {
self.display.print_info("No messages to display");
return Ok(());
}
// Format and display
self.display.print_conversation_history(&messages);
Ok(())
}
fn handle_export_command(&mut self, parts: &[&str]) -> Result<()> {
let format = if parts.len() > 1 {
parts[1].to_lowercase()
} else {
// Default to markdown
"markdown".to_string()
};
let valid_formats = ["markdown", "md", "json", "txt"];
if !valid_formats.contains(&format.as_str()) {
self.display.print_error(&format!(
"Invalid format '{}'. Supported formats: markdown, json, txt",
format
));
self.display.print_info("Usage: /export [format]");
return Ok(());
}
// Generate filename with timestamp
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let extension = match format.as_str() {
"markdown" | "md" => "md",
"json" => "json",
"txt" => "txt",
_ => "md",
};
// Create exports directory if it doesn't exist
let exports_dir = std::path::Path::new("exports");
if !exports_dir.exists() {
if let Err(e) = std::fs::create_dir(exports_dir) {
self.display
.print_error(&format!("Failed to create exports directory: {}", e));
return Ok(());
}
}
let filename = format!("{}_{}.{}", self.session.name, now, extension);
let file_path = exports_dir.join(&filename);
match self
.session
.export(&format, file_path.to_str().unwrap_or(&filename))
{
Ok(()) => {
self.display.print_command_result(&format!(
"Conversation exported to '{}'",
file_path.display()
));
}
Err(e) => {
self.display.print_error(&format!("Export failed: {}", e));
}
}
Ok(())
}
fn handle_save_command(&mut self, parts: &[&str]) -> Result<()> {
if parts.len() != 2 {
self.display.print_error("Usage: /save <new_session_name>");
self.display
.print_info("Example: /save my_important_conversation");
return Ok(());
}
let new_session_name = parts[1];
// Validate session name
if new_session_name.is_empty() {
self.display.print_error("Session name cannot be empty");
return Ok(());
}
// Check if session already exists
if let Ok(sessions) = Session::list_sessions() {
if sessions.iter().any(|(name, _)| name == new_session_name) {
self.display.print_error(&format!(
"Session '{}' already exists. Choose a different name.",
new_session_name
));
return Ok(());
}
}
// Save the current session first to ensure it's up to date
self.session.save()?;
// Copy the session with the new name
match self.session.save_as(new_session_name) {
Ok(()) => {
self.display.print_command_result(&format!(
"Current session saved as '{}' ({} messages copied)",
new_session_name,
self.session.messages.len().saturating_sub(1) // Exclude system prompt
));
self.display.print_info(&format!(
"Use '/switch' to switch to the new session, or continue with current session '{}'",
self.session.name
));
}
Err(e) => {
self.display
.print_error(&format!("Failed to save session: {}", e));
}
}
Ok(())
}
}