#[cfg(feature = "rig")]
use rig_core::completion::{GetTokenUsage, Usage};
use serde::{Deserialize, Serialize};
#[cfg(feature = "rig")]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Response {
pub text: String,
pub reasoning: Option<String>,
pub tool_calls: Vec<ToolCallData>,
pub usage: UsageData,
pub finish_reason: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub speculation: Option<SpeculationStatsData>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct SpeculationStatsData {
pub drafted: u64,
pub rounds: u64,
pub accepted_by_depth: Vec<u64>,
}
#[cfg(feature = "rig")]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCallData {
pub id: String,
pub name: String,
pub arguments: serde_json::Value,
}
#[cfg(feature = "rig")]
#[allow(
clippy::struct_field_names,
reason = "public token accounting follows established provider vocabulary"
)]
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
pub struct UsageData {
pub prompt_tokens: u64,
pub cached_tokens: u64,
pub completion_tokens: u64,
}
#[cfg(feature = "rig")]
impl UsageData {
pub(crate) fn to_rig(self) -> Usage {
Usage {
input_tokens: self.prompt_tokens,
output_tokens: self.completion_tokens,
total_tokens: self.prompt_tokens.saturating_add(self.completion_tokens),
cached_input_tokens: self.cached_tokens,
..Usage::new()
}
}
}
#[cfg(feature = "rig")]
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct StreamingResponse {
pub usage: UsageData,
pub finish_reason: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub speculation: Option<SpeculationStatsData>,
}
#[cfg(feature = "rig")]
pub const fn finish_reason_label(reason: crate::engine::generate::FinishReason) -> &'static str {
match reason {
crate::engine::generate::FinishReason::Stop => "stop",
crate::engine::generate::FinishReason::Length => "length",
crate::engine::generate::FinishReason::ToolCalls => "tool_calls",
crate::engine::generate::FinishReason::Aborted => "aborted",
}
}
#[cfg(feature = "rig")]
impl GetTokenUsage for StreamingResponse {
fn token_usage(&self) -> Usage {
self.usage.to_rig()
}
}