noc/tests/tool_call.rs

//! Integration test: verify tool call round-trip with Ollama's OpenAI-compatible API.
//! Requires Ollama running at OLLAMA_URL (default: http://100.84.7.49:11434).

use serde_json::json;

const OLLAMA_URL: &str = "http://100.84.7.49:11434/v1";
const MODEL: &str = "gemma4:31b";

fn tools() -> serde_json::Value {
    json!([{
        "type": "function",
        "function": {
            "name": "calculator",
            "description": "Calculate a math expression",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {"type": "string", "description": "Math expression to evaluate"}
                },
                "required": ["expression"]
            }
        }
    }])
}

/// Test non-streaming tool call round-trip
#[tokio::test]
async fn test_tool_call_roundtrip_non_streaming() {
    let client = reqwest::Client::new();
    let url = format!("{OLLAMA_URL}/chat/completions");

    // Round 1: ask the model to use the calculator
    let body = json!({
        "model": MODEL,
        "messages": [
            {"role": "user", "content": "What is 2+2? Use the calculator tool."}
        ],
        "tools": tools(),
    });

    let resp = client.post(&url).json(&body).send().await.unwrap();
    assert!(resp.status().is_success(), "Round 1 failed: {}", resp.status());

    let result: serde_json::Value = resp.json().await.unwrap();
    let choice = &result["choices"][0];
    assert_eq!(
        choice["finish_reason"].as_str().unwrap(),
        "tool_calls",
        "Expected tool_calls finish_reason, got: {choice}"
    );

    let tool_calls = choice["message"]["tool_calls"].as_array().unwrap();
    assert!(!tool_calls.is_empty(), "No tool calls returned");

    let tc = &tool_calls[0];
    let call_id = tc["id"].as_str().unwrap();
    let func_name = tc["function"]["name"].as_str().unwrap();
    assert_eq!(func_name, "calculator");

    // Round 2: send tool result back
    let body2 = json!({
        "model": MODEL,
        "messages": [
            {"role": "user", "content": "What is 2+2? Use the calculator tool."},
            {
                "role": "assistant",
                "content": "",
                "tool_calls": [{
                    "id": call_id,
                    "type": "function",
                    "function": {
                        "name": func_name,
                        "arguments": tc["function"]["arguments"].as_str().unwrap()
                    }
                }]
            },
            {
                "role": "tool",
                "tool_call_id": call_id,
                "content": "4"
            }
        ],
        "tools": tools(),
    });

    let resp2 = client.post(&url).json(&body2).send().await.unwrap();
    let status2 = resp2.status();
    let body2_text = resp2.text().await.unwrap();
    assert!(
        status2.is_success(),
        "Round 2 failed ({status2}): {body2_text}"
    );

    let result2: serde_json::Value = serde_json::from_str(&body2_text).unwrap();
    let content = result2["choices"][0]["message"]["content"]
        .as_str()
        .unwrap_or("");
    assert!(!content.is_empty(), "Expected content in round 2 response");
    println!("Round 2 response: {content}");
}

