mistralai-client-rs/tests/v1_client_chat_test.rs

use jrest::expect;
use mistralai_client::v1::{
    chat::{ChatMessage, ChatMessageRole, ChatParams, ChatResponseChoiceFinishReason},
    client::Client,
    constants::Model,
    tool::{Tool, ToolChoice},
};

mod setup;

#[test]
fn test_client_chat() {
    setup::setup();

    let client = Client::new(None, None, None, None).unwrap();

    let model = Model::mistral_small_latest();
    let messages = vec![ChatMessage::new_user_message(
        "Guess the next word: \"Eiffel ...\"?",
    )];
    let options = ChatParams {
        temperature: Some(0.0),
        random_seed: Some(42),
        ..Default::default()
    };

    let response = client.chat(model, messages, Some(options)).unwrap();

    expect!(response.object).to_be("chat.completion".to_string());
    expect!(response.choices.len()).to_be(1);
    expect!(response.choices[0].index).to_be(0);
    expect!(response.choices[0].message.role.clone()).to_be(ChatMessageRole::Assistant);
    expect!(response.choices[0]
        .message
        .content
        .text()
        .contains("Tower"))
    .to_be(true);
    expect!(response.choices[0].finish_reason.clone()).to_be(ChatResponseChoiceFinishReason::Stop);
    expect!(response.usage.prompt_tokens).to_be_greater_than(0);
    expect!(response.usage.completion_tokens).to_be_greater_than(0);
    expect!(response.usage.total_tokens).to_be_greater_than(0);
}

#[test]
fn test_client_chat_with_function_calling() {
    setup::setup();

    let tools = vec![Tool::new(
        "get_city_temperature".to_string(),
        "Get the current temperature in a city.".to_string(),
        serde_json::json!({
            "type": "object",
            "properties": {
                "city": {
                    "type": "string",
                    "description": "The name of the city."
                }
            },
            "required": ["city"]
        }),
    )];

    let client = Client::new(None, None, None, None).unwrap();

    let model = Model::mistral_small_latest();
    let messages = vec![ChatMessage::new_user_message(
        "What's the current temperature in Paris?",
    )];
    let options = ChatParams {
        temperature: Some(0.0),
        random_seed: Some(42),
        tool_choice: Some(ToolChoice::Auto),
        tools: Some(tools),
        ..Default::default()
    };

    let response = client.chat(model, messages, Some(options)).unwrap();

    expect!(response.object).to_be("chat.completion".to_string());
    expect!(response.choices.len()).to_be(1);
    expect!(response.choices[0].index).to_be(0);
    expect!(response.choices[0].message.role.clone()).to_be(ChatMessageRole::Assistant);
    expect!(response.choices[0].finish_reason.clone())
        .to_be(ChatResponseChoiceFinishReason::ToolCalls);
    expect!(response.usage.prompt_tokens).to_be_greater_than(0);
    expect!(response.usage.completion_tokens).to_be_greater_than(0);
    expect!(response.usage.total_tokens).to_be_greater_than(0);
}