Skip to main content

vtcode_core/llm/
rig_adapter.rs

1use crate::config::models::Provider;
2use crate::config::types::ReasoningEffortLevel;
3use anyhow::{Context, Result};
4use rig::client::CompletionClient;
5use rig::providers::gemini::completion::gemini_api_types::ThinkingConfig;
6use rig::providers::{anthropic, deepseek, gemini, openai, openrouter};
7use serde_json::{Value, json};
8
9/// Result of validating a provider/model combination through rig-core.
10#[derive(Debug, Clone)]
11pub struct RigValidationSummary {
12    pub provider: Provider,
13    pub model: String,
14}
15
16/// Internal bridge for Rig-backed provider/model capability checks.
17#[derive(Debug, Clone)]
18pub struct RigProviderCapabilities {
19    provider: Provider,
20    model: String,
21}
22
23impl RigProviderCapabilities {
24    #[must_use]
25    pub fn new(provider: Provider, model: impl Into<String>) -> Self {
26        Self {
27            provider,
28            model: model.into(),
29        }
30    }
31
32    /// Attempt to construct a rig-core client for the given provider and
33    /// instantiate the requested model. This performs a lightweight validation
34    /// without issuing a network request, ensuring that downstream calls can
35    /// reuse the rig client configuration paths.
36    pub fn validate_model(&self, api_key: &str) -> Result<RigValidationSummary> {
37        match self.provider {
38            Provider::Gemini => {
39                let client = gemini::Client::new(api_key)
40                    .context("failed to initialise Rig Gemini client")?;
41                let _ = client.completion_model(&self.model);
42            }
43            Provider::OpenAI => {
44                let client = openai::Client::new(api_key)
45                    .context("failed to initialise Rig OpenAI client")?;
46                let _ = client.completion_model(&self.model);
47            }
48            Provider::Anthropic => {
49                let client = anthropic::Client::new(api_key)
50                    .context("failed to initialise Rig Anthropic client")?;
51                let _ = client.completion_model(&self.model);
52            }
53            Provider::Copilot => {
54                // Copilot is authenticated through the official CLI, not rig.
55            }
56            Provider::Minimax => {
57                // MiniMax uses an Anthropic-compatible API; rig has no direct client.
58            }
59            Provider::DeepSeek => {
60                let client = deepseek::Client::new(api_key)
61                    .context("failed to initialise Rig DeepSeek client")?;
62                let _ = client.completion_model(&self.model);
63            }
64            Provider::HuggingFace => {
65                // Hugging Face exposes an OpenAI-compatible router; rig does not ship a dedicated client.
66            }
67            Provider::OpenRouter => {
68                let client = openrouter::Client::new(api_key)
69                    .context("failed to initialise Rig OpenRouter client")?;
70                let _ = client.completion_model(&self.model);
71            }
72            Provider::Ollama => {
73                // Rig does not provide an Ollama integration; validation is skipped.
74            }
75            Provider::LlamaCpp => {
76                // llama.cpp provides an OpenAI-compatible API; rig has no direct client.
77            }
78            Provider::LmStudio => {
79                // LM Studio provides an OpenAI-compatible API; rig has no direct client.
80            }
81            Provider::Moonshot => {
82                // Moonshot does not have a rig client integration yet.
83            }
84            Provider::Mistral => {
85                // Mistral exposes an OpenAI-compatible API; rig has no dedicated client.
86            }
87            Provider::ZAI => {
88                // The rig crate does not yet expose a dedicated Z.AI client.
89                // Skip instantiation while still marking the provider as verified.
90            }
91            Provider::OpenCodeZen | Provider::OpenCodeGo => {
92                // OpenCode Zen/Go are API gateways; rig has no direct client.
93            }
94            Provider::MiMo => {
95                // Xiaomi MiMo is an OpenAI-compatible API; rig has no dedicated client.
96            }
97            Provider::Qwen => {
98                // Alibaba Cloud Qwen uses an OpenAI-compatible API; rig has no dedicated client.
99            }
100            Provider::StepFun => {
101                // StepFun uses an OpenAI-compatible API; rig has no dedicated client.
102            }
103            Provider::Evolink => {
104                // Evolink is an OpenAI-compatible gateway; rig has no dedicated client.
105            }
106            Provider::Poolside => {
107                // Poolside uses an OpenAI-compatible API; rig has no dedicated client.
108            }
109        }
110
111        Ok(RigValidationSummary {
112            provider: self.provider,
113            model: self.model.clone(),
114        })
115    }
116
117    /// Convert a VT Code reasoning effort level to provider-specific parameters
118    /// using rig-core data structures. The resulting JSON payload can be merged
119    /// into provider requests when supported.
120    #[must_use]
121    pub fn reasoning_parameters(&self, effort: ReasoningEffortLevel) -> Option<Value> {
122        match self.provider {
123            Provider::OpenAI => {
124                let mut reasoning = openai::responses_api::Reasoning::new();
125                let mapped = match effort {
126                    ReasoningEffortLevel::None => return None,
127                    ReasoningEffortLevel::Minimal => {
128                        let effort = if is_gpt5_codex_model(&self.model) {
129                            "low"
130                        } else {
131                            "minimal"
132                        };
133                        return Some(json!({ "effort": effort }));
134                    }
135                    ReasoningEffortLevel::Low => openai::responses_api::ReasoningEffort::Low,
136                    ReasoningEffortLevel::Medium => openai::responses_api::ReasoningEffort::Medium,
137                    ReasoningEffortLevel::High => openai::responses_api::ReasoningEffort::High,
138                    ReasoningEffortLevel::XHigh => return Some(json!({ "effort": "xhigh" })),
139                    ReasoningEffortLevel::Max => return Some(json!({ "effort": "xhigh" })),
140                };
141                reasoning = reasoning.with_effort(mapped);
142                serde_json::to_value(reasoning).ok()
143            }
144            Provider::Gemini => {
145                let include_thoughts = matches!(
146                    effort,
147                    ReasoningEffortLevel::High
148                        | ReasoningEffortLevel::XHigh
149                        | ReasoningEffortLevel::Max
150                );
151                let budget = match effort {
152                    ReasoningEffortLevel::None => return None,
153                    ReasoningEffortLevel::Minimal => 16,
154                    ReasoningEffortLevel::Low => 64,
155                    ReasoningEffortLevel::Medium => 128,
156                    ReasoningEffortLevel::High
157                    | ReasoningEffortLevel::XHigh
158                    | ReasoningEffortLevel::Max => 256,
159                };
160                let config = ThinkingConfig {
161                    thinking_budget: Some(budget),
162                    thinking_level: None,
163                    include_thoughts: Some(include_thoughts),
164                };
165                serde_json::to_value(config)
166                    .ok()
167                    .map(|value| json!({ "thinking_config": value }))
168            }
169            Provider::HuggingFace => match effort {
170                ReasoningEffortLevel::None => None,
171                ReasoningEffortLevel::Minimal => Some(json!({ "reasoning_effort": "minimal" })),
172                ReasoningEffortLevel::Low => Some(json!({ "reasoning_effort": "low" })),
173                ReasoningEffortLevel::Medium => Some(json!({ "reasoning_effort": "medium" })),
174                ReasoningEffortLevel::High
175                | ReasoningEffortLevel::XHigh
176                | ReasoningEffortLevel::Max => Some(json!({ "reasoning_effort": "high" })),
177            },
178            // DeepSeek only accepts `high` and `max` for reasoning_effort.
179            // Per DeepSeek docs: low/medium → high, xhigh → max.
180            Provider::DeepSeek => match effort {
181                ReasoningEffortLevel::None => None,
182                ReasoningEffortLevel::Minimal
183                | ReasoningEffortLevel::Low
184                | ReasoningEffortLevel::Medium
185                | ReasoningEffortLevel::High => {
186                    Some(json!({"thinking": {"type": "enabled"}, "reasoning_effort": "high"}))
187                }
188                ReasoningEffortLevel::XHigh | ReasoningEffortLevel::Max => {
189                    Some(json!({"thinking": {"type": "enabled"}, "reasoning_effort": "max"}))
190                }
191            },
192            Provider::Minimax => None,
193            Provider::Ollama => None,
194            Provider::LlamaCpp => None,
195            Provider::ZAI => match effort {
196                ReasoningEffortLevel::None => None,
197                ReasoningEffortLevel::Minimal => Some(json!({
198                    "thinking": { "type": "enabled" },
199                    "thinking_effort": "minimal"
200                })),
201                ReasoningEffortLevel::Low => Some(json!({
202                    "thinking": { "type": "enabled" },
203                    "thinking_effort": "low"
204                })),
205                ReasoningEffortLevel::Medium => Some(json!({
206                    "thinking": { "type": "enabled" },
207                    "thinking_effort": "medium"
208                })),
209                ReasoningEffortLevel::High
210                | ReasoningEffortLevel::XHigh
211                | ReasoningEffortLevel::Max => Some(json!({
212                    "thinking": { "type": "enabled" },
213                    "thinking_effort": "high"
214                })),
215            },
216            Provider::StepFun => match effort {
217                ReasoningEffortLevel::None => None,
218                ReasoningEffortLevel::Minimal | ReasoningEffortLevel::Low => {
219                    Some(json!({ "reasoning_effort": "low" }))
220                }
221                ReasoningEffortLevel::Medium => Some(json!({ "reasoning_effort": "medium" })),
222                ReasoningEffortLevel::High
223                | ReasoningEffortLevel::XHigh
224                | ReasoningEffortLevel::Max => Some(json!({ "reasoning_effort": "high" })),
225            },
226            Provider::Evolink => match effort {
227                ReasoningEffortLevel::None => None,
228                ReasoningEffortLevel::Minimal | ReasoningEffortLevel::Low => {
229                    Some(json!({ "reasoning_effort": "low" }))
230                }
231                ReasoningEffortLevel::Medium => Some(json!({ "reasoning_effort": "medium" })),
232                ReasoningEffortLevel::High
233                | ReasoningEffortLevel::XHigh
234                | ReasoningEffortLevel::Max => Some(json!({ "reasoning_effort": "high" })),
235            },
236            _ => None,
237        }
238    }
239}
240
241fn is_gpt5_codex_model(model: &str) -> bool {
242    model == "gpt-5-codex" || (model.starts_with("gpt-5.") && model.contains("codex"))
243}
244
245#[cfg(test)]
246mod tests {
247    use super::RigProviderCapabilities;
248    use crate::config::models::Provider;
249    use crate::config::types::ReasoningEffortLevel;
250
251    #[test]
252    fn rig_capabilities_validate_rig_backed_and_noop_providers() {
253        let openai = RigProviderCapabilities::new(Provider::OpenAI, "gpt-5")
254            .validate_model("test-key")
255            .expect("openai validation");
256        assert_eq!(openai.provider, Provider::OpenAI);
257        assert_eq!(openai.model, "gpt-5");
258
259        let deepseek = RigProviderCapabilities::new(Provider::DeepSeek, "deepseek-chat")
260            .validate_model("test-key")
261            .expect("no-op validation");
262        assert_eq!(deepseek.provider, Provider::DeepSeek);
263        assert_eq!(deepseek.model, "deepseek-chat");
264    }
265
266    #[test]
267    fn rig_capabilities_generate_reasoning_payload_for_supported_provider() {
268        let payload = RigProviderCapabilities::new(Provider::ZAI, "glm-5")
269            .reasoning_parameters(ReasoningEffortLevel::Medium)
270            .expect("reasoning payload");
271
272        assert_eq!(payload["thinking"]["type"], "enabled");
273        assert_eq!(payload["thinking_effort"], "medium");
274    }
275
276    #[test]
277    fn rig_capabilities_preserve_reasoning_payload_for_openai() {
278        let payload = RigProviderCapabilities::new(Provider::OpenAI, "gpt-5")
279            .reasoning_parameters(ReasoningEffortLevel::Medium)
280            .expect("reasoning payload");
281
282        assert_eq!(payload["effort"], "medium");
283
284        let codex_payload = RigProviderCapabilities::new(Provider::OpenAI, "gpt-5-codex")
285            .reasoning_parameters(ReasoningEffortLevel::Minimal)
286            .expect("reasoning payload");
287
288        assert_eq!(codex_payload["effort"], "low");
289    }
290
291    #[test]
292    fn rig_capabilities_preserve_reasoning_payload_for_gemini() {
293        let payload = RigProviderCapabilities::new(Provider::Gemini, "gemini-2.5-pro")
294            .reasoning_parameters(ReasoningEffortLevel::High)
295            .expect("reasoning payload");
296
297        assert_eq!(payload["thinking_config"]["thinkingBudget"], 256);
298        assert_eq!(payload["thinking_config"]["includeThoughts"], true);
299    }
300
301    #[test]
302    fn rig_capabilities_preserve_reasoning_payload_for_anthropic() {
303        assert!(
304            RigProviderCapabilities::new(Provider::Anthropic, "claude-sonnet-4-5")
305                .reasoning_parameters(ReasoningEffortLevel::High)
306                .is_none()
307        );
308    }
309
310    #[test]
311    fn rig_capabilities_preserve_reasoning_payload_for_deepseek() {
312        let high_payload = RigProviderCapabilities::new(Provider::DeepSeek, "deepseek-chat")
313            .reasoning_parameters(ReasoningEffortLevel::Medium)
314            .expect("reasoning payload");
315
316        assert_eq!(high_payload["thinking"]["type"], "enabled");
317        assert_eq!(high_payload["reasoning_effort"], "high");
318
319        let max_payload = RigProviderCapabilities::new(Provider::DeepSeek, "deepseek-chat")
320            .reasoning_parameters(ReasoningEffortLevel::XHigh)
321            .expect("reasoning payload");
322
323        assert_eq!(max_payload["thinking"]["type"], "enabled");
324        assert_eq!(max_payload["reasoning_effort"], "max");
325    }
326
327    #[test]
328    fn rig_capabilities_preserve_reasoning_payload_for_openrouter() {
329        assert!(
330            RigProviderCapabilities::new(Provider::OpenRouter, "openai/gpt-5")
331                .reasoning_parameters(ReasoningEffortLevel::High)
332                .is_none()
333        );
334    }
335
336    #[test]
337    fn rig_capabilities_skip_reasoning_payload_for_unsupported_provider() {
338        assert!(
339            RigProviderCapabilities::new(Provider::Ollama, "qwen")
340                .reasoning_parameters(ReasoningEffortLevel::High)
341                .is_none()
342        );
343    }
344}