命令行里的智能体：Rust AI CLI 工具从架构设计到实现

原创于 2026-06-22 09:43:44 发布 · 15 阅读

0 ·

本内容遵循CC 4.0 BY-SA版权协议

GEO检测

标签

#rust #github #python

命令行里的智能体：Rust AI CLI 工具从架构设计到实现

cover

一、CLI 工具的"智能化"需求——为什么要在命令行里跑 AI

命令行工具是开发者的日常伴侣。git、cargo、npm——这些工具的核心逻辑是"接收指令、执行操作、返回结果"。但当操作涉及自然语言理解、代码生成、智能推荐时，传统的 if-else 逻辑就力不从心了。

实际场景中，这种需求很常见：一个 CLI 工具需要根据用户输入的自然语言描述搜索代码、根据错误信息推荐修复方案、根据项目结构生成配置文件。这些任务的核心是"理解意图 + 生成响应"，恰好是 AI 模型擅长的领域。

Rust 在 CLI 工具开发上有天然优势：编译为单二进制、启动快、跨平台。结合 AI 推理能力，可以构建出既高效又智能的命令行工具。但架构设计上有几个关键决策：模型放在本地还是远程？流式输出怎么做？上下文管理如何设计？

二、Rust AI CLI 工具的架构分层

2.1 整体架构

graph TB
    A[CLI 入口<br>clap 参数解析] --> B[命令路由<br>match 子命令]
    B --> C{推理模式}
    C -->|本地推理| D[本地模型加载<br>ONNX/Candle/GGUF]
    C -->|远程推理| E[HTTP 客户端<br>reqwest + 流式 SSE]

    D --> F[推理引擎<br>Token 生成循环]
    E --> G[API 调用<br>OpenAI/Claude/自定义]

    F --> H[流式输出<br>termcolor 渐进渲染]
    G --> H

    H --> I[上下文管理<br>对话历史持久化]

    subgraph 核心抽象层
        C
        F
        G
        I
    end

    style D fill:#f9f,stroke:#333
    style E fill:#bbf,stroke:#333
    style H fill:#bfb,stroke:#333

2.2 关键设计决策

本地 vs 远程推理：本地推理零延迟、无网络依赖，但模型大小受限于设备内存；远程推理模型能力强，但有网络延迟和 API 成本。对于 CLI 工具，推荐混合模式——简单任务本地推理，复杂任务远程推理。

流式输出：AI 生成文本是逐 token 产出的，用户等待完整响应的体验很差。必须实现流式输出：每生成一个 token 就立即渲染到终端。

上下文管理：多轮对话需要维护上下文历史。CLI 工具的上下文可以存储在本地文件（如 .chat_history.json），支持跨会话持久化。

三、生产级代码实现

3.1 CLI 框架与命令路由

use clap::{Parser, Subcommand};

#[derive(Parser)]
#[command(name = "ai-cli", about = "AI 驱动的命令行助手")]
struct Cli {
    #[command(subcommand)]
    command: Commands,

    /// 使用本地模型（默认使用远程 API）
    #[arg(long, global = true)]
    local: bool,

    /// 模型名称
    #[arg(long, default_value = "qwen2-1.5b-instruct")]
    model: String,
}

#[derive(Subcommand)]
enum Commands {
    /// 与 AI 对话
    Chat {
        /// 输入消息
        message: Vec<String>,
        /// 继续上一次对话
        #[arg(long)]
        continue_last: bool,
    },
    /// 根据错误信息推荐修复方案
    Debug {
        /// 错误信息（从 stdin 读取或作为参数传入）
        error: Option<String>,
    },
    /// 根据描述生成代码
    Code {
        /// 代码描述
        prompt: Vec<String>,
        /// 编程语言
        #[arg(long, default_value = "rust")]
        lang: String,
    },
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let cli = Cli::parse();

    match cli.command {
        Commands::Chat { message, continue_last } => {
            let prompt = message.join(" ");
            chat_command(&prompt, continue_last, cli.local, &cli.model).await
        }
        Commands::Debug { error } => {
            let error_msg = error.unwrap_or_else(|| {
                // 从 stdin 读取错误信息
                use std::io::Read;
                let mut buf = String::new();
                std::io::stdin().read_to_string(&mut buf).unwrap();
                buf.trim().to_string()
            });
            debug_command(&error_msg, cli.local, &cli.model).await
        }
        Commands::Code { prompt, lang } => {
            let desc = prompt.join(" ");
            code_command(&desc, &lang, cli.local, &cli.model).await
        }
    }
}

3.2 远程 API 调用与流式输出

use reqwest::Client;
use serde::{Deserialize, Serialize};
use tokio_stream::StreamExt;

#[derive(Serialize)]
struct ChatRequest {
    model: String,
    messages: Vec<Message>,
    stream: bool,
    max_tokens: u32,
    temperature: f32,
}

#[derive(Serialize, Deserialize, Clone)]
struct Message {
    role: String,
    content: String,
}

#[derive(Deserialize)]
struct ChatChunk {
    choices: Vec<Choice>,
}

#[derive(Deserialize)]
struct Choice {
    delta: Delta,
}

#[derive(Deserialize)]
struct Delta {
    content: Option<String>,
}

