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}