use std::{
collections::HashMap,
sync::Arc,
time::{Duration, Instant},
};
use futures::{Stream, StreamExt};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio::{
sync::mpsc::{self, Sender},
time,
};
use tokio_stream::wrappers::ReceiverStream;
use crate::{
_prelude::*,
http::{Auth, Client, Sse},
response::{Message, ResponseRequest, Role, StreamOptions},
tool::*,
};
#[derive(Clone)]
pub struct Agent {
client: Client,
options: AgentOptions,
custom_instructions: Option<String>,
tools: HashMap<String, Arc<dyn Tool>>,
}
impl Agent {
pub fn builder(auth: Auth) -> AgentBuilder {
AgentBuilder { auth, options: Default::default(), custom_instructions: Default::default() }
}
pub fn register_tool<T>(&mut self, tool: T)
where
T: 'static + Tool,
{
let name = tool.name().to_string();
tracing::info!("registering tool: {name}");
self.tools.insert(name, Arc::new(tool));
}
pub fn register_tools<I, T>(&mut self, tools: I)
where
I: IntoIterator<Item = T>,
T: 'static + Tool,
{
tools.into_iter().for_each(|tool| self.register_tool(tool));
}
pub fn find_tool(&self, name: &str) -> Option<Arc<dyn Tool>> {
self.tools.get(name).cloned()
}
pub fn list_tools(&self) -> Vec<String> {
self.tools.keys().cloned().collect()
}
pub async fn react_stream(&self, state: AgentState) -> impl Stream<Item = AgentEvent> {
let (tx, rx) = mpsc::channel(32);
let agent = self.clone();
tokio::spawn(async move {
if let Err(e) = run_agent_stream(agent, state, tx.clone()).await {
let _ = tx.send(AgentEvent::err(e.to_string())).await;
}
});
ReceiverStream::new(rx)
}
pub async fn reasoning_stream_with_prompt(
&self,
prompt: String,
) -> Result<impl Stream<Item = String> + '_> {
let system_message = Message { role: Role::Developer, content: self.system_prompt() };
let user_message = Message { role: Role::User, content: prompt };
let params = ResponseRequest {
messages: vec![system_message, user_message],
temperature: Some(self.options.temperature),
max_completion_tokens: Some(self.options.max_completion_tokens),
..Default::default()
};
let sse = self.completion_stream(params).await?;
Ok(Box::pin(sse.filter_map(|event| async move {
match event {
Ok(data) =>
if let Ok(json) = serde_json::from_str::<Value>(&data) {
json.get("choices")?
.get(0)?
.get("delta")?
.get("content")?
.as_str()
.map(|s| s.to_string())
} else {
None
},
Err(e) => {
tracing::warn!("error in reasoning stream: {e}");
None
},
}
})))
}
pub async fn reasoning_stream<'a>(
&'a self,
state: &'a AgentState,
) -> Result<impl Stream<Item = String> + 'a> {
let prompt = self.build_prompt(state);
let system_message = Message { role: Role::Developer, content: self.system_prompt() };
let user_message = Message { role: Role::User, content: prompt };
let params = ResponseRequest {
messages: vec![system_message, user_message],
temperature: Some(self.options.temperature),
max_completion_tokens: Some(self.options.max_completion_tokens),
..Default::default()
};
let sse = self.completion_stream(params).await?;
Ok(Box::pin(sse.filter_map(|event| async move {
match event {
Ok(data) =>
if let Ok(json) = serde_json::from_str::<Value>(&data) {
json.get("choices")?
.get(0)?
.get("delta")?
.get("content")?
.as_str()
.map(|s| s.to_string())
} else {
None
},
Err(e) => {
tracing::warn!("error in reasoning stream: {e}");
None
},
}
})))
}
fn build_prompt(&self, state: &AgentState) -> String {
let mut prompt = format!("Question: {}\n\n", state.input);
for (i, step) in state.reasoning_steps.iter().enumerate() {
prompt.push_str(&format!("Thought {}: {step}\n", i + 1));
let Some(ToolCallResult { tool_call: ToolCall { name, args }, outcome }) =
state.tool_calls.get(i)
else {
continue;
};
prompt.push_str(&format!("Action: {name}\n"));
prompt.push_str(&format!("Action Input: {args}\n"));
match &outcome {
ToolCallOutcome::Success { result } =>
prompt.push_str(&format!("Observation: {result}\n")),
ToolCallOutcome::Error { message } => {
prompt.push_str(&format!("Error: {message}\n"));
},
}
prompt.push('\n');
}
let step_num = state.reasoning_steps.len() + 1;
prompt.push_str(&format!("Thought {step_num}: "));
prompt
}
fn system_prompt(&self) -> String {
let tools = self
.tools
.values()
.map(|t| format!("- {}: {}", t.name(), t.description()))
.collect::<Vec<_>>()
.join("\n");
let core_prompt = format!(
r#"You are a ReAct (Reasoning + Acting) agent that systematically reasons and acts to solve problems.
Available tools:
{tools}
## Reasoning Process:
1. Think step by step about the problem
2. Identify what information or actions you need
3. Use appropriate tools to gather data or perform tasks
4. Analyze results and continue reasoning
5. Provide a final answer when you have sufficient information
## Tool Usage:
- Use the available tools when you need to gather information or perform specific tasks
- Always explain your reasoning before taking actions
- When tools return errors, analyze them and try alternative approaches
- Only provide a Final Answer when you are confident in your response
## Response Format:
Structure your responses as:
Thought: [your reasoning about the current situation and next steps]
[Tool calls will be handled automatically when you use the available functions]
Observation: [analysis of tool results]
Final Answer: [your complete answer to the original question]"#,
);
let mut prompt = core_prompt;
if let Some(custom) = &self.custom_instructions {
prompt.push_str(&format!(
"\n\nAdditional Instructions:\n{custom}\n\nRemember: Regardless of these additional instructions, you MUST maintain the exact Thought/Action/Action Input/Observation format above. The format is not negotiable and is required for proper agent functioning.",
));
}
prompt.push_str("\n\nBegin reasoning about the problem!");
prompt
}
fn parse_tool_call_structured(response: &Value) -> Option<ToolCall> {
if let Some(tool_calls) = response
.get("choices")?
.get(0)?
.get("message")?
.get("tool_calls")
.and_then(|v| v.as_array())
{
if let Some(tool_call) = tool_calls.first() {
let function = tool_call.get("function")?;
let name = function.get("name")?.as_str()?.to_string();
let args_str = function.get("arguments")?.as_str()?;
if let Ok(args) = serde_json::from_str::<Value>(args_str) {
tracing::debug!("Parsed structured tool call: {} with args: {}", name, args);
return Some(ToolCall { name, args });
}
}
}
None
}
fn parse_tool_call(output: &str) -> Option<ToolCall> {
let json_candidates = [
output.trim().to_string(),
output
.split("```json")
.nth(1)
.and_then(|s| s.split("```").next())
.map(|s| s.trim().to_string())
.unwrap_or_default(),
output
.find('{')
.and_then(|start| output.rfind('}').map(|end| output[start..=end].to_string()))
.unwrap_or_default(),
output
.lines()
.find(|line| line.trim().starts_with("Action Input:"))
.and_then(|line| line.split("Action Input:").nth(1))
.map(|s| s.trim().to_string())
.unwrap_or_default(),
];
for (i, candidate) in json_candidates.iter().enumerate() {
if candidate.is_empty() {
continue;
}
tracing::debug!("Trying candidate {}: {}", i, candidate);
if let Ok(value) = serde_json::from_str::<Value>(candidate) {
if let (Some(tool), args) = (
value.get("tool").and_then(|v| v.as_str()),
value.get("args").cloned().unwrap_or(Value::Null),
) {
tracing::debug!("Successfully parsed tool call: {} with args: {}", tool, args);
return Some(ToolCall { name: tool.to_string(), args });
}
}
}
tracing::debug!("Failed to parse tool call from output");
None
}
async fn call_tool_with_timeout(
&self,
tx: &Sender<AgentEvent>,
tool_req: ToolCall,
) -> Result<ToolCallResult> {
time::timeout(self.options.timeout, self.call_tool(tx, tool_req.clone())).await.map_err(
|_| {
let e = Error::Timeout(self.options.timeout);
tracing::error!("{e}");
e
},
)?
}
async fn call_tool(
&self,
tx: &Sender<AgentEvent>,
tool_req: ToolCall,
) -> Result<ToolCallResult> {
let ToolCall { name, args } = &tool_req;
tracing::debug!("calling tool '{name}' with args: {args}");
let Some(tool) = self.find_tool(name) else {
let e = ToolError::Unknown(name.to_owned());
tracing::error!("{e}");
let _ = tx.send(AgentEvent::err(e.to_string())).await;
Err(e)?
};
if tool.supports_stream() {
tracing::debug!("Using streaming execution for tool '{}'", name);
match tool.call_stream(args.clone()).await {
Ok(mut stream) => {
let mut acc = String::new();
while let Some(chunk) = stream.next().await {
let _ = tx
.send(AgentEvent::ToolResult {
name: name.to_string(),
result: Value::String(chunk.clone()),
is_streaming: Some(true),
})
.await;
acc.push_str(&chunk);
}
tracing::debug!("tool '{name}' streaming completed");
return Ok(ToolCallResult::success(
name.to_string(),
args.clone(),
Value::String(acc),
));
},
Err(e) => {
tracing::error!("{e}");
let _ = tx.send(AgentEvent::err(e.to_string())).await;
},
}
}
tracing::debug!("Using synchronous execution for tool '{}'", name);
match tool.call(args.clone()).await {
Ok(result) => {
tracing::debug!("Tool '{}' executed successfully", name);
let _ = tx
.send(AgentEvent::ToolResult {
name: name.to_string(),
result: result.clone(),
is_streaming: Some(false),
})
.await;
Ok(ToolCallResult::success(name.to_string(), args.clone(), result))
},
Err(err) => {
tracing::error!("Tool '{}' execution failed: {}", name, err);
Err(err)
},
}
}
async fn completion<P>(&self, params: P) -> Result<String>
where
P: Into<ResponseRequest>,
{
self.client.post(params.into()).await
}
async fn completion_stream<P>(&self, params: P) -> Result<Sse>
where
P: Into<ResponseRequest>,
{
let mut params = params.into();
params.stream = Some(true);
params.stream_options = Some(StreamOptions { include_usage: true });
self.client.sse_post(params).await
}
pub fn validate_system_prompt(&self) -> Result<()> {
let prompt = self.system_prompt();
let mut issues = Vec::new();
let required_keywords =
["Thought:", "Action:", "Action Input:", "Observation:", "Final Answer:"];
for keyword in &required_keywords {
if !prompt.contains(keyword) {
issues.push(format!("Missing required format keyword: '{}'", keyword));
}
}
if !prompt.contains(r#"{"tool":"#) && !prompt.contains(r#"{"tool": "#) {
issues.push("Missing required tool call JSON format specification".to_string());
}
if self.tools.is_empty() {
issues.push("No tools available to the agent".to_string());
} else {
let tools_mentioned = self.tools.keys().any(|tool_name| prompt.contains(tool_name));
if !tools_mentioned && !prompt.contains("Available tools:") {
issues.push("Tools may not be properly described in the prompt".to_string());
}
}
if issues.is_empty() {
Ok(())
} else {
let e = format!(
"System prompt validation failed with {} issues:\n{}",
issues.len(),
issues.join("\n- ")
);
Err(Error::any(e))
}
}
pub fn get_tool_definitions(&self) -> Vec<Value> {
self.tools
.values()
.map(|tool| {
serde_json::json!({
"type": "function",
"function": {
"name": tool.name(),
"description": tool.description(),
"parameters": tool.schema()
}
})
})
.collect()
}
async fn completion_with_tools<P>(&self, params: P) -> Result<(String, Vec<ToolCall>)>
where
P: Into<ResponseRequest>,
{
let mut params = params.into();
if !self.tools.is_empty() {
params.tools = Some(self.get_tool_definitions());
params.tool_choice = Some(serde_json::json!("auto"));
}
let response_text = self.client.post(params).await?;
if let Ok(response_json) = serde_json::from_str::<Value>(&response_text) {
if let Some(tool_call) = Self::parse_tool_call_structured(&response_json) {
return Ok((response_text, vec![tool_call]));
}
}
let tool_calls =
Self::parse_tool_call(&response_text).map(|tc| vec![tc]).unwrap_or_default();
Ok((response_text, tool_calls))
}
}
pub struct AgentBuilder {
pub auth: Auth,
pub options: AgentOptions,
pub custom_instructions: Option<String>,
}
impl AgentBuilder {
pub fn max_steps(mut self, steps: usize) -> Self {
self.options.max_steps = steps;
self
}
pub fn timeout(mut self, duration: Duration) -> Self {
self.options.timeout = duration;
self
}
pub fn temperature(mut self, temp: f32) -> Self {
self.options.temperature = temp.clamp(0.0, 2.0);
self
}
pub fn max_completion_tokens(mut self, tokens: u32) -> Self {
self.options.max_completion_tokens = tokens;
self
}
pub fn custom_instructions(mut self, instructions: String) -> Self {
self.custom_instructions = Some(instructions);
self
}
pub fn reasoning_effort(mut self, enabled: bool) -> Self {
self.options.reasoning_effort = enabled;
self
}
pub fn build(self, client: Client) -> Agent {
Agent {
client,
options: self.options,
tools: HashMap::new(),
custom_instructions: self.custom_instructions,
}
}
}
#[derive(Clone, Debug)]
pub struct AgentState {
pub input: String,
pub memory: HashMap<String, String>,
pub reasoning_steps: Vec<String>,
pub tool_calls: Vec<ToolCallResult>,
pub metadata: AgentMetadata,
}
impl AgentState {
pub fn new(input: String) -> Self {
tracing::info!("Creating new agent state for input: {}", input);
Self {
input,
reasoning_steps: Vec::new(),
tool_calls: Vec::new(),
memory: HashMap::new(),
metadata: AgentMetadata::new(),
}
}
pub fn add_step(&mut self, reasoning: String) {
tracing::debug!("Adding reasoning step: {}", reasoning);
self.reasoning_steps.push(reasoning);
self.metadata.total_steps += 1;
}
pub fn add_tool_call(&mut self, result: ToolCallResult) {
tracing::debug!("Adding tool call result: {:?}", result);
self.tool_calls.push(result);
self.metadata.tool_calls_count += 1;
}
pub fn remember(&mut self, key: String, value: String) {
tracing::debug!("Storing in memory: {} = {}", key, value);
self.memory.insert(key, value);
}
pub fn recall(&self, key: &str) -> Option<&String> {
self.memory.get(key)
}
pub fn current_step(&self) -> usize {
self.reasoning_steps.len()
}
pub fn total_steps(&self) -> usize {
self.metadata.total_steps
}
pub fn is_complete(&self) -> bool {
self.reasoning_steps.iter().any(|step| step.to_lowercase().contains("final answer:"))
}
}
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct AgentMetadata {
pub total_steps: usize,
pub tool_calls_count: usize,
#[serde(skip)]
pub start_time: Option<Instant>,
#[serde(skip)]
pub end_time: Option<Instant>,
pub duration_ms: Option<u64>,
}
impl AgentMetadata {
pub fn new() -> Self {
Self { start_time: Some(Instant::now()), ..Default::default() }
}
pub fn complete(&mut self) {
let now = Instant::now();
self.end_time = Some(now);
if let Some(start) = self.start_time {
let duration = now.duration_since(start);
self.duration_ms = Some(duration.as_millis() as u64);
}
}
pub fn get_duration(&self) -> Option<Duration> {
if let Some(ms) = self.duration_ms {
Some(Duration::from_millis(ms))
} else if let (Some(start), Some(end)) = (self.start_time, self.end_time) {
Some(end.duration_since(start))
} else {
self.start_time.map(|start| Instant::now().duration_since(start))
}
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum AgentEvent {
#[serde(rename = "reasoningToken")]
ReasoningToken { content: String },
#[serde(rename = "reasoningStepDone")]
ReasoningStepDone { content: String },
#[serde(rename = "toolCall")]
ToolCall { name: String, args: Value },
#[serde(rename = "toolResult")]
ToolResult {
name: String,
result: Value,
#[serde(skip_serializing_if = "Option::is_none")]
is_streaming: Option<bool>,
},
#[serde(rename = "finalAnswer")]
FinalAnswer { content: String },
#[serde(rename = "error")]
Error { message: String },
#[serde(rename = "metadata")]
Metadata {
#[serde(flatten)]
data: AgentMetadata,
},
#[serde(rename = "started")]
Started { max_steps: usize, tools: Vec<String> },
#[serde(rename = "completed")]
Completed { success: bool, total_steps: usize, duration: Option<Duration> },
}
impl AgentEvent {
pub fn err(message: impl Into<String>) -> Self {
Self::Error { message: message.into() }
}
pub fn reasoning_token(content: impl Into<String>) -> Self {
Self::ReasoningToken { content: content.into() }
}
pub fn started(max_steps: usize, tools: Vec<String>) -> Self {
Self::Started { max_steps, tools }
}
pub fn completed(success: bool, total_steps: usize, duration: Option<Duration>) -> Self {
Self::Completed { success, total_steps, duration }
}
}
#[derive(Clone, Debug)]
pub struct AgentOptions {
pub max_steps: usize,
pub timeout: Duration,
pub temperature: f32,
pub max_completion_tokens: u32,
pub reasoning_effort: bool,
}
impl Default for AgentOptions {
fn default() -> Self {
Self {
max_steps: 10,
timeout: Duration::from_secs(300),
temperature: 0.7,
max_completion_tokens: 4000,
reasoning_effort: false,
}
}
}
async fn run_agent_stream(
agent: Agent,
mut state: AgentState,
tx: Sender<AgentEvent>,
) -> Result<()> {
tracing::info!("Starting agent execution for input: {}", state.input);
let _ = tx.send(AgentEvent::started(agent.options.max_steps, agent.list_tools())).await;
for step in 0..agent.options.max_steps {
tracing::debug!("Starting step {} of {}", step + 1, agent.options.max_steps);
if state.is_complete() {
if let Some(answer) = extract_final_answer(state.reasoning_steps.last().unwrap()) {
tracing::info!("Agent found final answer: {}", answer);
let _ = tx.send(AgentEvent::FinalAnswer { content: answer }).await;
let _ = tx
.send(AgentEvent::completed(
true,
state.total_steps(),
state.metadata.get_duration(),
))
.await;
return Ok(());
}
}
let prompt = agent.build_prompt(&state);
let reasoning_res = agent.reasoning_stream_with_prompt(prompt).await;
match reasoning_res {
Ok(mut stream) => {
let mut full_reasoning = String::new();
let mut token_buffer = Vec::new();
while let Some(token) = stream.next().await {
token_buffer.push(token.clone());
full_reasoning.push_str(&token);
if token_buffer.len() >= 5 {
for chunk in token_buffer.drain(..) {
let _ = tx.send(AgentEvent::reasoning_token(chunk)).await;
}
}
}
for chunk in token_buffer {
let _ = tx.send(AgentEvent::reasoning_token(chunk)).await;
}
if full_reasoning.trim().is_empty() {
tracing::warn!("Empty reasoning generated at step {}", step + 1);
continue;
}
tracing::debug!("generated reasoning: {full_reasoning}");
let _ = tx
.send(AgentEvent::ReasoningStepDone { content: full_reasoning.clone() })
.await;
state.add_step(full_reasoning.clone());
if full_reasoning.to_lowercase().contains("final answer:") {
if let Some(answer) = extract_final_answer(&full_reasoning) {
tracing::info!("Agent provided final answer: {}", answer);
let _ = tx.send(AgentEvent::FinalAnswer { content: answer }).await;
state.metadata.complete();
let _ = tx
.send(AgentEvent::completed(
true,
state.total_steps(),
state.metadata.get_duration(),
))
.await;
return Ok(());
}
}
if let Some(tool_req) = Agent::parse_tool_call(&full_reasoning) {
tracing::info!(
"Parsed tool call: {} with args: {}",
tool_req.name,
tool_req.args
);
let _ = tx
.send(AgentEvent::ToolCall {
name: tool_req.name.clone(),
args: tool_req.args.clone(),
})
.await;
match agent.call_tool_with_timeout(&tx, tool_req.clone()).await {
Ok(result) => {
tracing::info!("tool call successful: {}", result.tool_call.name);
state.add_tool_call(result);
},
Err(e) => {
let _ = tx.send(AgentEvent::err(e.to_string())).await;
state.add_tool_call(ToolCallResult::err(
tool_req.name.clone(),
tool_req.args.clone(),
e.to_string(),
));
},
}
}
},
Err(e) => {
tracing::error!("reasoning failed at step {}: {e}", step + 1);
let _ = tx.send(AgentEvent::err(e.to_string())).await;
state.metadata.complete();
let _ = tx
.send(AgentEvent::completed(
false,
state.total_steps(),
state.metadata.get_duration(),
))
.await;
return Err(e);
},
}
let _ = tx.send(AgentEvent::Metadata { data: state.metadata.clone() }).await;
}
tracing::warn!(
"Agent reached maximum steps ({}) without final answer",
agent.options.max_steps
);
state.metadata.complete();
let e = AgentError::MaxStepsExceeded(agent.options.max_steps);
let _ = tx.send(AgentEvent::err(e.to_string())).await;
let _ = tx
.send(AgentEvent::completed(false, state.total_steps(), state.metadata.get_duration()))
.await;
Ok(())
}
fn extract_final_answer(text: &str) -> Option<String> {
let lower_text = text.to_lowercase();
if let Some(pos) = lower_text.find("final answer:") {
let after_marker = &text[pos + "final answer:".len()..];
let answer = after_marker.trim();
if !answer.is_empty() {
return Some(answer.to_string());
}
}
None
}