/// Test tool call with conversation history (simulates real scenario)
#[tokio::test]
async fn test_tool_call_with_history() {
    let client = reqwest::Client::new();
    let url = format!("{OLLAMA_URL}/chat/completions");

    // Simulate real message history with system prompt
    let body = json!({
        "model": MODEL,
        "stream": true,
        "messages": [
            {"role": "system", "content": "你是一个AI助手。你可以使用提供的工具来完成任务。当需要执行命令、运行代码或启动复杂子任务时，直接调用对应的工具，不要只是描述你会怎么做。"},
            {"role": "user", "content": "hi"},
            {"role": "assistant", "content": "Hello!"},
            {"role": "user", "content": "What is 3+4? Use the calculator."}
        ],
        "tools": tools(),
    });

    // Round 1: expect tool call
    let mut resp = client.post(&url).json(&body).send().await.unwrap();
    assert!(resp.status().is_success(), "Round 1 failed: {}", resp.status());

    let mut buffer = String::new();
    let mut tc_id = String::new();
    let mut tc_name = String::new();
    let mut tc_args = String::new();
    let mut has_tc = false;

    while let Some(chunk) = resp.chunk().await.unwrap() {
        buffer.push_str(&String::from_utf8_lossy(&chunk));
        while let Some(pos) = buffer.find('\n') {
            let line = buffer[..pos].to_string();
            buffer = buffer[pos + 1..].to_string();
            if let Some(data) = line.trim().strip_prefix("data: ") {
                if data.trim() == "[DONE]" { break; }
                if let Ok(j) = serde_json::from_str::<serde_json::Value>(data) {
                    if let Some(arr) = j["choices"][0]["delta"]["tool_calls"].as_array() {
                        has_tc = true;
                        for tc in arr {
                            if let Some(id) = tc["id"].as_str() { tc_id = id.into(); }
                            if let Some(n) = tc["function"]["name"].as_str() { tc_name = n.into(); }
                            if let Some(a) = tc["function"]["arguments"].as_str() { tc_args.push_str(a); }
                        }
                    }
                }
            }
        }
    }

    assert!(has_tc, "Expected tool call, got content only");
    println!("Tool: {tc_name}({tc_args}) id={tc_id}");

    // Round 2: tool result → expect content
    let body2 = json!({
        "model": MODEL,
        "stream": true,
        "messages": [
            {"role": "system", "content": "你是一个AI助手。"},
            {"role": "user", "content": "hi"},
            {"role": "assistant", "content": "Hello!"},
            {"role": "user", "content": "What is 3+4? Use the calculator."},
            {"role": "assistant", "content": "", "tool_calls": [{"id": tc_id, "type": "function", "function": {"name": tc_name, "arguments": tc_args}}]},
            {"role": "tool", "tool_call_id": tc_id, "content": "7"}
        ],
        "tools": tools(),
    });

    let resp2 = client.post(&url).json(&body2).send().await.unwrap();
    let status = resp2.status();
    if !status.is_success() {
        let err = resp2.text().await.unwrap();
        panic!("Round 2 failed ({status}): {err}");
    }

    let mut resp2 = client.post(&url).json(&body2).send().await.unwrap();
    let mut content = String::new();
    let mut buf2 = String::new();
    while let Some(chunk) = resp2.chunk().await.unwrap() {
        buf2.push_str(&String::from_utf8_lossy(&chunk));
        while let Some(pos) = buf2.find('\n') {
            let line = buf2[..pos].to_string();
            buf2 = buf2[pos + 1..].to_string();
            if let Some(data) = line.trim().strip_prefix("data: ") {
                if data.trim() == "[DONE]" { break; }
                if let Ok(j) = serde_json::from_str::<serde_json::Value>(data) {
                    if let Some(c) = j["choices"][0]["delta"]["content"].as_str() {
                        content.push_str(c);
                    }
                }
            }
        }
    }

    println!("Final response: {content}");
    assert!(!content.is_empty(), "Expected non-empty content in round 2");
}

/// Test multimodal image input
#[tokio::test]
async fn test_image_multimodal() {
    let client = reqwest::Client::new();
    let url = format!("{OLLAMA_URL}/chat/completions");

    // 2x2 red PNG generated by PIL
    let b64 = "iVBORw0KGgoAAAANSUhEUgAAAAIAAAACCAIAAAD91JpzAAAAFklEQVR4nGP8z8DAwMDAxMDAwMDAAAANHQEDasKb6QAAAABJRU5ErkJggg==";

    let body = json!({
        "model": MODEL,
        "messages": [{
            "role": "user",
            "content": [
                {"type": "text", "text": "What color is this image? Reply with just the color name."},
                {"type": "image_url", "image_url": {"url": format!("data:image/png;base64,{b64}")}}
            ]
        }],
    });

    let resp = client.post(&url).json(&body).send().await.unwrap();
    let status = resp.status();
    let text = resp.text().await.unwrap();
    assert!(status.is_success(), "Multimodal request failed ({status}): {text}");

    let result: serde_json::Value = serde_json::from_str(&text).unwrap();
    let content = result["choices"][0]["message"]["content"]
        .as_str()
        .unwrap_or("");
    println!("Image description: {content}");
    assert!(!content.is_empty(), "Expected non-empty response for image");
}

