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