Skip to main content

aimux_core/
result.rs

1//! Result types for `do_generate` and `do_stream`.
2
3use std::collections::HashMap;
4
5use futures::Stream;
6use serde::{Deserialize, Serialize};
7use std::pin::Pin;
8use ts_rs::TS;
9
10use crate::error::AiMuxError;
11use crate::shared::FileData;
12use crate::stream_part::StreamPart;
13use crate::types::{FinishReason, ProviderMetadata, ResponseMetadata, Usage, Warning};
14
15use serde_json::Value;
16
17/// Compatibility deserializer for `GenerateContent::ToolCall.input`.
18///
19/// The field used to be a `serde_json::Value` (the already-parsed arguments,
20/// since providers parsed their own input) and became a `String` (the raw,
21/// unparsed provider text) in the tool-input-parse-repair refactor, so a
22/// `GenerateResult` persisted or recorded before it carries an
23/// object/array/number/bool/null here. Accept both: a JSON string passes
24/// through unchanged, and any other JSON value is re-serialized to its
25/// compact JSON text so the field keeps meaning "the raw text a
26/// schema-validating parse would run against".
27fn deserialize_tool_call_input<'de, D>(deserializer: D) -> Result<String, D::Error>
28where
29    D: serde::Deserializer<'de>,
30{
31    Ok(match Value::deserialize(deserializer)? {
32        Value::String(raw) => raw,
33        legacy => legacy.to_string(),
34    })
35}
36
37/// A content item in the generation result.
38#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
39#[ts(export)]
40pub enum GenerateContent {
41    /// Generated text.
42    Text {
43        text: String,
44        #[serde(default, skip_serializing_if = "Option::is_none")]
45        provider_metadata: Option<ProviderMetadata>,
46    },
47    /// A tool call requested by the model.
48    ToolCall {
49        tool_call_id: String,
50        tool_name: String,
51        /// The model's raw argument text, exactly as the provider delivered
52        /// it (possibly malformed). Providers never parse it — `generate_text`
53        /// owns parsing, schema validation, and repair.
54        ///
55        /// Always serializes as a JSON string. Deserializes a JSON string
56        /// (the current wire shape) unchanged, and also re-serializes the
57        /// pre-refactor shape — an already-parsed JSON value — to its compact
58        /// JSON text, so a `GenerateResult` persisted before this field
59        /// became a `String` keeps loading. See docs/api/gaps.md §9.
60        #[serde(deserialize_with = "deserialize_tool_call_input")]
61        input: String,
62        /// Whether the tool call will be executed by the provider.
63        /// If false/unset, the tool call is executed by the client.
64        #[serde(default, skip_serializing_if = "Option::is_none")]
65        provider_executed: Option<bool>,
66        /// Whether the tool is dynamic (defined at runtime, e.g. MCP tools).
67        #[serde(default, skip_serializing_if = "Option::is_none")]
68        dynamic: Option<bool>,
69        /// Provider-assigned thought signature (e.g. Google Gemini
70        /// `thoughtSignature`). Must be echoed back verbatim on the follow-up
71        /// turn when the tool result is sent.
72        #[serde(default, skip_serializing_if = "Option::is_none")]
73        thought_signature: Option<String>,
74        /// Additional provider-specific metadata for the tool call.
75        #[serde(default, skip_serializing_if = "Option::is_none")]
76        provider_metadata: Option<ProviderMetadata>,
77    },
78    /// A source / citation (e.g. URL citation from search-preview models).
79    Source {
80        id: String,
81        source_type: String,
82        url: Option<String>,
83        title: Option<String>,
84        #[serde(default, skip_serializing_if = "Option::is_none")]
85        provider_metadata: Option<ProviderMetadata>,
86    },
87    /// A reasoning / thinking segment produced by the model.
88    ///
89    /// Mirrors the Vercel AI SDK `reasoning` content type. `provider_metadata`
90    /// carries provider-specific data such as the Anthropic thinking-block
91    /// signature (`providerMetadata.anthropic.signature`).
92    Reasoning {
93        text: String,
94        provider_metadata: Option<ProviderMetadata>,
95    },
96    /// A file generated by the model (e.g. an image or document).
97    ///
98    /// Mirrors the Vercel AI SDK `file` content type. `data` is a tagged union
99    /// covering base64 data, URL, or provider reference.
100    File {
101        data: FileData,
102        media_type: String,
103        #[serde(default, skip_serializing_if = "Option::is_none")]
104        provider_metadata: Option<ProviderMetadata>,
105    },
106    /// A tool result from a provider-executed tool (e.g. xAI file_search,
107    /// web_search). Emitted alongside the preceding `ToolCall` when the
108    /// provider executes the tool server-side.
109    ToolResult {
110        tool_call_id: String,
111        tool_name: String,
112        result: Value,
113        /// Whether the result is an error or error message.
114        #[serde(default, skip_serializing_if = "Option::is_none")]
115        is_error: Option<bool>,
116        /// Whether the result is preliminary (replaces prior, e.g. image previews).
117        #[serde(default, skip_serializing_if = "Option::is_none")]
118        preliminary: Option<bool>,
119        /// Whether the tool is dynamic (defined at runtime, e.g. MCP tools).
120        #[serde(default, skip_serializing_if = "Option::is_none")]
121        dynamic: Option<bool>,
122        /// Additional provider-specific metadata for the tool result.
123        #[serde(default, skip_serializing_if = "Option::is_none")]
124        provider_metadata: Option<ProviderMetadata>,
125    },
126}
127
128// ── M7: top-level aggregation types ─────────────────────────────────────────
129//
130// Simplified views extracted from `GenerateContent` for the top-level fields of
131// `GenerateTextResult`. They omit provider-internal metadata so the common case
132// (reading reasoning/sources/files) doesn't require digging through `raw.content`.
133
134/// A reasoning / thinking segment, extracted to the top level (M7).
135#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
136#[ts(export)]
137pub struct ReasoningPart {
138    /// The reasoning text.
139    pub text: String,
140}
141
142/// A source / citation, extracted to the top level (M7).
143#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
144#[ts(export)]
145pub struct SourcePart {
146    pub id: String,
147    pub source_type: String,
148    #[serde(default, skip_serializing_if = "Option::is_none")]
149    pub url: Option<String>,
150    #[serde(default, skip_serializing_if = "Option::is_none")]
151    pub title: Option<String>,
152}
153
154/// A file generated by the model, extracted to the top level (M7).
155#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
156#[ts(export)]
157pub struct FilePart {
158    pub data: FileData,
159    pub media_type: String,
160}
161
162/// Aggregated result of `stream_text().consume()` (M11). Mirrors
163/// `GenerateTextResult`'s user-facing fields (without `raw`, since streaming
164/// has no `GenerateResult` equivalent).
165#[derive(Debug, Clone, Serialize, Deserialize, TS)]
166#[ts(export)]
167pub struct StreamTextResultAggregated {
168    /// The generated text (concatenated `TextDelta`).
169    pub text: String,
170    /// Reasoning / thinking segments.
171    #[serde(default)]
172    pub reasoning: Vec<ReasoningPart>,
173    /// Concatenated reasoning text.
174    #[serde(default)]
175    pub reasoning_text: String,
176    /// Tool calls requested by the model.
177    #[serde(default)]
178    pub tool_calls: Vec<crate::tool::ToolCall>,
179    /// Sources / citations.
180    #[serde(default)]
181    pub sources: Vec<SourcePart>,
182    /// Files generated by the model.
183    #[serde(default)]
184    pub files: Vec<FilePart>,
185    /// Why generation stopped.
186    pub finish_reason: crate::types::FinishReason,
187    /// Raw provider-specific finish reason string.
188    #[serde(default)]
189    pub raw_finish_reason: Option<String>,
190    /// Token usage.
191    #[serde(default)]
192    pub usage: crate::types::Usage,
193    /// Total token usage across all steps. In single-step mode (aimux's
194    /// default), `total_usage` equals `usage`. Provided for AI SDK parity.
195    #[serde(default)]
196    pub total_usage: crate::types::Usage,
197    /// Warnings from the provider.
198    #[serde(default)]
199    pub warnings: Vec<crate::types::Warning>,
200    /// Provider-specific metadata from the Finish chunk.
201    #[serde(default)]
202    pub provider_metadata: Option<serde_json::Value>,
203    /// Response metadata (id, timestamp, model_id) if emitted by the stream.
204    #[serde(default)]
205    pub response: Option<crate::types::ResponseMetadata>,
206    /// Assistant messages ready to append for the next turn.
207    #[serde(default)]
208    pub response_messages: Vec<crate::message::ModelMessage>,
209}
210
211/// Result of `LanguageModel::do_generate` (non-streaming).
212#[derive(Debug, Serialize, Deserialize, TS)]
213#[ts(export)]
214pub struct GenerateResult {
215    /// Generated content items (text, tool calls, etc.).
216    pub content: Vec<GenerateContent>,
217    /// Why generation stopped.
218    pub finish_reason: FinishReason,
219    /// Token usage.
220    pub usage: Usage,
221    /// Warnings issued by the provider.
222    pub warnings: Vec<Warning>,
223    /// Provider-specific metadata.
224    pub provider_metadata: Option<ProviderMetadata>,
225    /// Response metadata (id, timestamp, model_id).
226    pub response: ResponseMetadata,
227    /// The request body that was sent (for debugging).
228    pub request_body: Option<serde_json::Value>,
229    /// Response headers.
230    pub response_headers: Option<HashMap<String, String>>,
231}
232
233/// Result of `LanguageModel::do_stream` (streaming).
234pub struct StreamResult {
235    /// The stream of `StreamPart` items.
236    pub stream: Pin<Box<dyn Stream<Item = Result<StreamPart, AiMuxError>> + Send>>,
237    /// The request body that was sent (for debugging).
238    pub request_body: Option<serde_json::Value>,
239    /// Response headers.
240    pub response_headers: Option<HashMap<String, String>>,
241}
242
243impl std::fmt::Debug for StreamResult {
244    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
245        f.debug_struct("StreamResult")
246            .field("request_body", &self.request_body)
247            .field("response_headers", &self.response_headers)
248            .field("stream", &"<stream>")
249            .finish()
250    }
251}
252
253#[cfg(test)]
254mod tests {
255    use super::*;
256
257    /// Regression test for PR #165 review finding: `GenerateContent::ToolCall`
258    /// went from carrying `input: Value` (the parsed arguments) to
259    /// `input: String` (the raw text), so a `GenerateResult` persisted before
260    /// the refactor carries an object here and must still load rather than
261    /// fail closed. Serialization stays one-way — always a plain string.
262    #[test]
263    fn tool_call_input_loads_the_legacy_object_shape_and_the_current_string_shape() {
264        let wire = |input: serde_json::Value| {
265            serde_json::json!({
266                "ToolCall": {
267                    "tool_call_id": "call_1",
268                    "tool_name": "get_weather",
269                    "input": input,
270                }
271            })
272        };
273        let from_legacy: GenerateContent =
274            serde_json::from_value(wire(serde_json::json!({ "city": "Tokyo" }))).unwrap();
275        let from_current: GenerateContent =
276            serde_json::from_value(wire(serde_json::json!(r#"{"city":"Tokyo"}"#))).unwrap();
277        assert_eq!(from_legacy, from_current);
278
279        let GenerateContent::ToolCall { input, .. } = &from_legacy else {
280            panic!("expected ToolCall, got {from_legacy:?}");
281        };
282        assert_eq!(input, r#"{"city":"Tokyo"}"#);
283        assert_eq!(
284            serde_json::to_value(&from_legacy).unwrap(),
285            wire(serde_json::json!(r#"{"city":"Tokyo"}"#))
286        );
287    }
288}