magi-openai 0.1.0

use std::{borrow::Cow, collections::HashMap};

use serde::{Deserialize, Serialize};
use serde_json::json;

/// Request body for the Responses API.
///
/// The structure mirrors the public OpenAI Responses endpoint while remaining
/// intentionally flexible. Unknown content types inside `input` are accepted as
/// raw JSON values so the SDK can continue to deserialize future additions
/// without requiring a breaking change.
#[derive(Clone, Serialize, Default, Debug, Deserialize, PartialEq)]
pub struct RequestBody {
    /// ID of the model to use.
    pub model: String,

    /// The primary input collection that drives the response.
    pub input: Vec<InputItem>,

    /// The unique ID of a previous response to continue a multi-turn conversation.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub previous_response_id: Option<String>,

    /// Specify additional output data to include in the model response.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub include: Option<Vec<Includable>>,

    /// Output types that you would like the model to generate for this request.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub modalities: Option<Vec<String>>,

    /// Maximum number of tokens that can be generated for a single response.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_output_tokens: Option<u32>,

    /// How many candidate responses to generate for each input.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub n: Option<u8>,

    /// Adjust the balance between concise and verbose answers.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub verbosity: Option<Verbosity>,

    /// Optional structured guidance for the model.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub instructions: Option<Instructions>,

    /// Configuration for a Predicted Output, which can improve response times
    /// when large parts of the model response are known ahead of time.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prediction: Option<PredictionConfig>,

    /// Parameters for audio output. Required when audio output is requested via
    /// the modalities field.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub audio: Option<AudioConfig>,

    /// Constrains model reasoning effort.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning: Option<ReasoningConfig>,

    /// Controls the sampling temperature, between 0 and 2.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,

    /// Alternative nucleus sampling configuration.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_p: Option<f32>,

    /// Whether to return log probabilities for output tokens.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub logprobs: Option<bool>,

    /// Number of top tokens to include when `logprobs` is enabled.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_logprobs: Option<u8>,

    /// Frequency penalty discouraging repeated tokens.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub frequency_penalty: Option<f32>,

    /// Presence penalty encouraging novel topics.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub presence_penalty: Option<f32>,

    /// Bias specific tokens during sampling.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub logit_bias: Option<HashMap<String, serde_json::Value>>,

    /// A unique identifier representing your end-user.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user: Option<String>,

    /// In beta. Attempts to make sampling deterministic.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub seed: Option<i64>,

    /// Early termination conditions for the generated text.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stop: Option<Stop>,

    /// Whether to return partial deltas via Server Sent Events.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream: Option<bool>,

    /// Options for streaming responses. Only set when `stream` is true.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream_options: Option<StreamOptions>,

    /// Specifies the output format expected from the model.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub response_format: Option<ResponseFormat>,

    /// Whether or not to store the output of this request.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub store: Option<bool>,

    /// Developer-defined tags and values used for filtering responses in the dashboard.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metadata: Option<HashMap<String, serde_json::Value>>,

    /// Whether to enable parallel function calling during tool use.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parallel_tool_calls: Option<bool>,

    /// A list of tools the model may call.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<Vec<Tool>>,

    /// Controls whether a tool must or may be used.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_choice: Option<ToolChoice>,

    /// Specifies the latency tier to use for processing the request.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub service_tier: Option<String>,

    /// Options for streaming output usage.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub web_search_options: Option<WebSearchOptions>,

    /// Open router compatible reasoning configuration.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub openrouter_reasoning: Option<OpenRouterReasoning>,

    /// Used by OpenAI to cache responses for similar requests.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prompt_cache_key: Option<String>,

    /// Truncation strategy for the model response.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub truncation: Option<Truncation>,

    /// Configuration options for a text response from the model.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub text: Option<serde_json::Value>,
}

impl RequestBody {
    /// Convenience constructor for a simple single-message text prompt.
    pub fn from_text(model: impl Into<String>, text: impl Into<String>) -> Self {
        let input = vec![InputItem::text(Role::User, text)];
        Self {
            model: model.into(),
            input,
            ..Default::default()
        }
    }

    /// Returns the first user message contained in the input collection, if any.
    pub fn first_user_input(&self) -> Option<&InputItem> {
        self.input.iter().find(|message| message.role == Role::User)
    }

    /// Returns the text from the first user message, if available.
    pub fn first_user_text(&self) -> Option<String> {
        self.first_user_input()?.first_text()
    }
}

/// One entry in the request `input` array.
#[derive(Clone, Serialize, Default, Debug, Deserialize, PartialEq)]
pub struct InputItem {
    /// The role of the author of this message.
    pub role: Role,

    /// Content parts comprising the message. Each part is stored as raw JSON to
    /// preserve compatibility with future content types.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub content: Vec<ContentPart>,

