crates\nfarms-core\src\chat.rs

源码路径:crates\nfarms-core\src\chat.rs · 520 行

//! 桌宠对话框的后端:调本地推理服务的 OpenAI 兼容接口。
//!
//! # 为什么不走 DSH
//!
//! DSH 是**会话式**的:它管历史、工具调用、权限、上下文裁剪。而桌宠这里要的是
//! 「问一句答一句」的一次性对话 —— 把它塞进 DSH 的会话模型里,两边都别扭。
//!
//! 直接打 llama-server 还少一个会坏的地方:**DSH 没启动时,桌宠照样能说话**。
//!
//! # 这个模块只做纯逻辑
//!
//! `ask` 会真的发请求,但它只依赖一个 `base_url` —— 谁传进来的、是哪个后端,
//! 它都不关心。所以测试可以直接对着正在跑的 `llama-server` 打
//! (见 `mod tests`),也可以对着一个不存在的端口断言错误路径。

use serde_json::json;
use std::time::{Duration, Instant};

/// 对话里的一条消息。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Msg {
    pub role: &'static str,
    pub content: String,
}

impl Msg {
    pub fn system(s: impl Into<String>) -> Self {
        Self {
            role: "system",
            content: s.into(),
        }
    }
    pub fn user(s: impl Into<String>) -> Self {
        Self {
            role: "user",
            content: s.into(),
        }
    }
    pub fn assistant(s: impl Into<String>) -> Self {
        Self {
            role: "assistant",
            content: s.into(),
        }
    }
}

/// 桌宠的人格。**写死在代码里**,不让用户配 —— 一个桌面宠物不需要人设编辑器,
/// 而留一堆可配项只会让「它怎么突然不这么说话了」变成一个要查的问题。
pub const PERSONA: &str = "你是一只住在 Windows 桌面上的鲸鱼娘,名字叫小鲸。\
回答要短(通常一两句)、口语化、带点可爱,但不要用颜文字堆砌。\
用户问技术问题时正常回答,不要卖萌耽误事。";

/// 默认最多生成多少 token。
///
/// 桌宠的回答本来就该短 —— 而且本地小模型生成慢,给太大只是让它在那儿打很久字。
pub const DEFAULT_MAX_TOKENS: u32 = 512;

/// 一次请求最多等多久。
///
/// ureq 的默认超时对本地的 30B MoE **不够**:实测出一段两百字的回答要十几秒,
/// 冷启动(模型还在加载)更久。给到 5 分钟,超时了就报错而不是永远转圈。
const TIMEOUT_SECS: u64 = 300;

/// `draw` 最多等 `sd-cli` 多久。
///
/// **和上面那个 HTTP 超时是两件事**:`TIMEOUT_SECS` 给的是 `ureq` 的
/// `timeout_global`,它管的是「HTTP 响应等多久」,而 `draw` 等的是一个**子进程**
/// ——HTTP 的超时对子进程一点用都没有(这正是出图会永久卡死的原因)。
///
/// 取 10 分钟,比 HTTP 那个还宽:出图实测 148 秒,冷启动时前面还要加加载
/// 8.8 GB 权重的时间。它**只当「卡死了」的兜底**,不是性能判据 —— 真按 5 分钟
/// 去砍,慢一点但正常的出图会被误杀,那比没有超时更糟。
const DRAW_TIMEOUT_SECS: u64 = 600;

fn agent() -> ureq::Agent {
    ureq::Agent::config_builder()
        .timeout_global(Some(Duration::from_secs(TIMEOUT_SECS)))
        .build()
        .into()
}

/// 只要端口上没人在听,`ureq` 报的就是这句(拿它区分「服务没起」和「别的问题」)。
///
/// 这个区分对用户很重要:前者要去 Model › llama.cpp 点「启动服务」,
/// 后者是配置或模型的问题 —— 而两件事的**处理方式完全不同**。
pub fn looks_like_connection_refused(e: &str) -> bool {
    let l = e.to_ascii_lowercase();
    l.contains("connection refused")
        || l.contains("actively refused")
        || l.contains("10061")
        || l.contains("connect error")
}

/// 组装请求体(纯函数,单独拎出来是为了能测)。
pub fn build_payload(model: &str, msgs: &[Msg], max_tokens: u32) -> serde_json::Value {
    json!({
        "model": model,
        "messages": msgs
            .iter()
            .map(|m| json!({ "role": m.role, "content": m.content }))
            .collect::<Vec<_>>(),
        "max_tokens": max_tokens,
        // 不用流式:桌宠的回答短,一次性拿回来比分块拼接简单得多
        "stream": false
    })
}

