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#[derive(Debug, Clone)]
11pub struct RigValidationSummary {
12 pub provider: Provider,
13 pub model: String,
14}
15
16#[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 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 }
56 Provider::Minimax => {
57 }
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 }
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 }
75 Provider::LlamaCpp => {
76 }
78 Provider::LmStudio => {
79 }
81 Provider::Moonshot => {
82 }
84 Provider::Mistral => {
85 }
87 Provider::ZAI => {
88 }
91 Provider::OpenCodeZen | Provider::OpenCodeGo => {
92 }
94 Provider::MiMo => {
95 }
97 Provider::Qwen => {
98 }
100 Provider::StepFun => {
101 }
103 Provider::Evolink => {
104 }
106 Provider::Poolside => {
107 }
109 }
110
111 Ok(RigValidationSummary {
112 provider: self.provider,
113 model: self.model.clone(),
114 })
115 }
116
117 #[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 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}