xai_rust/

chat.rs

//! Chat helper functions matching the Python SDK's ergonomic API.
//!
//! This module provides standalone functions for creating messages and content,
//! similar to the Python SDK's `xai_sdk.chat` module.
//!
//! # Example
//!
//! ```rust
//! use xai_rust::chat::{user, system, assistant, image, text, file};
//!
//! // Simple text messages
//! let msg1 = system("You are a helpful assistant.");
//! let msg2 = user("Hello!");
//! let msg3 = assistant("Hi there!");
//!
//! // Message with image
//! let msg4 = user((
//!     text("What's in this image?"),
//!     image("https://example.com/image.jpg"),
//! ));
//!
//! // Message with file
//! let msg5 = user((
//!     text("Summarize this document"),
//!     file("file-abc123"),
//! ));
//! ```

use crate::models::content::{ContentPart, FileContent, ImageDetail, ImageUrlContent};
use crate::models::message::{Message, MessageContent, Role};

// ============================================================================
// Message Constructors
// ============================================================================

/// Create a system message.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::system;
///
/// let msg = system("You are a helpful assistant.");
/// ```
pub fn system(content: impl Into<String>) -> Message {
    Message::new(Role::System, MessageContent::Text(content.into()))
}

/// Create a user message.
///
/// Accepts either a string or content parts.
///
/// # Examples
///
/// ```rust
/// use xai_rust::chat::{user, text, image};
///
/// // Simple text
/// let msg1 = user("Hello!");
///
/// // With image
/// let msg2 = user((text("What's this?"), image("https://example.com/img.jpg")));
/// ```
pub fn user(content: impl IntoMessageContent) -> Message {
    Message::new(Role::User, content.into_message_content())
}

/// Create an assistant message.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::assistant;
///
/// let msg = assistant("I'm here to help!");
/// ```
pub fn assistant(content: impl Into<String>) -> Message {
    Message::new(Role::Assistant, MessageContent::Text(content.into()))
}

/// Create a developer message (alias for system).
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::developer;
///
/// let msg = developer("You are a coding assistant.");
/// ```
pub fn developer(content: impl Into<String>) -> Message {
    Message::new(Role::Developer, MessageContent::Text(content.into()))
}

/// Create a tool result message.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::tool_result;
///
/// let msg = tool_result("call_123", r#"{"temperature": 72}"#);
/// ```
pub fn tool_result(tool_call_id: impl Into<String>, content: impl Into<String>) -> Message {
    Message::tool(tool_call_id, content)
}

// ============================================================================
// Content Constructors
// ============================================================================

/// Create text content.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::text;
///
/// let content = text("Hello, world!");
/// ```
pub fn text(content: impl Into<String>) -> ContentPart {
    ContentPart::Text {
        text: content.into(),
    }
}

/// Create image content from a URL.
///
/// # Examples
///
/// ```rust
/// use xai_rust::chat::image;
///
/// // From URL
/// let img = image("https://example.com/photo.jpg");
///
/// // From base64 data URL
/// let img = image("data:image/jpeg;base64,/9j/4AAQ...");
/// ```
pub fn image(url: impl Into<String>) -> ContentPart {
    ContentPart::ImageUrl {
        image_url: ImageUrlContent::new(url),
    }
}

/// Create image content with a specific detail level.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::image_with_detail;
/// use xai_rust::ImageDetail;
///
/// let img = image_with_detail("https://example.com/photo.jpg", ImageDetail::High);
/// ```
pub fn image_with_detail(url: impl Into<String>, detail: ImageDetail) -> ContentPart {
    ContentPart::ImageUrl {
        image_url: ImageUrlContent::new(url).with_detail(detail),
    }
}

/// Create image content from base64 data.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::image_base64;
///
/// let img = image_base64("/9j/4AAQ...", "image/jpeg");
/// ```
pub fn image_base64(data: impl Into<String>, mime_type: impl Into<String>) -> ContentPart {
    ContentPart::ImageUrl {
        image_url: ImageUrlContent::from_base64(&data.into(), &mime_type.into()),
    }
}

