Skip to main content

rig_core/providers/xai/
client.rs

1use crate::{
2    client::{
3        self, BearerAuth, Capabilities, Capable, DebugExt, Nothing, Provider, ProviderBuilder,
4        ProviderClient,
5    },
6    http_client,
7};
8
9#[derive(Debug, Default, Clone, Copy)]
10pub struct XAiExt;
11#[derive(Debug, Default, Clone, Copy)]
12pub struct XAiExtBuilder;
13
14type XAiApiKey = BearerAuth;
15
16pub type Client<H = reqwest::Client> = client::Client<XAiExt, H>;
17pub type ClientBuilder<H = crate::markers::Missing> =
18    client::ClientBuilder<XAiExtBuilder, XAiApiKey, H>;
19
20const XAI_BASE_URL: &str = "https://api.x.ai";
21
22impl Provider for XAiExt {
23    type Builder = XAiExtBuilder;
24
25    const VERIFY_PATH: &'static str = "/v1/api-key";
26}
27
28impl<H> Capabilities<H> for XAiExt {
29    type Completion = Capable<super::completion::CompletionModel<H>>;
30
31    type Embeddings = Nothing;
32    type Transcription = Nothing;
33    type ModelListing = Nothing;
34    #[cfg(feature = "image")]
35    type ImageGeneration = Capable<super::image_generation::ImageGenerationModel<H>>;
36    #[cfg(feature = "audio")]
37    type AudioGeneration = Capable<super::audio_generation::AudioGenerationModel<H>>;
38    type Rerank = Nothing;
39}
40
41impl DebugExt for XAiExt {}
42
43impl ProviderBuilder for XAiExtBuilder {
44    type Extension<H>
45        = XAiExt
46    where
47        H: http_client::HttpClientExt;
48    type ApiKey = XAiApiKey;
49
50    const BASE_URL: &'static str = XAI_BASE_URL;
51
52    fn build<H>(
53        _builder: &client::ClientBuilder<Self, Self::ApiKey, H>,
54    ) -> http_client::Result<Self::Extension<H>>
55    where
56        H: http_client::HttpClientExt,
57    {
58        Ok(XAiExt)
59    }
60}
61
62impl ProviderClient for Client {
63    type Input = String;
64    type Error = crate::client::ProviderClientError;
65
66    /// Create a new xAI client from the `XAI_API_KEY` environment variable.
67    fn from_env() -> Result<Self, Self::Error> {
68        let api_key = crate::client::required_env_var("XAI_API_KEY")?;
69        Self::new(&api_key).map_err(Into::into)
70    }
71
72    fn from_val(input: Self::Input) -> Result<Self, Self::Error> {
73        Self::new(&input).map_err(Into::into)
74    }
75}
76#[cfg(test)]
77mod tests {
78    #[test]
79    fn test_client_initialization() {
80        let _client_from_builder = crate::providers::xai::Client::builder()
81            .api_key("dummy-key")
82            .build()
83            .expect("Client::builder() failed");
84    }
85}