test

src/tools.rs

@@ -0,0 +1,129 @@
+use crate::api::{self, FunctionDefinition, Message, Tool, ToolCall};
+use std::collections::HashMap;
+
+pub fn get_tools() -> Vec<Tool> {
+    vec![
+        Tool {
+            tool_type: "function".to_string(),
+            function: FunctionDefinition {
+                name: "google_search".to_string(),
+                description: "Search the web for information".to_string(),
+                parameters: serde_json::json!({
+                    "type": "object",
+                    "properties": {
+                        "query": {
+                            "type": "string",
+                            "description": "The search query"
+                        }
+                    },
+                    "required": ["query"]
+                }),
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: FunctionDefinition {
+                name: "write_file".to_string(),
+                description: "Write content to a file. Ensure parent directories exist."
+                    .to_string(),
+                parameters: serde_json::json!({
+                    "type": "object",
+                    "properties": {
+                        "path": {
+                            "type": "string",
+                            "description": "The path to the file to write"
+                        },
+                        "content": {
+                            "type": "string",
+                            "description": "The content to write to the file"
+                        }
+                    },
+                    "required": ["path", "content"]
+                }),
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: FunctionDefinition {
+                name: "finish".to_string(),
+                description: "Finish the task".to_string(),
+                parameters: serde_json::json!({
+                    "type": "object",
+                    "properties": {
+                        "result": {
+                            "type": "string",
+                            "description": "The result of the task"
+                        }
+                    },
+                    "required": ["result"]
+                }),
+            },
+        },
+    ]
+}
+
+pub async fn handle_tool_call(
+    tool_call: &ToolCall,
+    tavily_api_key: &Option<String>,
+) -> Result<(Message, bool, Option<String>), Box<dyn std::error::Error>> {
+    let mut file_written = false;
+    let mut answer = None;
+    let name = &tool_call.function.name;
+
+    let (content, written) = if name == "google_search" {
+        let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
+        let query = args.get("query").ok_or("Missing query argument")?;
+
+        println!(
+            "--- Executing tool: google_search(query: \"{}\") ---",
+            query
+        );
+
+        let search_result = if let Some(key) = tavily_api_key {
+            match api::perform_search(query, key).await {
+                Ok(results) => results,
+                Err(e) => format!("Search error: {}", e),
+            }
+        } else {
+            "Error: TAVILY_API_KEY is not set. Cannot perform real search.".to_string()
+        };
+        (search_result, false)
+    } else if name == "write_file" {
+        let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
+        let path = args.get("path").ok_or("Missing path argument")?;
+        let content = args.get("content").ok_or("Missing content argument")?;
+
+        println!("--- Executing tool: write_file(path: \"{}\") ---", path);
+
+        let write_result = match std::fs::write(path, content) {
+            Ok(_) => {
+                file_written = true;
+                format!("Successfully wrote content to {}", path)
+            }
+            Err(e) => format!("Error writing to {}: {}", path, e),
+        };
+        (write_result, file_written)
+    } else if name == "finish" {
+        let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
+        let result = args.get("result").ok_or("Missing result argument")?;
+
+        println!("--- Finishing task: {}", result);
+
+        file_written = true;
+        answer = Some(result.clone());
+        (result.clone(), file_written)
+    } else {
+        (format!("Error: Unknown tool {}", name), false)
+    };
+
+    Ok((
+        Message {
+            role: "tool".to_string(),
+            content: Some(content),
+            tool_calls: None,
+            tool_call_id: Some(tool_call.id.clone()),
+        },
+        written,
+        answer,
+    ))
+}