rig_core/providers/zai/
extension.rs

1//! Z.AI's typed request options and reply extras. One entry serves its
2//! OpenAI-format API (<https://docs.z.ai/api-reference/llm/chat-completion>)
3//! and its Anthropic-format Messages API, which takes no Z.AI field of its
4//! own.
5//!
6//! ```
7//! use rig_core::completion::CompletionRequest;
8//! use rig_core::providers::zai::extension::ZaiOptions;
9//!
10//! let options = ZaiOptions::new().user_id("user-1");
11//! let request = CompletionRequest::new("hi").provider_option(options);
12//! # let _ = request;
13//! ```
14
15use serde::Serialize;
16use serde_json::Value;
17
18use crate::completion::provider_options::reply_field;
19use crate::completion::{ExtensionOptions, ProviderExtension, ReplyExtras};
20use crate::message::Api;
21use crate::providers::anthropic::extension::MessagesStop;
22use crate::providers::anthropic::wire::MESSAGES_API;
23
24/// Z.AI's extension marker.
25#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
26pub struct ZaiExt;
27
28impl ProviderExtension for ZaiExt {
29    const PROVIDER: &'static str = super::PROVIDER_NAME;
30    type Options = ZaiOptions;
31    type Extras = ZaiExtras;
32}
33
34/// Z.AI's request options, by route. The field setters here write the
35/// same field as [`chat`](Self::chat), which replaces every field set
36/// before it; the Messages API takes none of them.
37#[non_exhaustive]
38#[derive(Clone, Debug, Default, PartialEq, Serialize)]
39pub struct ZaiOptions {
40    /// Sent on the OpenAI-format API only.
41    #[serde(rename = "openai.chat")]
42    pub chat: ZaiChat,
43}
44
45impl ZaiOptions {
46    /// No option set.
47    pub fn new() -> Self {
48        Self::default()
49    }
50
51    /// Send `chat` on the OpenAI-format API.
52    pub fn chat(mut self, chat: ZaiChat) -> Self {
53        self.chat = chat;
54        self
55    }
56
57    /// Apply `set` to the OpenAI-format section.
58    fn with_chat(mut self, set: impl FnOnce(ZaiChat) -> ZaiChat) -> Self {
59        self.chat = set(std::mem::take(&mut self.chat));
60        self
61    }
62
63    /// Whether to sample, on the OpenAI-format API, as
64    /// [`ZaiChat::do_sample`].
65    pub fn do_sample(self, sample: bool) -> Self {
66        self.with_chat(|chat| chat.do_sample(sample))
67    }
68
69    /// Identify the request as `id`, on the OpenAI-format API, as
70    /// [`ZaiChat::request_id`].
71    pub fn request_id(self, id: impl Into<String>) -> Self {
72        self.with_chat(|chat| chat.request_id(id))
73    }
74
75    /// Name the end user, on the OpenAI-format API, as
76    /// [`ZaiChat::user_id`].
77    pub fn user_id(self, id: impl Into<String>) -> Self {
78        self.with_chat(|chat| chat.user_id(id))
79    }
80
81    /// Whether earlier turns' thinking is dropped from context, on the
82    /// OpenAI-format API, as [`ZaiChat::clear_thinking`].
83    pub fn clear_thinking(self, clear: bool) -> Self {
84        self.with_chat(|chat| chat.clear_thinking(clear))
85    }
86}
87
88impl ExtensionOptions for ZaiOptions {
89    type Ext = ZaiExt;
90}
91
92/// The fields Z.AI's OpenAI-format API takes.
93#[non_exhaustive]
94#[derive(Clone, Debug, Default, PartialEq, Serialize)]
95pub struct ZaiChat {
96    /// Whether to sample; `false` decodes greedily.
97    #[serde(skip_serializing_if = "Option::is_none")]
98    pub do_sample: Option<bool>,
99    /// The caller's id for the request.
100    #[serde(skip_serializing_if = "Option::is_none")]
101    pub request_id: Option<String>,
102    /// A stable id of the end user.
103    #[serde(skip_serializing_if = "Option::is_none")]
104    pub user_id: Option<String>,
105    /// Thinking fields sent beside the mapped thinking type.
106    #[serde(skip_serializing_if = "ZaiThinking::is_empty")]
107    pub thinking: ZaiThinking,
108}
109
110/// The `thinking` fields Z.AI takes beside the type the generation options
111/// map.
112#[non_exhaustive]
113#[derive(Clone, Debug, Default, PartialEq, Serialize)]
114pub struct ZaiThinking {
115    /// Whether earlier turns' thinking is dropped from context.
116    #[serde(skip_serializing_if = "Option::is_none")]
117    pub clear_thinking: Option<bool>,
118}
119
120impl ZaiThinking {
121    fn is_empty(&self) -> bool {
122        self.clear_thinking.is_none()
123    }
124}
125
126impl ZaiChat {
127    /// No option set.
128    pub fn new() -> Self {
129        Self::default()
130    }
131
132    /// Whether to sample.
133    pub fn do_sample(mut self, sample: bool) -> Self {
134        self.do_sample = Some(sample);
135        self
136    }
137
138    /// Identify the request as `id`.
139    pub fn request_id(mut self, id: impl Into<String>) -> Self {
140        self.request_id = Some(id.into());
141        self
142    }
143
144    /// Name the end user.
145    pub fn user_id(mut self, id: impl Into<String>) -> Self {
146        self.user_id = Some(id.into());
147        self
148    }
149
150    /// Whether earlier turns' thinking is dropped from context.
151    pub fn clear_thinking(mut self, clear: bool) -> Self {
152        self.thinking.clear_thinking = Some(clear);
153        self
154    }
155}
156
157/// Z.AI's reply fields. Each is `None` when the reply lacks it.
158#[non_exhaustive]
159#[derive(Clone, Debug, Default, PartialEq)]
160pub struct ZaiExtras {
161    /// The request's id. OpenAI-format API only.
162    pub request_id: Option<String>,
163    /// `stop_reason`, verbatim. Messages route.
164    pub stop_reason: Option<String>,
165    /// `stop_sequence`: the stop sequence the turn ended on. Messages route.
166    pub stop_sequence: Option<String>,
167}
168
169impl ReplyExtras for ZaiExtras {
170    fn from_reply(api: &Api, raw: &Value) -> Result<Self, serde_json::Error> {
171        if api.as_str() == MESSAGES_API {
172            let MessagesStop {
173                stop_reason,
174                stop_sequence,
175            } = MessagesStop::read("Z.AI", api, raw)?;
176            return Ok(Self {
177                stop_reason,
178                stop_sequence,
179                ..Self::default()
180            });
181        }
182        if api.as_str() != "openai.chat" {
183            return Ok(Self::default());
184        }
185        Ok(Self {
186            request_id: reply_field(raw, "/request_id")?,
187            ..Self::default()
188        })
189    }
190}
191
192#[cfg(test)]
193mod tests;
rig_core/providers/zai/extension.rs

rig_core/providers/zai/
extension.rs
