use std::sync::Arc;
use futures::stream::BoxStream;
use crate::error::Error;
use crate::provider::Provider;
use crate::stream_event::StreamEvent;
use crate::tool::Tool;
use crate::value_objects::{Message, UserMessage};
use super::response::TextResponse;
#[derive(Debug, Clone, Default)]
pub enum ToolChoice {
#[default]
Auto,
None,
Any,
Tool(String),
}
#[derive(Clone)]
pub struct TextRequest {
pub model: String,
pub system_prompts: Vec<String>,
pub messages: Vec<Message>,
pub max_tokens: Option<u32>,
pub temperature: Option<f32>,
pub top_p: Option<f32>,
pub max_steps: u32,
pub tools: Vec<Arc<dyn Tool>>,
pub provider_tools: Vec<serde_json::Value>,
pub tool_choice: ToolChoice,
pub cache_system_prompt: bool,
pub thinking_budget: Option<u32>,
pub reasoning_effort: Option<String>,
pub provider_options: serde_json::Value,
}
impl TextRequest {
pub fn new(model: impl Into<String>) -> Self {
Self {
model: model.into(),
system_prompts: Vec::new(),
messages: Vec::new(),
max_tokens: None,
temperature: None,
top_p: None,
max_steps: 1,
tools: Vec::new(),
provider_tools: Vec::new(),
tool_choice: ToolChoice::Auto,
cache_system_prompt: false,
thinking_budget: None,
reasoning_effort: None,
provider_options: serde_json::Value::Null,
}
}
}
pub struct PendingTextRequest {
provider: Arc<dyn Provider>,
request: TextRequest,
}
impl PendingTextRequest {
pub fn new(provider: Arc<dyn Provider>, model: impl Into<String>) -> Self {
Self {
provider,
request: TextRequest::new(model),
}
}
pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
self.request.system_prompts.push(prompt.into());
self
}
pub fn with_prompt(mut self, prompt: impl Into<String>) -> Self {
self.request
.messages
.push(Message::User(UserMessage::text(prompt)));
self
}
pub fn with_message(mut self, message: Message) -> Self {
self.request.messages.push(message);
self
}
pub fn with_messages(mut self, messages: impl IntoIterator<Item = Message>) -> Self {
self.request.messages.extend(messages);
self
}
pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
self.request.max_tokens = Some(max_tokens);
self
}
pub fn with_temperature(mut self, temperature: f32) -> Self {
self.request.temperature = Some(temperature);
self
}
pub fn with_top_p(mut self, top_p: f32) -> Self {
self.request.top_p = Some(top_p);
self
}
pub fn with_max_steps(mut self, max_steps: u32) -> Self {
self.request.max_steps = max_steps;
self
}
pub fn with_tool(mut self, tool: impl Tool + 'static) -> Self {
self.request.tools.push(Arc::new(tool));
self
}
pub fn with_tools(mut self, tools: impl IntoIterator<Item = Arc<dyn Tool>>) -> Self {
self.request.tools.extend(tools);
self
}
pub fn with_provider_tool(mut self, tool: serde_json::Value) -> Self {
self.request.provider_tools.push(tool);
self
}
pub fn with_tool_choice(mut self, choice: ToolChoice) -> Self {
self.request.tool_choice = choice;
self
}
pub fn with_prompt_caching(mut self) -> Self {
self.request.cache_system_prompt = true;
self
}
pub fn with_thinking_budget(mut self, budget_tokens: u32) -> Self {
self.request.thinking_budget = Some(budget_tokens);
self
}
pub fn with_reasoning_effort(mut self, effort: impl Into<String>) -> Self {
self.request.reasoning_effort = Some(effort.into());
self
}
pub fn to_request(&self) -> TextRequest {
self.request.clone()
}
pub async fn generate(self) -> Result<TextResponse, Error> {
crate::tool_loop::run_text(self.provider, self.request).await
}
pub fn stream(self) -> BoxStream<'static, Result<StreamEvent, Error>> {
crate::stream_loop::stream_text(self.provider, self.request)
}
}