aimux-core 0.5.0

//! Result types for `do_generate` and `do_stream`.

use std::collections::HashMap;

use futures::Stream;
use serde::{Deserialize, Serialize};
use std::pin::Pin;
use ts_rs::TS;

use crate::error::AiMuxError;
use crate::shared::FileData;
use crate::stream_part::StreamPart;
use crate::types::{FinishReason, ProviderMetadata, ResponseMetadata, Usage, Warning};

use serde_json::Value;

/// Compatibility deserializer for `GenerateContent::ToolCall.input`.
///
/// The field used to be a `serde_json::Value` (the already-parsed arguments,
/// since providers parsed their own input) and became a `String` (the raw,
/// unparsed provider text) in the tool-input-parse-repair refactor, so a
/// `GenerateResult` persisted or recorded before it carries an
/// object/array/number/bool/null here. Accept both: a JSON string passes
/// through unchanged, and any other JSON value is re-serialized to its
/// compact JSON text so the field keeps meaning "the raw text a
/// schema-validating parse would run against".
fn deserialize_tool_call_input<'de, D>(deserializer: D) -> Result<String, D::Error>
where
    D: serde::Deserializer<'de>,
{
    Ok(match Value::deserialize(deserializer)? {
        Value::String(raw) => raw,
        legacy => legacy.to_string(),
    })
}

/// A content item in the generation result.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
#[ts(export)]
pub enum GenerateContent {
    /// Generated text.
    Text {
        text: String,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        provider_metadata: Option<ProviderMetadata>,
    },
    /// A tool call requested by the model.
    ToolCall {
        tool_call_id: String,
        tool_name: String,
        /// The model's raw argument text, exactly as the provider delivered
        /// it (possibly malformed). Providers never parse it — `generate_text`
        /// owns parsing, schema validation, and repair.
        ///
        /// Always serializes as a JSON string. Deserializes a JSON string
        /// (the current wire shape) unchanged, and also re-serializes the
        /// pre-refactor shape — an already-parsed JSON value — to its compact
        /// JSON text, so a `GenerateResult` persisted before this field
        /// became a `String` keeps loading. See docs/api/gaps.md §9.
        #[serde(deserialize_with = "deserialize_tool_call_input")]
        input: String,
        /// Whether the tool call will be executed by the provider.
        /// If false/unset, the tool call is executed by the client.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        provider_executed: Option<bool>,
        /// Whether the tool is dynamic (defined at runtime, e.g. MCP tools).
        #[serde(default, skip_serializing_if = "Option::is_none")]
        dynamic: Option<bool>,
        /// Provider-assigned thought signature (e.g. Google Gemini
        /// `thoughtSignature`). Must be echoed back verbatim on the follow-up
        /// turn when the tool result is sent.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        thought_signature: Option<String>,
        /// Additional provider-specific metadata for the tool call.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        provider_metadata: Option<ProviderMetadata>,
    },
    /// A source / citation (e.g. URL citation from search-preview models).
    Source {
        id: String,
        source_type: String,
        url: Option<String>,
        title: Option<String>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        provider_metadata: Option<ProviderMetadata>,
    },
    /// A reasoning / thinking segment produced by the model.
    ///
    /// Mirrors the Vercel AI SDK `reasoning` content type. `provider_metadata`
    /// carries provider-specific data such as the Anthropic thinking-block
    /// signature (`providerMetadata.anthropic.signature`).
    Reasoning {
        text: String,
        provider_metadata: Option<ProviderMetadata>,
    },
    /// A file generated by the model (e.g. an image or document).
    ///
    /// Mirrors the Vercel AI SDK `file` content type. `data` is a tagged union
    /// covering base64 data, URL, or provider reference.
    File {
        data: FileData,
        media_type: String,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        provider_metadata: Option<ProviderMetadata>,
    },
    /// A tool result from a provider-executed tool (e.g. xAI file_search,
    /// web_search). Emitted alongside the preceding `ToolCall` when the
    /// provider executes the tool server-side.
    ToolResult {
        tool_call_id: String,
        tool_name: String,
        result: Value,
        /// Whether the result is an error or error message.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        is_error: Option<bool>,
        /// Whether the result is preliminary (replaces prior, e.g. image previews).
        #[serde(default, skip_serializing_if = "Option::is_none")]
        preliminary: Option<bool>,
        /// Whether the tool is dynamic (defined at runtime, e.g. MCP tools).
        #[serde(default, skip_serializing_if = "Option::is_none")]
        dynamic: Option<bool>,
        /// Additional provider-specific metadata for the tool result.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        provider_metadata: Option<ProviderMetadata>,
    },
}

// ── M7: top-level aggregation types ─────────────────────────────────────────
//
// Simplified views extracted from `GenerateContent` for the top-level fields of
// `GenerateTextResult`. They omit provider-internal metadata so the common case
// (reading reasoning/sources/files) doesn't require digging through `raw.content`.

/// A reasoning / thinking segment, extracted to the top level (M7).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct ReasoningPart {
    /// The reasoning text.
    pub text: String,
}