/// Create file content from a file ID.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::file;
///
/// let f = file("file-abc123");
/// ```
pub fn file(file_id: impl Into<String>) -> ContentPart {
    ContentPart::File {
        file: FileContent::new(file_id),
    }
}

/// Create inline file content.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::file_inline;
///
/// let f = file_inline("data.json", r#"{"key": "value"}"#);
/// ```
pub fn file_inline(filename: impl Into<String>, data: impl Into<String>) -> ContentPart {
    ContentPart::File {
        file: FileContent::inline(filename, data),
    }
}

// ============================================================================
// Trait for flexible content input
// ============================================================================

/// Trait for types that can be converted to message content.
///
/// This allows `user()` to accept both strings and content parts.
pub trait IntoMessageContent {
    /// Convert to message content.
    fn into_message_content(self) -> MessageContent;
}

impl IntoMessageContent for String {
    fn into_message_content(self) -> MessageContent {
        MessageContent::Text(self)
    }
}

impl IntoMessageContent for &str {
    fn into_message_content(self) -> MessageContent {
        MessageContent::Text(self.to_string())
    }
}

impl IntoMessageContent for ContentPart {
    fn into_message_content(self) -> MessageContent {
        MessageContent::Parts(vec![self])
    }
}

impl IntoMessageContent for Vec<ContentPart> {
    fn into_message_content(self) -> MessageContent {
        MessageContent::Parts(self)
    }
}

// Tuple implementations for combining content parts.
// Up to 4 content parts can be combined in a single tuple.
// For more parts, use `Vec<ContentPart>` instead.
impl IntoMessageContent for (ContentPart,) {
    fn into_message_content(self) -> MessageContent {
        MessageContent::Parts(vec![self.0])
    }
}

impl IntoMessageContent for (ContentPart, ContentPart) {
    fn into_message_content(self) -> MessageContent {
        MessageContent::Parts(vec![self.0, self.1])
    }
}

impl IntoMessageContent for (ContentPart, ContentPart, ContentPart) {
    fn into_message_content(self) -> MessageContent {
        MessageContent::Parts(vec![self.0, self.1, self.2])
    }
}

impl IntoMessageContent for (ContentPart, ContentPart, ContentPart, ContentPart) {
    fn into_message_content(self) -> MessageContent {
        MessageContent::Parts(vec![self.0, self.1, self.2, self.3])
    }
}

// ============================================================================
// Tool Helpers
// ============================================================================

/// Create a function tool definition.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::tool;
/// use serde_json::json;
///
/// let weather_tool = tool(
///     "get_weather",
///     "Get the current weather for a location",
///     json!({
///         "type": "object",
///         "properties": {
///             "location": {"type": "string", "description": "City name"}
///         },
///         "required": ["location"]
///     }),
/// );
/// ```
pub fn tool(
    name: impl Into<String>,
    description: impl Into<String>,
    parameters: serde_json::Value,
) -> crate::models::tool::Tool {
    crate::models::tool::Tool::function(name, description, parameters)
}

/// Create a function tool with strict schema enforcement.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::tool_strict;
/// use serde_json::json;
///
/// let tool = tool_strict(
///     "get_data",
///     "Fetch data",
///     json!({"type": "object", "properties": {}}),
/// );
/// ```
pub fn tool_strict(
    name: impl Into<String>,
    description: impl Into<String>,
    parameters: serde_json::Value,
) -> crate::models::tool::Tool {
    crate::models::tool::Tool::Function {
        function: crate::models::tool::FunctionDefinition {
            name: name.into(),
            description: Some(description.into()),
            parameters,
            strict: Some(true),
        },
    }
}

/// Create a web search tool.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::web_search;
///
/// let tool = web_search();
/// ```
pub fn web_search() -> crate::models::tool::Tool {
    crate::models::tool::Tool::web_search()
}

