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