/// 从响应体里抠出回复正文(纯函数)。
///
/// 分开测的理由和 `build_payload` 一样:HTTP 能不能通取决于环境,
/// 而**解析逻辑对不对是代码自己的事**,不该被环境的成败掩盖。
pub fn parse_reply(body: &str) -> Result<String, String> {
    let v: serde_json::Value =
        serde_json::from_str(body).map_err(|e| format!("响应不是合法 JSON: {e}"))?;
    // 服务端出错时返回的是 {"error": {...}},这条要先认出来 ——
    // 不然会被当成「choices 是空的」,报出一个和真实原因无关的错。
    if let Some(err) = v.get("error") {
        let m = err
            .get("message")
            .and_then(|x| x.as_str())
            .unwrap_or("服务端返回了错误");
        return Err(format!("服务端: {m}"));
    }
    let text = v
        .get("choices")
        .and_then(|c| c.get(0))
        .and_then(|c| c.get("message"))
        .and_then(|m| m.get("content"))
        .and_then(|c| c.as_str())
        .ok_or_else(|| "响应里没有 choices[0].message.content".to_string())?;
    let t = text.trim();
    if t.is_empty() {
        return Err("模型返回了空回复".to_string());
    }
    Ok(t.to_string())
}

/// 一轮回复:正文 + 用量。
///
/// **用量单独带出来**,不直接丢掉 —— 界面上有一个「Tokens」页要显示它。
/// 服务端本来就在响应里给了(`usage`),不取白不取;而一旦丢了,
/// 想再要就得自己数 token,那和真实值对不上。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Reply {
    pub text: String,
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
}

/// 从响应里取用量。缺字段就是 0 —— **不能因为服务端没给用量就判定请求失败**。
pub fn parse_usage(body: &str) -> (u32, u32) {
    let Ok(v) = serde_json::from_str::<serde_json::Value>(body) else {
        return (0, 0);
    };
    let g = |k: &str| v.get("usage").and_then(|u| u.get(k)).and_then(|x| x.as_u64()).unwrap_or(0) as u32;
    (g("prompt_tokens"), g("completion_tokens"))
}

/// 发一轮对话,返回回复正文。
///
/// `base_url` 形如 `http://127.0.0.1:8080/v1` —— 和 DSH 的 provider 用的是**同一个**
/// 地址,所以端口改了这里跟着走(唯一真值在 Setting › Proxy and Port)。
pub fn ask(
    base_url: &str,
    model: &str,
    msgs: &[Msg],
    max_tokens: u32,
) -> Result<Reply, String> {
    let url = format!("{}/chat/completions", base_url.trim_end_matches('/'));
    let payload = build_payload(model, msgs, max_tokens);
    let body = payload.to_string();
    let mut resp = agent()
        .post(&url)
        .header("Content-Type", "application/json")
        .send(body)
        .map_err(|e| {
            let s = e.to_string();
            if looks_like_connection_refused(&s) {
                format!("连不上 {url} —— 本地推理服务没在跑?(Model › llama.cpp › 启动服务)")
            } else {
                format!("请求失败: {s}")
            }
        })?;
    let body = resp
        .body_mut()
        .read_to_string()
        .map_err(|e| format!("读响应失败: {e}"))?;
    let text = parse_reply(&body)?;
    let (pt, ct) = parse_usage(&body);
    Ok(Reply {
        text,
        prompt_tokens: pt,
        completion_tokens: ct,
    })
}

/// 让 stable-diffusion.cpp 画一张图。**阻塞几分钟**,调用方必须放进 `Job`。
///
/// # 为什么走 `sd-cli` 而不是 `sd-server`
///
/// 桌宠画图是"要一张图"的一次性动作,不是常驻服务。`sd-cli` 起来→出图→退出,
/// 不用管生命周期;`sd-server` 要另起进程、管端口、管退出,而这一页并不需要
/// 连续出图。**等真需要连续出图(比如批量)时再换**,那时 server 的收益才成立。
///
/// 出图很慢:这个会话里实测 4 GB 扩散权重 + 4.8 GB text encoder 跑一次
/// **148 秒**。所以超时给得很宽(`DRAW_TIMEOUT_SECS`),而不是让它中途失败。
pub fn draw(
    prompt: &str,
    out: &std::path::Path,
    model_dir: &std::path::Path,
    sd_cli: &std::path::Path,
    w: u32,
    h: u32,
) -> Result<(), String> {
    draw_with_timeout(
        prompt,
        out,
        model_dir,
        sd_cli,
        w,
        h,
        Duration::from_secs(DRAW_TIMEOUT_SECS),
    )
}

