[WIP] Add ToolChoice for message requests

examples/conversation.rs

@@ -71,7 +71,7 @@ async fn main() -> anyhow::Result<()> {
     // 5. Store the first assistant message.
     request_body
         .messages
-        .push(response.crate_message());
+        .push(response.create_message());
 
     // 6. Add the second user message.
     request_body
@@ -96,7 +96,7 @@ async fn main() -> anyhow::Result<()> {
     // 9. Store the second assistant message.
     request_body
         .messages
-        .push(response.crate_message());
+        .push(response.create_message());
 
     // Continue the conversation...
 

src/messages.rs

@@ -68,6 +68,7 @@ pub use system_prompt::SystemPrompt;
 pub use temperature::Temperature;
 pub use tool::AsyncTool;
 pub use tool::Tool;
+pub use tool::ToolChoice;
 pub use tool::ToolDefinition;
 pub use tool::ToolList;
 pub use tool::ToolResult;

src/messages/messages_request_body.rs

@@ -1,7 +1,7 @@
 use crate::macros::impl_display_for_serialize;
 use crate::messages::{
     ClaudeModel, MaxTokens, Message, Metadata, StopSequence, StreamOption,
-    SystemPrompt, Temperature, ToolDefinition, TopK, TopP,
+    SystemPrompt, Temperature, ToolChoice, ToolDefinition, TopK, TopP,
 };
 use crate::ValidationError;
 
@@ -82,6 +82,17 @@ pub struct MessagesRequestBody {
     /// Recommended for advanced use cases only. You usually only need to use temperature.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub top_k: Option<TopK>,
+    /// How the model should use the provided tools. The model can use a specific tool, any available tool, or decide by itself.
+    ///
+    /// This field is an object with the following possible structures:
+    /// - `{"type": "auto"}`: Allows the model to decide whether to use tools.
+    /// - `{"type": "any"}`: Tells the model it must use one of the provided tools, but doesn't specify which.
+    /// - `{"type": "tool", "name": "<tool_name>"}`: Forces the model to use the specified tool.
+    ///
+    /// The `type` field is required and must be one of: "auto", "any", or "tool".
+    /// When `type` is "tool", the `name` field is also required.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_choice: Option<ToolChoice>,
 }
 
 impl_display_for_serialize!(MessagesRequestBody);
@@ -244,6 +255,15 @@ impl MessagesRequestBuilder {
         self
     }
 
+    /// Sets the tool choice.
+    pub fn tool_choice(
+        mut self,
+        tool_choice: ToolChoice,
+    ) -> Self {
+        self.request_body.tool_choice = Some(tool_choice);
+        self
+    }
+
     /// Builds the MessagesRequestBody.
     pub fn build(self) -> MessagesRequestBody {
         self.request_body
@@ -282,6 +302,7 @@ mod tests {
         assert_eq!(messages_request_body.temperature, None);
         assert_eq!(messages_request_body.top_p, None);
         assert_eq!(messages_request_body.top_k, None);
+        assert_eq!(messages_request_body.tool_choice, None);
     }
 
     #[test]
@@ -307,6 +328,7 @@ mod tests {
         assert_eq!(messages_request_body.tools, None);
         assert_eq!(messages_request_body.top_p, None);
         assert_eq!(messages_request_body.top_k, None);
+        assert_eq!(messages_request_body.tool_choice, None);
     }
 
     #[test]
@@ -343,10 +365,11 @@ mod tests {
             tools: None,
             top_p: Some(TopP::new(0.5).unwrap()),
             top_k: Some(TopK::new(50)),
+            tool_choice: Some(ToolChoice::any),
         };
         assert_eq!(
             serde_json::to_string(&messages_request_body).unwrap(),
-            "{\"model\":\"claude-3-sonnet-20240229\",\"messages\":[],\"system\":\"system-prompt\",\"max_tokens\":16,\"metadata\":{\"user_id\":\"metadata\"},\"stop_sequences\":[\"stop-sequence\"],\"stream\":false,\"temperature\":0.5,\"top_p\":0.5,\"top_k\":50}"
+            "{\"model\":\"claude-3-sonnet-20240229\",\"messages\":[],\"system\":\"system-prompt\",\"max_tokens\":16,\"metadata\":{\"user_id\":\"metadata\"},\"stop_sequences\":[\"stop-sequence\"],\"stream\":false,\"temperature\":0.5,\"top_p\":0.5,\"top_k\":50,\"tool_choice\":{\"type\":\"any\"}}"
         );
     }
 
@@ -373,11 +396,12 @@ mod tests {
             stream: Some(StreamOption::ReturnOnce),
             temperature: Some(Temperature::new(0.5).unwrap()),
             tools: None,
+            tool_choice: Some(ToolChoice::auto),
             top_p: Some(TopP::new(0.5).unwrap()),
             top_k: Some(TopK::new(50)),
         };
         assert_eq!(
-            serde_json::from_str::<MessagesRequestBody>("{\"model\":\"claude-3-sonnet-20240229\",\"messages\":[],\"system\":\"system-prompt\",\"max_tokens\":16,\"metadata\":{\"user_id\":\"metadata\"},\"stop_sequences\":[\"stop-sequence\"],\"stream\":false,\"temperature\":0.5,\"top_p\":0.5,\"top_k\":50}").unwrap(),
+            serde_json::from_str::<MessagesRequestBody>("{\"model\":\"claude-3-sonnet-20240229\",\"messages\":[],\"system\":\"system-prompt\",\"max_tokens\":16,\"metadata\":{\"user_id\":\"metadata\"},\"stop_sequences\":[\"stop-sequence\"],\"stream\":false,\"temperature\":0.5,\"top_p\":0.5,\"top_k\":50,\"tool_choice\":{\"type\":\"auto\"}}").unwrap(),
             messages_request_body
         );
     }
@@ -407,6 +431,9 @@ mod tests {
                 }])
                 .top_p(TopP::new(0.5).unwrap())
                 .top_k(TopK::new(50))