/// A source / citation, extracted to the top level (M7).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct SourcePart {
    pub id: String,
    pub source_type: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub url: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub title: Option<String>,
}

/// A file generated by the model, extracted to the top level (M7).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct FilePart {
    pub data: FileData,
    pub media_type: String,
}

/// Aggregated result of `stream_text().consume()` (M11). Mirrors
/// `GenerateTextResult`'s user-facing fields (without `raw`, since streaming
/// has no `GenerateResult` equivalent).
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct StreamTextResultAggregated {
    /// The generated text (concatenated `TextDelta`).
    pub text: String,
    /// Reasoning / thinking segments.
    #[serde(default)]
    pub reasoning: Vec<ReasoningPart>,
    /// Concatenated reasoning text.
    #[serde(default)]
    pub reasoning_text: String,
    /// Tool calls requested by the model.
    #[serde(default)]
    pub tool_calls: Vec<crate::tool::ToolCall>,
    /// Sources / citations.
    #[serde(default)]
    pub sources: Vec<SourcePart>,
    /// Files generated by the model.
    #[serde(default)]
    pub files: Vec<FilePart>,
    /// Why generation stopped.
    pub finish_reason: crate::types::FinishReason,
    /// Raw provider-specific finish reason string.
    #[serde(default)]
    pub raw_finish_reason: Option<String>,
    /// Token usage.
    #[serde(default)]
    pub usage: crate::types::Usage,
    /// Total token usage across all steps. In single-step mode (aimux's
    /// default), `total_usage` equals `usage`. Provided for AI SDK parity.
    #[serde(default)]
    pub total_usage: crate::types::Usage,
    /// Warnings from the provider.
    #[serde(default)]
    pub warnings: Vec<crate::types::Warning>,
    /// Provider-specific metadata from the Finish chunk.
    #[serde(default)]
    pub provider_metadata: Option<serde_json::Value>,
    /// Response metadata (id, timestamp, model_id) if emitted by the stream.
    #[serde(default)]
    pub response: Option<crate::types::ResponseMetadata>,
    /// Assistant messages ready to append for the next turn.
    #[serde(default)]
    pub response_messages: Vec<crate::message::ModelMessage>,
}

/// Result of `LanguageModel::do_generate` (non-streaming).
#[derive(Debug, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct GenerateResult {
    /// Generated content items (text, tool calls, etc.).
    pub content: Vec<GenerateContent>,
    /// Why generation stopped.
    pub finish_reason: FinishReason,
    /// Token usage.
    pub usage: Usage,
    /// Warnings issued by the provider.
    pub warnings: Vec<Warning>,
    /// Provider-specific metadata.
    pub provider_metadata: Option<ProviderMetadata>,
    /// Response metadata (id, timestamp, model_id).
    pub response: ResponseMetadata,
    /// The request body that was sent (for debugging).
    pub request_body: Option<serde_json::Value>,
    /// Response headers.
    pub response_headers: Option<HashMap<String, String>>,
}

/// Result of `LanguageModel::do_stream` (streaming).
pub struct StreamResult {
    /// The stream of `StreamPart` items.
    pub stream: Pin<Box<dyn Stream<Item = Result<StreamPart, AiMuxError>> + Send>>,
    /// The request body that was sent (for debugging).
    pub request_body: Option<serde_json::Value>,
    /// Response headers.
    pub response_headers: Option<HashMap<String, String>>,
}

impl std::fmt::Debug for StreamResult {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StreamResult")
            .field("request_body", &self.request_body)
            .field("response_headers", &self.response_headers)
            .field("stream", &"<stream>")
            .finish()
    }
}

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

    /// Regression test for PR #165 review finding: `GenerateContent::ToolCall`
    /// went from carrying `input: Value` (the parsed arguments) to
    /// `input: String` (the raw text), so a `GenerateResult` persisted before
    /// the refactor carries an object here and must still load rather than
    /// fail closed. Serialization stays one-way — always a plain string.
    #[test]
    fn tool_call_input_loads_the_legacy_object_shape_and_the_current_string_shape() {
        let wire = |input: serde_json::Value| {
            serde_json::json!({
                "ToolCall": {
                    "tool_call_id": "call_1",
                    "tool_name": "get_weather",
                    "input": input,
                }
            })
        };
        let from_legacy: GenerateContent =
            serde_json::from_value(wire(serde_json::json!({ "city": "Tokyo" }))).unwrap();
        let from_current: GenerateContent =
            serde_json::from_value(wire(serde_json::json!(r#"{"city":"Tokyo"}"#))).unwrap();
        assert_eq!(from_legacy, from_current);

        let GenerateContent::ToolCall { input, .. } = &from_legacy else {
            panic!("expected ToolCall, got {from_legacy:?}");
        };
        assert_eq!(input, r#"{"city":"Tokyo"}"#);
        assert_eq!(
            serde_json::to_value(&from_legacy).unwrap(),
            wire(serde_json::json!(r#"{"city":"Tokyo"}"#))
        );
    }
}