/// `draw` 的本体。超时值单独当参数,只为了让测试能用几百毫秒**真的**构造出
/// 「`sd-cli` 不退出」的场景 —— 拿真的 10 分钟去测,测试就没法跑了。
fn draw_with_timeout(
    prompt: &str,
    out: &std::path::Path,
    model_dir: &std::path::Path,
    sd_cli: &std::path::Path,
    w: u32,
    h: u32,
    timeout: Duration,
) -> Result<(), String> {
    let args = crate::libs::sd_cli_args(prompt, out, model_dir, w, h);
    let mut cmd = std::process::Command::new(sd_cli);
    cmd.args(&args);
    // `output()` 会默认把两路管道接走,`spawn()` 不会:不显式给 `piped`,
    // 子进程会继承本进程的 stdout/stderr,失败时想看的报错就抓不到了。
    cmd.stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped());
    // **不弹控制台**:GUI 里起子进程不加这个会闪一个黑框出来
    #[cfg(windows)]
    {
        use std::os::windows::process::CommandExt;
        const CREATE_NO_WINDOW: u32 = 0x0800_0000;
        cmd.creation_flags(CREATE_NO_WINDOW);
    }
    let child = cmd
        .spawn()
        .map_err(|e| format!("起不了 sd-cli({}): {e}", sd_cli.display()))?;
    let r = wait_with_timeout(child, timeout).map_err(|e| match e {
        // 报错要指明**是超时**。写成「退出码 xxx」的话用户会去查参数,
        // 而真实原因通常是模型太大 / 显存不足 / 后端挂住。
        WaitError::Timeout => format!(
            "出图超时(>{} 秒),已终止 sd-cli —— 卡在加载或采样上了;\
换更小的模型,或看 Model 页的日志再试",
            timeout.as_secs()
        ),
        WaitError::Failed(e) => format!("等 sd-cli 结束失败: {e}"),
    })?;
    if !r.status.success() {
        let err = String::from_utf8_lossy(&r.stderr);
        let tail: String = err.chars().rev().take(800).collect::<String>().chars().rev().collect();
        return Err(format!("sd-cli 退出码 {:?}:{tail}", r.status.code()));
    }
    if !out.exists() {
        return Err("sd-cli 报成功,但输出文件不在 —— 参数可能没被它认".to_string());
    }
    Ok(())
}

/// `wait_with_timeout` 的两条失败路径。
///
/// **分开**是因为它们对用户的意义完全不同:超时要能说一句「我把它杀了」,
/// 而 `try_wait` 自己出错是系统层面的事(句柄没了之类),混成一句话会把人带偏。
enum WaitError {
    Timeout,
    Failed(std::io::Error),
}

