Skip to main content

vtcode_commons/
llm.rs

1//! Core LLM types shared across the project
2
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
6pub enum BackendKind {
7    Gemini,
8    OpenAI,
9    Anthropic,
10    DeepSeek,
11    Mistral,
12    OpenRouter,
13    Ollama,
14    LlamaCpp,
15    ZAI,
16    Moonshot,
17    HuggingFace,
18    Minimax,
19    MiMo,
20    OpenCodeZen,
21    OpenCodeGo,
22    Qwen,
23    StepFun,
24    Evolink,
25    Poolside,
26}
27
28#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
29pub struct Usage {
30    pub prompt_tokens: u32,
31    pub completion_tokens: u32,
32    pub total_tokens: u32,
33    pub cached_prompt_tokens: Option<u32>,
34    pub cache_creation_tokens: Option<u32>,
35    pub cache_read_tokens: Option<u32>,
36    /// Per-iteration token usage for Anthropic server-side fallback and compaction.
37    /// Each entry represents one sampling pass (message, fallback_message, or compaction).
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub iterations: Option<Vec<serde_json::Value>>,
40}
41
42impl Usage {
43    #[inline]
44    fn has_cache_read_metric(&self) -> bool {
45        self.cache_read_tokens.is_some() || self.cached_prompt_tokens.is_some()
46    }
47
48    #[inline]
49    fn has_any_cache_metrics(&self) -> bool {
50        self.has_cache_read_metric() || self.cache_creation_tokens.is_some()
51    }
52
53    #[inline]
54    pub fn cache_read_tokens_or_fallback(&self) -> u32 {
55        self.cache_read_tokens.or(self.cached_prompt_tokens).unwrap_or(0)
56    }
57
58    #[inline]
59    pub fn cache_creation_tokens_or_zero(&self) -> u32 {
60        self.cache_creation_tokens.unwrap_or(0)
61    }
62
63    #[inline]
64    pub fn cache_hit_rate(&self) -> Option<f64> {
65        if !self.has_any_cache_metrics() {
66            return None;
67        }
68        let read = self.cache_read_tokens_or_fallback() as f64;
69        let creation = self.cache_creation_tokens_or_zero() as f64;
70        let total = read + creation;
71        if total > 0.0 {
72            Some((read / total) * 100.0)
73        } else {
74            None
75        }
76    }
77
78    #[inline]
79    pub fn is_cache_hit(&self) -> Option<bool> {
80        self.has_any_cache_metrics().then(|| self.cache_read_tokens_or_fallback() > 0)
81    }
82
83    #[inline]
84    pub fn is_cache_miss(&self) -> Option<bool> {
85        self.has_any_cache_metrics()
86            .then(|| self.cache_creation_tokens_or_zero() > 0 && self.cache_read_tokens_or_fallback() == 0)
87    }
88
89    #[inline]
90    pub fn total_cache_tokens(&self) -> u32 {
91        let read = self.cache_read_tokens_or_fallback();
92        let creation = self.cache_creation_tokens_or_zero();
93        read + creation
94    }
95
96    #[inline]
97    pub fn cache_savings_ratio(&self) -> Option<f64> {
98        if !self.has_cache_read_metric() {
99            return None;
100        }
101        let read = self.cache_read_tokens_or_fallback() as f64;
102        let prompt = self.prompt_tokens as f64;
103        if prompt > 0.0 { Some(read / prompt) } else { None }
104    }
105}
106
107/// Provider-agnostic balance information for account status display.
108#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
109pub struct BalanceInfo {
110    /// Human-readable balance string (e.g. "100.00¥", "$50.00").
111    pub display: String,
112    /// Whether the account has sufficient balance for API calls.
113    pub is_available: bool,
114}
115
116/// DeepSeek-specific balance info from GET /user/balance
117#[derive(Debug, Clone, Serialize, Deserialize)]
118pub struct DeepSeekBalanceResponse {
119    pub is_available: bool,
120    pub balance_infos: Vec<DeepSeekCurrencyBalance>,
121}
122
123#[derive(Debug, Clone, Serialize, Deserialize)]
124pub struct DeepSeekCurrencyBalance {
125    pub currency: String,
126    pub total_balance: String,
127    #[serde(default)]
128    pub granted_balance: String,
129    #[serde(default)]
130    pub topped_up_balance: String,
131}
132
133impl From<DeepSeekBalanceResponse> for BalanceInfo {
134    fn from(resp: DeepSeekBalanceResponse) -> Self {
135        let display = resp
136            .balance_infos
137            .first()
138            .map(|b| {
139                let symbol = match b.currency.as_str() {
140                    "CNY" => "¥",
141                    "USD" => "$",
142                    _ => &b.currency,
143                };
144                format!("{}{}", b.total_balance, symbol)
145            })
146            .unwrap_or_else(|| "N/A".to_string());
147        BalanceInfo { display, is_available: resp.is_available }
148    }
149}
150
151#[cfg(test)]
152mod usage_tests {
153    use super::Usage;
154
155    #[test]
156    fn cache_helpers_fall_back_to_cached_prompt_tokens() {
157        let usage = Usage {
158            prompt_tokens: 1_000,
159            completion_tokens: 200,
160            total_tokens: 1_200,
161            cached_prompt_tokens: Some(600),
162            cache_creation_tokens: Some(150),
163            cache_read_tokens: None,
164            iterations: None,
165        };
166
167        assert_eq!(usage.cache_read_tokens_or_fallback(), 600);
168        assert_eq!(usage.cache_creation_tokens_or_zero(), 150);
169        assert_eq!(usage.total_cache_tokens(), 750);
170        assert_eq!(usage.is_cache_hit(), Some(true));
171        assert_eq!(usage.is_cache_miss(), Some(false));
172        assert_eq!(usage.cache_savings_ratio(), Some(0.6));
173        assert_eq!(usage.cache_hit_rate(), Some(80.0));
174    }
175
176    #[test]
177    fn cache_helpers_preserve_unknown_without_metrics() {
178        let usage = Usage {
179            prompt_tokens: 1_000,
180            completion_tokens: 200,
181            total_tokens: 1_200,
182            cached_prompt_tokens: None,
183            cache_creation_tokens: None,
184            cache_read_tokens: None,
185            iterations: None,
186        };
187
188        assert_eq!(usage.total_cache_tokens(), 0);
189        assert_eq!(usage.is_cache_hit(), None);
190        assert_eq!(usage.is_cache_miss(), None);
191        assert_eq!(usage.cache_savings_ratio(), None);
192        assert_eq!(usage.cache_hit_rate(), None);
193    }
194}
195
196#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
197pub enum FinishReason {
198    #[default]
199    Stop,
200    Length,
201    ToolCalls,
202    ContentFilter,
203    Pause,
204    Refusal,
205    Error(String),
206}
207
208/// Universal tool call that matches OpenAI/Anthropic/Gemini specifications
209#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
210pub struct ToolCall {
211    /// Unique identifier for this tool call (e.g., "call_123")
212    pub id: String,
213
214    /// The type of tool call: "function", "custom" (GPT-5 freeform), or other
215    #[serde(rename = "type")]
216    pub call_type: String,
217
218    /// Function call details (for function-type tools)
219    #[serde(skip_serializing_if = "Option::is_none")]
220    pub function: Option<FunctionCall>,
221
222    /// Raw text payload (for custom freeform tools in GPT-5)
223    #[serde(skip_serializing_if = "Option::is_none")]
224    pub text: Option<String>,
225
226    /// Gemini-specific thought signature for maintaining reasoning context
227    #[serde(skip_serializing_if = "Option::is_none")]
228    pub thought_signature: Option<String>,
229}
230
231/// Function call within a tool call
232#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
233pub struct FunctionCall {
234    /// Optional namespace for grouped or deferred tools.
235    #[serde(default, skip_serializing_if = "Option::is_none")]
236    pub namespace: Option<String>,
237
238    /// The name of the function to call
239    pub name: String,
240
241    /// The arguments to pass to the function, as a JSON string
242    pub arguments: String,
243}
244
245impl ToolCall {
246    /// Create a new function tool call
247    pub fn function(id: String, name: String, arguments: String) -> Self {
248        Self::function_with_namespace(id, None, name, arguments)
249    }
250
251    /// Create a new function tool call with an optional namespace.
252    pub fn function_with_namespace(id: String, namespace: Option<String>, name: String, arguments: String) -> Self {
253        Self {
254            id,
255            call_type: "function".to_owned(),
256            function: Some(FunctionCall { namespace, name, arguments }),
257            text: None,
258            thought_signature: None,
259        }
260    }
261
262    /// Create a new custom tool call with raw text payload (GPT-5 freeform)
263    pub fn custom(id: String, name: String, text: String) -> Self {
264        Self {
265            id,
266            call_type: "custom".to_owned(),
267            function: Some(FunctionCall { namespace: None, name, arguments: text.clone() }),
268            text: Some(text),
269            thought_signature: None,
270        }
271    }
272
273    /// Returns true when this tool call uses GPT-5 custom/freeform semantics.
274    pub fn is_custom(&self) -> bool {
275        self.call_type == "custom"
276    }
277
278    /// Returns the tool name when the call includes function details.
279    pub fn tool_name(&self) -> Option<&str> {
280        self.function.as_ref().map(|function| function.name.as_str())
281    }
282
283    /// Returns the raw payload text exactly as emitted by the model.
284    pub fn raw_input(&self) -> Option<&str> {
285        self.text
286            .as_deref()
287            .or_else(|| self.function.as_ref().map(|function| function.arguments.as_str()))
288    }
289
290    /// Parse the arguments as JSON Value (for function-type tools)
291    pub fn parsed_arguments(&self) -> Result<serde_json::Value, serde_json::Error> {
292        if let Some(ref func) = self.function {
293            parse_tool_arguments(&func.arguments)
294        } else {
295            // Return an error by trying to parse invalid JSON
296            serde_json::from_str("")
297        }
298    }
299
300    /// Returns the execution payload for this tool call.
301    ///
302    /// Function tools keep their JSON semantics. Custom tools execute with their
303    /// raw text payload wrapped as a JSON string value so freeform inputs can
304    /// flow through the existing tool pipeline.
305    pub fn execution_arguments(&self) -> Result<serde_json::Value, serde_json::Error> {
306        if self.is_custom() {
307            return Ok(serde_json::Value::String(self.raw_input().unwrap_or_default().to_string()));
308        }
309
310        self.parsed_arguments()
311    }
312
313    /// Validate that this tool call is properly formed
314    pub fn validate(&self) -> Result<(), String> {
315        if self.id.is_empty() {
316            return Err("Tool call ID cannot be empty".to_owned());
317        }
318
319        match self.call_type.as_str() {
320            "function" => {
321                if let Some(func) = &self.function {
322                    if func.name.is_empty() {
323                        return Err("Function name cannot be empty".to_owned());
324                    }
325                    // Validate that arguments is valid JSON for function tools
326                    if let Err(e) = self.parsed_arguments() {
327                        return Err(format!("Invalid JSON in function arguments: {e}"));
328                    }
329                } else {
330                    return Err("Function tool call missing function details".to_owned());
331                }
332            }
333            "custom" => {
334                // For custom tools, we allow raw text payload without JSON validation
335                if let Some(func) = &self.function {
336                    if func.name.is_empty() {
337                        return Err("Custom tool name cannot be empty".to_owned());
338                    }
339                } else {
340                    return Err("Custom tool call missing function details".to_owned());
341                }
342            }
343            _ => return Err(format!("Unsupported tool call type: {}", self.call_type)),
344        }
345
346        Ok(())
347    }
348}
349
350fn parse_tool_arguments(raw_arguments: &str) -> Result<serde_json::Value, serde_json::Error> {
351    let trimmed = raw_arguments.trim();
352    match serde_json::from_str(trimmed) {
353        Ok(parsed) => Ok(parsed),
354        Err(primary_error) => {
355            if let Some(candidate) = extract_balanced_json(trimmed)
356                && let Ok(parsed) = serde_json::from_str(candidate)
357            {
358                return Ok(parsed);
359            }
360            if let Some(candidate) = repair_tag_polluted_json(trimmed)
361                && let Ok(parsed) = serde_json::from_str(&candidate)
362            {
363                return Ok(parsed);
364            }
365            if let Some(repaired) = close_incomplete_json_prefix(trimmed)
366                && let Ok(parsed) = serde_json::from_str(&repaired)
367            {
368                return Ok(parsed);
369            }
370            Err(primary_error)
371        }
372    }
373}
374
375fn extract_balanced_json(input: &str) -> Option<&str> {
376    let start = input.find(['{', '['])?;
377    let opening = input.as_bytes().get(start).copied()?;
378    let closing = match opening {
379        b'{' => b'}',
380        b'[' => b']',
381        _ => return None,
382    };
383
384    let mut depth = 0usize;
385    let mut in_string = false;
386    let mut escaped = false;
387
388    for (offset, ch) in input[start..].char_indices() {
389        if in_string {
390            if escaped {
391                escaped = false;
392                continue;
393            }
394            if ch == '\\' {
395                escaped = true;
396                continue;
397            }
398            if ch == '"' {
399                in_string = false;
400            }
401            continue;
402        }
403
404        match ch {
405            '"' => in_string = true,
406            _ if ch as u32 == opening as u32 => depth += 1,
407            _ if ch as u32 == closing as u32 => {
408                depth = depth.saturating_sub(1);
409                if depth == 0 {
410                    let end = start + offset + ch.len_utf8();
411                    return input.get(start..end);
412                }
413            }
414            _ => {}
415        }
416    }
417
418    None
419}
420
421fn repair_tag_polluted_json(input: &str) -> Option<String> {
422    let start = input.find(['{', '['])?;
423    let candidate = input.get(start..)?;
424    let boundary = find_provider_markup_boundary(candidate)?;
425    if boundary == 0 {
426        return None;
427    }
428
429    close_incomplete_json_prefix(candidate[..boundary].trim_end())
430}
431
432fn find_provider_markup_boundary(input: &str) -> Option<usize> {
433    const PROVIDER_MARKERS: &[&str] = &[
434        "<</",
435        "</parameter>",
436        "</invoke>",
437        "</minimax:tool_call>",
438        "<minimax:tool_call>",
439        "<parameter name=\"",
440        "<invoke name=\"",
441        "<tool_call>",
442        "</tool_call>",
443    ];
444
445    input.char_indices().find_map(|(offset, _)| {
446        let rest = input.get(offset..)?;
447        PROVIDER_MARKERS.iter().any(|marker| rest.starts_with(marker)).then_some(offset)
448    })
449}
450
451fn close_incomplete_json_prefix(prefix: &str) -> Option<String> {
452    if prefix.is_empty() {
453        return None;
454    }
455
456    let mut repaired = String::with_capacity(prefix.len() + 8);
457    let mut expected_closers = Vec::new();
458    let mut in_string = false;
459    let mut escaped = false;
460
461    for ch in prefix.chars() {
462        repaired.push(ch);
463
464        if in_string {
465            if escaped {
466                escaped = false;
467                continue;
468            }
469
470            match ch {
471                '\\' => escaped = true,
472                '"' => in_string = false,
473                _ => {}
474            }
475            continue;
476        }
477
478        match ch {
479            '"' => in_string = true,
480            '{' => expected_closers.push('}'),
481            '[' => expected_closers.push(']'),
482            '}' | ']' if expected_closers.pop() != Some(ch) => return None,
483            '}' | ']' => {}
484            _ => {}
485        }
486    }
487
488    if in_string {
489        repaired.push('"');
490    }
491    for closer in expected_closers.drain(..) {
492        repaired.push(closer);
493    }
494
495    Some(repaired)
496}
497
498/// Universal LLM response structure
499#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
500pub struct LLMResponse {
501    /// The response content text
502    pub content: Option<String>,
503
504    /// Tool calls made by the model
505    pub tool_calls: Option<Vec<ToolCall>>,
506
507    /// The model that generated this response
508    pub model: String,
509
510    /// Token usage statistics
511    pub usage: Option<Usage>,
512
513    /// Why the response finished
514    pub finish_reason: FinishReason,
515
516    /// Reasoning content (for models that support it)
517    pub reasoning: Option<String>,
518
519    /// Detailed reasoning traces (for models that support it)
520    pub reasoning_details: Option<Vec<String>>,
521
522    /// Tool references for context
523    pub tool_references: Vec<String>,
524
525    /// Request ID from the provider
526    pub request_id: Option<String>,
527
528    /// Organization ID from the provider
529    pub organization_id: Option<String>,
530
531    /// Compaction summary content from Anthropic's server-side compaction.
532    /// Populated when `stop_reason` is `Pause` (from `"compaction"`).
533    /// The caller should pass this back in subsequent requests so the API
534    /// can drop prior messages before the compaction block.
535    pub compaction: Option<String>,
536}
537
538impl LLMResponse {
539    /// Create a new LLM response with mandatory fields
540    pub fn new(model: impl Into<String>, content: impl Into<String>) -> Self {
541        Self {
542            content: Some(content.into()),
543            tool_calls: None,
544            model: model.into(),
545            usage: None,
546            finish_reason: FinishReason::Stop,
547            reasoning: None,
548            reasoning_details: None,
549            tool_references: Vec::new(),
550            request_id: None,
551            organization_id: None,
552            compaction: None,
553        }
554    }
555
556    /// Get content or empty string
557    pub fn content_text(&self) -> &str {
558        self.content.as_deref().unwrap_or("")
559    }
560
561    /// Get content as String (clone)
562    pub fn content_string(&self) -> String {
563        self.content.clone().unwrap_or_default()
564    }
565}
566
567#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
568pub struct LLMErrorMetadata {
569    pub provider: Option<String>,
570    pub status: Option<u16>,
571    pub code: Option<String>,
572    pub request_id: Option<String>,
573    pub organization_id: Option<String>,
574    pub retry_after: Option<String>,
575    pub message: Option<String>,
576}
577
578impl LLMErrorMetadata {
579    /// Boxed constructor because metadata is always stored inside `Option<Box<LLMErrorMetadata>>`
580    /// in the LLMError enum variants.
581    #[must_use]
582    pub fn new(
583        provider: impl Into<String>,
584        status: Option<u16>,
585        code: Option<String>,
586        request_id: Option<String>,
587        organization_id: Option<String>,
588        retry_after: Option<String>,
589        message: Option<String>,
590    ) -> Box<Self> {
591        Box::new(Self {
592            provider: Some(provider.into()),
593            status,
594            code,
595            request_id,
596            organization_id,
597            retry_after,
598            message,
599        })
600    }
601}
602
603/// LLM error types with optional provider metadata
604#[derive(Debug, thiserror::Error, Serialize, Deserialize, Clone)]
605#[serde(tag = "type", rename_all = "snake_case")]
606pub enum LLMError {
607    #[error("Authentication failed: {message}")]
608    Authentication {
609        message: String,
610        metadata: Option<Box<LLMErrorMetadata>>,
611    },
612    #[error("Rate limit exceeded")]
613    RateLimit { metadata: Option<Box<LLMErrorMetadata>> },
614    #[error("Invalid request: {message}")]
615    InvalidRequest {
616        message: String,
617        metadata: Option<Box<LLMErrorMetadata>>,
618    },
619    #[error("Network error: {message}")]
620    Network {
621        message: String,
622        metadata: Option<Box<LLMErrorMetadata>>,
623    },
624    #[error("Provider error: {message}")]
625    Provider {
626        message: String,
627        metadata: Option<Box<LLMErrorMetadata>>,
628    },
629}
630
631#[cfg(test)]
632mod tests {
633    use super::ToolCall;
634    use serde_json::json;
635
636    #[test]
637    fn parsed_arguments_accepts_trailing_characters() {
638        let call = ToolCall::function(
639            "call_read".to_string(),
640            "exec_command".to_string(),
641            r#"{"path":"src/main.rs"} trailing text"#.to_string(),
642        );
643
644        let parsed = call.parsed_arguments().expect("arguments with trailing text should recover");
645        assert_eq!(parsed, json!({"path":"src/main.rs"}));
646    }
647
648    #[test]
649    fn parsed_arguments_accepts_code_fenced_json() {
650        let call = ToolCall::function(
651            "call_read".to_string(),
652            "exec_command".to_string(),
653            "```json\n{\"path\":\"src/lib.rs\",\"limit\":25}\n```".to_string(),
654        );
655
656        let parsed = call.parsed_arguments().expect("code-fenced arguments should recover");
657        assert_eq!(parsed, json!({"path":"src/lib.rs","limit":25}));
658    }
659
660    #[test]
661    fn parsed_arguments_recovers_truncated_json_missing_closing_brace() {
662        let call = ToolCall::function(
663            "call_search".to_string(),
664            "code_search".to_string(),
665            r#"{"query":"context","path":".","file_types":["rust"],"result_types":["definition"],"max_results":20"#
666                .to_string(),
667        );
668
669        let parsed = call
670            .parsed_arguments()
671            .expect("truncated JSON missing closing brace should recover");
672        assert_eq!(
673            parsed,
674            json!({
675                "query": "context",
676                "path": ".",
677                "file_types": ["rust"],
678                "result_types": ["definition"],
679                "max_results": 20
680            })
681        );
682    }
683
684    #[test]
685    fn parsed_arguments_rejects_incomplete_json() {
686        let call = ToolCall::function(
687            "call_read".to_string(),
688            "exec_command".to_string(),
689            r#"{"path":"src/main.rs","limit""#.to_string(),
690        );
691
692        assert!(call.parsed_arguments().is_err());
693    }
694
695    #[test]
696    fn parsed_arguments_recovers_truncated_minimax_markup() {
697        let call = ToolCall::function(
698            "call_search".to_string(),
699            "code_search".to_string(),
700            "{\"query\":\"persistent_memory\",\"file_types\":[\"rust\"],\"result_types\":[\"text\"],\"max_results\":20,\"path\":\"crates/codegen/vtcode-core/src</parameter>\n<</invoke>\n</minimax:tool_call>".to_string(),
701        );
702
703        let parsed = call.parsed_arguments().expect("minimax markup spillover should recover");
704        assert_eq!(
705            parsed,
706            json!({
707                "query": "persistent_memory",
708                "path": "crates/codegen/vtcode-core/src",
709                "file_types": ["rust"],
710                "result_types": ["text"],
711                "max_results": 20
712            })
713        );
714    }
715
716    #[test]
717    fn function_call_serializes_optional_namespace() {
718        let call = ToolCall::function_with_namespace(
719            "call_read".to_string(),
720            Some("workspace".to_string()),
721            "exec_command".to_string(),
722            r#"{"path":"src/main.rs"}"#.to_string(),
723        );
724
725        let json = serde_json::to_value(&call).expect("tool call should serialize");
726        assert_eq!(json["function"]["namespace"], "workspace");
727        assert_eq!(json["function"]["name"], "exec_command");
728    }
729
730    #[test]
731    fn custom_tool_call_exposes_raw_execution_arguments() {
732        let patch = "*** Begin Patch\n*** End Patch\n".to_string();
733        let call = ToolCall::custom("call_patch".to_string(), "apply_patch".to_string(), patch.clone());
734
735        assert!(call.is_custom());
736        assert_eq!(call.tool_name(), Some("apply_patch"));
737        assert_eq!(call.raw_input(), Some(patch.as_str()));
738        assert_eq!(call.execution_arguments().expect("custom arguments"), json!(patch));
739        assert!(call.parsed_arguments().is_err(), "custom tool payload should stay freeform rather than JSON");
740    }
741}