/// 流式调用远程 AI API
/// 逐 token 接收并实时输出到终端
async fn stream_chat(
    client: &Client,
    api_url: &str,
    api_key: &str,
    messages: &[Message],
    model: &str,
) -> anyhow::Result<String> {
    let request = ChatRequest {
        model: model.to_string(),
        messages: messages.to_vec(),
        stream: true,
        max_tokens: 2048,
        temperature: 0.7,
    };

    let response = client
        .post(api_url)
        .header("Authorization", format!("Bearer {}", api_key))
        .header("Content-Type", "application/json")
        .json(&request)
        .send()
        .await?;

    // 检查 HTTP 状态码
    if !response.status().is_success() {
        let status = response.status();
        let body = response.text().await?;
        anyhow::bail!("API 请求失败: {} - {}", status, body);
    }

    // 解析 SSE 流
    let mut stream = response.bytes_stream();
    let mut full_response = String::new();

    use termcolor::{ColorChoice, StandardStream, WriteColor};
    use termcolor::ColorSpec;
    let mut stdout = StandardStream::stdout(ColorChoice::Always);

    while let Some(chunk) = stream.next().await {
        let chunk = chunk?;
        let text = String::from_utf8_lossy(&chunk);

        for line in text.lines() {
            if let Some(data) = line.strip_prefix("data: ") {
                if data == "[DONE]" {
                    break;
                }
                if let Ok(parsed) = serde_json::from_str::<ChatChunk>(data) {
                    if let Some(choice) = parsed.choices.first() {
                        if let Some(content) = &choice.delta.content {
                            // 实时输出到终端
                            print!("{}", content);
                            full_response.push_str(content);
                            use std::io::Write;
                            stdout.flush()?;
                        }
                    }
                }
            }
        }
    }

    println!(); // 换行
    Ok(full_response)
}

3.3 上下文管理与持久化

use std::path::PathBuf;
use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize, Default)]
struct ChatHistory {
    sessions: Vec<ChatSession>,
}

#[derive(Serialize, Deserialize, Clone)]
struct ChatSession {
    id: String,
    messages: Vec<Message>,
    created_at: String,
}

impl ChatHistory {
    /// 获取历史文件路径
    fn history_path() -> PathBuf {
        let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
        home.join(".ai-cli").join("history.json")
    }

    /// 加载历史记录
    fn load() -> anyhow::Result<Self> {
        let path = Self::history_path();
        if path.exists() {
            let content = std::fs::read_to_string(&path)?;
            Ok(serde_json::from_str(&content)?)
        } else {
            Ok(ChatHistory::default())
        }
    }

    /// 保存历史记录
    fn save(&self) -> anyhow::Result<()> {
        let path = Self::history_path();
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        let content = serde_json::to_string_pretty(self)?;
        std::fs::write(&path, content)?;
        Ok(())
    }

    /// 获取最后一次会话
    fn last_session(&self) -> Option<&ChatSession> {
        self.sessions.last()
    }

    /// 添加消息到当前会话
    fn add_message(&mut self, session_id: &str, message: Message) {
        if let Some(session) = self.sessions.iter_mut()
            .find(|s| s.id == session_id)
        {
            session.messages.push(message);
        }
    }

    /// 创建新会话
    fn new_session(&mut self) -> &ChatSession {
        let session = ChatSession {
            id: uuid::Uuid::new_v4().to_string(),
            messages: vec![],
            created_at: chrono::Utc::now().to_rfc3339(),
        };
        self.sessions.push(session);
        self.sessions.last().unwrap()
    }
}

3.4 完整的 Chat 命令实现

async fn chat_command(
    prompt: &str,
    continue_last: bool,
    use_local: bool,
    model: &str,
) -> anyhow::Result<()> {
    let mut history = ChatHistory::load()?;

    // 确定会话
    let session_id = if continue_last {
        history.last_session()
            .map(|s| s.id.clone())
            .ok_or_else(|| anyhow::anyhow!("没有历史会话"))?
    } else {
        history.new_session().id.clone()
    };

    // 构建消息列表
    let user_message = Message {
        role: "user".to_string(),
        content: prompt.to_string(),
    };
    history.add_message(&session_id, user_message.clone());

    // 收集上下文消息
    let messages: Vec<Message> = history.sessions
        .iter()
        .find(|s| s.id == session_id)
        .map(|s| s.messages.clone())
        .unwrap_or_default();

    // 调用推理
    let response = if use_local {
        // 本地推理路径（简化示例）
        local_inference(&messages, model).await?
    } else {
        // 远程 API 路径
        let client = Client::new();
        let api_url = std::env::var("AI_API_URL")
            .unwrap_or_else(|_| "https://api.openai.com/v1/chat/completions".into());
        let api_key = std::env::var("AI_API_KEY")
            .map_err(|_| anyhow::anyhow!("请设置 AI_API_KEY 环境变量"))?;
        stream_chat(&client, &api_url, &api_key, &messages, model).await?
    };

    // 保存助手回复到历史
    let assistant_message = Message {
        role: "assistant".to_string(),
        content: response,
    };
    history.add_message(&session_id, assistant_message);
    history.save()?;

    Ok(())
}

async fn local_inference(
    _messages: &[Message],
    _model: &str,
) -> anyhow::Result<String> {
    // 本地推理实现——使用 Candle 或 ONNX Runtime
    // 此处为占位，实际实现参考第 2 篇文章
    println!("[本地推理模式] 正在加载模型...");
    Ok("本地推理功能开发中".to_string())
}