Skip to main content

rune_chain_core/
generate.rs

1use serde::{Deserialize, Serialize};
2
3/// Token consumption reported by a single LLM generation call.
4#[derive(Debug, Clone, Default, Serialize, Deserialize)]
5pub struct TokenUsage {
6    /// Tokens consumed by the prompt (input side).
7    pub prompt_tokens: u32,
8    /// Tokens produced in the completion (output side).
9    pub completion_tokens: u32,
10    /// `prompt_tokens + completion_tokens`.
11    pub total_tokens: u32,
12}
13
14impl TokenUsage {
15    /// Combine two [`TokenUsage`] records by summing each field.
16    ///
17    /// Useful when aggregating usage across multiple LLM calls in one chain run.
18    ///
19    /// # Example
20    ///
21    /// ```rust
22    /// use rune_chain_core::TokenUsage;
23    ///
24    /// let a = TokenUsage { prompt_tokens: 10, completion_tokens: 20, total_tokens: 30 };
25    /// let b = TokenUsage { prompt_tokens: 5,  completion_tokens: 8,  total_tokens: 13 };
26    /// let combined = a.combine(&b);
27    /// assert_eq!(combined.total_tokens, 43);
28    /// ```
29    pub fn combine(&self, other: &TokenUsage) -> TokenUsage {
30        TokenUsage {
31            prompt_tokens: self.prompt_tokens + other.prompt_tokens,
32            completion_tokens: self.completion_tokens + other.completion_tokens,
33            total_tokens: self.total_tokens + other.total_tokens,
34        }
35    }
36}
37
38/// The result of a single LLM generation call.
39#[derive(Debug, Clone, Default, Serialize, Deserialize)]
40pub struct GenerateResult {
41    /// The text produced by the model.
42    pub generation: String,
43    /// Token usage, if reported by the provider.
44    pub tokens: Option<TokenUsage>,
45}
46
47impl GenerateResult {
48    /// Create a [`GenerateResult`] with just the generated text and no token data.
49    ///
50    /// # Example
51    ///
52    /// ```rust
53    /// use rune_chain_core::GenerateResult;
54    ///
55    /// let result = GenerateResult::from_text("4");
56    /// assert_eq!(result.generation, "4");
57    /// assert!(result.tokens.is_none());
58    /// ```
59    pub fn from_text(text: impl Into<String>) -> Self {
60        Self {
61            generation: text.into(),
62            tokens: None,
63        }
64    }
65}