rig-core 0.44.0

An opinionated library for building LLM powered applications.
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

//! xAI's typed request options and reply extras, on Chat Completions and
//! Responses alike (<https://docs.x.ai/developers/cost-tracking>).
//!
//! ```
//! use rig_core::completion::CompletionRequest;
//! use rig_core::providers::xai::extension::{XaiOptions};
//!
//! let options = XaiOptions::new().prompt_cache_key("conversation-42");
//! let request = CompletionRequest::new("hi").provider_option(options);
//! # let _ = request;
//! ```

use serde::Serialize;
use serde_json::Value;

use crate::completion::provider_options::reply_field;
use crate::completion::{ExtensionOptions, ProviderExtension, ReplyExtras};
use crate::message::Api;

/// xAI's extension marker.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct XaiExt;

impl ProviderExtension for XaiExt {
    const PROVIDER: &'static str = super::PROVIDER_NAME;
    type Options = XaiOptions;
    type Extras = XaiExtras;
}

/// xAI's request options.
#[non_exhaustive]
#[derive(Clone, Debug, Default, PartialEq, Serialize)]
pub struct XaiOptions {
    /// The fields both routes take.
    #[serde(rename = "*")]
    pub shared: XaiShared,
}

/// The fields xAI takes on both routes.
#[non_exhaustive]
#[derive(Clone, Debug, Default, PartialEq, Serialize)]
pub struct XaiShared {
    /// Routes requests that share a prefix to the same cache.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prompt_cache_key: Option<String>,
}

impl XaiOptions {
    /// No option set.
    pub fn new() -> Self {
        Self::default()
    }

    /// Route the prompt cache by `key`.
    pub fn prompt_cache_key(mut self, key: impl Into<String>) -> Self {
        self.shared.prompt_cache_key = Some(key.into());
        self
    }
}

impl ExtensionOptions for XaiOptions {
    type Ext = XaiExt;
}

/// xAI's reply fields, from `usage` on both routes. Each is `None` when
/// the reply lacks it.
#[non_exhaustive]
#[derive(Clone, Debug, Default, PartialEq)]
pub struct XaiExtras {
    /// The request's cost in ticks of 10^-10 USD.
    pub cost_in_usd_ticks: Option<u64>,
    /// Search sources the request read.
    pub num_sources_used: Option<u64>,
    /// Server-side tool calls the request made.
    pub num_server_side_tools_used: Option<u64>,
    /// Server-side tool calls per tool, such as `web_search_calls`.
    pub server_side_tool_usage_details: Option<Value>,
}

impl ReplyExtras for XaiExtras {
    fn from_reply(_api: &Api, raw: &Value) -> Result<Self, serde_json::Error> {
        Ok(Self {
            cost_in_usd_ticks: reply_field(raw, "/usage/cost_in_usd_ticks")?,
            num_sources_used: reply_field(raw, "/usage/num_sources_used")?,
            num_server_side_tools_used: reply_field(raw, "/usage/num_server_side_tools_used")?,
            server_side_tool_usage_details: reply_field(
                raw,
                "/usage/server_side_tool_usage_details",
            )?,
        })
    }
}

#[cfg(test)]
mod tests;