Skip to main content

rig_core/providers/
moonshot.rs

1//! Moonshot AI (Kimi) API client and Rig integration
2//!
3//! # Example
4//! ```no_run
5//! use rig_core::providers::moonshot;
6//! use rig_core::client::CompletionClient;
7//!
8//! let client = moonshot::Client::new("YOUR_API_KEY").expect("Failed to build client");
9//!
10//! let kimi_model = client.completion_model(moonshot::KIMI_K2_5);
11//! ```
12//!
13//! # Custom base URL
14//! The default base URL is `https://api.moonshot.ai/v1`. For China access,
15//! use `https://api.moonshot.cn/v1`:
16//! ```no_run
17//! use rig_core::providers::moonshot;
18//!
19//! let client = moonshot::Client::builder()
20//!     .api_key("YOUR_API_KEY")
21//!     .base_url("https://api.moonshot.ai/v1")
22//!     .build()
23//!     .expect("Failed to build Moonshot client");
24//! ```
25use crate::client::{
26    self, BearerAuth, Capabilities, Capable, DebugExt, Nothing, Provider, ProviderBuilder,
27    ProviderClient,
28};
29use crate::http_client;
30use crate::http_client::HttpClientExt;
31use crate::providers::anthropic::client::{
32    AnthropicBuilder as AnthropicCompatBuilder, AnthropicKey, finish_anthropic_builder,
33};
34use crate::{completion::CompletionError, providers::openai};
35
36// ================================================================
37// Main Moonshot Client
38// ================================================================
39/// Global OpenAI-compatible base URL.
40pub const GLOBAL_API_BASE_URL: &str = "https://api.moonshot.ai/v1";
41/// China OpenAI-compatible base URL.
42pub const CHINA_API_BASE_URL: &str = "https://api.moonshot.cn/v1";
43/// Anthropic-compatible base URL.
44pub const ANTHROPIC_API_BASE_URL: &str = "https://api.moonshot.ai/anthropic";
45
46#[derive(Debug, Default, Clone, Copy)]
47pub struct MoonshotExt;
48#[derive(Debug, Default, Clone, Copy)]
49pub struct MoonshotBuilder;
50#[derive(Debug, Default, Clone)]
51pub struct MoonshotAnthropicBuilder {
52    anthropic: AnthropicCompatBuilder,
53}
54#[derive(Debug, Default, Clone, Copy)]
55pub struct MoonshotAnthropicExt;
56
57type MoonshotApiKey = BearerAuth;
58
59impl Provider for MoonshotExt {
60    type Builder = MoonshotBuilder;
61
62    const VERIFY_PATH: &'static str = "/models";
63}
64
65impl Provider for MoonshotAnthropicExt {
66    type Builder = MoonshotAnthropicBuilder;
67
68    const VERIFY_PATH: &'static str = "/v1/models";
69}
70
71impl DebugExt for MoonshotExt {}
72impl DebugExt for MoonshotAnthropicExt {}
73
74impl ProviderBuilder for MoonshotBuilder {
75    type Extension<H>
76        = MoonshotExt
77    where
78        H: HttpClientExt;
79    type ApiKey = MoonshotApiKey;
80
81    const BASE_URL: &'static str = GLOBAL_API_BASE_URL;
82
83    fn build<H>(
84        _builder: &crate::client::ClientBuilder<Self, Self::ApiKey, H>,
85    ) -> http_client::Result<Self::Extension<H>>
86    where
87        H: HttpClientExt,
88    {
89        Ok(MoonshotExt)
90    }
91}
92
93impl ProviderBuilder for MoonshotAnthropicBuilder {
94    type Extension<H>
95        = MoonshotAnthropicExt
96    where
97        H: HttpClientExt;
98    type ApiKey = AnthropicKey;
99
100    const BASE_URL: &'static str = ANTHROPIC_API_BASE_URL;
101
102    fn build<H>(
103        _builder: &crate::client::ClientBuilder<Self, Self::ApiKey, H>,
104    ) -> http_client::Result<Self::Extension<H>>
105    where
106        H: HttpClientExt,
107    {
108        Ok(MoonshotAnthropicExt)
109    }
110
111    fn finish<H>(
112        &self,
113        builder: client::ClientBuilder<Self, AnthropicKey, H>,
114    ) -> http_client::Result<client::ClientBuilder<Self, AnthropicKey, H>> {
115        finish_anthropic_builder(&self.anthropic, builder)
116    }
117}
118
119impl<H> Capabilities<H> for MoonshotExt {
120    type Completion = Capable<CompletionModel<H>>;
121    type Embeddings = Nothing;
122    type Transcription = Nothing;
123    type ModelListing = Nothing;
124    #[cfg(feature = "image")]
125    type ImageGeneration = Nothing;
126    #[cfg(feature = "audio")]
127    type AudioGeneration = Nothing;
128    type Rerank = Nothing;
129}
130
131impl<H> Capabilities<H> for MoonshotAnthropicExt {
132    type Completion =
133        Capable<super::anthropic::completion::GenericCompletionModel<MoonshotAnthropicExt, H>>;
134    type Embeddings = Nothing;
135    type Transcription = Nothing;
136    type ModelListing = Nothing;
137    #[cfg(feature = "image")]
138    type ImageGeneration = Nothing;
139    #[cfg(feature = "audio")]
140    type AudioGeneration = Nothing;
141    type Rerank = Nothing;
142}
143
144pub type Client<H = reqwest::Client> = client::Client<MoonshotExt, H>;
145pub type ClientBuilder<H = crate::markers::Missing> =
146    client::ClientBuilder<MoonshotBuilder, MoonshotApiKey, H>;
147pub type AnthropicClient<H = reqwest::Client> = client::Client<MoonshotAnthropicExt, H>;
148pub type AnthropicClientBuilder<H = crate::markers::Missing> =
149    client::ClientBuilder<MoonshotAnthropicBuilder, AnthropicKey, H>;
150
151impl ProviderClient for Client {
152    type Input = String;
153    type Error = crate::client::ProviderClientError;
154
155    /// Create a new Moonshot client from the `MOONSHOT_API_KEY` environment variable.
156    fn from_env() -> Result<Self, Self::Error> {
157        let api_key = crate::client::required_env_var("MOONSHOT_API_KEY")?;
158        let mut builder = Self::builder().api_key(&api_key);
159        if let Some(base_url) = crate::client::optional_env_var("MOONSHOT_API_BASE")? {
160            builder = builder.base_url(base_url);
161        }
162        builder.build().map_err(Into::into)
163    }
164
165    fn from_val(input: Self::Input) -> Result<Self, Self::Error> {
166        Self::new(&input).map_err(Into::into)
167    }
168}
169
170impl ProviderClient for AnthropicClient {
171    type Input = String;
172    type Error = crate::client::ProviderClientError;
173
174    fn from_env() -> Result<Self, Self::Error> {
175        let api_key = crate::client::required_env_var("MOONSHOT_API_KEY")?;
176        let mut builder = Self::builder().api_key(api_key);
177        if let Some(base_url) =
178            anthropic_base_override("MOONSHOT_ANTHROPIC_API_BASE", "MOONSHOT_API_BASE")?
179        {
180            builder = builder.base_url(base_url);
181        }
182        builder.build().map_err(Into::into)
183    }
184
185    fn from_val(input: Self::Input) -> Result<Self, Self::Error> {
186        Self::builder().api_key(input).build().map_err(Into::into)
187    }
188}
189
190impl<H> ClientBuilder<H> {
191    pub fn global(self) -> Self {
192        self.base_url(GLOBAL_API_BASE_URL)
193    }
194
195    pub fn china(self) -> Self {
196        self.base_url(CHINA_API_BASE_URL)
197    }
198}
199
200impl<H> AnthropicClientBuilder<H> {
201    pub fn global(self) -> Self {
202        self.base_url(ANTHROPIC_API_BASE_URL)
203    }
204
205    pub fn anthropic_version(self, anthropic_version: &str) -> Self {
206        self.over_ext(|mut ext| {
207            ext.anthropic.anthropic_version = anthropic_version.into();
208            ext
209        })
210    }
211
212    pub fn anthropic_betas(self, anthropic_betas: &[&str]) -> Self {
213        self.over_ext(|mut ext| {
214            ext.anthropic
215                .anthropic_betas
216                .extend(anthropic_betas.iter().copied().map(String::from));
217            ext
218        })
219    }
220
221    pub fn anthropic_beta(self, anthropic_beta: &str) -> Self {
222        self.over_ext(|mut ext| {
223            ext.anthropic.anthropic_betas.push(anthropic_beta.into());
224            ext
225        })
226    }
227}
228
229impl super::anthropic::completion::AnthropicCompatibleProvider for MoonshotAnthropicExt {
230    const PROVIDER_NAME: &'static str = "moonshot";
231
232    fn default_max_tokens(_model: &str) -> Option<u64> {
233        Some(4096)
234    }
235}
236
237fn anthropic_base_override(
238    primary_env: &'static str,
239    fallback_env: &'static str,
240) -> crate::client::ProviderClientResult<Option<String>> {
241    let primary = crate::client::optional_env_var(primary_env)?;
242    let fallback = crate::client::optional_env_var(fallback_env)?;
243
244    Ok(resolve_anthropic_base_override(
245        primary.as_deref(),
246        fallback.as_deref(),
247    ))
248}
249
250fn resolve_anthropic_base_override(
251    primary: Option<&str>,
252    fallback: Option<&str>,
253) -> Option<String> {
254    primary
255        .map(str::to_owned)
256        .or_else(|| fallback.and_then(normalize_anthropic_base_url))
257}
258
259fn normalize_anthropic_base_url(base_url: &str) -> Option<String> {
260    if base_url.contains("/anthropic") {
261        return Some(base_url.to_owned());
262    }
263
264    let mut url = url::Url::parse(base_url).ok()?;
265    if !matches!(url.path(), "/v1" | "/v1/") {
266        return None;
267    }
268    url.set_path("/anthropic");
269    Some(url.to_string())
270}
271
272// ================================================================
273// Moonshot Completion API
274// ================================================================
275
276/// Moonshot v1 128K context model (legacy)
277pub const MOONSHOT_CHAT: &str = "moonshot-v1-128k";
278
279/// Kimi K2 — Mixture-of-Experts model (1T total params, 32B active)
280pub const KIMI_K2: &str = "kimi-k2";
281
282/// Kimi K2.5 — Native multimodal agentic model with 256K context
283pub const KIMI_K2_5: &str = "kimi-k2.5";
284
285/// Moonshot completion model, driven by the shared OpenAI Chat Completions path.
286pub type CompletionModel<H = reqwest::Client> =
287    openai::completion::GenericCompletionModel<MoonshotExt, H>;
288
289impl openai::completion::OpenAICompatibleProvider for MoonshotExt {
290    const PROVIDER_NAME: &'static str = "moonshot";
291
292    type StreamingUsage = openai::Usage;
293
294    // Moonshot's API rejects the `json_schema` response format; keep the
295    // pre-migration behavior of dropping `output_schema` with a warning.
296    const SUPPORTS_RESPONSE_FORMAT: bool = false;
297
298    type Response = openai::CompletionResponse;
299
300    fn prepare_request(
301        &self,
302        request: &mut openai::completion::CompletionRequest,
303    ) -> Result<(), CompletionError> {
304        // Moonshot only supports `auto`/`none` tool choices. Forcing one
305        // specific tool has no workaround; fail fast like the pre-migration
306        // conversion did (on main, `openai::ToolChoice::try_from` returned
307        // "Provider doesn't support only using specific tools" for every
308        // `ToolChoice::Specific`, single- or multi-name).
309        if matches!(
310            request.tool_choice,
311            Some(openai::completion::ToolChoice::Function { .. })
312        ) {
313            return Err(CompletionError::ProviderError(
314                "Moonshot does not support forcing a specific tool".to_string(),
315            ));
316        }
317
318        // Moonshot does not support `tool_choice: "required"`; coerce it to
319        // `auto` and steer the model with an extra user message instead.
320        if matches!(
321            request.tool_choice,
322            Some(openai::completion::ToolChoice::Required)
323        ) {
324            tracing::warn!(
325                "Moonshot does not support tool_choice=required; coercing to auto with an additional steering message"
326            );
327            request.tool_choice = Some(openai::completion::ToolChoice::Auto);
328            request.messages.push(openai::Message::User {
329                content: crate::OneOrMany::one(openai::UserContent::Text {
330                    text: "Please select a tool to handle the current issue.".to_string(),
331                }),
332                name: None,
333            });
334        }
335
336        Ok(())
337    }
338}
339
340#[cfg(test)]
341mod tests {
342    use super::{MoonshotExt, normalize_anthropic_base_url, resolve_anthropic_base_override};
343    use crate::completion::CompletionRequest;
344    use crate::message::{
345        AssistantContent, Message, Reasoning, ToolCall, ToolChoice, ToolFunction,
346    };
347    use crate::providers::openai::completion::{
348        CompletionRequest as OpenAICompletionRequest, OpenAICompatibleProvider, OpenAIRequestParams,
349    };
350
351    fn prepared_body(request: CompletionRequest, model: &str) -> serde_json::Value {
352        let mut request = OpenAICompletionRequest::try_from(OpenAIRequestParams {
353            model: model.to_string(),
354            request,
355            strict_tools: false,
356            tool_result_array_content: false,
357            supports_response_format: MoonshotExt::SUPPORTS_RESPONSE_FORMAT,
358            supports_tools: true,
359        })
360        .expect("request should convert");
361        MoonshotExt
362            .prepare_request(&mut request)
363            .expect("prepare_request should succeed");
364        serde_json::to_value(request).expect("request should serialize")
365    }
366
367    #[test]
368    fn test_client_initialization() {
369        let _client =
370            crate::providers::moonshot::Client::new("dummy-key").expect("Client::new() failed");
371        let _client_from_builder = crate::providers::moonshot::Client::builder()
372            .api_key("dummy-key")
373            .build()
374            .expect("Client::builder() failed");
375        let _anthropic_client = crate::providers::moonshot::AnthropicClient::new("dummy-key")
376            .expect("AnthropicClient::new() failed");
377        let _anthropic_client_from_builder = crate::providers::moonshot::AnthropicClient::builder()
378            .api_key("dummy-key")
379            .build()
380            .expect("AnthropicClient::builder() failed");
381    }
382
383    #[test]
384    fn moonshot_preserves_reasoning_content_in_assistant_history() {
385        let assistant = Message::Assistant {
386            id: None,
387            content: crate::OneOrMany::many(vec![
388                AssistantContent::Reasoning(Reasoning::new("tool planning")),
389                AssistantContent::ToolCall(ToolCall {
390                    id: "call_1".to_string(),
391                    call_id: None,
392                    function: ToolFunction {
393                        name: "lookup".to_string(),
394                        arguments: serde_json::json!({}),
395                    },
396                    signature: None,
397                    additional_params: None,
398                }),
399            ])
400            .expect("assistant content"),
401        };
402
403        let request = CompletionRequest {
404            model: Some("kimi-k2-thinking".to_string()),
405            preamble: None,
406            chat_history: crate::OneOrMany::one(assistant),
407            documents: vec![],
408            tools: vec![],
409            temperature: None,
410            max_tokens: None,
411            tool_choice: None,
412            additional_params: None,
413            output_schema: None,
414        };
415
416        let body = prepared_body(request, "kimi-k2-thinking");
417        assert_eq!(
418            body["messages"][0]["reasoning_content"],
419            serde_json::json!("tool planning")
420        );
421    }
422
423    #[test]
424    fn moonshot_joins_multiple_reasoning_blocks_with_newline() {
425        // A replayed assistant turn carrying two distinct reasoning blocks must
426        // keep them newline-separated on the wire, not glued together.
427        let assistant = Message::Assistant {
428            id: None,
429            content: crate::OneOrMany::many(vec![
430                AssistantContent::Reasoning(Reasoning::new("first thought")),
431                AssistantContent::Reasoning(Reasoning::new("second thought")),
432                AssistantContent::Text("done".into()),
433            ])
434            .expect("assistant content"),
435        };
436
437        let request = CompletionRequest {
438            model: Some("kimi-k2-thinking".to_string()),
439            preamble: None,
440            chat_history: crate::OneOrMany::one(assistant),
441            documents: vec![],
442            tools: vec![],
443            temperature: None,
444            max_tokens: None,
445            tool_choice: None,
446            additional_params: None,
447            output_schema: None,
448        };
449
450        let body = prepared_body(request, "kimi-k2-thinking");
451        assert_eq!(
452            body["messages"][0]["reasoning_content"],
453            serde_json::json!("first thought\nsecond thought")
454        );
455    }
456
457    #[test]
458    fn moonshot_specific_tool_choice_is_rejected() {
459        let request = CompletionRequest {
460            model: Some("kimi-k2.5".to_string()),
461            preamble: None,
462            chat_history: crate::OneOrMany::one(Message::user("Use a tool.")),
463            documents: vec![],
464            tools: vec![],
465            temperature: None,
466            max_tokens: None,
467            tool_choice: Some(ToolChoice::Specific {
468                function_names: vec!["lookup".to_string()],
469            }),
470            additional_params: None,
471            output_schema: None,
472        };
473
474        let mut request = OpenAICompletionRequest::try_from(OpenAIRequestParams {
475            model: "kimi-k2.5".to_string(),
476            request,
477            strict_tools: false,
478            tool_result_array_content: false,
479            supports_response_format: MoonshotExt::SUPPORTS_RESPONSE_FORMAT,
480            supports_tools: true,
481        })
482        .expect("request should convert");
483
484        let error = MoonshotExt
485            .prepare_request(&mut request)
486            .expect_err("specific tool choice should be rejected");
487        assert!(error.to_string().contains("specific tool"));
488    }
489
490    #[test]
491    fn moonshot_required_tool_choice_is_coerced() {
492        let request = CompletionRequest {
493            model: Some("kimi-k2.5".to_string()),
494            preamble: None,
495            chat_history: crate::OneOrMany::one(Message::user("Use a tool.")),
496            documents: vec![],
497            tools: vec![],
498            temperature: None,
499            max_tokens: None,
500            tool_choice: Some(ToolChoice::Required),
501            additional_params: None,
502            output_schema: None,
503        };
504
505        let body = prepared_body(request, "kimi-k2.5");
506        assert_eq!(body["tool_choice"], "auto");
507        assert_eq!(
508            body["messages"]
509                .as_array()
510                .and_then(|messages| messages.last())
511                .and_then(|message| message.get("content"))
512                .and_then(|content| content.as_str()),
513            Some("Please select a tool to handle the current issue.")
514        );
515    }
516
517    #[test]
518    fn normalize_openai_style_base_to_anthropic_base() {
519        assert_eq!(
520            normalize_anthropic_base_url("https://api.moonshot.ai/v1").as_deref(),
521            Some("https://api.moonshot.ai/anthropic")
522        );
523        assert_eq!(
524            normalize_anthropic_base_url("https://api.moonshot.cn/v1").as_deref(),
525            Some("https://api.moonshot.cn/anthropic")
526        );
527        assert_eq!(
528            normalize_anthropic_base_url("https://proxy.example.com/v1").as_deref(),
529            Some("https://proxy.example.com/anthropic")
530        );
531    }
532
533    #[test]
534    fn normalize_preserves_existing_anthropic_base() {
535        assert_eq!(
536            normalize_anthropic_base_url("https://proxy.example.com/anthropic").as_deref(),
537            Some("https://proxy.example.com/anthropic")
538        );
539    }
540
541    #[test]
542    fn anthropic_primary_override_wins() {
543        let override_url = resolve_anthropic_base_override(
544            Some("https://primary.example.com/anthropic"),
545            Some("https://api.moonshot.cn/v1"),
546        );
547
548        assert_eq!(
549            override_url.as_deref(),
550            Some("https://primary.example.com/anthropic")
551        );
552    }
553}