rig_core/providers/venice/
completion.rs1use serde::{Deserialize, Serialize};
14
15use crate::completion::{self, CompletionError, NormalizeCompletionResponse};
16use crate::providers::openai;
17use crate::telemetry::ProviderResponseExt;
18
19pub const GLM_4_7: &str = "zai-org-glm-4.7";
28pub const GLM_5_2: &str = "zai-org-glm-5-2";
30pub const QWEN3_5_9B: &str = "qwen3-5-9b";
32pub const QWEN3_5_397B_A17B: &str = "qwen3-5-397b-a17b";
34pub const QWEN3_235B_A22B_THINKING: &str = "qwen3-235b-a22b-thinking-2507";
36pub const QWEN3_VL_235B_A22B: &str = "qwen3-vl-235b-a22b";
38pub const QWEN3_CODER_480B: &str = "qwen3-coder-480b-a35b-instruct-turbo";
40pub const VENICE_UNCENSORED_1_2: &str = "venice-uncensored-1-2";
42pub const GEMINI_3_6_FLASH: &str = "gemini-3-6-flash";
44pub const GROK_4_6: &str = "grok-4-6";
46pub const MISTRAL_SMALL_2603: &str = "mistral-small-2603";
48pub const MISTRAL_SMALL_3_2_24B: &str = "mistral-small-3-2-24b-instruct";
50
51pub type CompletionModel<H = reqwest::Client> =
55 openai::completion::GenericCompletionModel<super::client::VeniceExt, H>;
56
57#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
63#[serde(rename_all = "lowercase")]
64pub enum WebSearchMode {
65 Off,
67 On,
69 Auto,
71}
72
73#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
106pub struct VeniceParameters {
107 #[serde(default, skip_serializing_if = "Option::is_none")]
109 pub character_slug: Option<String>,
110 #[serde(default, skip_serializing_if = "Option::is_none")]
112 pub strip_thinking_response: Option<bool>,
113 #[serde(default, skip_serializing_if = "Option::is_none")]
115 pub disable_thinking: Option<bool>,
116 #[serde(default, skip_serializing_if = "Option::is_none")]
118 pub enable_web_search: Option<WebSearchMode>,
119 #[serde(default, skip_serializing_if = "Option::is_none")]
121 pub enable_web_scraping: Option<bool>,
122 #[serde(default, skip_serializing_if = "Option::is_none")]
124 pub enable_x_search: Option<bool>,
125 #[serde(default, skip_serializing_if = "Option::is_none")]
127 pub enable_web_citations: Option<bool>,
128 #[serde(default, skip_serializing_if = "Option::is_none")]
130 pub include_search_results_in_stream: Option<bool>,
131 #[serde(default, skip_serializing_if = "Option::is_none")]
133 pub return_search_results_as_documents: Option<bool>,
134 #[serde(default, skip_serializing_if = "Option::is_none")]
136 pub include_venice_system_prompt: Option<bool>,
137 #[serde(default, skip_serializing_if = "Option::is_none")]
139 pub prompt_cache_key: Option<String>,
140}
141
142impl VeniceParameters {
143 pub fn new() -> Self {
145 Self::default()
146 }
147
148 pub fn character_slug(mut self, slug: impl Into<String>) -> Self {
150 self.character_slug = Some(slug.into());
151 self
152 }
153
154 pub fn strip_thinking_response(mut self, strip: bool) -> Self {
156 self.strip_thinking_response = Some(strip);
157 self
158 }
159
160 pub fn disable_thinking(mut self, disable: bool) -> Self {
162 self.disable_thinking = Some(disable);
163 self
164 }
165
166 pub fn enable_web_search(mut self, mode: WebSearchMode) -> Self {
168 self.enable_web_search = Some(mode);
169 self
170 }
171
172 pub fn enable_web_scraping(mut self, enable: bool) -> Self {
174 self.enable_web_scraping = Some(enable);
175 self
176 }
177
178 pub fn enable_x_search(mut self, enable: bool) -> Self {
180 self.enable_x_search = Some(enable);
181 self
182 }
183
184 pub fn enable_web_citations(mut self, enable: bool) -> Self {
186 self.enable_web_citations = Some(enable);
187 self
188 }
189
190 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 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 pub fn include_venice_system_prompt(mut self, include: bool) -> Self {
204 self.include_venice_system_prompt = Some(include);
205 self
206 }
207
208 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 pub fn into_additional_params(self) -> serde_json::Value {
217 serde_json::json!({ "venice_parameters": self })
218 }
219}
220
221#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
227pub struct WebSearchCitation {
228 #[serde(default)]
230 pub title: String,
231 #[serde(default)]
233 pub url: String,
234 #[serde(default)]
236 pub content: String,
237 #[serde(default)]
239 pub date: String,
240}
241
242#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
247pub struct VeniceParametersEcho {
248 #[serde(flatten)]
250 pub parameters: VeniceParameters,
251 #[serde(default, skip_serializing_if = "Option::is_none")]
253 pub enable_e2ee: Option<bool>,
254 #[serde(default, skip_serializing_if = "Vec::is_empty")]
256 pub web_search_citations: Vec<WebSearchCitation>,
257}
258
259#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
261pub struct Cost {
262 #[serde(default)]
264 pub usd: f64,
265 #[serde(default)]
267 pub diem: f64,
268}
269
270#[derive(Debug, Deserialize, Serialize)]
277pub struct CompletionResponse {
278 #[serde(flatten)]
280 pub openai: openai::CompletionResponse,
281 #[serde(default, skip_serializing_if = "Option::is_none")]
283 pub venice_parameters: Option<VeniceParametersEcho>,
284 #[serde(default, skip_serializing_if = "Option::is_none")]
286 pub cost: Option<Cost>,
287}
288
289impl CompletionResponse {
290 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 #[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(¶ms).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 #[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}