From 8af48390a9b89ba8d0f7d7fb20932f13322dd8d8 Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Tue, 3 Feb 2026 10:50:20 -0800 Subject: [PATCH] Tool use --- src/agent/agent_loop.rs | 184 +++++++++++++++++++++++++++++++-- src/llm/mod.rs | 2 +- src/llm/nearai.rs | 147 +++++++++++++++++++++++++-- src/llm/reasoning.rs | 218 +++++++++++++++++++++------------------- 4 files changed, 431 insertions(+), 120 deletions(-) diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index cf6a6492..9c61b671 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -19,9 +19,10 @@ use crate::agent::{ use crate::channels::{ChannelManager, IncomingMessage, OutgoingResponse, StatusUpdate}; use crate::config::{AgentConfig, HeartbeatConfig}; use crate::context::ContextManager; +use crate::context::JobContext; use crate::error::Error; use crate::history::Store; -use crate::llm::{LlmProvider, Reasoning, ReasoningContext}; +use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult}; use crate::safety::SafetyLayer; use crate::tools::ToolRegistry; use crate::workspace::Workspace; @@ -368,13 +369,10 @@ impl Agent { ) .await; - // Call LLM with thread context and available tools - let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone()); - let tool_defs = self.tools.tool_definitions().await; - let context = ReasoningContext::new() - .with_messages(turn_messages) - .with_tools(tool_defs); - let llm_result = reasoning.respond(&context).await; + // Run the agentic tool execution loop + let result = self + .run_agentic_loop(message, session.clone(), thread_id, turn_messages) + .await; // Re-acquire lock and check if interrupted let mut sess = session.lock().await; @@ -392,7 +390,7 @@ impl Agent { } // Complete or fail the turn - match llm_result { + match result { Ok(response) => { thread.complete_turn(&response); let _ = self @@ -408,6 +406,174 @@ impl Agent { } } + /// Run the agentic loop: call LLM, execute tools, repeat until text response. + async fn run_agentic_loop( + &self, + message: &IncomingMessage, + session: Arc>, + thread_id: Uuid, + initial_messages: Vec, + ) -> Result { + let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone()); + let tool_defs = self.tools.tool_definitions().await; + + // Build context with messages that we'll mutate during the loop + let mut context_messages = initial_messages; + + // Create a JobContext for tool execution (chat doesn't have a real job) + let job_ctx = JobContext::with_user(&message.user_id, "chat", "Interactive chat session"); + + const MAX_TOOL_ITERATIONS: usize = 10; + let mut iteration = 0; + + loop { + iteration += 1; + if iteration > MAX_TOOL_ITERATIONS { + return Err(crate::error::LlmError::InvalidResponse { + provider: "nearai".to_string(), + reason: format!("Exceeded maximum tool iterations ({})", MAX_TOOL_ITERATIONS), + } + .into()); + } + + // Check if interrupted + { + let sess = session.lock().await; + if let Some(thread) = sess.threads.get(&thread_id) { + if thread.state == ThreadState::Interrupted { + return Err(crate::error::JobError::ContextError { + id: thread_id, + reason: "Interrupted".to_string(), + } + .into()); + } + } + } + + // Call LLM with current context + let context = ReasoningContext::new() + .with_messages(context_messages.clone()) + .with_tools(tool_defs.clone()); + + let result = reasoning.respond_with_tools(&context).await?; + + match result { + RespondResult::Text(text) => { + // Final response, return it + return Ok(text); + } + RespondResult::ToolCalls(tool_calls) => { + // Execute tools and add results to context + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Thinking(format!( + "Executing {} tool(s)...", + tool_calls.len() + )), + ) + .await; + + // Record tool calls in the thread + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + if let Some(turn) = thread.last_turn_mut() { + for tc in &tool_calls { + turn.record_tool_call(&tc.name, tc.arguments.clone()); + } + } + } + } + + // Execute each tool + for tc in tool_calls { + let tool_result = self + .execute_chat_tool(&tc.name, &tc.arguments, &job_ctx) + .await; + + // Record result in thread + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + if let Some(turn) = thread.last_turn_mut() { + match &tool_result { + Ok(output) => { + turn.record_tool_result(serde_json::json!(output)); + } + Err(e) => { + turn.record_tool_error(e.to_string()); + } + } + } + } + } + + // Add tool result to context for next LLM call + let result_content = match tool_result { + Ok(output) => { + // Sanitize output before showing to LLM + let sanitized = self.safety.sanitize_tool_output(&tc.name, &output); + self.safety.wrap_for_llm( + &tc.name, + &sanitized.content, + sanitized.was_modified, + ) + } + Err(e) => format!("Error: {}", e), + }; + + context_messages.push(ChatMessage::tool_result( + &tc.id, + &tc.name, + result_content, + )); + } + } + } + } + } + + /// Execute a tool for chat (without full job context). + async fn execute_chat_tool( + &self, + tool_name: &str, + params: &serde_json::Value, + job_ctx: &JobContext, + ) -> Result { + let tool = + self.tools + .get(tool_name) + .await + .ok_or_else(|| crate::error::ToolError::NotFound { + name: tool_name.to_string(), + })?; + + // Execute with timeout + let result = tokio::time::timeout(std::time::Duration::from_secs(60), async { + tool.execute(params.clone(), job_ctx).await + }) + .await + .map_err(|_| crate::error::ToolError::Timeout { + name: tool_name.to_string(), + timeout: std::time::Duration::from_secs(60), + })? + .map_err(|e| crate::error::ToolError::ExecutionFailed { + name: tool_name.to_string(), + reason: e.to_string(), + })?; + + // Convert result to string + serde_json::to_string_pretty(&result.result).map_err(|e| { + crate::error::ToolError::ExecutionFailed { + name: tool_name.to_string(), + reason: format!("Failed to serialize result: {}", e), + } + .into() + }) + } + /// Handle job-related intents without turn tracking. async fn handle_job_or_command( &self, diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 528119c3..69d85d42 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -12,7 +12,7 @@ pub use provider::{ ChatMessage, CompletionRequest, CompletionResponse, LlmProvider, Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult, }; -pub use reasoning::{ActionPlan, Reasoning, ReasoningContext, ToolSelection}; +pub use reasoning::{ActionPlan, Reasoning, ReasoningContext, RespondResult, ToolSelection}; pub use session::{SessionConfig, SessionManager, create_session_manager}; use std::sync::Arc; diff --git a/src/llm/nearai.rs b/src/llm/nearai.rs index 5e34f1f0..0acf93ba 100644 --- a/src/llm/nearai.rs +++ b/src/llm/nearai.rs @@ -153,14 +153,39 @@ impl NearAiProvider { } } +/// Split messages into system instructions and non-system input messages. +/// The OpenAI Responses API expects system prompts in an `instructions` field, +/// not as a message with role "system" in the input array. +fn split_messages(messages: Vec) -> (Option, Vec) { + let mut instructions: Vec = Vec::new(); + let mut input: Vec = Vec::new(); + + for msg in messages { + if msg.role == Role::System { + instructions.push(msg.content); + } else { + input.push(msg.into()); + } + } + + let instructions = if instructions.is_empty() { + None + } else { + Some(instructions.join("\n\n")) + }; + + (instructions, input) +} + #[async_trait] impl LlmProvider for NearAiProvider { async fn complete(&self, req: CompletionRequest) -> Result { - let messages: Vec = req.messages.into_iter().map(Into::into).collect(); + let (instructions, input) = split_messages(req.messages); let request = NearAiRequest { model: self.config.model.clone(), - input: messages, + instructions, + input, temperature: req.temperature, max_output_tokens: req.max_tokens, stream: Some(false), @@ -272,7 +297,7 @@ impl LlmProvider for NearAiProvider { &self, req: ToolCompletionRequest, ) -> Result { - let messages: Vec = req.messages.into_iter().map(Into::into).collect(); + let (instructions, input) = split_messages(req.messages); let tools: Vec = req .tools @@ -287,7 +312,8 @@ impl LlmProvider for NearAiProvider { let request = NearAiRequest { model: self.config.model.clone(), - input: messages, + instructions, + input, temperature: req.temperature, max_output_tokens: req.max_tokens, stream: Some(false), @@ -302,16 +328,28 @@ impl LlmProvider for NearAiProvider { // Try parsing as alternative response format if let Ok(alt) = serde_json::from_str::(raw_text) { - tracing::info!("NEAR AI returned alternative response format (tool request)"); let text = extract_text_from_output(&alt.output); + let tool_calls = extract_tool_calls_from_output(&alt.output); let usage = alt.usage.unwrap_or(NearAiUsage { input_tokens: 0, output_tokens: 0, }); + + let finish_reason = if tool_calls.is_empty() { + FinishReason::Stop + } else { + FinishReason::ToolUse + }; + + tracing::info!( + "NEAR AI returned alternative response format ({} tool calls)", + tool_calls.len() + ); + return Ok(ToolCompletionResponse { content: if text.is_empty() { None } else { Some(text) }, - tool_calls: vec![], - finish_reason: FinishReason::Stop, + tool_calls, + finish_reason, input_tokens: usage.input_tokens, output_tokens: usage.output_tokens, }); @@ -403,7 +441,10 @@ impl LlmProvider for NearAiProvider { struct NearAiRequest { /// Model identifier (e.g., "fireworks::accounts/fireworks/models/llama-v3p1-405b-instruct") model: String, - /// Input messages - can be a string or array of message objects + /// System instructions (replaces sending system role in input) + #[serde(skip_serializing_if = "Option::is_none")] + instructions: Option, + /// Input messages (user/assistant/tool only, NOT system) input: Vec, #[serde(skip_serializing_if = "Option::is_none")] temperature: Option, @@ -524,6 +565,10 @@ fn extract_text_from_output(output: &Option) -> String { let texts: Vec = arr .iter() .filter_map(|item| { + // Skip function_call items + if item.get("type").and_then(|t| t.as_str()) == Some("function_call") { + return None; + } // Check for direct text field if let Some(text) = item.get("text").and_then(|t| t.as_str()) { return Some(text.to_string()); @@ -566,6 +611,49 @@ fn extract_text_from_output(output: &Option) -> String { output.to_string() } +/// Extract tool calls from alternative output format. +fn extract_tool_calls_from_output(output: &Option) -> Vec { + let Some(output) = output else { + return vec![]; + }; + + let Some(arr) = output.as_array() else { + return vec![]; + }; + + arr.iter() + .filter_map(|item| { + // Look for function_call type items + let item_type = item.get("type").and_then(|t| t.as_str())?; + if item_type != "function_call" { + return None; + } + + let name = item.get("name").and_then(|n| n.as_str())?; + let call_id = item + .get("call_id") + .and_then(|c| c.as_str()) + .unwrap_or("unknown"); + + // Arguments can be a string (JSON) or already an object + let arguments = if let Some(args_str) = item.get("arguments").and_then(|a| a.as_str()) { + serde_json::from_str(args_str) + .unwrap_or(serde_json::Value::Object(Default::default())) + } else if let Some(args_obj) = item.get("arguments") { + args_obj.clone() + } else { + serde_json::Value::Object(Default::default()) + }; + + Some(ToolCall { + id: call_id.to_string(), + name: name.to_string(), + arguments, + }) + }) + .collect() +} + #[cfg(test)] mod tests { use super::*; @@ -584,4 +672,47 @@ mod tests { let nearai_msg: NearAiMessage = msg.into(); assert_eq!(nearai_msg.role, "system"); } + + #[test] + fn test_split_messages_with_system() { + let messages = vec![ + ChatMessage::system("You are a helpful assistant"), + ChatMessage::user("Hello"), + ChatMessage::assistant("Hi there!"), + ]; + let (instructions, input) = split_messages(messages); + assert_eq!( + instructions, + Some("You are a helpful assistant".to_string()) + ); + assert_eq!(input.len(), 2); + assert_eq!(input[0].role, "user"); + assert_eq!(input[1].role, "assistant"); + } + + #[test] + fn test_split_messages_no_system() { + let messages = vec![ + ChatMessage::user("Hello"), + ChatMessage::assistant("Hi there!"), + ]; + let (instructions, input) = split_messages(messages); + assert!(instructions.is_none()); + assert_eq!(input.len(), 2); + } + + #[test] + fn test_split_messages_multiple_system() { + let messages = vec![ + ChatMessage::system("First instruction"), + ChatMessage::system("Second instruction"), + ChatMessage::user("Hello"), + ]; + let (instructions, input) = split_messages(messages); + assert_eq!( + instructions, + Some("First instruction\n\nSecond instruction".to_string()) + ); + assert_eq!(input.len(), 1); + } } diff --git a/src/llm/reasoning.rs b/src/llm/reasoning.rs index 46035594..2206b17e 100644 --- a/src/llm/reasoning.rs +++ b/src/llm/reasoning.rs @@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize}; use crate::error::LlmError; use crate::llm::{ - ChatMessage, CompletionRequest, LlmProvider, ToolCompletionRequest, ToolDefinition, + ChatMessage, CompletionRequest, LlmProvider, ToolCall, ToolCompletionRequest, ToolDefinition, }; use crate::safety::SafetyLayer; @@ -105,6 +105,17 @@ pub struct ToolSelection { pub alternatives: Vec, } +/// Result of a response with potential tool calls. +/// +/// Used by the agent loop to handle tool execution before returning a final response. +#[derive(Debug, Clone)] +pub enum RespondResult { + /// A text response (no tools needed). + Text(String), + /// The model wants to call tools. Caller should execute them and call back. + ToolCalls(Vec), +} + /// Reasoning engine for the agent. pub struct Reasoning { llm: Arc, @@ -236,7 +247,32 @@ Respond in JSON format: /// Generate a response to a user message. /// /// If tools are available in the context, uses tool completion mode. + /// This is a convenience wrapper around `respond_with_tools()` that formats + /// tool calls as text for simple cases. Use `respond_with_tools()` when you + /// need to actually execute tool calls in an agentic loop. pub async fn respond(&self, context: &ReasoningContext) -> Result { + match self.respond_with_tools(context).await? { + RespondResult::Text(text) => Ok(text), + RespondResult::ToolCalls(calls) => { + // Format tool calls as text (legacy behavior for non-agentic callers) + let tool_info: Vec = calls + .iter() + .map(|tc| format!("`{}({})`", tc.name, tc.arguments)) + .collect(); + Ok(format!("[Calling tools: {}]", tool_info.join(", "))) + } + } + } + + /// Generate a response that may include tool calls. + /// + /// Returns `RespondResult::ToolCalls` if the model wants to call tools, + /// allowing the caller to execute them and continue the conversation. + /// Returns `RespondResult::Text` when the model has a final text response. + pub async fn respond_with_tools( + &self, + context: &ReasoningContext, + ) -> Result { let system_prompt = self.build_conversation_prompt(context); let mut messages = vec![ChatMessage::system(system_prompt)]; @@ -251,22 +287,16 @@ Respond in JSON format: let response = self.llm.complete_with_tools(request).await?; - // If there were tool calls, the content is usually just internal reasoning - // Don't show it - just acknowledge the tool calls (actual execution handled by caller) + // If there were tool calls, return them for execution if !response.tool_calls.is_empty() { - let tool_info: Vec = response - .tool_calls - .iter() - .map(|tc| format!("`{}({})`", tc.name, tc.arguments)) - .collect(); - return Ok(format!("[Calling tools: {}]", tool_info.join(", "))); + return Ok(RespondResult::ToolCalls(response.tool_calls)); } // No tool calls - clean up the response let content = response .content .unwrap_or_else(|| "I'm not sure how to respond to that.".to_string()); - Ok(clean_response(&content)) + Ok(RespondResult::Text(clean_response(&content))) } else { // No tools, use simple completion let request = CompletionRequest::new(messages) @@ -274,7 +304,7 @@ Respond in JSON format: .with_temperature(0.7); let response = self.llm.complete(request).await?; - Ok(clean_response(&response.content)) + Ok(RespondResult::Text(clean_response(&response.content))) } } @@ -331,7 +361,7 @@ Respond with a JSON plan in this format: .map(|t| format!(" - {}: {}", t.name, t.description)) .collect(); format!( - "\n\n## Available Tools\nYou have access to these tools:\n{}\n\nCall tools directly when needed - don't announce what you're going to do.", + "\n\n## Available Tools\nYou have access to these tools:\n{}\n\nCall tools directly when needed.", tool_list.join("\n") ) }; @@ -339,25 +369,26 @@ Respond with a JSON plan in this format: format!( r#"You are NEAR AI Agent, an autonomous assistant. -CRITICAL: Never output your internal reasoning or thinking process. Your response must contain ONLY the final answer or action. +## Response Format -FORBIDDEN patterns (never start with these): -- "The user wants..." / "The user is asking..." -- "I need to..." / "I should..." / "I will..." -- "Let me think..." / "Let me first..." -- "This is a request to..." -- Any self-narration about what you're doing +If you need to think through a problem, wrap your thinking in tags. Everything outside these tags goes directly to the user. -CORRECT behavior: -- Answer questions directly -- Call tools without announcing it -- Ask clarifying questions if genuinely needed -- Provide code/content without preamble{} +Example: + +Let me consider the options... +Option 1: ... +Option 2: ... +I'll go with option 1. + +Here's the solution: [actual response to user] -## Format -- Be concise -- Use markdown where helpful -- Code blocks with language tags"#, +## Guidelines +- Be concise and direct +- Use markdown formatting where helpful +- For code, use appropriate code blocks with language tags +- Call tools when they would help accomplish the task{} + +The user sees ONLY content outside tags."#, tools_section ) } @@ -449,68 +480,44 @@ fn strip_thinking_tags(text: &str) -> String { cleaned } -/// Strip common reasoning/thinking patterns from the start of responses. +/// Strip any remaining reasoning that wasn't in proper tags. /// -/// Models sometimes output their thinking process as plain text despite -/// instructions not to. This strips common patterns like "The user wants...", -/// "Let me think...", "I need to...", etc. +/// This is a simple fallback for models that don't follow the tag +/// instruction. It looks for paragraph breaks where actual content follows. fn strip_reasoning_patterns(text: &str) -> String { let text = text.trim(); + if text.is_empty() { + return text.to_string(); + } - // Patterns that indicate internal reasoning (case-insensitive check) - let reasoning_prefixes = [ - "the user wants", - "the user is asking", - "the user would like", - "i need to", - "i should", - "i will", - "i'll", - "let me think", - "let me first", - "let me check", - "let me look", - "let me explore", - "let me search", - "this is a request", - "this request", - "to answer this", - "to help with this", - "first, i", - "okay, so", - "alright, ", - ]; + // If text already looks clean (starts with actual content), return as-is + // Actual content often starts with: markdown, code blocks, direct statements + let first_char = text.chars().next().unwrap_or(' '); + if first_char == '#' || first_char == '`' || first_char == '*' || first_char == '-' { + return text.to_string(); + } - // Find where reasoning ends and actual content begins - // Look for paragraph breaks or sentences that start the actual response - let lines: Vec<&str> = text.lines().collect(); - let mut skip_until = 0; - - for (i, line) in lines.iter().enumerate() { - let lower = line.to_lowercase(); - - // Check if this line starts with a reasoning pattern - let is_reasoning = reasoning_prefixes - .iter() - .any(|p| lower.trim_start().starts_with(p)); - - if is_reasoning { - skip_until = i + 1; - } else if !line.trim().is_empty() && skip_until <= i { - // Found non-reasoning content, stop looking - break; + // Look for paragraph break followed by actual content + // Often models output: "thinking...\n\nActual response" + if let Some(idx) = text.find("\n\n") { + let after_break = text[idx + 2..].trim(); + if !after_break.is_empty() { + let first_after = after_break.chars().next().unwrap_or(' '); + // If it starts with typical response markers, use content after break + if first_after == '#' + || first_after == '`' + || first_after == '*' + || first_after == '-' + || after_break.to_lowercase().starts_with("here") + || after_break.to_lowercase().starts_with("i'd") + || after_break.to_lowercase().starts_with("sure") + { + return after_break.to_string(); + } } } - if skip_until > 0 && skip_until < lines.len() { - // Skip the reasoning lines and return the rest - let result = lines[skip_until..].join("\n").trim().to_string(); - if !result.is_empty() { - return result; - } - } - - // If we'd strip everything, just return the original (better than empty) + // Return original if no clear split found text.to_string() } @@ -582,43 +589,50 @@ Here is my response to your question."#; } #[test] - fn test_strip_reasoning_patterns_basic() { - let input = "The user wants me to implement something.\n\nHere's the implementation:"; + fn test_strip_reasoning_paragraph_break() { + // Content after paragraph break with "here" marker + let input = "Some thinking here.\n\nHere's the answer:"; let output = strip_reasoning_patterns(input); - assert_eq!(output, "Here's the implementation:"); + assert_eq!(output, "Here's the answer:"); } #[test] - fn test_strip_reasoning_patterns_multiline() { - let input = r#"The user is asking about Telegram. -I need to think about what this involves. -Let me first check the existing code. - -Here's what I found in the codebase."#; + fn test_strip_reasoning_markdown_after_break() { + // Content after paragraph break starting with markdown + let input = "Some reasoning.\n\n**The Solution**\n- Item 1"; let output = strip_reasoning_patterns(input); - assert_eq!(output, "Here's what I found in the codebase."); + assert_eq!(output, "**The Solution**\n- Item 1"); } #[test] - fn test_strip_reasoning_no_patterns() { - let input = "Here's a direct answer to your question."; + fn test_strip_reasoning_preserves_markdown_start() { + // If response starts with markdown, keep as-is + let input = "**What type of tool?**\n- Option 1\n- Option 2"; let output = strip_reasoning_patterns(input); - assert_eq!(output, "Here's a direct answer to your question."); + assert_eq!(output, "**What type of tool?**\n- Option 1\n- Option 2"); } #[test] - fn test_strip_reasoning_preserves_all_if_only_reasoning() { - // If stripping would leave nothing, keep the original - let input = "The user wants to know X."; + fn test_strip_reasoning_preserves_code_start() { + // If response starts with code block, keep as-is + let input = "```rust\nfn main() {}\n```"; let output = strip_reasoning_patterns(input); - assert_eq!(output, "The user wants to know X."); + assert_eq!(output, "```rust\nfn main() {}\n```"); + } + + #[test] + fn test_strip_reasoning_no_paragraph_break() { + // Without clear paragraph break, return original + let input = "Some text without clear separation."; + let output = strip_reasoning_patterns(input); + assert_eq!(output, "Some text without clear separation."); } #[test] fn test_clean_response_combined() { - let input = - "Internal thoughtI need to check this.\n\nActual response here."; + // Combines thinking tags + paragraph break fallback + let input = "Internal thoughtSome text.\n\nHere's the answer."; let output = clean_response(input); - assert_eq!(output, "Actual response here."); + assert_eq!(output, "Here's the answer."); } }