    /// Identifier referencing a previous tool call, when applicable.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,

    /// Optional name for the participant.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

impl InputItem {
    /// Creates a new message with a collection of content parts.
    pub fn new(role: Role, content: Vec<ContentPart>) -> Self {
        Self {
            role,
            content,
            ..Default::default()
        }
    }

    /// Creates a message with a single text part.
    pub fn text(role: Role, text: impl Into<String>) -> Self {
        let content = vec![json!({
            "type": "input_text",
            "text": text.into()
        })];
        Self::new(role, content)
    }

    /// Attempts to extract the first textual content part as a string.
    pub fn first_text(&self) -> Option<String> {
        self.content.iter().find_map(|part| {
            part.get("text")
                .and_then(|value| value.as_str())
                .map(|text| text.to_string())
        })
    }
}

/// Individual content parts are represented as raw JSON values to remain
/// forward compatible with future OpenAI additions.
pub type ContentPart = serde_json::Value;

/// Includable output data options for responses.
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq, Eq)]
pub enum Includable {
    /// Include file search call results in the response.
    #[serde(rename = "file_search_call.results")]
    FileSearchCallResults,
    /// Include image URLs from input messages.
    #[serde(rename = "message.input_image.image_url")]
    MessageInputImageUrl,
    /// Include image URLs from computer call output.
    #[serde(rename = "computer_call_output.output.image_url")]
    ComputerCallOutputImageUrl,
}

/// Role of a message author.
#[derive(Clone, Copy, Serialize, Default, Debug, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    System,
    Developer,
    #[default]
    User,
    Assistant,
    Tool,
}

/// Instructions accepted by the Responses API can be expressed as a string or
/// as an array of structured content parts.
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum Instructions {
    Text(String),
    Content(Vec<ContentPart>),
}

/// Modify the likelihood of specified tokens appearing in the completion.
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum Stop {
    String(String),
    Array(Vec<String>),
}

/// A tool that the model may call. Supports function calls, file search, web search, and computer use.
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Tool {
    /// A custom function the model can call.
    Function(FunctionTool),
    /// A file search tool for searching uploaded files.
    FileSearch(FileSearchTool),
    /// A web search tool for searching the internet.
    #[serde(rename = "web_search_preview")]
    WebSearchPreview(WebSearchTool),
    /// A versioned web search tool.
    #[serde(rename = "web_search_preview_2025_03_11")]
    WebSearchPreview20250311(WebSearchTool),
    /// A tool for computer use (controlling a virtual computer).
    #[serde(rename = "computer_use_preview")]
    ComputerUsePreview(ComputerUseTool),
}

impl Default for Tool {
    fn default() -> Self {
        Tool::Function(FunctionTool::default())
    }
}

/// A function tool definition for custom function calls.
#[derive(Clone, Serialize, Default, Debug, Deserialize, PartialEq)]
pub struct FunctionTool {
    /// The name of the function to be called.
    pub name: String,
    /// Optional description of the function.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<Cow<'static, str>>,
    /// JSON schema describing the arguments the function expects.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parameters: Option<serde_json::Value>,
    /// Whether to enforce strict parameter validation.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub strict: Option<bool>,
}

/// A file search tool for searching uploaded files.
#[derive(Clone, Serialize, Default, Debug, Deserialize, PartialEq)]
pub struct FileSearchTool {
    /// The IDs of the vector stores to search.
    pub vector_store_ids: Vec<String>,
    /// The maximum number of results to return (1-50).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_num_results: Option<u32>,
    /// Ranking options for search.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ranking_options: Option<serde_json::Value>,
    /// A filter to apply to the search.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub filters: Option<serde_json::Value>,
}

/// A web search tool for searching the internet.
#[derive(Clone, Serialize, Default, Debug, Deserialize, PartialEq)]
pub struct WebSearchTool {
    /// The user's location for search context.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user_location: Option<WebSearchUserLocation>,
    /// High level guidance for the amount of context window space to use for the search.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub search_context_size: Option<SearchContextSize>,
}

/// Search context size options for web search.
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum SearchContextSize {
    Low,
    Medium,
    High,
}

/// A computer use tool for controlling a virtual computer.
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
pub struct ComputerUseTool {
    /// The type of computer environment to control.
    pub environment: ComputerEnvironment,
    /// The width of the computer display.
    pub display_width: u32,
    /// The height of the computer display.
    pub display_height: u32,
}

/// Computer environment types for computer use tool.
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ComputerEnvironment {
    Windows,
    Mac,
    Linux,
    Ubuntu,
    Browser,
}

/// Specifies a tool type the model should use.
#[derive(Clone, Serialize, Default, Debug, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum ToolType {
    #[default]
    Function,
    FileSearch,
    #[serde(rename = "web_search_preview")]
    WebSearchPreview,
    #[serde(rename = "computer_use_preview")]
    ComputerUsePreview,
}

