use rig_core::{
OneOrMany,
completion::AssistantContent,
message::{Reasoning, ToolCall, ToolFunction},
};
use crate::{
engine::generate::{GenerateReply, Usage},
model::UsageData,
};
pub fn choice(reply: &GenerateReply) -> OneOrMany<AssistantContent> {
let mut items: Vec<AssistantContent> = Vec::new();
if let Some(reasoning) = &reply.reasoning
&& !reasoning.is_empty()
{
items.push(AssistantContent::Reasoning(Reasoning::new(reasoning)));
}
if !reply.text.is_empty() {
items.push(AssistantContent::text(&reply.text));
}
for call in &reply.tool_calls {
items.push(AssistantContent::ToolCall(ToolCall::new(
call.id.clone(),
ToolFunction::new(call.name.clone(), call.arguments.clone()),
)));
}
OneOrMany::many(items).unwrap_or_else(|_| OneOrMany::one(AssistantContent::text("")))
}
pub const fn usage_data(usage: Usage) -> UsageData {
UsageData {
prompt_tokens: usage.prompt_tokens as u64,
cached_tokens: usage.cached_tokens as u64,
completion_tokens: usage.completion_tokens as u64,
}
}
pub fn speculation_data(
stats: Option<&crate::engine::generate::SpeculationStats>,
) -> Option<crate::model::SpeculationStatsData> {
let stats = stats?;
Some(crate::model::SpeculationStatsData {
drafted: stats.drafted,
rounds: stats.rounds,
accepted_by_depth: stats.accepted_by_depth.clone(),
})
}