grok_api 0.1.6

//! Tool / function calling example.
//!
//! Shows how to:
//!   1. Define tools the model can invoke
//!   2. Send a request with those tools attached
//!   3. Receive and dispatch tool calls
//!   4. Feed results back to get the final answer
//!
//! Run with:
//!   cargo run --example tools_example

use grok_api::{ChatMessage, GrokClient, Result};
use serde_json::json;

/// Model to use — grok-4-1-fast-reasoning is the recommended default for tool use.
const MODEL: &str = "grok-4-1-fast-reasoning";

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt::init();

    let api_key = std::env::var("GROK_API_KEY").expect("GROK_API_KEY environment variable not set");

    let client = GrokClient::builder().api_key(api_key).build()?;

    println!("🔧 Grok API — Tool / Function Calling\n");
    println!("Model: {MODEL}\n");

    // ── Define tools ──────────────────────────────────────────────────────────
    let tools = vec![
        json!({
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "Get the current weather for a location",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "City and country, e.g. London, UK"
                        },
                        "unit": {
                            "type": "string",
                            "enum": ["celsius", "fahrenheit"],
                            "description": "Temperature unit"
                        }
                    },
                    "required": ["location"]
                }
            }
        }),
        json!({
            "type": "function",
            "function": {
                "name": "get_time",
                "description": "Get the current local time for a timezone",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "timezone": {
                            "type": "string",
                            "description": "IANA timezone, e.g. America/New_York"
                        }
                    },
                    "required": ["timezone"]
                }
            }
        }),
    ];

    // ── Initial user query ────────────────────────────────────────────────────
    let user_query = "What is the weather in San Francisco right now, and what time is it there?";

    let mut messages = vec![
        ChatMessage::system(
            "You are a helpful assistant. Use the available tools to answer the user's question.",
        ),
        ChatMessage::user(user_query),
    ];

    println!("👤 User: {user_query}\n");

    // ── First request — model may call tools ──────────────────────────────────
    let response = client
        .chat_with_history(&messages)
        .model(MODEL)
        .tools(tools.clone())
        .send()
        .await?;

    if !response.has_tool_calls() {
        // Model answered directly — no tools needed.
        println!("🤖 Assistant: {}\n", response.content().unwrap_or(""));
        return Ok(());
    }

    // ── Dispatch tool calls ───────────────────────────────────────────────────
    println!("🔧 Model is calling tools...\n");

    // Add the assistant turn (with its tool_calls) to history.
    let assistant_msg = response.message().unwrap();
    messages.push(ChatMessage {
        role: assistant_msg.role.clone(),
        content: assistant_msg.content.clone(),
        tool_calls: assistant_msg.tool_calls.clone(),
        tool_call_id: None,
    });

    for call in response.tool_calls().unwrap() {
        let args = call.function.parse_arguments()?;
        println!("  → {}({})", call.function.name, call.function.arguments);

        let result = match call.function.name.as_str() {
            "get_weather" => {
                let loc = args["location"].as_str().unwrap_or("unknown");
                let unit = args
                    .get("unit")
                    .and_then(|u| u.as_str())
                    .unwrap_or("celsius");
                get_weather(loc, unit)
            }
            "get_time" => {
                let tz = args["timezone"].as_str().unwrap_or("UTC");
                get_time(tz)
            }
            other => format!("{{\"error\": \"unknown function {other}\"}}"),
        };

        println!("     result: {result}\n");
        messages.push(ChatMessage::tool(result, &call.id));
    }

    // ── Second request — model summarises tool results ────────────────────────
    println!("📤 Sending tool results back…\n");

    let final_response = client
        .chat_with_history(&messages)
        .model(MODEL)
        .tools(tools)
        .send()
        .await?;

    println!("🤖 Assistant: {}\n", final_response.content().unwrap_or(""));

    // ── Token usage ───────────────────────────────────────────────────────────
    let u = &final_response.usage;
    print!(
        "📊 Tokens — prompt: {}, completion: {}, total: {}",
        u.prompt_tokens, u.completion_tokens, u.total_tokens
    );
    if let Some(c) = u.cached_prompt_tokens {
        print!(", cached: {c} 💰");
    }
    println!();

    println!("\n✨ Done!");
    Ok(())
}

// ── Stub implementations (replace with real API calls in production) ──────────

fn get_weather(location: &str, unit: &str) -> String {
    let (temp, symbol) = if unit == "celsius" {
        (18, "°C")
    } else {
        (64, "°F")
    };
    format!(
        r#"{{"location":"{location}","temperature":"{temp}{symbol}","condition":"Partly cloudy","humidity":"72%"}}"#
    )
}

fn get_time(timezone: &str) -> String {
    format!(r#"{{"timezone":"{timezone}","time":"14:32","date":"2026-04-17"}}"#)
}