mistralai-client-rs/tests/v1_agents_api_test.rs

use mistralai_client::v1::{
    agents::*,
    client::Client,
};

mod setup;

fn make_client() -> Client {
    Client::new(None, None, None, None).unwrap()
}

// ---------------------------------------------------------------------------
// Sync tests
// ---------------------------------------------------------------------------

#[test]
fn test_create_and_delete_agent() {
    setup::setup();
    let client = make_client();

    let req = CreateAgentRequest {
        model: "mistral-medium-latest".to_string(),
        name: "test-create-delete".to_string(),
        description: Some("Integration test agent".to_string()),
        instructions: Some("You are a test agent. Respond briefly.".to_string()),
        tools: None,
        handoffs: None,
        completion_args: None,
        metadata: None,
    };

    let agent = client.create_agent(&req).unwrap();
    assert!(!agent.id.is_empty());
    assert_eq!(agent.name, "test-create-delete");
    assert_eq!(agent.model, "mistral-medium-latest");
    assert_eq!(agent.object, "agent");
    // Version starts at 0 in the API
    assert!(agent.description.as_deref() == Some("Integration test agent"));
    assert!(agent.instructions.as_deref() == Some("You are a test agent. Respond briefly."));

    // Cleanup
    let del = client.delete_agent(&agent.id).unwrap();
    assert!(del.deleted);
}

#[test]
fn test_create_agent_with_tools() {
    setup::setup();
    let client = make_client();

    let req = CreateAgentRequest {
        model: "mistral-medium-latest".to_string(),
        name: "test-agent-tools".to_string(),
        description: None,
        instructions: Some("You can search.".to_string()),
        tools: Some(vec![
            AgentTool::function(
                "search".to_string(),
                "Search for things".to_string(),
                serde_json::json!({
                    "type": "object",
                    "properties": {
                        "query": {"type": "string", "description": "Search query"}
                    },
                    "required": ["query"]
                }),
            ),
            AgentTool::web_search(),
        ]),
        handoffs: None,
        completion_args: Some(CompletionArgs {
            temperature: Some(0.3),
            ..Default::default()
        }),
        metadata: None,
    };

    let agent = client.create_agent(&req).unwrap();
    assert_eq!(agent.tools.len(), 2);
    assert!(matches!(&agent.tools[0], AgentTool::Function(_)));
    assert!(matches!(&agent.tools[1], AgentTool::WebSearch {}));

    // Verify completion_args round-tripped
    let args = agent.completion_args.as_ref().unwrap();
    assert!((args.temperature.unwrap() - 0.3).abs() < 0.01);

    client.delete_agent(&agent.id).unwrap();
}

#[test]
fn test_get_agent() {
    setup::setup();
    let client = make_client();

    let req = CreateAgentRequest {
        model: "mistral-medium-latest".to_string(),
        name: "test-get-agent".to_string(),
        description: Some("Get test".to_string()),
        instructions: None,
        tools: None,
        handoffs: None,
        completion_args: None,
        metadata: None,
    };

    let created = client.create_agent(&req).unwrap();
    let fetched = client.get_agent(&created.id).unwrap();

    assert_eq!(fetched.id, created.id);
    assert_eq!(fetched.name, "test-get-agent");
    assert_eq!(fetched.model, "mistral-medium-latest");
    assert_eq!(fetched.description.as_deref(), Some("Get test"));

    client.delete_agent(&created.id).unwrap();
}

#[test]
fn test_update_agent() {
    setup::setup();
    let client = make_client();

    let req = CreateAgentRequest {
        model: "mistral-medium-latest".to_string(),
        name: "test-update-agent".to_string(),
        description: Some("Before update".to_string()),
        instructions: Some("Original instructions".to_string()),
        tools: None,
        handoffs: None,
        completion_args: None,
        metadata: None,
    };

    let created = client.create_agent(&req).unwrap();

    let update = UpdateAgentRequest {
        name: Some("test-update-agent-renamed".to_string()),
        description: Some("After update".to_string()),
        instructions: Some("Updated instructions".to_string()),
        ..Default::default()
    };

    let updated = client.update_agent(&created.id, &update).unwrap();
    assert_eq!(updated.id, created.id);
    assert_eq!(updated.name, "test-update-agent-renamed");
    assert_eq!(updated.description.as_deref(), Some("After update"));
    assert_eq!(updated.instructions.as_deref(), Some("Updated instructions"));
    // Version should have incremented
    assert!(updated.version >= created.version);

    client.delete_agent(&created.id).unwrap();
}

#[test]
fn test_list_agents() {
    setup::setup();
    let client = make_client();

    // Create two agents
    let req1 = CreateAgentRequest {
        model: "mistral-medium-latest".to_string(),
        name: "test-list-agent-1".to_string(),
        description: None,
        instructions: None,
        tools: None,
        handoffs: None,
        completion_args: None,
        metadata: None,
    };
    let req2 = CreateAgentRequest {
        model: "mistral-medium-latest".to_string(),
        name: "test-list-agent-2".to_string(),
        description: None,
        instructions: None,
        tools: None,
        handoffs: None,
        completion_args: None,
        metadata: None,
    };

    let a1 = client.create_agent(&req1).unwrap();
    let a2 = client.create_agent(&req2).unwrap();

    let list = client.list_agents().unwrap();
    assert!(list.data.len() >= 2);

    // Our two agents should be in the list
    let ids: Vec<&str> = list.data.iter().map(|a| a.id.as_str()).collect();
    assert!(ids.contains(&a1.id.as_str()));
    assert!(ids.contains(&a2.id.as_str()));

    // Cleanup
    client.delete_agent(&a1.id).unwrap();
    client.delete_agent(&a2.id).unwrap();
}

