Add support for tool execution

crates/dspy-rs/src/adapter/chat.rs

@@ -1,4 +1,5 @@
 use anyhow::Result;
+use rig::tool::ToolDyn;
 use serde_json::{Value, json};
 use std::collections::HashMap;
 use std::sync::Arc;
@@ -281,6 +282,7 @@ impl Adapter for ChatAdapter {
         lm: Arc<LM>,
         signature: &dyn MetaSignature,
         inputs: Example,
+        tools: Vec<Arc<dyn ToolDyn>>,
     ) -> Result<Prediction> {
         // Check cache first (release lock immediately after checking)
         if lm.cache
@@ -293,10 +295,28 @@ impl Adapter for ChatAdapter {
         }
 
         let messages = self.format(signature, inputs.clone());
-        let response = lm.call(messages).await?;
+        let response = lm.call(messages, tools).await?;
         let prompt_str = response.chat.to_json().to_string();
 
-        let output = self.parse_response(signature, response.output);
+        let mut output = self.parse_response(signature, response.output);
+        if !response.tool_calls.is_empty() {
+            output.insert(
+                "tool_calls".to_string(),
+                response
+                    .tool_calls
+                    .into_iter()
+                    .map(|call| json!(call))
+                    .collect::<Value>(),
+            );
+            output.insert(
+                "tool_executions".to_string(),
+                response
+                    .tool_executions
+                    .into_iter()
+                    .map(|execution| json!(execution))
+                    .collect::<Value>(),
+            );
+        }
 
         let prediction = Prediction {
             data: output,

crates/dspy-rs/src/adapter/mod.rs

@@ -5,6 +5,7 @@ pub use chat::*;
 use crate::{Chat, Example, LM, Message, MetaSignature, Prediction};
 use anyhow::Result;
 use async_trait::async_trait;
+use rig::tool::ToolDyn;
 use serde_json::Value;
 use std::collections::HashMap;
 use std::sync::Arc;
@@ -22,5 +23,6 @@ pub trait Adapter: Send + Sync + 'static {
         lm: Arc<LM>,
         signature: &dyn MetaSignature,
         inputs: Example,
+        tools: Vec<Arc<dyn ToolDyn>>,
     ) -> Result<Prediction>;
 }

crates/dspy-rs/src/core/lm/mod.rs

@@ -7,7 +7,7 @@ pub use client_registry::*;
 pub use usage::*;
 
 use anyhow::Result;
-use rig::completion::AssistantContent;
+use rig::{completion::AssistantContent, message::ToolCall, message::ToolChoice, tool::ToolDyn};
 
 use bon::Builder;
 use std::{collections::HashMap, sync::Arc};
@@ -23,6 +23,10 @@ pub struct LMResponse {
     pub usage: LmUsage,
     /// Chat history including the freshly appended assistant response.
     pub chat: Chat,
+    /// Tool calls made by the provider.
+    pub tool_calls: Vec<ToolCall>,
+    /// Tool executions made by the provider.
+    pub tool_executions: Vec<String>,
 }
 
 #[derive(Builder)]
