Skip to main content

ferrin_core/batch/
request.rs

1//! Batch request types and their validation.
2
3use std::collections::HashMap;
4
5use ferrin_message::Message;
6use ferrin_spec::AspectRatio;
7use ferrin_spec::ImageSize;
8use ferrin_spec::ModelId;
9use ferrin_spec::ProviderOptions;
10use ferrin_spec::ResponseFormat;
11use ferrin_spec::ToolChoice;
12use ferrin_spec::ToolDefinition;
13use ferrin_spec::ToolName;
14use ferrin_spec::image_model::ImageFile;
15use ferrin_tool::ToolSet;
16
17use crate::error::Error;
18use crate::prompt::CallSettings;
19use crate::prompt::Instructions;
20
21/// A text generation request of a batch (the settings part of
22/// `generate_text`, without callbacks or cancellation).
23#[derive(Debug, Clone)]
24pub struct TextBatchRequest {
25    /// Request id, unique within the batch.
26    pub id: String,
27    /// Model id at the provider.
28    pub model_id: ModelId,
29    /// System instructions.
30    pub system: Option<Instructions>,
31    /// Single user prompt (exclusive with `messages`).
32    pub prompt: Option<String>,
33    /// Conversation (exclusive with `prompt`).
34    pub messages: Option<Vec<Message>>,
35    /// Whether system messages may appear inside `messages`.
36    pub allow_system_in_messages: bool,
37    /// Tools offered to the model (the batch provider cannot execute them).
38    pub tools: ToolSet,
39    /// Tool choice.
40    pub tool_choice: Option<ToolChoice>,
41    /// Tools sent to the model (all when `None`).
42    pub active_tools: Option<Vec<ToolName>>,
43    /// Tool order.
44    pub tool_order: Vec<ToolName>,
45    /// Shared tool context used to resolve dynamic descriptions.
46    pub tools_context: Option<ferrin_spec::JsonValue>,
47    /// Sampling settings and provider options.
48    pub settings: CallSettings,
49    /// Response format.
50    pub response_format: Option<ResponseFormat>,
51}
52
53impl TextBatchRequest {
54    /// Creates a request for `model_id`.
55    #[must_use]
56    pub fn new(id: impl Into<String>, model_id: impl Into<ModelId>) -> Self {
57        Self {
58            id: id.into(),
59            model_id: model_id.into(),
60            system: None,
61            prompt: None,
62            messages: None,
63            allow_system_in_messages: false,
64            tools: ToolSet::new(),
65            tool_choice: None,
66            active_tools: None,
67            tool_order: Vec::new(),
68            tools_context: None,
69            settings: CallSettings::default(),
70            response_format: None,
71        }
72    }
73
74    /// Sets the system instructions.
75    #[must_use]
76    pub fn system(mut self, system: impl Into<Instructions>) -> Self {
77        self.system = Some(system.into());
78        self
79    }
80
81    /// Sets a single user prompt.
82    #[must_use]
83    pub fn prompt(mut self, prompt: impl Into<String>) -> Self {
84        self.prompt = Some(prompt.into());
85        self
86    }
87
88    /// Sets the conversation.
89    #[must_use]
90    pub fn messages(mut self, messages: impl IntoIterator<Item = Message>) -> Self {
91        self.messages = Some(messages.into_iter().collect());
92        self
93    }
94
95    /// Allows system messages inside the conversation.
96    #[must_use]
97    pub fn allow_system_in_messages(mut self, allow: bool) -> Self {
98        self.allow_system_in_messages = allow;
99        self
100    }
101
102    /// Sets the tools.
103    #[must_use]
104    pub fn tools(mut self, tools: ToolSet) -> Self {
105        self.tools = tools;
106        self
107    }
108
109    /// Sets the tool choice.
110    #[must_use]
111    pub fn tool_choice(mut self, tool_choice: ToolChoice) -> Self {
112        self.tool_choice = Some(tool_choice);
113        self
114    }
115
116    /// Restricts the tools sent to the model.
117    #[must_use]
118    pub fn active_tools(
119        mut self,
120        active_tools: impl IntoIterator<Item = impl Into<ToolName>>,
121    ) -> Self {
122        self.active_tools = Some(active_tools.into_iter().map(Into::into).collect());
123        self
124    }
125
126    /// Sets the tool order.
127    #[must_use]
128    pub fn tool_order(mut self, tool_order: impl IntoIterator<Item = impl Into<ToolName>>) -> Self {
129        self.tool_order = tool_order.into_iter().map(Into::into).collect();
130        self
131    }
132
133    /// Sets the shared tool context.
134    #[must_use]
135    pub fn tools_context(mut self, context: ferrin_spec::JsonValue) -> Self {
136        self.tools_context = Some(context);
137        self
138    }
139
140    /// Sets the sampling settings and provider options.
141    #[must_use]
142    pub fn settings(mut self, settings: CallSettings) -> Self {
143        self.settings = settings;
144        self
145    }
146
147    /// Sets the response format.
148    #[must_use]
149    pub fn response_format(mut self, response_format: ResponseFormat) -> Self {
150        self.response_format = Some(response_format);
151        self
152    }
153
154    /// Sets the provider options.
155    #[must_use]
156    pub fn provider_options(mut self, provider_options: ProviderOptions) -> Self {
157        self.settings.provider_options = provider_options;
158        self
159    }
160}
161
162/// An image generation request of a batch.
163#[derive(Debug, Clone)]
164pub struct ImageBatchRequest {
165    /// Request id, unique within the batch.
166    pub id: String,
167    /// Model id at the provider.
168    pub model_id: ModelId,
169    /// Text prompt.
170    pub prompt: Option<String>,
171    /// Number of images.
172    pub n: u32,
173    /// Image size.
174    pub size: Option<ImageSize>,
175    /// Aspect ratio.
176    pub aspect_ratio: Option<AspectRatio>,
177    /// Seed.
178    pub seed: Option<u64>,
179    /// Reference images.
180    pub files: Vec<ImageFile>,
181    /// Mask.
182    pub mask: Option<ImageFile>,
183    /// Provider options.
184    pub provider_options: ProviderOptions,
185}
186
187impl ImageBatchRequest {
188    /// Creates a request for `model_id`.
189    #[must_use]
190    pub fn new(
191        id: impl Into<String>,
192        model_id: impl Into<ModelId>,
193        prompt: impl Into<String>,
194    ) -> Self {
195        Self {
196            id: id.into(),
197            model_id: model_id.into(),
198            prompt: Some(prompt.into()),
199            n: 1,
200            size: None,
201            aspect_ratio: None,
202            seed: None,
203            files: Vec::new(),
204            mask: None,
205            provider_options: ProviderOptions::new(),
206        }
207    }
208
209    /// Number of images (default 1).
210    #[must_use]
211    pub fn n(mut self, n: u32) -> Self {
212        self.n = n;
213        self
214    }
215
216    /// Image size.
217    #[must_use]
218    pub fn size(mut self, size: ImageSize) -> Self {
219        self.size = Some(size);
220        self
221    }
222
223    /// Aspect ratio.
224    #[must_use]
225    pub fn aspect_ratio(mut self, aspect_ratio: AspectRatio) -> Self {
226        self.aspect_ratio = Some(aspect_ratio);
227        self
228    }
229
230    /// Seed.
231    #[must_use]
232    pub fn seed(mut self, seed: u64) -> Self {
233        self.seed = Some(seed);
234        self
235    }
236
237    /// Reference images.
238    #[must_use]
239    pub fn files(mut self, files: Vec<ImageFile>) -> Self {
240        self.files = files;
241        self
242    }
243
244    /// Mask.
245    #[must_use]
246    pub fn mask(mut self, mask: ImageFile) -> Self {
247        self.mask = Some(mask);
248        self
249    }
250
251    /// Provider options.
252    #[must_use]
253    pub fn provider_options(mut self, provider_options: ProviderOptions) -> Self {
254        self.provider_options = provider_options;
255        self
256    }
257}
258
259/// A request of a batch.
260#[derive(Debug, Clone)]
261#[non_exhaustive]
262pub enum BatchRequest {
263    /// Text generation.
264    Text(Box<TextBatchRequest>),
265    /// Image generation.
266    Image(Box<ImageBatchRequest>),
267}
268
269impl BatchRequest {
270    /// The request id.
271    #[must_use]
272    pub fn id(&self) -> &str {
273        match self {
274            Self::Text(request) => &request.id,
275            Self::Image(request) => &request.id,
276        }
277    }
278}
279
280impl From<TextBatchRequest> for BatchRequest {
281    fn from(request: TextBatchRequest) -> Self {
282        Self::Text(Box::new(request))
283    }
284}
285
286impl From<ImageBatchRequest> for BatchRequest {
287    fn from(request: ImageBatchRequest) -> Self {
288        Self::Image(Box::new(request))
289    }
290}
291
292pub(super) fn validate_requests(requests: &[BatchRequest]) -> Result<(), Error> {
293    if requests.is_empty() {
294        return Err(Error::invalid_argument("requests", "must not be empty"));
295    }
296    let mut ids: std::collections::HashSet<&str> = std::collections::HashSet::new();
297    for request in requests {
298        let id = request.id();
299        if id.trim().is_empty() {
300            return Err(Error::invalid_argument(
301                "requests",
302                "request ids must not be empty",
303            ));
304        }
305        if !ids.insert(id) {
306            return Err(Error::invalid_argument(
307                "requests",
308                format!("request ids must be unique; duplicate id `{id}`"),
309            ));
310        }
311    }
312    Ok(())
313}
314
315pub(super) fn validate_compatible_tools(
316    request_id: &str,
317    definitions: &[ToolDefinition],
318    seen: &mut HashMap<ToolName, ToolDefinition>,
319) -> Result<(), Error> {
320    for definition in definitions {
321        let name = match definition {
322            ToolDefinition::Function { name, .. } | ToolDefinition::Provider { name, .. } => {
323                name.clone()
324            }
325            #[allow(unreachable_patterns, reason = "ToolDefinition is non-exhaustive")]
326            _ => continue,
327        };
328        if let Some(previous) = seen.get(&name)
329            && previous != definition
330        {
331            return Err(Error::invalid_argument(
332                "requests",
333                format!(
334                    "tool `{name}` must have the same definition in every batch request \
335                     (request `{request_id}` differs)"
336                ),
337            ));
338        }
339        seen.insert(name, definition.clone());
340    }
341    Ok(())
342}