// ---------------------------------------------------------------------------
// Async tests
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_create_and_delete_agent_async() {
    setup::setup();
    let client = make_client();

    let req = CreateAgentRequest {
        model: "mistral-medium-latest".to_string(),
        name: "test-async-create-delete".to_string(),
        description: Some("Async integration test".to_string()),
        instructions: Some("Respond briefly.".to_string()),
        tools: None,
        handoffs: None,
        completion_args: None,
        metadata: None,
    };

    let agent = client.create_agent_async(&req).await.unwrap();
    assert!(!agent.id.is_empty());
    assert_eq!(agent.name, "test-async-create-delete");
    assert_eq!(agent.object, "agent");

    let del = client.delete_agent_async(&agent.id).await.unwrap();
    assert!(del.deleted);
}

#[tokio::test]
async fn test_get_agent_async() {
    setup::setup();
    let client = make_client();

    let req = CreateAgentRequest {
        model: "mistral-medium-latest".to_string(),
        name: "test-async-get".to_string(),
        description: None,
        instructions: None,
        tools: None,
        handoffs: None,
        completion_args: None,
        metadata: None,
    };

    let created = client.create_agent_async(&req).await.unwrap();
    let fetched = client.get_agent_async(&created.id).await.unwrap();
    assert_eq!(fetched.id, created.id);
    assert_eq!(fetched.name, "test-async-get");

    client.delete_agent_async(&created.id).await.unwrap();
}

#[tokio::test]
async fn test_update_agent_async() {
    setup::setup();
    let client = make_client();

    let req = CreateAgentRequest {
        model: "mistral-medium-latest".to_string(),
        name: "test-async-update".to_string(),
        description: Some("Before".to_string()),
        instructions: None,
        tools: None,
        handoffs: None,
        completion_args: None,
        metadata: None,
    };

    let created = client.create_agent_async(&req).await.unwrap();

    let update = UpdateAgentRequest {
        description: Some("After".to_string()),
        ..Default::default()
    };
    let updated = client.update_agent_async(&created.id, &update).await.unwrap();
    assert_eq!(updated.description.as_deref(), Some("After"));

    client.delete_agent_async(&created.id).await.unwrap();
}

#[tokio::test]
async fn test_list_agents_async() {
    setup::setup();
    let client = make_client();

    let req = CreateAgentRequest {
        model: "mistral-medium-latest".to_string(),
        name: "test-async-list".to_string(),
        description: None,
        instructions: None,
        tools: None,
        handoffs: None,
        completion_args: None,
        metadata: None,
    };

    let agent = client.create_agent_async(&req).await.unwrap();
    let list = client.list_agents_async().await.unwrap();
    assert!(list.data.iter().any(|a| a.id == agent.id));

    client.delete_agent_async(&agent.id).await.unwrap();
}

#[test]
fn test_create_agent_with_handoffs() {
    setup::setup();
    let client = make_client();

    // Create a target agent first
    let target_req = CreateAgentRequest {
        model: "mistral-medium-latest".to_string(),
        name: "test-handoff-target".to_string(),
        description: Some("Target agent for handoff".to_string()),
        instructions: Some("You handle math questions.".to_string()),
        tools: None,
        handoffs: None,
        completion_args: None,
        metadata: None,
    };
    let target = client.create_agent(&target_req).unwrap();

    // Create orchestrator with handoff to target
    let orch_req = CreateAgentRequest {
        model: "mistral-medium-latest".to_string(),
        name: "test-handoff-orchestrator".to_string(),
        description: Some("Orchestrator with handoffs".to_string()),
        instructions: Some("Delegate math questions.".to_string()),
        tools: None,
        handoffs: Some(vec![target.id.clone()]),
        completion_args: None,
        metadata: None,
    };
    let orch = client.create_agent(&orch_req).unwrap();
    assert_eq!(orch.handoffs.as_ref().unwrap().len(), 1);
    assert_eq!(orch.handoffs.as_ref().unwrap()[0], target.id);

    // Cleanup
    client.delete_agent(&orch.id).unwrap();
    client.delete_agent(&target.id).unwrap();
}

#[test]
fn test_agent_completion_with_created_agent() {
    setup::setup();
    let client = make_client();

    let req = CreateAgentRequest {
        model: "mistral-medium-latest".to_string(),
        name: "test-completion-agent".to_string(),
        description: None,
        instructions: Some("Always respond with exactly the word 'pong'.".to_string()),
        tools: None,
        handoffs: None,
        completion_args: Some(CompletionArgs {
            temperature: Some(0.0),
            ..Default::default()
        }),
        metadata: None,
    };

    let agent = client.create_agent(&req).unwrap();

    // Use the existing agent_completion method with the created agent
    use mistralai_client::v1::chat::ChatMessage;
    let messages = vec![ChatMessage::new_user_message("ping")];
    let response = client
        .agent_completion(agent.id.clone(), messages, None)
        .unwrap();

    assert!(!response.choices.is_empty());
    let text = response.choices[0].message.content.text().to_lowercase();
    assert!(text.contains("pong"), "Expected 'pong', got: {text}");
    assert!(response.usage.total_tokens > 0);

    client.delete_agent(&agent.id).unwrap();
}