just-deepseek 0.2.0

//! Full tool-calling loop with DeepSeek thinking mode.
//!
//! Demonstrates: define a tool -> model calls it -> execute locally -> send result back -> final answer.
//! The assistant message is reconstructed manually to preserve `reasoning_content`.
//!
//! ```bash
//! JUST_LLM_DEEPSEEK_API_KEY=your-key JUST_LLM_DEEPSEEK_MODEL=deepseek-chat \
//!   cargo run -p just-deepseek --example tool_calling
//! ```

use just_deepseek::types::chat::{
    ChatCompletionRequest, ChatMessage, FunctionDefinition, ReasoningEffort, ThinkingConfig,
    ThinkingMode, ToolCallsMessage, ToolDefinition, ToolType,
};
use just_deepseek::{DeepSeekClient, Error};

/// A mock tool implementation.
fn add(args: &serde_json::Value) -> serde_json::Value {
    let x = args["x"].as_f64().expect("x is not a number");
    let y = args["y"].as_f64().expect("y is not a number");
    serde_json::json!({ "result": x + y })
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    dotenvy::dotenv().ok();

    let api_key =
        std::env::var("JUST_LLM_DEEPSEEK_API_KEY").expect("JUST_LLM_DEEPSEEK_API_KEY must be set");
    let base_url = std::env::var("JUST_LLM_DEEPSEEK_BASE_URL").ok();
    let model =
        std::env::var("JUST_LLM_DEEPSEEK_MODEL").expect("JUST_LLM_DEEPSEEK_MODEL must be set");

    let mut builder = DeepSeekClient::builder().api_key(&api_key);
    if let Some(url) = base_url {
        builder = builder.base_url(&url);
    }
    let client = builder.build()?;

    let add_tool = ToolDefinition {
        kind: ToolType::Function,
        function: FunctionDefinition {
            name: "add".to_owned(),
            description: Some("Add two numbers together.".to_owned()),
            parameters: Some(serde_json::json!({
                "type": "object",
                "properties": {
                    "x": { "type": "number", "description": "The first number." },
                    "y": { "type": "number", "description": "The second number." }
                },
                "required": ["x", "y"]
            })),
            strict: None,
        },
    };

    let system_prompt = "You are a helpful math assistant. Use the provided tools.";
    let user_prompt = "What is 12345 + 67890?";

    // --- Request 1: ask with tools + thinking ---
    let mut request = ChatCompletionRequest::new(
        model,
        vec![
            ChatMessage::system(system_prompt),
            ChatMessage::user(user_prompt),
        ],
    );
    request.tools = Some(vec![add_tool]);
    request.thinking = Some(ThinkingConfig {
        kind: ThinkingMode::Enabled,
    });
    request.reasoning_effort = Some(ReasoningEffort::High);

    println!("--- request 1 ---");
    println!("  [system] {system_prompt}");
    println!("  [user] {user_prompt}");

    let completion = client.chat_completion(request).await?;

    let choice = completion
        .choices
        .first()
        .expect("expected at least one choice");
    let reasoning = choice.message.reasoning_content.clone();
    let tool_calls = choice
        .message
        .tool_calls
        .as_ref()
        .expect("expected tool calls in response");

    println!("\n--- response 1 ---");
    if let Some(r) = &reasoning {
        println!("  [reasoning] {r}");
    }
    let call = &tool_calls[0];
    println!(
        "  [tool call] {}({})",
        call.function.name, call.function.arguments
    );

    // --- Execute the tool locally ---
    let args: serde_json::Value = serde_json::from_str(&call.function.arguments)?;
    let tool_result = add(&args);
    println!("\n--- tool result ---");
    println!("  {tool_result}");

    // --- Request 2: replay conversation with tool result ---
    // Manually construct the assistant ToolCallsMessage to preserve reasoning_content.
    // ChatMessage::assistant_tool_calls() always sets reasoning_content to None, so direct
    // struct construction is needed when the model's thinking-mode response must be replayed.
    let assistant_msg = ChatMessage::ToolCalls(ToolCallsMessage {
        role: "assistant".to_owned(),
        content: None,
        name: None,
        tool_calls: tool_calls.clone(),
        reasoning_content: reasoning,
    });

    let request2 = ChatCompletionRequest::new(
        completion.model,
        vec![
            ChatMessage::system(system_prompt),
            ChatMessage::user(user_prompt),
            assistant_msg,
            ChatMessage::tool_result(tool_result.to_string(), &call.id),
        ],
    );

    let final_completion = client.chat_completion(request2).await?;

    println!("\n--- response 2 ---");
    if let Some(choice) = final_completion.choices.first() {
        if let Some(r) = &choice.message.reasoning_content {
            println!("  [reasoning] {r}");
        }
        println!(
            "  [assistant] {}",
            choice.message.content.as_deref().unwrap_or_default()
        );
    }
    if let Some(usage) = &final_completion.usage {
        println!(
            "  [usage] prompt={} completion={} total={}",
            usage.prompt_tokens, usage.completion_tokens, usage.total_tokens
        );
    }

    Ok(())
}