/// Create a web search tool with allowed domains.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::web_search_allow;
///
/// let tool = web_search_allow(vec!["wikipedia.org".into(), "docs.rs".into()]);
/// ```
pub fn web_search_allow(domains: Vec<String>) -> crate::models::tool::Tool {
    crate::models::tool::Tool::web_search_filtered(
        crate::models::tool::WebSearchFilters::allow_domains(domains),
    )
}

/// Create a web search tool with excluded domains.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::web_search_exclude;
///
/// let tool = web_search_exclude(vec!["example.com".into()]);
/// ```
pub fn web_search_exclude(domains: Vec<String>) -> crate::models::tool::Tool {
    crate::models::tool::Tool::web_search_filtered(
        crate::models::tool::WebSearchFilters::exclude_domains(domains),
    )
}

/// Create an X (Twitter) search tool.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::x_search;
///
/// let tool = x_search();
/// ```
pub fn x_search() -> crate::models::tool::Tool {
    crate::models::tool::Tool::x_search()
}

/// Create a code interpreter tool.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::code_interpreter;
///
/// let tool = code_interpreter();
/// ```
pub fn code_interpreter() -> crate::models::tool::Tool {
    crate::models::tool::Tool::code_interpreter()
}

/// Create a collections search tool.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::collections_search;
///
/// let tool = collections_search(vec!["collection-123".into()]);
/// ```
pub fn collections_search(collection_ids: Vec<String>) -> crate::models::tool::Tool {
    crate::models::tool::Tool::collections_search(collection_ids)
}

/// Create an MCP (Model Context Protocol) tool.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::mcp;
///
/// let tool = mcp("https://mcp-server.example.com");
/// ```
pub fn mcp(server_url: impl Into<String>) -> crate::models::tool::Tool {
    crate::models::tool::Tool::mcp(server_url)
}

// ============================================================================
// Tool Choice Helpers
// ============================================================================

/// Specify that the model should automatically decide whether to use tools.
pub fn tool_choice_auto() -> crate::models::tool::ToolChoice {
    crate::models::tool::ToolChoice::auto()
}

/// Specify that the model must use at least one tool.
pub fn tool_choice_required() -> crate::models::tool::ToolChoice {
    crate::models::tool::ToolChoice::required()
}

/// Specify that the model should not use any tools.
pub fn tool_choice_none() -> crate::models::tool::ToolChoice {
    crate::models::tool::ToolChoice::none()
}

