rune-chain-core 0.1.2

use serde::{Deserialize, Serialize};

use crate::ToolCall;

/// 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.
    ///
    /// # 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 };
    /// assert_eq!(a.combine(&b).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.
///
/// When the model requests tool calls (native function calling), `tool_calls`
/// will be non-empty and `generation` may be empty or contain the model's
/// reasoning text. After executing the tools and feeding results back, call the
/// LLM again until `tool_calls` is empty.
#[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>,
    /// Tool calls the model wants to make (native function calling).
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub tool_calls: Vec<ToolCall>,
}

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

    /// Return `true` if the model is requesting tool calls rather than giving a
    /// final text answer.
    pub fn wants_tool_calls(&self) -> bool {
        !self.tool_calls.is_empty()
    }
}