/// Controls whether a tool must be called.
#[derive(Clone, Serialize, Default, Debug, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum ToolChoice {
    #[default]
    Auto,
    None,
    Tool(ToolChoiceFunction),
}

/// Specifies the function that must be invoked when `tool_choice` is used.
#[derive(Clone, Serialize, Default, Debug, Deserialize, PartialEq)]
pub struct ToolChoiceFunction {
    /// The type of the tool. Only `function` is supported.
    pub r#type: ToolType,
    /// Function metadata specifying which function must run.
    pub function: FunctionName,
}

/// Name wrapper used inside `tool_choice`.
#[derive(Clone, Serialize, Default, Debug, Deserialize, PartialEq)]
pub struct FunctionName {
    /// The name of the function to call.
    pub name: String,
}

/// Specifies the format that the model must output.
#[derive(Debug, Deserialize, Default, Serialize, Clone, PartialEq)]
#[serde(rename_all = "snake_case")]
#[serde(tag = "type")]
pub enum ResponseFormat {
    #[default]
    Text,
    JsonObject,
    JsonSchema {
        /// Optional JSON schema description.
        description: Option<String>,
        /// JSON schema properties.
        properties: Option<serde_json::Value>,
        /// JSON schema name.
        name: String,
        /// Whether to enforce strict JSON schema compliance.
        strict: Option<bool>,
    },
}

/// Options for streaming responses.
#[derive(Debug, Deserialize, Default, Serialize, Clone, PartialEq)]
pub struct StreamOptions {
    /// If set, an additional chunk containing usage statistics will be streamed
    /// before the final `[DONE]` message.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub include_usage: Option<bool>,
}

/// Configuration for predicted outputs.
#[derive(Debug, Deserialize, Default, Serialize, Clone, PartialEq)]
pub struct PredictionConfig {
    /// The predicted output text that you expect the model to generate.
    pub text: String,
    /// Optional list of predicted log probabilities for each token.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub logprobs: Option<Vec<f32>>,
}

/// Audio generation configuration.
#[derive(Debug, Deserialize, Default, Serialize, Clone, PartialEq)]
pub struct AudioConfig {
    /// The voice to use for text-to-speech.
    pub voice: String,
    /// The audio output format.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub format: Option<String>,
    /// Speed multiplier for the generated audio.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub speed: Option<f32>,
}

/// Web search configuration used by OpenAI's web search tool.
#[derive(Debug, Deserialize, Default, Serialize, Clone, PartialEq)]
pub struct WebSearchOptions {
    /// High level guidance for the amount of context window space to use for the search.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub search_context_size: Option<String>,
    /// Approximate location parameters for the search.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user_location: Option<WebSearchUserLocation>,
}

/// Location configuration for web search. Uses approximate location.
#[derive(Debug, Deserialize, Default, Serialize, Clone, PartialEq)]
pub struct WebSearchUserLocation {
    /// The type of location approximation. Always `approximate`.
    #[serde(default = "default_approximate_type")]
    pub r#type: String,
    /// City of the user.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub city: Option<String>,
    /// Country in ISO two-letter format.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub country: Option<String>,
    /// Region of the user.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub region: Option<String>,
    /// IANA timezone identifier.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub timezone: Option<String>,
}

fn default_approximate_type() -> String {
    "approximate".to_string()
}

/// Effort options for reasoning models.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ReasoningEffort {
    High,
    Medium,
    Low,
}

/// Summary options for reasoning output.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ReasoningSummary {
    Auto,
    Concise,
    Detailed,
}

/// Truncation strategy for responses.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum Truncation {
    Auto,
    Disabled,
}

/// Controls verbosity of responses.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum Verbosity {
    Low,
    Medium,
    High,
}

/// Reasoning configuration for OpenRouter compatibility.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
pub struct OpenRouterReasoning {
    /// Requested reasoning effort.
    pub effort: ReasoningEffort,
    /// Whether to exclude reasoning tokens from the response.
    pub exclude: bool,
}

/// Configuration for the OpenAI reasoning beta controls.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
pub struct ReasoningConfig {
    /// How much reasoning effort the model should apply.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub effort: Option<ReasoningEffort>,
    /// A summary of the reasoning performed by the model. One of `auto`, `concise`, or `detailed`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub summary: Option<ReasoningSummary>,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn create_simple_text_request() {
        let request = RequestBody::from_text("gpt-4.1-mini", "Hello world");
        assert_eq!(request.model, "gpt-4.1-mini");
        assert_eq!(request.input.len(), 1);
        let first = request.first_user_input().unwrap();
        assert_eq!(first.role, Role::User);
        assert_eq!(first.first_text().as_deref(), Some("Hello world"));
    }
}