/// Force the model to use a specific function.
///
/// # Example
///
/// ```rust
/// use xai_rust::chat::required_tool;
///
/// let choice = required_tool("get_weather");
/// ```
pub fn required_tool(function_name: impl Into<String>) -> crate::models::tool::ToolChoice {
    crate::models::tool::ToolChoice::function(function_name)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::models::content::ContentPart;
    use crate::models::message::{MessageContent, Role};
    use crate::models::tool::{Tool, ToolChoice};

    // ── Message constructor helpers ────────────────────────────────────

    #[test]
    fn system_creates_system_message() {
        let msg = system("You are helpful");
        assert_eq!(msg.role, Role::System);
        assert_eq!(msg.text(), "You are helpful");
        assert!(msg.name.is_none());
        assert!(msg.tool_call_id.is_none());
    }

    #[test]
    fn user_creates_user_message_from_str() {
        let msg = user("Hello!");
        assert_eq!(msg.role, Role::User);
        assert_eq!(msg.text(), "Hello!");
    }

    #[test]
    fn user_creates_user_message_from_string() {
        let msg = user(String::from("Hello!"));
        assert_eq!(msg.role, Role::User);
        assert_eq!(msg.text(), "Hello!");
    }

    #[test]
    fn assistant_creates_assistant_message() {
        let msg = assistant("Hi there!");
        assert_eq!(msg.role, Role::Assistant);
        assert_eq!(msg.text(), "Hi there!");
    }

    #[test]
    fn developer_creates_developer_message() {
        let msg = developer("You are a coder");
        assert_eq!(msg.role, Role::Developer);
        assert_eq!(msg.text(), "You are a coder");
    }

    #[test]
    fn tool_result_creates_tool_message() {
        let msg = tool_result("call_123", r#"{"temp":72}"#);
        assert_eq!(msg.role, Role::Tool);
        assert_eq!(msg.text(), r#"{"temp":72}"#);
        assert_eq!(msg.tool_call_id.as_deref(), Some("call_123"));
    }

    // ── Content constructor helpers ────────────────────────────────────

    #[test]
    fn text_creates_text_part() {
        let part = text("Hello");
        assert_eq!(part.as_text().unwrap(), "Hello");
    }

    #[test]
    fn image_creates_image_url_part() {
        let part = image("https://example.com/img.jpg");
        if let ContentPart::ImageUrl { image_url } = part {
            assert_eq!(image_url.url, "https://example.com/img.jpg");
            assert!(image_url.detail.is_none());
        } else {
            panic!("Expected ImageUrl variant");
        }
    }

    #[test]
    fn image_with_detail_creates_image_part() {
        let part = image_with_detail(
            "https://example.com/img.jpg",
            crate::models::content::ImageDetail::High,
        );
        if let ContentPart::ImageUrl { image_url } = part {
            assert_eq!(
                image_url.detail,
                Some(crate::models::content::ImageDetail::High)
            );
        } else {
            panic!("Expected ImageUrl variant");
        }
    }

    #[test]
    fn image_base64_creates_data_url() {
        let part = image_base64("abc123", "image/png");
        if let ContentPart::ImageUrl { image_url } = part {
            assert!(image_url.url.starts_with("data:image/png;base64,"));
            assert!(image_url.url.ends_with("abc123"));
        } else {
            panic!("Expected ImageUrl variant");
        }
    }

    #[test]
    fn file_creates_file_part() {
        let part = file("file-abc123");
        if let ContentPart::File { file: fc } = part {
            assert_eq!(fc.file_id.as_deref(), Some("file-abc123"));
        } else {
            panic!("Expected File variant");
        }
    }

    #[test]
    fn file_inline_creates_inline_file_part() {
        let part = file_inline("test.json", r#"{"x": 1}"#);
        if let ContentPart::File { file: fc } = part {
            assert!(fc.file_id.is_none());
            assert_eq!(fc.filename.as_deref(), Some("test.json"));
            assert_eq!(fc.file_data.as_deref(), Some(r#"{"x": 1}"#));
        } else {
            panic!("Expected File variant");
        }
    }

    // ── IntoMessageContent trait ──────────────────────────────────────

    #[test]
    fn into_message_content_str() {
        let content: MessageContent = "hello".into_message_content();
        assert_eq!(content.as_text().unwrap(), "hello");
    }

    #[test]
    fn into_message_content_string() {
        let content: MessageContent = String::from("hello").into_message_content();
        assert_eq!(content.as_text().unwrap(), "hello");
    }

    #[test]
    fn into_message_content_single_part() {
        let part = ContentPart::text("part");
        let content = part.into_message_content();
        if let MessageContent::Parts(parts) = content {
            assert_eq!(parts.len(), 1);
            assert_eq!(parts[0].as_text().unwrap(), "part");
        } else {
            panic!("Expected Parts variant");
        }
    }

    #[test]
    fn into_message_content_vec_parts() {
        let parts = vec![ContentPart::text("a"), ContentPart::text("b")];
        let content = parts.into_message_content();
        if let MessageContent::Parts(p) = content {
            assert_eq!(p.len(), 2);
        } else {
            panic!("Expected Parts variant");
        }
    }

    #[test]
    fn into_message_content_tuple_1() {
        let content = (ContentPart::text("a"),).into_message_content();
        if let MessageContent::Parts(p) = content {
            assert_eq!(p.len(), 1);
        } else {
            panic!("Expected Parts variant");
        }
    }

    #[test]
    fn into_message_content_tuple_2() {
        let content =
            (text("What's this?"), image("https://example.com/img.jpg")).into_message_content();
        if let MessageContent::Parts(parts) = content {
            assert_eq!(parts.len(), 2);
            assert_eq!(parts[0].as_text().unwrap(), "What's this?");
            assert!(matches!(parts[1], ContentPart::ImageUrl { .. }));
        } else {
            panic!("Expected Parts variant");
        }
    }

    #[test]
    fn into_message_content_tuple_3() {
        let content = (text("a"), text("b"), text("c")).into_message_content();
        if let MessageContent::Parts(p) = content {
            assert_eq!(p.len(), 3);
        } else {
            panic!("Expected Parts variant");
        }
    }

    #[test]
    fn into_message_content_tuple_4() {
        let content = (text("a"), text("b"), text("c"), text("d")).into_message_content();
        if let MessageContent::Parts(p) = content {
            assert_eq!(p.len(), 4);
        } else {
            panic!("Expected Parts variant");
        }
    }

    #[test]
    fn user_accepts_tuple_content() {
        let msg = user((text("Look at this"), image("https://example.com/photo.jpg")));
        assert_eq!(msg.role, Role::User);
        if let MessageContent::Parts(parts) = &msg.content {
            assert_eq!(parts.len(), 2);
        } else {
            panic!("Expected Parts content");
        }
    }

    // ── Tool helpers ──────────────────────────────────────────────────

    #[test]
    fn tool_helper_creates_function_tool() {
        let t = tool(
            "test_fn",
            "A test function",
            serde_json::json!({"type": "object"}),
        );
        assert!(matches!(t, Tool::Function { .. }));
    }

    #[test]
    fn tool_strict_creates_strict_function_tool() {
        let t = tool_strict("test_fn", "strict function", serde_json::json!({}));
        if let Tool::Function { function } = t {
            assert_eq!(function.strict, Some(true));
        } else {
            panic!("Expected Function variant");
        }
    }

    #[test]
    fn web_search_helper() {
        let t = web_search();
        assert!(matches!(t, Tool::WebSearch { .. }));
    }

    #[test]
    fn web_search_allow_helper() {
        let t = web_search_allow(vec!["docs.rs".to_string()]);
        if let Tool::WebSearch { filters, .. } = t {
            assert!(filters.is_some());
            let f = filters.unwrap();
            assert_eq!(f.allowed_domains.unwrap(), vec!["docs.rs"]);
        } else {
            panic!("Expected WebSearch variant");
        }
    }

    #[test]
    fn web_search_exclude_helper() {
        let t = web_search_exclude(vec!["example.com".to_string()]);
        if let Tool::WebSearch { filters, .. } = t {
            let f = filters.unwrap();
            assert_eq!(f.excluded_domains.unwrap(), vec!["example.com"]);
        } else {
            panic!("Expected WebSearch variant");
        }
    }

    #[test]
    fn x_search_helper() {
        let t = x_search();
        assert!(matches!(t, Tool::XSearch { .. }));
    }

    #[test]
    fn code_interpreter_helper() {
        let t = code_interpreter();
        assert!(matches!(t, Tool::CodeInterpreter {}));
    }

    #[test]
    fn collections_search_helper() {
        let t = collections_search(vec!["col-1".to_string()]);
        assert!(matches!(t, Tool::CollectionsSearch { .. }));
    }

    #[test]
    fn mcp_helper() {
        let t = mcp("https://mcp.example.com");
        assert!(matches!(t, Tool::Mcp { .. }));
    }

    // ── Tool choice helpers ──────────────────────────────────────────

    #[test]
    fn tool_choice_auto_helper() {
        let tc = tool_choice_auto();
        assert!(matches!(tc, ToolChoice::Auto(_)));
    }

    #[test]
    fn tool_choice_required_helper() {
        let tc = tool_choice_required();
        assert!(matches!(tc, ToolChoice::Auto(_)));
    }

    #[test]
    fn tool_choice_none_helper() {
        let tc = tool_choice_none();
        assert!(matches!(tc, ToolChoice::Auto(_)));
    }

    #[test]
    fn required_tool_helper() {
        let tc = required_tool("get_weather");
        if let ToolChoice::Specific(spec) = tc {
            assert_eq!(spec.tool_type, "function");
            assert_eq!(spec.function.unwrap().name, "get_weather");
        } else {
            panic!("Expected Specific variant");
        }
    }
}