Skip to main content

rig_core/providers/venice/
completion.rs

1//! Venice completion models.
2//!
3//! Completions run through the shared OpenAI-compatible
4//! [`GenericCompletionModel`](openai::completion::GenericCompletionModel); the
5//! dialect is declared by the `OpenAICompatibleProvider` impl on
6//! [`VeniceExt`](super::client::VeniceExt) in `client.rs`.
7//!
8//! Venice's chat payload is OpenAI's plus two blocks it adds itself: the
9//! resolved [`VeniceParameters`] echo (which is where web-search citations
10//! arrive) and a per-request [`Cost`]. [`CompletionResponse`] preserves both,
11//! so `raw_completion` callers keep everything Venice sent.
12
13use serde::{Deserialize, Serialize};
14
15use crate::completion::{self, CompletionError, NormalizeCompletionResponse};
16use crate::providers::openai;
17use crate::telemetry::ProviderResponseExt;
18
19// ================================================================
20// Venice Completion Models
21// ================================================================
22// A non-exhaustive selection; the authoritative list is `GET /models`, which
23// also reports per-model capabilities (`supportsFunctionCalling`,
24// `supportsVision`, `supportsReasoning`, `supportsResponseSchema`, …).
25
26/// `zai-org-glm-4.7` — Venice's `default` and `function_calling_default` model.
27pub const GLM_4_7: &str = "zai-org-glm-4.7";
28/// `zai-org-glm-5-2`
29pub const GLM_5_2: &str = "zai-org-glm-5-2";
30/// `qwen3-5-9b` — small, tool-capable, and vision-capable.
31pub const QWEN3_5_9B: &str = "qwen3-5-9b";
32/// `qwen3-5-397b-a17b`
33pub const QWEN3_5_397B_A17B: &str = "qwen3-5-397b-a17b";
34/// `qwen3-235b-a22b-thinking-2507` — Venice's `default_reasoning` model.
35pub const QWEN3_235B_A22B_THINKING: &str = "qwen3-235b-a22b-thinking-2507";
36/// `qwen3-vl-235b-a22b` — Venice's `default_vision` model.
37pub const QWEN3_VL_235B_A22B: &str = "qwen3-vl-235b-a22b";
38/// `qwen3-coder-480b-a35b-instruct-turbo` — Venice's `default_code` model.
39pub const QWEN3_CODER_480B: &str = "qwen3-coder-480b-a35b-instruct-turbo";
40/// `venice-uncensored-1-2` — Venice's `most_uncensored` model.
41pub const VENICE_UNCENSORED_1_2: &str = "venice-uncensored-1-2";
42/// `gemini-3-6-flash`
43pub const GEMINI_3_6_FLASH: &str = "gemini-3-6-flash";
44/// `grok-4-6`
45pub const GROK_4_6: &str = "grok-4-6";
46/// `mistral-small-2603`
47pub const MISTRAL_SMALL_2603: &str = "mistral-small-2603";
48/// `mistral-small-3-2-24b-instruct`
49pub const MISTRAL_SMALL_3_2_24B: &str = "mistral-small-3-2-24b-instruct";
50
51/// Venice completion model — the shared OpenAI-compatible
52/// [`GenericCompletionModel`](openai::completion::GenericCompletionModel)
53/// specialized to Venice.
54pub type CompletionModel<H = reqwest::Client> =
55    openai::completion::GenericCompletionModel<super::client::VeniceExt, H>;
56
57// ================================================================
58// Venice-specific request parameters
59// ================================================================
60
61/// How Venice's web search behaves for a request.
62#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
63#[serde(rename_all = "lowercase")]
64pub enum WebSearchMode {
65    /// Never search.
66    Off,
67    /// Always search.
68    On,
69    /// Let the model decide.
70    Auto,
71}
72
73/// Venice's `venice_parameters` request block.
74///
75/// Venice accepts this alongside the OpenAI chat-completions body. Rig passes
76/// it through [`additional_params`](crate::completion::CompletionRequest),
77/// which is the same merge path every other provider's dialect extras use, so
78/// there is no separate request abstraction to keep in sync:
79///
80/// ```no_run
81/// use rig_core::client::{CompletionClient, ProviderClient};
82/// use rig_core::completion::CompletionModel;
83/// use rig_core::providers::venice::{self, VeniceParameters, WebSearchMode};
84///
85/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
86/// let client = venice::Client::from_env()?;
87/// let model = client.completion_model(venice::QWEN3_5_9B);
88/// let request = model
89///     .completion_request("Summarize today's Rust news.")
90///     .additional_params(
91///         VeniceParameters::new()
92///             .enable_web_search(WebSearchMode::On)
93///             .enable_web_citations(true)
94///             .into_additional_params(),
95///     )
96///     .build();
97/// let response = model.completion(request).await?;
98/// # let _ = response;
99/// # Ok(())
100/// # }
101/// ```
102///
103/// Every field is optional; omitted fields are left to Venice's own defaults
104/// (notably `include_venice_system_prompt`, which Venice defaults to `true`).
105#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
106pub struct VeniceParameters {
107    /// Public character slug to converse with.
108    #[serde(default, skip_serializing_if = "Option::is_none")]
109    pub character_slug: Option<String>,
110    /// Strip `<think>` blocks from the response.
111    #[serde(default, skip_serializing_if = "Option::is_none")]
112    pub strip_thinking_response: Option<bool>,
113    /// Disable reasoning on models that support it.
114    #[serde(default, skip_serializing_if = "Option::is_none")]
115    pub disable_thinking: Option<bool>,
116    /// Web-search mode.
117    #[serde(default, skip_serializing_if = "Option::is_none")]
118    pub enable_web_search: Option<WebSearchMode>,
119    /// Scrape URLs found in the prompt.
120    #[serde(default, skip_serializing_if = "Option::is_none")]
121    pub enable_web_scraping: Option<bool>,
122    /// Use xAI's native search on Grok models.
123    #[serde(default, skip_serializing_if = "Option::is_none")]
124    pub enable_x_search: Option<bool>,
125    /// Emit `[REF]`-style source citations.
126    #[serde(default, skip_serializing_if = "Option::is_none")]
127    pub enable_web_citations: Option<bool>,
128    /// Include search results in the stream (experimental).
129    #[serde(default, skip_serializing_if = "Option::is_none")]
130    pub include_search_results_in_stream: Option<bool>,
131    /// Return search results as tool-call documents.
132    #[serde(default, skip_serializing_if = "Option::is_none")]
133    pub return_search_results_as_documents: Option<bool>,
134    /// Include Venice's default system prompt (Venice defaults to `true`).
135    #[serde(default, skip_serializing_if = "Option::is_none")]
136    pub include_venice_system_prompt: Option<bool>,
137    /// Prompt-cache routing hint.
138    #[serde(default, skip_serializing_if = "Option::is_none")]
139    pub prompt_cache_key: Option<String>,
140}
141
142impl VeniceParameters {
143    /// An empty parameter block; every field falls back to Venice's default.
144    pub fn new() -> Self {
145        Self::default()
146    }
147
148    /// Converse with a public Venice character.
149    pub fn character_slug(mut self, slug: impl Into<String>) -> Self {
150        self.character_slug = Some(slug.into());
151        self
152    }
153
154    /// Strip `<think>` blocks from the response.
155    pub fn strip_thinking_response(mut self, strip: bool) -> Self {
156        self.strip_thinking_response = Some(strip);
157        self
158    }
159
160    /// Disable reasoning on models that support it.
161    pub fn disable_thinking(mut self, disable: bool) -> Self {
162        self.disable_thinking = Some(disable);
163        self
164    }
165
166    /// Set the web-search mode.
167    pub fn enable_web_search(mut self, mode: WebSearchMode) -> Self {
168        self.enable_web_search = Some(mode);
169        self
170    }
171
172    /// Scrape URLs found in the prompt.
173    pub fn enable_web_scraping(mut self, enable: bool) -> Self {
174        self.enable_web_scraping = Some(enable);
175        self
176    }
177
178    /// Use xAI's native search on Grok models.
179    pub fn enable_x_search(mut self, enable: bool) -> Self {
180        self.enable_x_search = Some(enable);
181        self
182    }
183
184    /// Emit `[REF]`-style source citations.
185    pub fn enable_web_citations(mut self, enable: bool) -> Self {
186        self.enable_web_citations = Some(enable);
187        self
188    }
189
190    /// Include search results in the stream (experimental).
191    pub fn include_search_results_in_stream(mut self, include: bool) -> Self {
192        self.include_search_results_in_stream = Some(include);
193        self
194    }
195
196    /// Return search results as tool-call documents.
197    pub fn return_search_results_as_documents(mut self, as_documents: bool) -> Self {
198        self.return_search_results_as_documents = Some(as_documents);
199        self
200    }
201
202    /// Include Venice's default system prompt.
203    pub fn include_venice_system_prompt(mut self, include: bool) -> Self {
204        self.include_venice_system_prompt = Some(include);
205        self
206    }
207
208    /// Set the prompt-cache routing hint.
209    pub fn prompt_cache_key(mut self, key: impl Into<String>) -> Self {
210        self.prompt_cache_key = Some(key.into());
211        self
212    }
213
214    /// Wrap this block in the `{"venice_parameters": …}` object Rig merges
215    /// into the request body through `additional_params`.
216    pub fn into_additional_params(self) -> serde_json::Value {
217        serde_json::json!({ "venice_parameters": self })
218    }
219}
220
221// ================================================================
222// Venice completion response
223// ================================================================
224
225/// A web-search source Venice consulted for a completion.
226#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
227pub struct WebSearchCitation {
228    /// Page title.
229    #[serde(default)]
230    pub title: String,
231    /// Source URL.
232    #[serde(default)]
233    pub url: String,
234    /// Extracted page content, as Venice returned it.
235    #[serde(default)]
236    pub content: String,
237    /// Publication date, empty when Venice could not determine one.
238    #[serde(default)]
239    pub date: String,
240}
241
242/// Venice's resolved `venice_parameters` block, echoed on every response.
243///
244/// The requested fields come back with the values Venice actually applied,
245/// alongside the response-only fields below.
246#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
247pub struct VeniceParametersEcho {
248    /// The parameters Venice resolved for this request.
249    #[serde(flatten)]
250    pub parameters: VeniceParameters,
251    /// Whether end-to-end encryption applied to this request.
252    #[serde(default, skip_serializing_if = "Option::is_none")]
253    pub enable_e2ee: Option<bool>,
254    /// Sources consulted when web search ran; empty otherwise.
255    #[serde(default, skip_serializing_if = "Vec::is_empty")]
256    pub web_search_citations: Vec<WebSearchCitation>,
257}
258
259/// What Venice charged for a request.
260#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
261pub struct Cost {
262    /// Cost in USD credits.
263    #[serde(default)]
264    pub usd: f64,
265    /// Cost in DIEM.
266    #[serde(default)]
267    pub diem: f64,
268}
269
270/// Venice's chat-completions payload: OpenAI's response plus the
271/// `venice_parameters` echo and the request's `cost`.
272///
273/// Normalization and telemetry delegate to the OpenAI payload — the wire
274/// shape of `choices`/`usage` is OpenAI's — so the Venice-only blocks are
275/// preserved for `raw_completion` callers without forking the conversion.
276#[derive(Debug, Deserialize, Serialize)]
277pub struct CompletionResponse {
278    /// The OpenAI-compatible portion of the payload.
279    #[serde(flatten)]
280    pub openai: openai::CompletionResponse,
281    /// Venice's resolved parameter block, including web-search citations.
282    #[serde(default, skip_serializing_if = "Option::is_none")]
283    pub venice_parameters: Option<VeniceParametersEcho>,
284    /// What Venice charged for this request.
285    #[serde(default, skip_serializing_if = "Option::is_none")]
286    pub cost: Option<Cost>,
287}
288
289impl CompletionResponse {
290    /// The web-search sources Venice consulted, empty when search did not run.
291    pub fn web_search_citations(&self) -> &[WebSearchCitation] {
292        self.venice_parameters
293            .as_ref()
294            .map(|parameters| parameters.web_search_citations.as_slice())
295            .unwrap_or_default()
296    }
297}
298
299impl NormalizeCompletionResponse for CompletionResponse {
300    fn normalize(self, provider: &str) -> Result<completion::CompletionResponse, CompletionError> {
301        self.openai.normalize(provider)
302    }
303}
304
305impl ProviderResponseExt for CompletionResponse {
306    type Usage = <openai::CompletionResponse as ProviderResponseExt>::Usage;
307
308    fn get_response_id(&self) -> Option<String> {
309        self.openai.get_response_id()
310    }
311
312    fn get_response_model_name(&self) -> Option<String> {
313        self.openai.get_response_model_name()
314    }
315
316    fn get_text_response(&self) -> Option<String> {
317        self.openai.get_text_response()
318    }
319
320    fn get_usage(&self) -> Option<Self::Usage> {
321        self.openai.get_usage()
322    }
323}
324
325#[cfg(test)]
326mod tests {
327    use super::*;
328
329    /// Serialization shape of the request block is definitory, not observed:
330    /// the cassette suite pins that Venice *accepts* it, this pins that
331    /// unset fields stay off the wire entirely rather than being sent null.
332    #[test]
333    fn venice_parameters_only_serialize_set_fields() {
334        let params = VeniceParameters::new()
335            .enable_web_search(WebSearchMode::Auto)
336            .disable_thinking(true);
337
338        let json = serde_json::to_value(&params).expect("parameters should serialize");
339
340        assert_eq!(
341            json,
342            serde_json::json!({
343                "enable_web_search": "auto",
344                "disable_thinking": true,
345            })
346        );
347    }
348
349    #[test]
350    fn venice_parameters_wrap_into_additional_params() {
351        let json = VeniceParameters::new()
352            .character_slug("venice")
353            .into_additional_params();
354
355        assert_eq!(
356            json,
357            serde_json::json!({ "venice_parameters": { "character_slug": "venice" } })
358        );
359    }
360
361    /// Response decoding is pinned by cassettes; this asserts the flattened
362    /// wrapper keeps *both* halves — an OpenAI-only decode would silently
363    /// drop citations and cost.
364    #[test]
365    fn completion_response_preserves_venice_blocks() {
366        let body = serde_json::json!({
367            "id": "chatcmpl-1",
368            "object": "chat.completion",
369            "created": 0,
370            "model": "qwen3-5-9b",
371            "choices": [{
372                "index": 0,
373                "message": {"role": "assistant", "content": "hi"},
374                "finish_reason": "stop"
375            }],
376            "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
377            "cost": {"usd": 0.000_002_65, "diem": 0.0},
378            "venice_parameters": {
379                "enable_web_search": "on",
380                "enable_e2ee": true,
381                "web_search_citations": [{
382                    "title": "Rust",
383                    "url": "https://example.com",
384                    "content": "text",
385                    "date": ""
386                }]
387            }
388        });
389
390        let response: CompletionResponse =
391            serde_json::from_value(body).expect("response should decode");
392
393        assert_eq!(response.openai.id, "chatcmpl-1");
394        assert_eq!(response.get_text_response().as_deref(), Some("hi"));
395        assert_eq!(response.cost.expect("cost").diem, 0.0);
396        assert_eq!(response.web_search_citations().len(), 1);
397        assert_eq!(response.web_search_citations()[0].title, "Rust");
398        assert_eq!(
399            response
400                .venice_parameters
401                .expect("venice parameters")
402                .parameters
403                .enable_web_search,
404            Some(WebSearchMode::On)
405        );
406    }
407}