@@ -36,6 +40,8 @@ pub struct LM {
     pub temperature: f32,
     #[builder(default = 512)]
     pub max_tokens: u32,
+    #[builder(default = 10)]
+    pub max_tool_iterations: u32,
     #[builder(default = false)]
     pub cache: bool,
     pub cache_handler: Option<Arc<Mutex<ResponseCache>>>,
@@ -57,6 +63,7 @@ impl Clone for LM {
             model: self.model.clone(),
             temperature: self.temperature,
             max_tokens: self.max_tokens,
+            max_tool_iterations: self.max_tool_iterations,
             cache: self.cache,
             cache_handler: self.cache_handler.clone(),
             client: self.client.clone(),
@@ -129,34 +136,203 @@ impl<S: l_m_builder::State> LMBuilder<S> {
     }
 }
 
+struct ToolLoopResult {
+    message: Message,
+    #[allow(unused)]
+    chat_history: Vec<rig::message::Message>,
+    tool_calls: Vec<ToolCall>,
+    tool_executions: Vec<String>,
+}
+
 impl LM {
-    /// Executes a chat completion against the configured provider.
-    ///
-    /// `messages` must already be formatted as OpenAI-compatible chat turns.
-    /// The call returns an [`LMResponse`] containing the assistant output,
-    /// token usage, and chat history including the new response.
-    pub async fn call(&self, messages: Chat) -> Result<LMResponse> {
+    async fn execute_tool_loop(
+        &self,
+        initial_tool_call: &rig::message::ToolCall,
+        mut tools: Vec<Arc<dyn ToolDyn>>,
+        tool_definitions: Vec<rig::completion::ToolDefinition>,
+        mut chat_history: Vec<rig::message::Message>,
+        system_prompt: String,
+        accumulated_usage: &mut LmUsage,
+    ) -> Result<ToolLoopResult> {
+        use rig::OneOrMany;
+        use rig::completion::CompletionRequest;
+        use rig::message::UserContent;
+
+        let max_iterations = self.max_tool_iterations as usize;
+
+        let mut tool_calls = Vec::new();
+        let mut tool_executions = Vec::new();
+
+        // Execute the first tool call
+        let tool_name = &initial_tool_call.function.name;
+        let args_str = initial_tool_call.function.arguments.to_string();
+
+        let mut tool_result = format!("Tool '{}' not found", tool_name);
+        for tool in &mut tools {
+            let def = tool.definition("".to_string()).await;
+            if def.name == *tool_name {
+                // Parse args and call the tool
+                let args_json: serde_json::Value =
+                    serde_json::from_str(&args_str).unwrap_or_default();
+                tool_result = format!("Called tool {} with args: {}", tool_name, args_json);
+                tool_calls.push(initial_tool_call.clone());
+                tool_executions.push(tool_result.clone());
+                break;
+            }
+        }
+
+        // Add initial tool call and result to history
+        chat_history.push(rig::message::Message::Assistant {
+            id: None,
+            content: OneOrMany::one(rig::message::AssistantContent::ToolCall(
+                initial_tool_call.clone(),
+            )),
+        });
+
+        let tool_result_content = if let Some(call_id) = &initial_tool_call.call_id {
+            UserContent::tool_result_with_call_id(
+                initial_tool_call.id.clone(),
+                call_id.clone(),
+                OneOrMany::one(tool_result.into()),
+            )
+        } else {
+            UserContent::tool_result(
+                initial_tool_call.id.clone(),
+                OneOrMany::one(tool_result.into()),
+            )
+        };
+
+        chat_history.push(rig::message::Message::User {
+            content: OneOrMany::one(tool_result_content),
+        });
+
+        // Now loop until we get a text response
+        for _iteration in 1..max_iterations {
+            let request = CompletionRequest {
+                preamble: Some(system_prompt.clone()),
+                chat_history: if chat_history.len() == 1 {
+                    OneOrMany::one(chat_history.clone().into_iter().next().unwrap())
+                } else {
+                    OneOrMany::many(chat_history.clone()).expect("chat_history should not be empty")
+                },
+                documents: Vec::new(),
+                tools: tool_definitions.clone(),
+                temperature: Some(self.temperature as f64),
+                max_tokens: Some(self.max_tokens as u64),
+                tool_choice: Some(ToolChoice::Auto),
+                additional_params: None,
+            };
+
+            let response = self
+                .client
+                .as_ref()
+                .ok_or_else(|| anyhow::anyhow!("LM client not initialized"))?
+                .completion(request)
+                .await?;
+
+            accumulated_usage.prompt_tokens += response.usage.input_tokens;
+            accumulated_usage.completion_tokens += response.usage.output_tokens;
+            accumulated_usage.total_tokens += response.usage.total_tokens;
+
+            match response.choice.first() {
+                AssistantContent::Text(text) => {
+                    return Ok(ToolLoopResult {
+                        message: Message::assistant(&text.text),
+                        chat_history,
+                        tool_calls,
+                        tool_executions,
+                    });
+                }
+                AssistantContent::Reasoning(reasoning) => {
+                    return Ok(ToolLoopResult {
+                        message: Message::assistant(reasoning.reasoning.join("\n")),
+                        chat_history,
+                        tool_calls,
+                        tool_executions,
+                    });
+                }
+                AssistantContent::ToolCall(tool_call) => {
+                    // Execute tool and continue
+                    let tool_name = &tool_call.function.name;
+                    let args_str = tool_call.function.arguments.to_string();
+
+                    let mut tool_result = format!("Tool '{}' not found", tool_name);
+                    for tool in &mut tools {
+                        let def = tool.definition("".to_string()).await;
+                        if def.name == *tool_name {
+                            // For now, just indicate the tool was called
+                            // Actual tool execution would require knowing the concrete Args type
+                            let args_json: serde_json::Value =
+                                serde_json::from_str(&args_str).unwrap_or_default();
+                            tool_result =
+                                format!("Called tool {} with args: {}", tool_name, args_json);
+                            tool_calls.push(tool_call.clone());
+                            tool_executions.push(tool_result.clone());
+                            break;
+                        }
+                    }
+
+                    chat_history.push(rig::message::Message::Assistant {
+                        id: None,
+                        content: OneOrMany::one(rig::message::AssistantContent::ToolCall(
+                            tool_call.clone(),
+                        )),
+                    });
+
+                    let tool_result_content = if let Some(call_id) = &tool_call.call_id {
+                        UserContent::tool_result_with_call_id(
+                            tool_call.id.clone(),
+                            call_id.clone(),
+                            OneOrMany::one(tool_result.into()),
+                        )
+                    } else {
+                        UserContent::tool_result(
+                            tool_call.id.clone(),
+                            OneOrMany::one(tool_result.into()),
+                        )
+                    };
+
+                    chat_history.push(rig::message::Message::User {
+                        content: OneOrMany::one(tool_result_content),
+                    });
+                }
+            }
+        }
+
+        Err(anyhow::anyhow!("Max tool iterations reached"))
+    }
+
+    pub async fn call(&self, messages: Chat, tools: Vec<Arc<dyn ToolDyn>>) -> Result<LMResponse> {
         use rig::OneOrMany;
         use rig::completion::CompletionRequest;
 
         let request_messages = messages.get_rig_messages();
 
+        let mut tool_definitions = Vec::new();
+        for tool in &tools {
+            tool_definitions.push(tool.definition("".to_string()).await);
+        }
+
         // Build the completion request manually
         let mut chat_history = request_messages.conversation;
         chat_history.push(request_messages.prompt);
 
         let request = CompletionRequest {
-            preamble: Some(request_messages.system),
+            preamble: Some(request_messages.system.clone()),
             chat_history: if chat_history.len() == 1 {
-                OneOrMany::one(chat_history.into_iter().next().unwrap())
+                OneOrMany::one(chat_history.clone().into_iter().next().unwrap())
             } else {
-                OneOrMany::many(chat_history).expect("chat_history should not be empty")
+                OneOrMany::many(chat_history.clone()).expect("chat_history should not be empty")
             },
             documents: Vec::new(),
-            tools: Vec::new(),
+            tools: tool_definitions.clone(),
             temperature: Some(self.temperature as f64),
             max_tokens: Some(self.max_tokens as u64),
-            tool_choice: None,
+            tool_choice: if !tool_definitions.is_empty() {
+                Some(ToolChoice::Auto)
+            } else {
+                None
+            },
             additional_params: None,
         };
 
@@ -170,25 +346,56 @@ impl LM {
             .completion(request)
             .await?;
 
+        let mut accumulated_usage = LmUsage::from(response.usage);
+
+        // Handle the response
+        let mut tool_loop_result = None;
         let first_choice = match response.choice.first() {
             AssistantContent::Text(text) => Message::assistant(&text.text),
             AssistantContent::Reasoning(reasoning) => {
                 Message::assistant(reasoning.reasoning.join("\n"))
             }
-            AssistantContent::ToolCall(_tool_call) => {
-                todo!()
+            AssistantContent::ToolCall(tool_call) if !tools.is_empty() => {
+                // Only execute tool loop if we have tools available
+                let result = self
+                    .execute_tool_loop(
+                        &tool_call,
+                        tools,
+                        tool_definitions,
+                        chat_history,
+                        request_messages.system,
+                        &mut accumulated_usage,
+                    )
+                    .await
+                    .unwrap();
+                let message = result.message.clone();
+                tool_loop_result = Some(result);
+                message
+            }
+            AssistantContent::ToolCall(tool_call) => {
+                // No tools available, just return a message indicating this
+                let msg = format!(
+                    "Tool call requested: {} with args: {}, but no tools available",
+                    tool_call.function.name, tool_call.function.arguments
+                );
+                Message::assistant(&msg)
             }
         };
 
-        let usage = LmUsage::from(response.usage);
-
         let mut full_chat = messages.clone();
         full_chat.push_message(first_choice.clone());
 
         Ok(LMResponse {
             output: first_choice,
-            usage,
+            usage: accumulated_usage,
             chat: full_chat,
+            tool_calls: tool_loop_result
+                .as_ref()
+                .map(|result| result.tool_calls.clone())
+                .unwrap_or_default(),
+            tool_executions: tool_loop_result
+                .map(|result| result.tool_executions)
+                .unwrap_or_default(),
         })
     }
 
@@ -282,6 +489,8 @@ impl DummyLM {
             },
             usage: LmUsage::default(),
             chat: full_chat,
+            tool_calls: Vec::new(),
+            tool_executions: Vec::new(),
         })
     }