/// Test streaming tool call round-trip (matches our actual code path)
#[tokio::test]
async fn test_tool_call_roundtrip_streaming() {
    let client = reqwest::Client::new();
    let url = format!("{OLLAMA_URL}/chat/completions");

    // Round 1: streaming, get tool calls
    let body = json!({
        "model": MODEL,
        "stream": true,
        "messages": [
            {"role": "user", "content": "What is 7*6? Use the calculator tool."}
        ],
        "tools": tools(),
    });

    let mut resp = client.post(&url).json(&body).send().await.unwrap();
    assert!(resp.status().is_success(), "Round 1 streaming failed");

    // Parse SSE to extract tool calls
    let mut buffer = String::new();
    let mut tool_call_id = String::new();
    let mut tool_call_name = String::new();
    let mut tool_call_args = String::new();
    let mut has_tool_calls = false;

    while let Some(chunk) = resp.chunk().await.unwrap() {
        buffer.push_str(&String::from_utf8_lossy(&chunk));

        while let Some(pos) = buffer.find('\n') {
            let line = buffer[..pos].to_string();
            buffer = buffer[pos + 1..].to_string();

            let trimmed = line.trim();
            if let Some(data) = trimmed.strip_prefix("data: ") {
                if data.trim() == "[DONE]" {
                    break;
                }
                if let Ok(json) = serde_json::from_str::<serde_json::Value>(data) {
                    let delta = &json["choices"][0]["delta"];
                    if let Some(tc_arr) = delta["tool_calls"].as_array() {
                        has_tool_calls = true;
                        for tc in tc_arr {
                            if let Some(id) = tc["id"].as_str() {
                                tool_call_id = id.to_string();
                            }
                            if let Some(name) = tc["function"]["name"].as_str() {
                                tool_call_name = name.to_string();
                            }
                            if let Some(args) = tc["function"]["arguments"].as_str() {
                                tool_call_args.push_str(args);
                            }
                        }
                    }
                }
            }
        }
    }

    assert!(has_tool_calls, "No tool calls in streaming response");
    assert_eq!(tool_call_name, "calculator");
    println!("Tool call: {tool_call_name}({tool_call_args}) id={tool_call_id}");

    // Round 2: send tool result, streaming
    let body2 = json!({
        "model": MODEL,
        "stream": true,
        "messages": [
            {"role": "user", "content": "What is 7*6? Use the calculator tool."},
            {
                "role": "assistant",
                "content": "",
                "tool_calls": [{
                    "id": tool_call_id,
                    "type": "function",
                    "function": {
                        "name": tool_call_name,
                        "arguments": tool_call_args
                    }
                }]
            },
            {
                "role": "tool",
                "tool_call_id": tool_call_id,
                "content": "42"
            }
        ],
        "tools": tools(),
    });

    let resp2 = client.post(&url).json(&body2).send().await.unwrap();
    let status2 = resp2.status();
    if !status2.is_success() {
        let err = resp2.text().await.unwrap();
        panic!("Round 2 streaming failed ({status2}): {err}");
    }

    // Collect content from streaming response
    let mut resp2 = client
        .post(&url)
        .json(&body2)
        .send()
        .await
        .unwrap();
    let mut content = String::new();
    let mut buffer2 = String::new();

    while let Some(chunk) = resp2.chunk().await.unwrap() {
        buffer2.push_str(&String::from_utf8_lossy(&chunk));
        while let Some(pos) = buffer2.find('\n') {
            let line = buffer2[..pos].to_string();
            buffer2 = buffer2[pos + 1..].to_string();
            let trimmed = line.trim();
            if let Some(data) = trimmed.strip_prefix("data: ") {
                if data.trim() == "[DONE]" {
                    break;
                }
                if let Ok(json) = serde_json::from_str::<serde_json::Value>(data) {
                    if let Some(c) = json["choices"][0]["delta"]["content"].as_str() {
                        content.push_str(c);
                    }
                }
            }
        }
    }

    assert!(!content.is_empty(), "Expected content in round 2 streaming");
    println!("Round 2 streaming content: {content}");
}