/// 等一个**已经 `spawn` 出来**的子进程结束,最多等 `timeout`;到点就把它杀掉。
///
/// # 为什么不用 `Command::output()`
///
/// `output()` 会一直阻塞到子进程退出并读完管道 —— **std 里没有任何给子进程设
/// 超时的 API**。`sd-cli` 一旦卡住(显存不足、后端挂住),`draw` 就永远不返回:
/// 界面上的转圈永不停,而 `Job::start` 的 `if self.running { return }` 还会把
/// 之后每一次出图请求**静默丢掉**,于是这一次运行里的出图功能等于废了,只能重启。
/// 所以这里自己 `try_wait` 轮询(写法照 `tools.rs::capture` 已有的那份)。
///
/// # 为什么两路管道各起一个线程读
///
/// 子进程的输出超过管道缓冲区(Windows 上 64 KB)时会阻塞在写管道上,而阻塞中的
/// 子进程不会退出 —— 只轮询 `try_wait` 会把「输出太多」误判成「超时」,还会把已经
/// 写出来的输出一起丢掉。读线程一挂上,子进程就能一直写完。
fn wait_with_timeout(
    mut child: std::process::Child,
    timeout: Duration,
) -> Result<std::process::Output, WaitError> {
    let out_handle = child.stdout.take().map(|mut p| {
        std::thread::spawn(move || {
            use std::io::Read;
            let mut buf = Vec::new();
            let _ = p.read_to_end(&mut buf);
            buf
        })
    });
    let err_handle = child.stderr.take().map(|mut p| {
        std::thread::spawn(move || {
            use std::io::Read;
            let mut buf = Vec::new();
            let _ = p.read_to_end(&mut buf);
            buf
        })
    });

    let deadline = Instant::now() + timeout;
    let status = loop {
        match child.try_wait() {
            Ok(Some(st)) => break Ok(st),
            Ok(None) => {
                if Instant::now() >= deadline {
                    let _ = child.kill();
                    // kill 之后照样 wait 一次:把进程句柄收掉,免得留在自己手里。
                    let _ = child.wait();
                    break Err(WaitError::Timeout);
                }
                // 20 ms:出图以分钟计,这个粒度足够,又不会把 CPU 占住
                std::thread::sleep(Duration::from_millis(20));
            }
            Err(e) => break Err(WaitError::Failed(e)),
        }
    };

    match status {
        Ok(st) => {
            let stdout = out_handle.and_then(|h| h.join().ok()).unwrap_or_default();
            let stderr = err_handle.and_then(|h| h.join().ok()).unwrap_or_default();
            Ok(std::process::Output {
                status: st,
                stdout,
                stderr,
            })
        }
        // 超时这条路上**故意不 join 读线程**:`kill` 只杀直接子进程,它要是留下了
        // 继承了管道句柄的后代(比如批处理里的 `ping`),join 会白等到那个后代自己
        // 退出 —— 那就等于没超时。句柄关掉之后这两个线程自己会结束,放它们走。
        Err(e) => Err(e),
    }
}

