Skip to main content

ferrin_core/batch/
result.rs

1//! Batch result types and the conversion of provider items.
2
3use 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/// Successful text item of a batch: the shape of one `generate_text` step.
25#[derive(Debug, Clone, PartialEq)]
26pub struct TextBatchResult {
27    /// Content parts (tool calls parsed against the tools given to
28    /// [`GetBatchResults::tools`](super::GetBatchResults::tools)).
29    pub content: Vec<StepContent>,
30    /// Finish reason.
31    pub finish_reason: FinishReason,
32    /// Token usage.
33    pub usage: Usage,
34    /// Adapter warnings.
35    pub warnings: Vec<Warning>,
36    /// Request metadata.
37    pub request: RequestMetadata,
38    /// Response metadata.
39    pub response: ResponseMetadata,
40    /// Provider metadata.
41    pub provider_metadata: Option<ProviderMetadata>,
42}
43
44impl TextBatchResult {
45    /// Concatenated text parts.
46    #[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/// Successful image item of a batch.
59#[derive(Debug, Clone, PartialEq)]
60pub struct ImageBatchResult {
61    /// Generated images.
62    pub images: Vec<GeneratedImage>,
63    /// Adapter warnings.
64    pub warnings: Vec<Warning>,
65    /// Response metadata.
66    pub response: ResponseMetadata,
67    /// Provider metadata.
68    pub provider_metadata: Option<ProviderMetadata>,
69    /// Usage.
70    pub usage: Option<ImageUsage>,
71}
72
73/// One item of [`get_batch_results`](super::get_batch_results).
74#[derive(Debug, Clone, PartialEq)]
75#[non_exhaustive]
76pub enum BatchResultItem {
77    /// A text request.
78    Text(Box<BatchItem<TextBatchResult>>),
79    /// An image request.
80    Image(Box<BatchItem<ImageBatchResult>>),
81}
82
83impl BatchResultItem {
84    /// The request id.
85    #[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
94/// The stream returned by [`get_batch_results`](super::get_batch_results).
95pub 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}