1use ferrin_spec::BoxStream;
4use ferrin_spec::FinishReason;
5use ferrin_spec::ProviderMetadata;
6use ferrin_spec::RequestMetadata;
7use ferrin_spec::ResponseMetadata;
8use ferrin_spec::Usage;
9use ferrin_spec::Warning;
10use ferrin_spec::batch::BatchError;
11use ferrin_spec::batch::BatchItem;
12use ferrin_spec::batch::BatchItemResult;
13use ferrin_spec::image_model::ImageUsage;
14use ferrin_tool::ToolSet;
15
16use crate::error::Error;
17use crate::generate_text::RefineToolInputs;
18use crate::generate_text::StepContent;
19use crate::generate_text::parse_tool_call::ParseContext;
20use crate::generate_text::run::convert_content;
21use crate::image::GeneratedImage;
22use crate::image::convert_images;
23
24#[derive(Debug, Clone, PartialEq)]
26pub struct TextBatchResult {
27 pub content: Vec<StepContent>,
30 pub finish_reason: FinishReason,
32 pub usage: Usage,
34 pub warnings: Vec<Warning>,
36 pub request: RequestMetadata,
38 pub response: ResponseMetadata,
40 pub provider_metadata: Option<ProviderMetadata>,
42}
43
44impl TextBatchResult {
45 #[must_use]
47 pub fn text(&self) -> String {
48 self.content
49 .iter()
50 .filter_map(|part| match part {
51 StepContent::Text { text, .. } => Some(text.as_str()),
52 _ => None,
53 })
54 .collect()
55 }
56}
57
58#[derive(Debug, Clone, PartialEq)]
60pub struct ImageBatchResult {
61 pub images: Vec<GeneratedImage>,
63 pub warnings: Vec<Warning>,
65 pub response: ResponseMetadata,
67 pub provider_metadata: Option<ProviderMetadata>,
69 pub usage: Option<ImageUsage>,
71}
72
73#[derive(Debug, Clone, PartialEq)]
75#[non_exhaustive]
76pub enum BatchResultItem {
77 Text(Box<BatchItem<TextBatchResult>>),
79 Image(Box<BatchItem<ImageBatchResult>>),
81}
82
83impl BatchResultItem {
84 #[must_use]
86 pub fn id(&self) -> &str {
87 match self {
88 Self::Text(item) => item.id(),
89 Self::Image(item) => item.id(),
90 }
91 }
92}
93
94pub type BatchResults = BoxStream<'static, Result<BatchResultItem, Error>>;
96
97pub(super) fn unsupported_status<R>(id: String) -> BatchItem<R> {
98 BatchItem::Failed {
99 id,
100 error: BatchError {
101 message: "unsupported batch item status".to_owned(),
102 error_type: None,
103 code: None,
104 status_code: None,
105 },
106 provider_metadata: None,
107 }
108}
109
110pub(super) async fn convert_item(item: BatchItemResult, tools: &ToolSet) -> BatchResultItem {
111 let item_id = item.id().to_owned();
112 match item {
113 BatchItemResult::Text(item) => BatchResultItem::Text(Box::new(match *item {
114 BatchItem::Succeeded { id, result } => {
115 let refine = RefineToolInputs::default();
116 let parse_ctx = ParseContext {
117 tools,
118 tool_choice: None,
119 repair: None,
120 refine: &refine,
121 system: None,
122 messages: &[],
123 };
124 let (content, _calls) = convert_content(&result.content, &parse_ctx, tools).await;
125 BatchItem::Succeeded {
126 id,
127 result: TextBatchResult {
128 content,
129 finish_reason: result.finish_reason,
130 usage: result.usage,
131 warnings: result.warnings,
132 request: result.request,
133 response: result.response,
134 provider_metadata: result.provider_metadata,
135 },
136 }
137 }
138 BatchItem::Failed {
139 id,
140 error,
141 provider_metadata,
142 } => BatchItem::Failed {
143 id,
144 error,
145 provider_metadata,
146 },
147 BatchItem::Cancelled {
148 id,
149 error,
150 provider_metadata,
151 } => BatchItem::Cancelled {
152 id,
153 error,
154 provider_metadata,
155 },
156 BatchItem::Expired {
157 id,
158 error,
159 provider_metadata,
160 } => BatchItem::Expired {
161 id,
162 error,
163 provider_metadata,
164 },
165 #[allow(unreachable_patterns, reason = "BatchItem is non-exhaustive")]
166 _ => unsupported_status(item_id),
167 })),
168 BatchItemResult::Image(item) => BatchResultItem::Image(Box::new(match *item {
169 BatchItem::Succeeded { id, result } => BatchItem::Succeeded {
170 id,
171 result: ImageBatchResult {
172 images: convert_images(&result),
173 warnings: result.warnings,
174 response: result.response,
175 provider_metadata: result.provider_metadata,
176 usage: result.usage,
177 },
178 },
179 BatchItem::Failed {
180 id,
181 error,
182 provider_metadata,
183 } => BatchItem::Failed {
184 id,
185 error,
186 provider_metadata,
187 },
188 BatchItem::Cancelled {
189 id,
190 error,
191 provider_metadata,
192 } => BatchItem::Cancelled {
193 id,
194 error,
195 provider_metadata,
196 },
197 BatchItem::Expired {
198 id,
199 error,
200 provider_metadata,
201 } => BatchItem::Expired {
202 id,
203 error,
204 provider_metadata,
205 },
206 #[allow(unreachable_patterns, reason = "BatchItem is non-exhaustive")]
207 _ => unsupported_status(item_id),
208 })),
209 #[allow(unreachable_patterns, reason = "BatchItemResult is non-exhaustive")]
210 _ => BatchResultItem::Text(Box::new(unsupported_status(item_id))),
211 }
212}