aisdk/core/language_model/request.rs
1//! Defines the `LanguageModelRequest` struct and its builder for configuring text generation requests.
2//!
3//! This module provides the `LanguageModelRequest` type, which encapsulates a language model
4//! and options for generating text or streaming responses. It includes a type-state builder
5//! pattern to ensure requests are constructed correctly and safely.
6
7use crate::core::Messages;
8use crate::core::capabilities::*;
9use crate::core::language_model::{LanguageModel, LanguageModelOptions};
10use crate::core::tools::Tool;
11use schemars::{JsonSchema, schema_for};
12use std::fmt::Debug;
13use std::ops::{Deref, DerefMut};
14use std::sync::Arc;
15
16/// Options for text generation requests such as `generate_text` and `stream_text`.
17#[derive(Debug)]
18pub struct LanguageModelRequest<M: LanguageModel> {
19 /// The language model to use for text generation.
20 pub model: M,
21
22 /// An optional simple text prompt for the request.
23 ///
24 /// This should not be set if `messages` are provided in the options.
25 pub prompt: Option<String>,
26
27 /// Configuration options for the language model request.
28 pub(crate) options: LanguageModelOptions,
29}
30
31impl<M: LanguageModel> LanguageModelRequest<M> {
32 /// Creates a new builder for constructing a `LanguageModelRequest`.
33 ///
34 /// This method initiates the type-state builder pattern, starting with the
35 /// [`ModelStage`] where you must specify the language model.
36 pub fn builder() -> LanguageModelRequestBuilder<M> {
37 LanguageModelRequestBuilder::default()
38 }
39}
40
41impl<M: LanguageModel> Deref for LanguageModelRequest<M> {
42 type Target = LanguageModelOptions;
43
44 fn deref(&self) -> &Self::Target {
45 &self.options
46 }
47}
48
49impl<M: LanguageModel> DerefMut for LanguageModelRequest<M> {
50 fn deref_mut(&mut self) -> &mut Self::Target {
51 &mut self.options
52 }
53}
54
55/// Type-state markers for the `LanguageModelRequestBuilder`.
56///
57/// These zero-sized types ensure the builder is used in the correct order,
58/// preventing invalid request configurations at compile time.
59///
60/// The initial builder state where the language model must be set.
61///
62/// Transitions to [`SystemStage`] after calling [`model`](LanguageModelRequestBuilder::model).
63pub struct ModelStage {}
64
65/// The state after setting the model, where a system prompt can be optionally added.
66///
67/// Transitions to [`ConversationStage`] after calling [`system`](LanguageModelRequestBuilder::system),
68/// or directly to [`OptionsStage`] after calling [`prompt`](LanguageModelRequestBuilder::prompt) or [`messages`](LanguageModelRequestBuilder::messages).
69pub struct SystemStage {}
70
71/// The state after optionally setting a system prompt, where conversation input must be provided.
72///
73/// Transitions to [`OptionsStage`] after calling [`prompt`](LanguageModelRequestBuilder::prompt) or [`messages`](LanguageModelRequestBuilder::messages).
74pub struct ConversationStage {}
75
76/// The final state where additional options can be configured before building.
77///
78/// Transitions to the completed `LanguageModelRequest` after calling [`build`](LanguageModelRequestBuilder::build).
79pub struct OptionsStage {}
80
81/// A type-state builder for constructing `LanguageModelRequest` instances.
82///
83/// This builder uses phantom types to enforce a specific construction order,
84/// ensuring that required fields (like the model) are set before optional ones.
85///
86/// # Type Parameters
87///
88/// * `M` - The language model type.
89/// * `State` - The current builder state, determining available methods.
90pub struct LanguageModelRequestBuilder<M: LanguageModel, State = ModelStage> {
91 model: Option<M>,
92 prompt: Option<String>,
93 options: LanguageModelOptions,
94 state: std::marker::PhantomData<State>,
95}
96
97impl<M: LanguageModel, State> Deref for LanguageModelRequestBuilder<M, State> {
98 type Target = LanguageModelOptions;
99
100 /// Dereferences to the underlying `LanguageModelOptions`.
101 ///
102 /// This allows direct access to the options fields during building.
103 fn deref(&self) -> &Self::Target {
104 &self.options
105 }
106}
107
108impl<M: LanguageModel, State> DerefMut for LanguageModelRequestBuilder<M, State> {
109 /// Mutably dereferences to the underlying `LanguageModelOptions`.
110 ///
111 /// This allows direct mutation of the options fields during building.
112 fn deref_mut(&mut self) -> &mut Self::Target {
113 &mut self.options
114 }
115}
116
117impl<M: LanguageModel> LanguageModelRequestBuilder<M> {
118 fn default() -> Self {
119 LanguageModelRequestBuilder {
120 model: None,
121 prompt: None,
122 options: LanguageModelOptions::default(),
123 state: std::marker::PhantomData,
124 }
125 }
126}
127
128/// Methods available in the [`ModelStage`] state.
129impl<M: LanguageModel> LanguageModelRequestBuilder<M, ModelStage> {
130 /// Sets the language model for the request.
131 ///
132 /// This is the first required step in building a request.
133 ///
134 /// # Parameters
135 ///
136 /// * `model` - The language model instance to use.
137 ///
138 /// # Returns
139 ///
140 /// The builder in the [`SystemStage`] state.
141 pub fn model(self, model: M) -> LanguageModelRequestBuilder<M, SystemStage> {
142 LanguageModelRequestBuilder {
143 model: Some(model),
144 prompt: self.prompt,
145 options: self.options,
146 state: std::marker::PhantomData,
147 }
148 }
149}
150
151/// Methods available in the [`SystemStage`] state.
152impl<M: LanguageModel> LanguageModelRequestBuilder<M, SystemStage> {
153 /// Sets an optional system prompt for the request.
154 ///
155 /// The system prompt provides context or instructions to the model.
156 ///
157 /// # Parameters
158 ///
159 /// * `system` - The system prompt text.
160 ///
161 /// # Returns
162 ///
163 /// The builder in the [`ConversationStage`] state.
164 pub fn system(
165 self,
166 system: impl Into<String>,
167 ) -> LanguageModelRequestBuilder<M, ConversationStage> {
168 LanguageModelRequestBuilder {
169 model: self.model,
170 prompt: self.prompt,
171 options: LanguageModelOptions {
172 system: Some(system.into()),
173 ..self.options
174 },
175 state: std::marker::PhantomData,
176 }
177 }
178
179 /// Sets a simple text prompt for the request.
180 ///
181 /// This skips the system prompt and goes directly to options.
182 ///
183 /// # Parameters
184 ///
185 /// * `prompt` - The user prompt text.
186 ///
187 /// # Returns
188 ///
189 /// The builder in the [`OptionsStage`] state.
190 pub fn prompt(self, prompt: impl Into<String>) -> LanguageModelRequestBuilder<M, OptionsStage> {
191 LanguageModelRequestBuilder {
192 model: self.model,
193 prompt: Some(prompt.into()),
194 options: self.options,
195 state: std::marker::PhantomData,
196 }
197 }
198
199 /// Sets conversation messages for the request.
200 ///
201 /// This allows for multi-turn conversations with the model.
202 ///
203 /// # Parameters
204 ///
205 /// * `messages` - `Messages` instances representing the conversation.
206 ///
207 /// # Returns
208 ///
209 /// The builder in the [`OptionsStage`] state.
210 pub fn messages(self, messages: Messages) -> LanguageModelRequestBuilder<M, OptionsStage> {
211 LanguageModelRequestBuilder {
212 model: self.model,
213 prompt: self.prompt,
214 options: LanguageModelOptions {
215 messages: messages.into_iter().map(|msg| msg.into()).collect(),
216 ..self.options
217 },
218 state: std::marker::PhantomData,
219 }
220 }
221}
222
223/// Methods available in the [`ConversationStage`] state.
224impl<M: LanguageModel> LanguageModelRequestBuilder<M, ConversationStage> {
225 /// Sets a simple text prompt for the request.
226 ///
227 /// This method allows setting a user prompt.
228 /// The prompt represents the user's input to the language model.
229 ///
230 /// # Parameters
231 ///
232 /// * `prompt` - The user prompt text.
233 ///
234 /// # Returns
235 ///
236 /// The builder in the [`OptionsStage`] state.
237 pub fn prompt(self, prompt: impl Into<String>) -> LanguageModelRequestBuilder<M, OptionsStage>
238 where
239 M: TextInputSupport,
240 {
241 LanguageModelRequestBuilder {
242 model: self.model,
243 prompt: Some(prompt.into()),
244 options: self.options,
245 state: std::marker::PhantomData,
246 }
247 }
248
249 /// Sets conversation messages for the request.
250 ///
251 /// This method allows providing a full conversation history as a vector of messages,
252 /// enabling multi-turn conversations with the language model.
253 ///
254 /// # Parameters
255 ///
256 /// * `messages` - `Messages` instances representing the conversation.
257 ///
258 /// # Returns
259 ///
260 /// The builder in the [`OptionsStage`] state.
261 pub fn messages(self, messages: Messages) -> LanguageModelRequestBuilder<M, OptionsStage>
262 where
263 M: TextInputSupport,
264 {
265 LanguageModelRequestBuilder {
266 model: self.model,
267 prompt: self.prompt,
268 options: LanguageModelOptions {
269 messages: messages.into_iter().map(|msg| msg.into()).collect(),
270 ..self.options
271 },
272 state: std::marker::PhantomData,
273 }
274 }
275}
276
277/// Methods available in the [`OptionsStage`] state.
278impl<M: LanguageModel> LanguageModelRequestBuilder<M, OptionsStage> {
279 /// Sets the output schema for structured generation.
280 ///
281 /// This method configures the language model to generate output that conforms
282 /// to the provided JSON schema. The schema is derived from the given type `T`.
283 ///
284 /// # Type Parameters
285 ///
286 /// * `T` - A type that implements [`JsonSchema`], used to generate the output schema.
287 ///
288 /// # Returns
289 ///
290 /// The builder with the schema configured.
291 pub fn schema<T: JsonSchema>(mut self) -> Self
292 where
293 M: StructuredOutputSupport,
294 {
295 self.schema = Some(schema_for!(T));
296 self
297 }
298
299 /// Sets a seed for deterministic generation.
300 ///
301 /// # Parameters
302 ///
303 /// * `seed` - The random seed value.
304 ///
305 /// # Returns
306 ///
307 /// The builder with the seed set.
308 pub fn seed(mut self, seed: impl Into<u32>) -> Self {
309 self.seed = Some(seed.into());
310 self
311 }
312
313 /// Sets the temperature for generation randomness (0-100, scaled to 0.0-1.0).
314 ///
315 /// Higher values increase creativity, lower values increase determinism.
316 ///
317 /// # Parameters
318 ///
319 /// * `temperature` - The temperature value (0-100).
320 ///
321 /// # Returns
322 ///
323 /// The builder with the temperature set.
324 pub fn temperature(mut self, temperature: impl Into<u32>) -> Self {
325 self.temperature = Some(temperature.into());
326 self
327 }
328
329 /// Sets the top-p (nucleus) sampling parameter (0-100, scaled to 0.0-1.0).
330 ///
331 /// # Parameters
332 ///
333 /// * `top_p` - The top-p value (0-100).
334 ///
335 /// # Returns
336 ///
337 /// The builder with top-p set.
338 pub fn top_p(mut self, top_p: impl Into<u32>) -> Self {
339 self.top_p = Some(top_p.into());
340 self
341 }
342
343 /// Sets the top-k sampling parameter.
344 ///
345 /// # Parameters
346 ///
347 /// * `top_k` - The top-k value.
348 ///
349 /// # Returns
350 ///
351 /// The builder with top-k set.
352 pub fn top_k(mut self, top_k: impl Into<u32>) -> Self {
353 self.top_k = Some(top_k.into());
354 self
355 }
356
357 /// Sets stop sequences that halt generation.
358 ///
359 /// # Parameters
360 ///
361 /// * `stop_sequences` - A list of strings that stop generation when encountered.
362 ///
363 /// # Returns
364 ///
365 /// The builder with stop sequences set.
366 pub fn stop_sequences(mut self, stop_sequences: impl Into<Vec<String>>) -> Self {
367 self.stop_sequences = Some(stop_sequences.into());
368 self
369 }
370
371 /// Sets the maximum number of retries for failed requests.
372 ///
373 /// # Parameters
374 ///
375 /// * `max_retries` - The maximum retry count.
376 ///
377 /// # Returns
378 ///
379 /// The builder with max retries set.
380 pub fn max_retries(mut self, max_retries: impl Into<u32>) -> Self {
381 self.max_retries = Some(max_retries.into());
382 self
383 }
384
385 /// Sets the frequency penalty to reduce repetition.
386 ///
387 /// # Parameters
388 ///
389 /// * `frequency_penalty` - The penalty value.
390 ///
391 /// # Returns
392 ///
393 /// The builder with frequency penalty set.
394 pub fn frequency_penalty(mut self, frequency_penalty: impl Into<f32>) -> Self {
395 self.frequency_penalty = Some(frequency_penalty.into());
396 self
397 }
398
399 /// Adds a tool to the request.
400 ///
401 /// # Arguments
402 ///
403 /// * `tool` - The tool to add.
404 ///
405 /// # Returns
406 ///
407 /// The builder with the tool added.
408 pub fn with_tool(mut self, tool: Tool) -> Self
409 where
410 M: ToolCallSupport,
411 {
412 self.tools.get_or_insert_default().add_tool(tool);
413 self
414 }
415
416 /// Sets a condition to stop the generation loop.
417 ///
418 /// # Parameters
419 ///
420 /// * `hook` - A function that returns `true` when generation should stop.
421 ///
422 /// # Returns
423 ///
424 /// The builder with the stop condition set.
425 pub fn stop_when<F>(mut self, hook: F) -> Self
426 where
427 F: Fn(&LanguageModelOptions) -> bool + Send + Sync + 'static,
428 {
429 self.stop_when = Some(Arc::new(hook));
430 self
431 }
432
433 /// Sets a hook to run at the start of each generation step.
434 ///
435 /// # Parameters
436 ///
437 /// * `hook` - A function called before each step.
438 ///
439 /// # Returns
440 ///
441 /// The builder with the hook set.
442 pub fn on_step_start<F>(mut self, hook: F) -> Self
443 where
444 F: Fn(&mut LanguageModelOptions) + Send + Sync + 'static,
445 {
446 self.on_step_start = Some(Arc::new(hook));
447 self
448 }
449
450 /// Sets a hook to run at the end of each generation step.
451 ///
452 /// # Parameters
453 ///
454 /// * `hook` - A function called after each step.
455 ///
456 /// # Returns
457 ///
458 /// The builder with the hook set.
459 pub fn on_step_finish<F>(mut self, hook: F) -> Self
460 where
461 F: Fn(&LanguageModelOptions) + Send + Sync + 'static,
462 {
463 self.on_step_finish = Some(Arc::new(hook));
464 self
465 }
466
467 /// Sets the reasoning effort level.
468 ///
469 /// # Parameters
470 ///
471 /// * `reasoning_effort` - The effort level.
472 ///
473 /// # Returns
474 ///
475 /// The builder with reasoning effort set.
476 pub fn reasoning_effort(
477 mut self,
478 reasoning_effort: impl Into<crate::core::language_model::ReasoningEffort>,
479 ) -> Self
480 where
481 M: ReasoningSupport,
482 {
483 self.reasoning_effort = Some(reasoning_effort.into());
484 self
485 }
486
487 /// Builds the `LanguageModelRequest`.
488 ///
489 /// This method consumes the builder and returns the configured request.
490 ///
491 /// # Returns
492 ///
493 /// The constructed `LanguageModelRequest`.
494 pub fn build(self) -> LanguageModelRequest<M> {
495 let model = self
496 .model
497 .unwrap_or_else(|| unreachable!("Model must be set"));
498
499 LanguageModelRequest {
500 model,
501 prompt: self.prompt,
502 options: self.options,
503 }
504 }
505}