qai-sdk 0.1.6

Universal Rust SDK for AI Providers
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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

//! # QAI xAI
//!
//! xAI Grok provider for the QAI SDK. Provides access to Grok models
//! through an OpenAI-compatible API adapter, with additional support for
//! image generation and the Responses API.
//!
//! ## Usage
//!
//! ```rust,no_run
//! use qai_sdk::xai::create_xai;
//! use qai_sdk::core::types::ProviderSettings;
//!
//! let provider = create_xai(ProviderSettings {
//!     api_key: Some("xai-...".to_string()),
//!     ..Default::default()
//! });
//!
//! let model = provider.chat("grok-2");
//! ```

pub mod error;
pub mod image;
pub mod responses;
pub mod tools;
pub mod types;

use crate::core::types::{GenerateOptions, GenerateResult, Prompt, ProviderSettings, StreamPart};
use crate::openai::OpenAIModel;
use async_trait::async_trait;
use futures::stream::BoxStream;
use reqwest::Client;

pub struct XAIModel {
    pub inner: OpenAIModel,
}

impl XAIModel {
    #[must_use]
    pub fn new(api_key: String) -> Self {
        Self {
            inner: OpenAIModel {
                api_key,
                base_url: "https://api.x.ai/v1".to_string(),
                client: Client::new(),
            },
        }
    }
}

#[async_trait]
impl crate::core::LanguageModel for XAIModel {
    #[tracing::instrument(skip(self, prompt), fields(model = options.model_id))]
    async fn generate(
        &self,
        prompt: Prompt,
        options: GenerateOptions,
    ) -> crate::core::Result<GenerateResult> {
        self.inner.generate(prompt, options).await
    }

    async fn generate_stream(
        &self,
        prompt: Prompt,
        options: GenerateOptions,
    ) -> crate::core::Result<BoxStream<'static, StreamPart>> {
        self.inner.generate_stream(prompt, options).await
    }
}

// --- Provider Factory ---

/// xAI provider with configurable settings.
pub struct XAIProvider {
    settings: ProviderSettings,
}

impl XAIProvider {
    /// Creates a chat language model.
    #[must_use]
    pub fn chat(&self, _model_id: &str) -> XAIModel {
        let api_key = self
            .settings
            .api_key
            .clone()
            .or_else(|| std::env::var("XAI_API_KEY").ok())
            .unwrap_or_default();
        let base_url = self
            .settings
            .base_url
            .clone()
            .unwrap_or_else(|| "https://api.x.ai/v1".to_string());
        XAIModel {
            inner: OpenAIModel {
                api_key,
                base_url,
                client: Client::new(),
            },
        }
    }

    /// Alias for `chat`.
    #[must_use]
    pub fn language_model(&self, model_id: &str) -> XAIModel {
        self.chat(model_id)
    }

    /// Creates an image generation model.
    #[must_use]
    pub fn image(&self, _model_id: &str) -> image::XaiImageModel {
        let api_key = self
            .settings
            .api_key
            .clone()
            .or_else(|| std::env::var("XAI_API_KEY").ok())
            .unwrap_or_default();
        let base_url = self
            .settings
            .base_url
            .clone()
            .unwrap_or_else(|| "https://api.x.ai/v1".to_string());
        image::XaiImageModel {
            api_key,
            base_url,
            client: Client::new(),
        }
    }

    /// Creates a Responses API model.
    #[must_use]
    pub fn responses(&self, _model_id: &str) -> responses::XaiResponsesModel {
        let api_key = self
            .settings
            .api_key
            .clone()
            .or_else(|| std::env::var("XAI_API_KEY").ok())
            .unwrap_or_default();
        responses::XaiResponsesModel::new(api_key)
    }
}

/// Create an xAI provider instance with the given settings.
#[must_use]
pub fn create_xai(settings: ProviderSettings) -> XAIProvider {
    XAIProvider { settings }
}

impl crate::core::registry::Provider for XAIProvider {
    fn language_model(&self, model_id: &str) -> Option<Box<dyn crate::core::LanguageModel>> {
        Some(Box::new(self.chat(model_id)))
    }

    fn image_model(&self, model_id: &str) -> Option<Box<dyn crate::core::ImageModel>> {
        Some(Box::new(self.image(model_id)))
    }
}