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