use serde::Serialize;
#[derive(Debug, Clone)]
pub enum Provider {
Gemini,
Anthropic,
}
#[derive(Debug, Clone)]
pub struct GenerateRequest {
pub provider: Provider,
pub model: String,
pub messages: Vec<Message>,
pub system: Option<String>,
pub max_tokens: Option<u32>,
pub temperature: Option<f32>,
pub top_p: Option<f32>,
pub top_k: Option<u32>,
pub tools: Option<Vec<ToolDefinition>>,
}
#[derive(Debug, Clone)]
pub struct Message {
pub role: Role,
pub content: Vec<ContentPart>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Role {
User,
Assistant,
}
#[derive(Debug, Clone)]
pub enum ContentPart {
Text(String),
ToolCall {
id: String,
name: String,
arguments: serde_json::Value,
},
ToolResult {
tool_call_id: String,
name: String,
content: serde_json::Value,
},
}
#[derive(Debug, Clone, Serialize)]
pub struct ToolDefinition {
pub name: String,
pub description: String,
pub parameters: serde_json::Value,
}
#[derive(Debug, Clone)]
pub struct GenerateResponse {
pub content: Vec<ContentPart>,
pub stop_reason: Option<String>,
pub usage: Option<Usage>,
}
#[derive(Debug, Clone)]
pub struct Usage {
pub input_tokens: Option<u32>,
pub output_tokens: Option<u32>,
pub total_tokens: Option<u32>,
}
#[derive(Debug, Clone)]
pub enum StreamChunk {
Text(String),
Done { finish_reason: String },
}
impl Message {
pub fn user(text: &str) -> Self {
Self {
role: Role::User,
content: vec![ContentPart::Text(text.into())],
}
}
pub fn assistant(text: &str) -> Self {
Self {
role: Role::Assistant,
content: vec![ContentPart::Text(text.into())],
}
}
pub fn tool_result(tool_call_id: &str, name: &str, content: serde_json::Value) -> Self {
Self {
role: Role::User,
content: vec![ContentPart::ToolResult {
tool_call_id: tool_call_id.into(),
name: name.into(),
content,
}],
}
}
}
impl GenerateRequest {
pub fn new(provider: Provider, model: &str, messages: Vec<Message>) -> Self {
Self {
provider,
model: model.into(),
messages,
system: None,
max_tokens: None,
temperature: None,
top_p: None,
top_k: None,
tools: None,
}
}
pub fn with_system(mut self, system: &str) -> Self {
self.system = Some(system.into());
self
}
pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
self.max_tokens = Some(max_tokens);
self
}
pub fn with_temperature(mut self, temperature: f32) -> Self {
self.temperature = Some(temperature);
self
}
pub fn with_top_p(mut self, top_p: f32) -> Self {
self.top_p = Some(top_p);
self
}
pub fn with_top_k(mut self, top_k: u32) -> Self {
self.top_k = Some(top_k);
self
}
pub fn with_tools(mut self, tools: Vec<ToolDefinition>) -> Self {
self.tools = Some(tools);
self
}
}
impl GenerateResponse {
pub fn text(&self) -> Option<String> {
let text: String = self
.content
.iter()
.filter_map(|p| match p {
ContentPart::Text(t) => Some(t.as_str()),
_ => None,
})
.collect();
if text.is_empty() { None } else { Some(text) }
}
pub fn tool_calls(&self) -> Vec<&ContentPart> {
self.content
.iter()
.filter(|p| matches!(p, ContentPart::ToolCall { .. }))
.collect()
}
pub fn stop_reason(&self) -> Option<&str> {
self.stop_reason.as_deref()
}
}