/// 读 llama-server 的 Prometheus 指标。
///
/// # 为什么不用它替代本地累计
///
/// 这是**服务端进程**的累计值,所以它同时包含 DSH 那边的调用 —— 比只统计
/// 桌宠对话框完整。但它是"服务器起来之后的总量",分不清是谁用的,
/// 也和对话的上下文窗口没关系。两个视角各有各的用,所以**两个都留**。
///
/// # 为什么不解析成数字
///
/// 我**没有核实过**这个版本 llama.cpp 的指标名。猜一个 `llamacpp:tokens_*`
/// 出来,如果名字变了就会显示成 0 —— 而那看起来像"没用过 token",
/// 是个会骗人的错。所以这里把**含 token 的原始行**交给界面照显,
/// 名字对不对用户一眼就能看见。
pub fn fetch_metrics(base_url: &str) -> Result<Vec<String>, String> {
    // **指标在根路径,不在 /v1 下面。** 传进来的 base 形如
    // `http://127.0.0.1:8080/v1`(和对话同一个地址),直接接 `/metrics`
    // 会请求成 `/v1/metrics` —— 服务在跑却回 404。
    // 404 和「连不上」长得完全不同,但都能让人以为指标功能坏了。
    let root = base_url
        .trim_end_matches('/')
        .trim_end_matches("/v1");
    let url = format!("{root}/metrics");
    let mut r = agent()
        .get(&url)
        .call()
        .map_err(|e| {
            let s = e.to_string();
            if looks_like_connection_refused(&s) {
                format!("连不上 {url} —— 本地推理服务没在跑?")
            }
            else if s.contains("501") {
                // 501 = 接口在,但没启用。这个区分很重要:
                // 「连不上」要去启动服务,「501」是要带 --metrics 重启 ——
                // 两件事的处理方式完全不同,而报错长得一样近。
                "服务在跑,但指标接口没开(501)。\
用本程序 Model › llama.cpp 启动的服务会自动带上 --metrics;\
若是手动起的,要加这个参数再重启。"
                    .to_string()
            }
            else {
                format!("取指标失败: {s}")
            }
        })?;
    let body = r
        .body_mut()
        .read_to_string()
        .map_err(|e| format!("读响应失败: {e}"))?;
    let lines: Vec<String> = body
        .lines()
        .filter(|l| !l.starts_with('#') && l.to_ascii_lowercase().contains("token"))
        .map(|l| l.trim().to_string())
        .collect();
    Ok(lines)
}
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn payload_carries_role_and_content() {
        let msgs = vec![Msg::system(PERSONA), Msg::user("你好")];
        let p = build_payload("nfarms-local", &msgs, 128);
        assert_eq!(p["model"], "nfarms-local");
        assert_eq!(p["messages"][0]["role"], "system");
        assert_eq!(p["messages"][1]["content"], "你好");
        assert_eq!(p["stream"], false);
        assert_eq!(p["max_tokens"], 128);
    }

    #[test]
    fn parses_a_normal_reply() {
        let body = r#"{"choices":[{"message":{"role":"assistant","content":"  你好呀  "}}]}"#;
        assert_eq!(parse_reply(body).unwrap(), "你好呀");
    }

    /// 服务端的 `error` 字段要被认出来,而不是掉进「choices 为空」那条错里 ——
    /// 后者的报错和真实原因无关,会把人带偏。
    #[test]
    fn surfaces_the_server_side_error_message() {
        let body = r#"{"error":{"message":"request exceeds context size","type":"exceed_context_size_error"}}"#;
        let e = parse_reply(body).unwrap_err();
        assert!(e.contains("exceeds context"), "实际: {e}");
    }

    #[test]
    fn rejects_a_body_without_choices() {
        assert!(parse_reply("{}").is_err());
        assert!(parse_reply("not json").is_err());
    }

    /// 空回复不算成功 —— 界面上显示一个空气泡比显示一句错误更让人困惑。
    #[test]
    fn empty_content_is_an_error_not_an_empty_answer() {
        let body = r#"{"choices":[{"message":{"content":"   "}}]}"#;
        assert!(parse_reply(body).is_err());
    }

    #[test]
    fn recognises_connection_refused_in_the_various_spellings() {
        assert!(looks_like_connection_refused("Connection refused (os error 10061)"));
        assert!(looks_like_connection_refused(
            "No connection could be made because the target machine actively refused it"
        ));
        assert!(!looks_like_connection_refused("timed out"));
    }

    /// 对着一个**确定没人监听**的端口打一次,断言错误里带的是「服务没在跑」那句。
    ///
    /// 这条不依赖外部服务:它要验的就是**连不上时的提示对不对**。
    #[test]
    fn asking_a_dead_port_explains_that_the_server_is_not_running() {
        // 59999 不在常规分配范围,基本可以确定没人监听
        let r = ask("http://127.0.0.1:59999/v1", "m", &[Msg::user("hi")], 16);
        let e = r.unwrap_err();
        assert!(
            e.contains("没在跑"),
            "连不上时应当提示去启动服务,实际: {e}"
        );
    }

    /// 守的是**「`sd-cli` 卡住时 `draw` 会自己收场」**这件事 —— 出图线程
    /// 一旦在里面无限等下去,界面上就是永久转圈,而且之后的出图请求会被
    /// 静默丢掉,只能重启程序。
    ///
    /// 假 `sd-cli` 是个批处理:参数一个不看,自己 `ping` 约 7 秒再退出。故意用
    /// **有限**的时长,是为了万一超时被拿掉(退回 `output()` / `wait()`),
    /// 这条测试会在约 7 秒后**失败**,而不是把整个测试进程挂死。
    #[cfg(windows)]
    #[test]
    fn draw_gives_up_on_a_hung_sd_cli_instead_of_waiting_forever() {
        let dir = std::env::temp_dir().join(format!("nfarms-draw-hang-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let fake = dir.join("fake-sd-cli.cmd");
        // 先 `cd` 到 %TEMP%:被留了一下的 `ping` 会把工作目录占住,否则这个目录删不掉
        std::fs::write(
            &fake,
            "@echo off\r\ncd /d \"%TEMP%\"\r\nping -n 8 127.0.0.1 >nul\r\n",
        )
        .unwrap();

        let out = dir.join("draw.png");
        let t0 = std::time::Instant::now();
        let e =
            draw_with_timeout("卡住", &out, &dir, &fake, 64, 64, Duration::from_millis(500))
                .unwrap_err();

        assert!(e.contains("超时"), "子进程不退出时应当报「出图超时」,实际: {e}");
        // 超时给的是 500 ms,却等了 5 秒以上 —— 说明等的其实是子进程自己退出,
        // 也就是超时根本没生效(上面的断言此时也会因为报的是别的错而失败)。
        assert!(
            t0.elapsed() < Duration::from_secs(5),
            "没有在超时后立刻收场,等了 {:?}",
            t0.elapsed()
        );
        // 收拾现场而已:杀完进程到目录句柄真的放掉之间有一点时间差,删不掉就
        // 重试两下 —— 删不掉也不该让这条测试失败(失败信息会指向完全无关的地方)。
        for _ in 0..5 {
            if std::fs::remove_dir_all(&dir).is_ok() {
                break;
            }
            std::thread::sleep(Duration::from_millis(100));
        }
    }
}