Skip to main content

omni_dev/cli/
ai.rs

1//! AI commands.
2
3use std::io::{self, Write};
4
5use anyhow::{Context, Result};
6use clap::{Parser, Subcommand};
7use crossterm::{
8    event::{self, Event, KeyCode, KeyModifiers},
9    terminal::{disable_raw_mode, enable_raw_mode},
10};
11
12/// AI operations.
13#[derive(Parser)]
14pub struct AiCommand {
15    /// The AI subcommand to execute.
16    #[command(subcommand)]
17    pub command: AiSubcommand,
18}
19
20/// AI subcommands.
21#[derive(Subcommand)]
22pub enum AiSubcommand {
23    /// Interactive AI chat session.
24    Chat(ChatCommand),
25}
26
27impl AiCommand {
28    /// Executes the AI command.
29    pub fn execute(self) -> Result<()> {
30        match self.command {
31            AiSubcommand::Chat(cmd) => cmd.execute(),
32        }
33    }
34}
35
36/// Interactive AI chat session.
37#[derive(Parser)]
38pub struct ChatCommand {
39    /// AI model to use (overrides environment configuration).
40    #[arg(long)]
41    pub model: Option<String>,
42}
43
44impl ChatCommand {
45    /// Executes the chat command.
46    pub fn execute(self) -> Result<()> {
47        let ai_info = crate::utils::preflight::check_ai_credentials(self.model.as_deref())?;
48        eprintln!(
49            "Connected to {} (model: {})",
50            ai_info.provider, ai_info.model
51        );
52        eprintln!("Enter to send, Shift+Enter for newline, Ctrl+D to exit.\n");
53
54        let client = crate::claude::create_default_claude_client(self.model.clone(), None)?;
55
56        let rt = tokio::runtime::Runtime::new().context("Failed to create tokio runtime")?;
57        rt.block_on(chat_loop(&client))
58    }
59}
60
61async fn chat_loop(client: &crate::claude::client::ClaudeClient) -> Result<()> {
62    let system_prompt = "You are a helpful assistant.";
63
64    loop {
65        let input = match read_user_input() {
66            Ok(Some(text)) => text,
67            Ok(None) => {
68                eprintln!("\nGoodbye!");
69                break;
70            }
71            Err(e) => {
72                eprintln!("\nInput error: {e}");
73                break;
74            }
75        };
76
77        let trimmed = input.trim();
78        if trimmed.is_empty() {
79            continue;
80        }
81
82        let response = client.send_message(system_prompt, trimmed).await?;
83        println!("{response}\n");
84    }
85
86    Ok(())
87}
88
89/// Guard that disables raw mode on drop.
90struct RawModeGuard;
91
92impl Drop for RawModeGuard {
93    fn drop(&mut self) {
94        let _ = disable_raw_mode();
95    }
96}
97
98/// Reads multiline user input with "> " prompt.
99///
100/// Returns `Ok(Some(text))` on Enter, `Ok(None)` on Ctrl+D/Ctrl+C.
101fn read_user_input() -> Result<Option<String>> {
102    eprint!("> ");
103    io::stderr().flush()?;
104
105    enable_raw_mode()?;
106    let _guard = RawModeGuard;
107
108    let mut buffer = String::new();
109
110    loop {
111        if let Event::Key(key_event) = event::read()? {
112            match key_event.code {
113                KeyCode::Enter => {
114                    if key_event.modifiers.contains(KeyModifiers::SHIFT) {
115                        buffer.push('\n');
116                        eprint!("\r\n... ");
117                        io::stderr().flush()?;
118                    } else {
119                        eprint!("\r\n");
120                        io::stderr().flush()?;
121                        return Ok(Some(buffer));
122                    }
123                }
124                KeyCode::Char('d') if key_event.modifiers.contains(KeyModifiers::CONTROL) => {
125                    if buffer.is_empty() {
126                        return Ok(None);
127                    }
128                    eprint!("\r\n");
129                    io::stderr().flush()?;
130                    return Ok(Some(buffer));
131                }
132                KeyCode::Char('c') if key_event.modifiers.contains(KeyModifiers::CONTROL) => {
133                    return Ok(None);
134                }
135                KeyCode::Char(c) => {
136                    buffer.push(c);
137                    eprint!("{c}");
138                    io::stderr().flush()?;
139                }
140                KeyCode::Backspace => {
141                    if buffer.pop().is_some() {
142                        eprint!("\x08 \x08");
143                        io::stderr().flush()?;
144                    }
145                }
146                _ => {}
147            }
148        }
149    }
150}