crabmate 0.5.0

Rust AI agent: OpenAI-compatible chat/completions, function calling, HTTP serve, ops CLI
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

//! P 步：向模型要本轮输出（含重试）。

use std::sync::atomic::AtomicBool;

use tokio::sync::mpsc;

use crate::config::AgentConfig;
use crate::llm::{
    CompleteChatRetryingParams, LlmCompleteError, LlmRetryingTransportOpts, complete_chat_retrying,
    tool_chat_request,
};
use crate::types::{LlmSeedOverride, Message};

/// P：构造请求并调用模型（`no_stream` 为 true 时走 `stream: false`），**不**修改 `messages`。
pub(crate) struct PerPlanCallModelParams<'a> {
    pub llm_backend: &'a (dyn crate::llm::ChatCompletionsBackend + 'static),
    pub client: &'a reqwest::Client,
    /// 当前 api_key（executor 阶段可能被覆盖）
    pub api_key: &'a str,
    pub cfg: &'a AgentConfig,
    /// 默认全量工具；分阶段步级子代理时传入收窄后的切片。
    pub tools_defs: &'a [crate::types::Tool],
    pub messages: &'a [Message],
    pub out: Option<&'a mpsc::Sender<String>>,
    pub no_stream: bool,
    pub cancel: Option<&'a AtomicBool>,
    pub temperature_override: Option<f32>,
    pub model_override: Option<&'a str>,
    pub seed_override: LlmSeedOverride,
    pub request_chrome_trace: Option<std::sync::Arc<crate::request_chrome_trace::RequestTurnTrace>>,
    /// 当 use_executor_model 为 true 时，优先使用此 api_base。
    pub executor_api_base: Option<&'a str>,
    /// 当 use_executor_model 为 true 时，优先使用此 api_key。
    pub executor_api_key: Option<&'a str>,
    pub turn_budget: Option<&'a std::sync::Arc<crate::agent::turn_budget::TurnBudgetCounter>>,
    pub provider_usage_sink:
        Option<&'a std::sync::Arc<std::sync::Mutex<Option<crate::cm_types::Usage>>>>,
}

pub(crate) async fn per_plan_call_model_retrying(
    p: PerPlanCallModelParams<'_>,
) -> Result<(Message, String), LlmCompleteError> {
    let PerPlanCallModelParams {
        llm_backend,
        client,
        api_key,
        cfg,
        tools_defs,
        messages,
        out,
        no_stream,
        cancel,
        temperature_override,
        model_override,
        seed_override,
        request_chrome_trace,
        executor_api_base,
        executor_api_key,
        turn_budget,
        provider_usage_sink,
    } = p;

    // 确定 effective api_base 和 api_key
    let (effective_cfg, effective_api_key) =
        if executor_api_base.is_some() || executor_api_key.is_some() {
            let mut c = (*cfg).clone();
            let mut key = api_key.to_string();
            if let Some(base) = executor_api_base {
                c.llm.api_base = base.to_string();
            }
            if let Some(key_override) = executor_api_key {
                key = key_override.to_string();
            }
            (std::sync::Arc::new(c), key)
        } else {
            (std::sync::Arc::new(cfg.clone()), api_key.to_string())
        };

    let llm_cfg = crate::cm_types::llm_config::LlmConfig {
        llm: effective_cfg.llm.clone(),
        sampling: effective_cfg.llm_sampling.clone(),
        vendor_flags: effective_cfg.llm_vendor_flags.clone(),
        http_retry: effective_cfg.llm_http_retry.clone(),
    };
    let req = tool_chat_request(
        &llm_cfg,
        messages,
        tools_defs,
        temperature_override,
        model_override,
        seed_override,
    );
    let cc = CompleteChatRetryingParams::new(
        llm_backend,
        client,
        &effective_api_key,
        &effective_cfg,
        LlmRetryingTransportOpts {
            out,
            no_stream,
            cancel,
        },
        request_chrome_trace,
        model_override,
    )
    .with_turn_budget(turn_budget)
    .with_provider_usage_sink(provider_usage_sink);
    let (msg, finish_reason) = complete_chat_retrying(&cc, &req).await?;
    Ok((msg, finish_reason))
}