+                .tool_choice(ToolChoice::tool {
+                    name: "tool".into(),
+                })
                 .build();
 
         assert_eq!(
@@ -458,6 +485,12 @@ mod tests {
             messages_request_body.top_k,
             Some(TopK::new(50))
         );
+        assert_eq!(
+            messages_request_body.tool_choice,
+            Some(ToolChoice::tool {
+                name: "tool".into(),
+            })
+        );
     }
 
     #[test]
@@ -536,5 +569,49 @@ mod tests {
             messages_request_body.top_k,
             Some(TopK::new(50))
         );
+        assert_eq!(messages_request_body.tool_choice, None);
+    }
+
+    #[test]
+    fn test_messages_request_body_with_tool_choice() {
+        let request_body =
+            MessagesRequestBuilder::new(ClaudeModel::Claude3Sonnet20240229)
+                .messages(vec![Message::user(
+                    "Hello, Claude!",
+                )])
+                .max_tokens(
+                    MaxTokens::new(100, ClaudeModel::Claude3Sonnet20240229)
+                        .unwrap(),
+                )
+                .tool_choice(ToolChoice::auto)
+                .build();
+
+        assert_eq!(
+            request_body.tool_choice,
+            Some(ToolChoice::auto)
+        );
+
+        let json = serde_json::to_string(&request_body).unwrap();
+        assert!(json.contains(r#""tool_choice":{"type":"auto"}"#));
+    }
+
+    #[test]
+    fn test_messages_request_body_serialization_with_tool_choice() {
+        let request_body = MessagesRequestBody {
+            model: ClaudeModel::Claude3Sonnet20240229,
+            messages: vec![Message::user(
+                "Hello",
+            )],
+            max_tokens: MaxTokens::new(100, ClaudeModel::Claude3Sonnet20240229)
+                .unwrap(),
+            tool_choice: Some(ToolChoice::tool {
+                name: "get_weather".to_string(),
+            }),
+            ..Default::default()
+        };
+
+        let json = serde_json::to_string(&request_body).unwrap();
+        assert!(json
+            .contains(r#""tool_choice":{"type":"tool","name":"get_weather"}"#));
     }
 }

src/messages/messages_response_body.rs

@@ -4,8 +4,8 @@ use crate::macros::{
     impl_display_for_serialize, impl_enum_string_serialization,
 };
 use crate::messages::{
-    ClaudeModel, Content, Message, Role, StopReason,
-    StopSequence, Usage,
+    ClaudeModel, Content, ContentBlock, Message, Role, StopReason,
+    StopSequence, ToolUse, Usage,
 };
 
 /// The response body for the Messages API.
@@ -76,12 +76,35 @@ impl_display_for_serialize!(MessagesResponseBody);
 
 impl MessagesResponseBody {
     /// Creates `Message` from the response body.
-    pub fn crate_message(self) -> Message {
+    pub fn create_message(self) -> Message {
         Message {
             role: self.role,
             content: self.content,
         }
     }
+
+    /// Extracts tool use information from the response body.
+    pub fn extract_tool_use(&self) -> Option<&ToolUse> {
+        if let Content::MultipleBlocks(blocks) = &self.content {
+            for block in blocks {
+                if let ContentBlock::ToolUse(tool_use_block) = block {
+                    return Some(&tool_use_block.tool_use);
+                }
+            }
+        }
+        None
+    }
+
+    /// Extracts tool name and input from the response body.
+    pub fn extract_tool_fields(&self) -> Option<(String, serde_json::Value)> {
+        self.extract_tool_use()
+            .map(|tool_use| {
+                (
+                    tool_use.name.clone(),
+                    tool_use.input.clone(),
+                )
+            })
+    }
 }
 
 /// The object type for message.
@@ -215,7 +238,7 @@ mod tests {
             MessageObjectType::Message
         );
     }
-    
+
     #[test]
     fn create_message() {
         let response = MessagesResponseBody {
@@ -224,8 +247,152 @@ mod tests {
         };
 
         assert_eq!(
-            response.crate_message(),
+            response.create_message(),
             Message::assistant("content")
         );
     }
+
+    #[test]
+    fn test_extract_tool_use() {
+        // Test case 1: Valid tool use
+        let response = MessagesResponseBody {
+            content: Content::MultipleBlocks(vec![
+                ContentBlock::ToolUse(ToolUseContentBlock {
+                    _type: ContentType::ToolUse,
+                    tool_use: ToolUse {
+                        id: "tool_id".to_string(),
+                        name: "test_tool".to_string(),
+                        input: serde_json::json!({"key": "value"}),
+                    },
+                }),
+            ]),
+            ..Default::default()
+        };
+        let tool_use = response.extract_tool_use();
+        assert!(tool_use.is_some());
+        assert_eq!(tool_use.unwrap().name, "test_tool");
+
+        // Test case 2: No tool use
+        let response = MessagesResponseBody {
+            content: Content::MultipleBlocks(vec![ContentBlock::Text(
+                TextContentBlock::new("text"),
+            )]),
+            ..Default::default()
+        };
+        let tool_use = response.extract_tool_use();
+        assert!(tool_use.is_none());
+
+        // Test case 3: Empty content
+        let response = MessagesResponseBody {
+            content: Content::MultipleBlocks(vec![]),
+            ..Default::default()
+        };
+        let tool_use = response.extract_tool_use();
+        assert!(tool_use.is_none());
+
+        // Test case 4: Multiple tool uses (should return the first one)
+        let response = MessagesResponseBody {
+            content: Content::MultipleBlocks(vec![
+                ContentBlock::ToolUse(ToolUseContentBlock {
+                    _type: ContentType::ToolUse,
+                    tool_use: ToolUse {
+                        id: "tool_id_1".to_string(),
+                        name: "test_tool_1".to_string(),
+                        input: serde_json::json!({"key": "value1"}),
+                    },
+                }),
+                ContentBlock::ToolUse(ToolUseContentBlock {
+                    _type: ContentType::ToolUse,
+                    tool_use: ToolUse {
+                        id: "tool_id_2".to_string(),
+                        name: "test_tool_2".to_string(),
+                        input: serde_json::json!({"key": "value2"}),
+                    },
+                }),
+            ]),
+            ..Default::default()
+        };
+        let tool_use = response.extract_tool_use();
+        assert!(tool_use.is_some());
+        assert_eq!(tool_use.unwrap().name, "test_tool_1");
+    }
+
+    #[test]
+    fn test_extract_tool_fields() {
+        // Test case 1: Valid tool use
+        let response = MessagesResponseBody {
+            content: Content::MultipleBlocks(vec![
+                ContentBlock::ToolUse(ToolUseContentBlock {
+                    _type: ContentType::ToolUse,
+                    tool_use: ToolUse {
+                        id: "tool_id".to_string(),
+                        name: "test_tool".to_string(),
+                        input: serde_json::json!({"key": "value"}),
+                    },
+                }),
+            ]),
+            ..Default::default()
+        };
+        let fields = response.extract_tool_fields();
+        assert!(fields.is_some());
+        let (name, input) = fields.unwrap();
+        assert_eq!(name, "test_tool");
+        assert_eq!(
+            input,
+            serde_json::json!({"key": "value"})
+        );
+
+        // Test case 2: No tool use
+        let response = MessagesResponseBody {
+            content: Content::MultipleBlocks(vec![ContentBlock::Text(
+                TextContentBlock::new("text"),
+            )]),
+            ..Default::default()
+        };
+        let fields = response.extract_tool_fields();
+        assert!(fields.is_none());
+
+        // Test case 3: Empty content
+        let response = MessagesResponseBody {
+            content: Content::MultipleBlocks(vec![]),
+            ..Default::default()
+        };
+        let fields = response.extract_tool_fields();
+        assert!(fields.is_none());
+
+        // Test case 4: Complex input
+        let response = MessagesResponseBody {
+            content: Content::MultipleBlocks(vec![
+                ContentBlock::ToolUse(ToolUseContentBlock {
+                    _type: ContentType::ToolUse,
+                    tool_use: ToolUse {
+                        id: "tool_id".to_string(),
+                        name: "complex_tool".to_string(),
+                        input: serde_json::json!({
+                            "nested": {
+                                "array": [1, 2, 3],
+                                "object": {"a": "b"}
+                            },
+                            "boolean": true
+                        }),
+                    },
+                }),
+            ]),
+            ..Default::default()
+        };
+        let fields = response.extract_tool_fields();
+        assert!(fields.is_some());
+        let (name, input) = fields.unwrap();
+        assert_eq!(name, "complex_tool");
+        assert_eq!(
+            input,
+            serde_json::json!({
+                "nested": {
+                    "array": [1, 2, 3],
+                    "object": {"a": "b"}
+                },
+                "boolean": true
+            })
+        );
+    }
 }