rune-chain-core 0.1.0

use serde::{Deserialize, Serialize};

/// Token consumption reported by a single LLM generation call.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TokenUsage {
    /// Tokens consumed by the prompt (input side).
    pub prompt_tokens: u32,
    /// Tokens produced in the completion (output side).
    pub completion_tokens: u32,
    /// `prompt_tokens + completion_tokens`.
    pub total_tokens: u32,
}

impl TokenUsage {
    /// Combine two [`TokenUsage`] records by summing each field.
    ///
    /// Useful when aggregating usage across multiple LLM calls in one chain run.
    ///
    /// # Example
    ///
    /// ```rust
    /// use rune_chain_core::TokenUsage;
    ///
    /// let a = TokenUsage { prompt_tokens: 10, completion_tokens: 20, total_tokens: 30 };
    /// let b = TokenUsage { prompt_tokens: 5,  completion_tokens: 8,  total_tokens: 13 };
    /// let combined = a.combine(&b);
    /// assert_eq!(combined.total_tokens, 43);
    /// ```
    pub fn combine(&self, other: &TokenUsage) -> TokenUsage {
        TokenUsage {
            prompt_tokens: self.prompt_tokens + other.prompt_tokens,
            completion_tokens: self.completion_tokens + other.completion_tokens,
            total_tokens: self.total_tokens + other.total_tokens,
        }
    }
}

/// The result of a single LLM generation call.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct GenerateResult {
    /// The text produced by the model.
    pub generation: String,
    /// Token usage, if reported by the provider.
    pub tokens: Option<TokenUsage>,
}

impl GenerateResult {
    /// Create a [`GenerateResult`] with just the generated text and no token data.
    ///
    /// # Example
    ///
    /// ```rust
    /// use rune_chain_core::GenerateResult;
    ///
    /// let result = GenerateResult::from_text("4");
    /// assert_eq!(result.generation, "4");
    /// assert!(result.tokens.is_none());
    /// ```
    pub fn from_text(text: impl Into<String>) -> Self {
        Self {
            generation: text.into(),
            tokens: None,
        }
    }
}