llm-universal 0.2.0

//! The same tool loop runs against Anthropic and OpenAI (or any compatible server).
//!
//! The provider is picked from environment variables:
//!   ANTHROPIC_API_KEY=...  -> Anthropic
//!   OPENAI_API_KEY=...     -> OpenAI (set OPENAI_BASE_URL for a compatible server)
//!   LLM_MODEL=...          -> optional, defaults below
//!
//! Run with: `cargo run --example tools`

use llm_universal::{
    AnthropicProvider, ChatRequest, LlmError, LlmProvider, Message, OpenAiProvider, Part, Result,
    Role, StopReason, Tool,
};
use serde_json::{json, Value};

/// Guards against an endless loop if the model keeps calling tools.
const MAX_ROUNDS: usize = 5;

fn get_weather(args: &Value) -> String {
    let city = args["city"].as_str().unwrap_or("?");
    json!({ "city": city, "temp_c": 18, "sky": "cloudy" }).to_string()
}

fn make_provider() -> (Box<dyn LlmProvider>, String) {
    if let Ok(key) = std::env::var("ANTHROPIC_API_KEY") {
        let model = std::env::var("LLM_MODEL").unwrap_or_else(|_| "claude-opus-5-5".into());
        return (Box::new(AnthropicProvider::new(key)), model);
    }
    let key = std::env::var("OPENAI_API_KEY").unwrap_or_default();
    let mut provider = OpenAiProvider::new(key);
    if let Ok(url) = std::env::var("OPENAI_BASE_URL") {
        provider = provider.with_base_url(url);
    }
    let model = std::env::var("LLM_MODEL").unwrap_or_else(|_| "gpt-4o".into());
    (Box::new(provider), model)
}

#[tokio::main]
async fn main() -> Result<()> {
    let (provider, model) = make_provider();

    let weather = Tool::new(
        "get_weather",
        "Get the current weather in a city",
        json!({
            "type": "object",
            "properties": { "city": { "type": "string" } },
            "required": ["city"],
            "additionalProperties": false
        }),
    );

    let mut req = ChatRequest::new(model)
        .system("Answer briefly.")
        .max_tokens(1024)
        .tool(weather)
        .message(Message::user("What's the weather in Berlin?"));

    for _ in 0..MAX_ROUNDS {
        let resp = provider.chat(&req).await?;
        println!(
            "[{}] stop={:?} in={} out={}",
            resp.model, resp.stop_reason, resp.usage.input_tokens, resp.usage.output_tokens
        );

        // pause_turn: send the response back as-is and the model continues.
        if resp.stop_reason == StopReason::Paused {
            req.messages.push(resp.message.clone());
            continue;
        }

        let calls = resp.tool_calls();
        if calls.is_empty() {
            println!("{}", resp.text());
            return Ok(());
        }

        // Keep the assistant turn (with its tool calls) in the history unchanged.
        req.messages.push(resp.message);

        let results: Vec<Part> = calls
            .iter()
            .map(|(id, name, args)| {
                let (content, is_error) = match name.as_str() {
                    "get_weather" => (get_weather(args), false),
                    other => (format!("unknown tool `{other}`"), true),
                };
                Part::ToolResult {
                    call_id: id.clone(),
                    content,
                    is_error,
                }
            })
            .collect();
        req.messages.push(Message::new(Role::Tool, results));
    }
    Err(LlmError::InvalidResponse(format!(
        "model did not finish within {MAX_ROUNDS} tool rounds"
    )))
}