Skip to main content

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}