oxide-agent 0.1.0

//! Tool / function calling with `ToolRegistry` and `ToolBuilder`.
//!
//! Demonstrates:
//!   - Building a tool definition via `ToolBuilder`
//!   - Registering an async handler in `ToolRegistry`
//!   - Passing tool definitions to a chat request
//!   - Detecting and dispatching tool calls from the model response
//!   - Feeding the tool result back into the conversation
//!
//! Run with:
//!   cargo run --example tool_calling

use oxide_agent::client::{HttpOllamaClient, OllamaClient};
use oxide_agent::tools::{ToolBuilder, ToolRegistry};
use oxide_agent::types::{ChatRequest, Message, Role};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let client = HttpOllamaClient::new("http://localhost:11434");

    // ── 1. Define and register tools ──────────────────────────────────────────
    let mut registry = ToolRegistry::new();

    // A mock weather tool.
    let weather_def = ToolBuilder::new("get_weather", "Get the current weather for a city")
        .string_param("city", "The city name", true)
        .string_param(
            "unit",
            "Temperature unit: 'celsius' or 'fahrenheit'",
            false,
        )
        .build();

    registry.register(weather_def, |args| async move {
        let city = args["city"].as_str().unwrap_or("unknown").to_string();
        let unit = args["unit"].as_str().unwrap_or("celsius");
        // In a real implementation you'd call a weather API here.
        Ok(serde_json::json!({
            "city": city,
            "temperature": 22,
            "unit": unit,
            "condition": "sunny"
        }))
    });

    // A simple arithmetic tool.
    let calc_def = ToolBuilder::new("calculate", "Evaluate a basic arithmetic expression")
        .number_param("a", "Left operand", true)
        .number_param("b", "Right operand", true)
        .string_param("op", "Operator: add, sub, mul, div", true)
        .build();

    registry.register(calc_def, |args| async move {
        let a = args["a"].as_f64().unwrap_or(0.0);
        let b = args["b"].as_f64().unwrap_or(0.0);
        let result = match args["op"].as_str().unwrap_or("add") {
            "sub" => a - b,
            "mul" => a * b,
            "div" if b != 0.0 => a / b,
            _ => a + b,
        };
        Ok(serde_json::json!({ "result": result }))
    });

    // ── 2. Send a chat request with tool definitions attached ─────────────────
    let mut messages = vec![Message {
        role: Role::User,
        content: "What's the weather in Berlin and what is 42 times 7?".into(),
        tool_calls: None,
    }];

    let req = ChatRequest {
        model: "llama3".into(),
        messages: messages.clone(),
        tools: Some(registry.definitions()),
        stream: false,
    };

    let resp = client.chat(req).await?;
    messages.push(resp.message.clone());

    // ── 3. Dispatch any tool calls the model requested ────────────────────────
    if let Some(calls) = &resp.message.tool_calls {
        for call in calls {
            println!("Model called tool: {}", call.function.name);

            let result = registry
                .dispatch(&call.function.name, call.function.arguments.clone())
                .await?;

            println!("Tool result: {result}");

            // Feed the result back as a Tool-role message.
            messages.push(Message {
                role: Role::Tool,
                content: result.to_string(),
                tool_calls: None,
            });
        }

        // ── 4. Continue the conversation with the tool results ────────────────
        let follow_up = ChatRequest {
            model: "llama3".into(),
            messages: messages.clone(),
            tools: Some(registry.definitions()),
            stream: false,
        };

        let final_resp = client.chat(follow_up).await?;
        println!("\nFinal answer: {}", final_resp.message.content);
    } else {
        println!("Direct answer: {}", resp.message.content);
    }

    Ok(())
}