rig_core/providers/xai/
extension.rs

1//! xAI's typed request options and reply extras, on Chat Completions and
2//! Responses alike (<https://docs.x.ai/developers/cost-tracking>).
3//!
4//! ```
5//! use rig_core::completion::CompletionRequest;
6//! use rig_core::providers::xai::extension::{XaiOptions};
7//!
8//! let options = XaiOptions::new().prompt_cache_key("conversation-42");
9//! let request = CompletionRequest::new("hi").provider_option(options);
10//! # let _ = request;
11//! ```
12
13use serde::Serialize;
14use serde_json::Value;
15
16use crate::completion::provider_options::reply_field;
17use crate::completion::{ExtensionOptions, ProviderExtension, ReplyExtras};
18use crate::message::Api;
19
20/// xAI's extension marker.
21#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
22pub struct XaiExt;
23
24impl ProviderExtension for XaiExt {
25    const PROVIDER: &'static str = super::PROVIDER_NAME;
26    type Options = XaiOptions;
27    type Extras = XaiExtras;
28}
29
30/// xAI's request options.
31#[non_exhaustive]
32#[derive(Clone, Debug, Default, PartialEq, Serialize)]
33pub struct XaiOptions {
34    /// The fields both routes take.
35    #[serde(rename = "*")]
36    pub shared: XaiShared,
37}
38
39/// The fields xAI takes on both routes.
40#[non_exhaustive]
41#[derive(Clone, Debug, Default, PartialEq, Serialize)]
42pub struct XaiShared {
43    /// Routes requests that share a prefix to the same cache.
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub prompt_cache_key: Option<String>,
46}
47
48impl XaiOptions {
49    /// No option set.
50    pub fn new() -> Self {
51        Self::default()
52    }
53
54    /// Route the prompt cache by `key`.
55    pub fn prompt_cache_key(mut self, key: impl Into<String>) -> Self {
56        self.shared.prompt_cache_key = Some(key.into());
57        self
58    }
59}
60
61impl ExtensionOptions for XaiOptions {
62    type Ext = XaiExt;
63}
64
65/// xAI's reply fields, from `usage` on both routes. Each is `None` when
66/// the reply lacks it.
67#[non_exhaustive]
68#[derive(Clone, Debug, Default, PartialEq)]
69pub struct XaiExtras {
70    /// The request's cost in ticks of 10^-10 USD.
71    pub cost_in_usd_ticks: Option<u64>,
72    /// Search sources the request read.
73    pub num_sources_used: Option<u64>,
74    /// Server-side tool calls the request made.
75    pub num_server_side_tools_used: Option<u64>,
76    /// Server-side tool calls per tool, such as `web_search_calls`.
77    pub server_side_tool_usage_details: Option<Value>,
78}
79
80impl ReplyExtras for XaiExtras {
81    fn from_reply(_api: &Api, raw: &Value) -> Result<Self, serde_json::Error> {
82        Ok(Self {
83            cost_in_usd_ticks: reply_field(raw, "/usage/cost_in_usd_ticks")?,
84            num_sources_used: reply_field(raw, "/usage/num_sources_used")?,
85            num_server_side_tools_used: reply_field(raw, "/usage/num_server_side_tools_used")?,
86            server_side_tool_usage_details: reply_field(
87                raw,
88                "/usage/server_side_tool_usage_details",
89            )?,
90        })
91    }
92}
93
94#[cfg(test)]
95mod tests;
rig_core/providers/xai/extension.rs

rig_core/providers/xai/
extension.rs
