rig_core/providers/venice/
extension.rs

1//! Venice's typed request options and reply extras. Turning reasoning off
2//! is [`Reasoning::Off`](crate::completion::Reasoning::Off), not a Venice
3//! parameter.
4//!
5//! ```
6//! use rig_core::completion::CompletionRequest;
7//! use rig_core::providers::venice::extension::{VeniceOptions, WebSearchMode};
8//!
9//! let options = VeniceOptions::new()
10//!     .enable_web_search(WebSearchMode::On)
11//!     .enable_web_citations(true);
12//! let request = CompletionRequest::new("Summarize today's Rust news.")
13//!     .provider_option(options);
14//! # let _ = request;
15//! ```
16
17use serde::{Deserialize, Serialize};
18use serde_json::Value;
19
20use crate::completion::provider_options::reply_field;
21use crate::completion::{ExtensionOptions, ProviderExtension, ReplyExtras};
22use crate::message::Api;
23
24/// Venice's extension marker.
25#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
26pub struct VeniceExt;
27
28impl ProviderExtension for VeniceExt {
29    const PROVIDER: &'static str = super::PROVIDER_NAME;
30    type Options = VeniceOptions;
31    type Extras = VeniceExtras;
32}
33
34/// Venice's request options. The parameter setters here write the same
35/// `venice_parameters` entry as
36/// [`venice_parameters`](Self::venice_parameters), which replaces every
37/// entry set before it.
38#[non_exhaustive]
39#[derive(Clone, Debug, Default, PartialEq, Serialize)]
40pub struct VeniceOptions {
41    /// The fields every route takes.
42    #[serde(rename = "*")]
43    pub shared: VeniceShared,
44}
45
46/// The fields Venice takes.
47#[non_exhaustive]
48#[derive(Clone, Debug, Default, PartialEq, Serialize)]
49pub struct VeniceShared {
50    /// Venice's own parameter block.
51    #[serde(skip_serializing_if = "Option::is_none")]
52    pub venice_parameters: Option<VeniceParameters>,
53    /// The prompt-cache routing key.
54    #[serde(skip_serializing_if = "Option::is_none")]
55    pub prompt_cache_key: Option<String>,
56}
57
58impl VeniceOptions {
59    /// No option set.
60    pub fn new() -> Self {
61        Self::default()
62    }
63
64    /// Send Venice's parameter block `parameters`.
65    pub fn venice_parameters(mut self, parameters: VeniceParameters) -> Self {
66        self.shared.venice_parameters = Some(parameters);
67        self
68    }
69
70    /// Route the prompt cache by `key`.
71    pub fn prompt_cache_key(mut self, key: impl Into<String>) -> Self {
72        self.shared.prompt_cache_key = Some(key.into());
73        self
74    }
75
76    /// Apply `set` to the parameter block, starting from an empty one.
77    fn with_parameters(mut self, set: impl FnOnce(VeniceParameters) -> VeniceParameters) -> Self {
78        let parameters = self.shared.venice_parameters.take().unwrap_or_default();
79        self.shared.venice_parameters = Some(set(parameters));
80        self
81    }
82
83    /// Converse with the public character `slug`, as
84    /// [`VeniceParameters::character_slug`].
85    pub fn character_slug(self, slug: impl Into<String>) -> Self {
86        self.with_parameters(|parameters| parameters.character_slug(slug))
87    }
88
89    /// Strip `<think>` blocks from the reply, as
90    /// [`VeniceParameters::strip_thinking_response`].
91    pub fn strip_thinking_response(self, strip: bool) -> Self {
92        self.with_parameters(|parameters| parameters.strip_thinking_response(strip))
93    }
94
95    /// Set the web-search mode, as [`VeniceParameters::enable_web_search`].
96    pub fn enable_web_search(self, mode: WebSearchMode) -> Self {
97        self.with_parameters(|parameters| parameters.enable_web_search(mode))
98    }
99
100    /// Cite sources as `[REF]` markers, as
101    /// [`VeniceParameters::enable_web_citations`].
102    pub fn enable_web_citations(self, enable: bool) -> Self {
103        self.with_parameters(|parameters| parameters.enable_web_citations(enable))
104    }
105
106    /// Include Venice's default system prompt, as
107    /// [`VeniceParameters::include_venice_system_prompt`].
108    pub fn include_venice_system_prompt(self, include: bool) -> Self {
109        self.with_parameters(|parameters| parameters.include_venice_system_prompt(include))
110    }
111}
112
113impl ExtensionOptions for VeniceOptions {
114    type Ext = VeniceExt;
115}
116
117/// How Venice's web search behaves for a request.
118#[non_exhaustive]
119#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
120#[serde(rename_all = "lowercase")]
121pub enum WebSearchMode {
122    /// Never search.
123    Off,
124    /// Always search.
125    On,
126    /// Let the model decide.
127    Auto,
128}
129
130/// Venice's `venice_parameters` block. Unset fields keep Venice's defaults
131/// (`include_venice_system_prompt` defaults to `true`).
132#[non_exhaustive]
133#[derive(Clone, Debug, Default, PartialEq, Serialize)]
134pub struct VeniceParameters {
135    /// A public character to converse with, by slug.
136    #[serde(skip_serializing_if = "Option::is_none")]
137    pub character_slug: Option<String>,
138    /// Strip `<think>` blocks from the reply.
139    #[serde(skip_serializing_if = "Option::is_none")]
140    pub strip_thinking_response: Option<bool>,
141    /// The web-search mode.
142    #[serde(skip_serializing_if = "Option::is_none")]
143    pub enable_web_search: Option<WebSearchMode>,
144    /// Scrape URLs found in the prompt.
145    #[serde(skip_serializing_if = "Option::is_none")]
146    pub enable_web_scraping: Option<bool>,
147    /// Use xAI's own search on Grok models.
148    #[serde(skip_serializing_if = "Option::is_none")]
149    pub enable_x_search: Option<bool>,
150    /// Cite sources as `[REF]` markers.
151    #[serde(skip_serializing_if = "Option::is_none")]
152    pub enable_web_citations: Option<bool>,
153    /// Stream the search results too.
154    #[serde(skip_serializing_if = "Option::is_none")]
155    pub include_search_results_in_stream: Option<bool>,
156    /// Return the search results as tool-call documents.
157    #[serde(skip_serializing_if = "Option::is_none")]
158    pub return_search_results_as_documents: Option<bool>,
159    /// Include Venice's default system prompt.
160    #[serde(skip_serializing_if = "Option::is_none")]
161    pub include_venice_system_prompt: Option<bool>,
162}
163
164impl VeniceParameters {
165    /// No parameter set.
166    pub fn new() -> Self {
167        Self::default()
168    }
169
170    /// Converse with the public character `slug`.
171    pub fn character_slug(mut self, slug: impl Into<String>) -> Self {
172        self.character_slug = Some(slug.into());
173        self
174    }
175
176    /// Strip `<think>` blocks from the reply.
177    pub fn strip_thinking_response(mut self, strip: bool) -> Self {
178        self.strip_thinking_response = Some(strip);
179        self
180    }
181
182    /// Set the web-search mode.
183    pub fn enable_web_search(mut self, mode: WebSearchMode) -> Self {
184        self.enable_web_search = Some(mode);
185        self
186    }
187
188    /// Scrape URLs found in the prompt.
189    pub fn enable_web_scraping(mut self, enable: bool) -> Self {
190        self.enable_web_scraping = Some(enable);
191        self
192    }
193
194    /// Use xAI's own search on Grok models.
195    pub fn enable_x_search(mut self, enable: bool) -> Self {
196        self.enable_x_search = Some(enable);
197        self
198    }
199
200    /// Cite sources as `[REF]` markers.
201    pub fn enable_web_citations(mut self, enable: bool) -> Self {
202        self.enable_web_citations = Some(enable);
203        self
204    }
205
206    /// Stream the search results too.
207    pub fn include_search_results_in_stream(mut self, include: bool) -> Self {
208        self.include_search_results_in_stream = Some(include);
209        self
210    }
211
212    /// Return the search results as tool-call documents.
213    pub fn return_search_results_as_documents(mut self, as_documents: bool) -> Self {
214        self.return_search_results_as_documents = Some(as_documents);
215        self
216    }
217
218    /// Include Venice's default system prompt.
219    pub fn include_venice_system_prompt(mut self, include: bool) -> Self {
220        self.include_venice_system_prompt = Some(include);
221        self
222    }
223}
224
225/// Venice's reply fields. Each is `None` when the reply lacks it.
226#[non_exhaustive]
227#[derive(Clone, Debug, Default, PartialEq)]
228pub struct VeniceExtras {
229    /// The parameters Venice applied, echoed.
230    pub venice_parameters: Option<VeniceParametersEcho>,
231    /// What the request cost.
232    pub cost: Option<VeniceCost>,
233}
234
235/// The `venice_parameters` block a reply echoes: the values Venice applied
236/// and its reply-only fields.
237#[non_exhaustive]
238#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
239pub struct VeniceParametersEcho {
240    /// The character conversed with.
241    #[serde(default)]
242    pub character_slug: Option<String>,
243    /// Whether `<think>` blocks were stripped.
244    #[serde(default)]
245    pub strip_thinking_response: Option<bool>,
246    /// Whether reasoning was off.
247    #[serde(default)]
248    pub disable_thinking: Option<bool>,
249    /// The web-search mode.
250    #[serde(default)]
251    pub enable_web_search: Option<WebSearchMode>,
252    /// Whether prompt URLs were scraped.
253    #[serde(default)]
254    pub enable_web_scraping: Option<bool>,
255    /// Whether xAI's search was used.
256    #[serde(default)]
257    pub enable_x_search: Option<bool>,
258    /// Whether sources were cited.
259    #[serde(default)]
260    pub enable_web_citations: Option<bool>,
261    /// Whether search results were streamed.
262    #[serde(default)]
263    pub include_search_results_in_stream: Option<bool>,
264    /// Whether search results came back as documents.
265    #[serde(default)]
266    pub return_search_results_as_documents: Option<bool>,
267    /// Whether Venice's system prompt was included.
268    #[serde(default)]
269    pub include_venice_system_prompt: Option<bool>,
270    /// Whether end-to-end encryption applied.
271    #[serde(default)]
272    pub enable_e2ee: Option<bool>,
273    /// The sources web search consulted.
274    #[serde(default)]
275    pub web_search_citations: Option<Vec<WebSearchCitation>>,
276}
277
278/// A source Venice's web search consulted.
279#[non_exhaustive]
280#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
281pub struct WebSearchCitation {
282    /// The page title.
283    #[serde(default)]
284    pub title: Option<String>,
285    /// The source URL.
286    #[serde(default)]
287    pub url: Option<String>,
288    /// The page content Venice extracted.
289    #[serde(default)]
290    pub content: Option<String>,
291    /// The publication date, when Venice found one.
292    #[serde(default)]
293    pub date: Option<String>,
294}
295
296/// What Venice charged for a request.
297#[non_exhaustive]
298#[derive(Clone, Copy, Debug, Default, PartialEq, Deserialize)]
299pub struct VeniceCost {
300    /// In USD credits.
301    #[serde(default)]
302    pub usd: Option<f64>,
303    /// In DIEM.
304    #[serde(default)]
305    pub diem: Option<f64>,
306}
307
308impl ReplyExtras for VeniceExtras {
309    fn from_reply(_api: &Api, raw: &Value) -> Result<Self, serde_json::Error> {
310        Ok(Self {
311            venice_parameters: reply_field(raw, "/venice_parameters")?,
312            cost: reply_field(raw, "/cost")?,
313        })
314    }
315}
316
317#[cfg(test)]
318mod tests;
rig_core/providers/venice/extension.rs

rig_core/providers/venice/
extension.rs
