mistralrs/
messages.rs

1use std::{collections::HashMap, fmt::Display, sync::Arc};
2
3use super::*;
4use either::Either;
5use image::DynamicImage;
6use indexmap::IndexMap;
7use serde_json::{json, Value};
8
9/// A type which can be used as a chat request.
10pub trait RequestLike {
11    fn messages_ref(&self) -> &[IndexMap<String, MessageContent>];
12    fn images_ref(&self) -> &[DynamicImage];
13    fn take_messages(&mut self) -> RequestMessage;
14    fn take_logits_processors(&mut self) -> Option<Vec<Arc<dyn CustomLogitsProcessor>>>;
15    fn take_adapters(&mut self) -> Option<Vec<String>>;
16    fn return_logprobs(&self) -> bool;
17    fn enable_search(&self) -> Option<bool>;
18    fn take_constraint(&mut self) -> Constraint;
19    fn take_tools(&mut self) -> Option<(Vec<Tool>, ToolChoice)>;
20    fn take_sampling_params(&mut self) -> SamplingParams;
21    fn take_web_search_options(&mut self) -> Option<WebSearchOptions>;
22    fn truncate_sequence(&self) -> bool {
23        false
24    }
25}
26
27#[derive(Debug, Clone, PartialEq)]
28/// Plain text (chat) messages.
29///
30/// No constraints, logits processors, logprobs, tools, or adapters.
31///
32/// Sampling is deterministic.
33pub struct TextMessages {
34    messages: Vec<IndexMap<String, MessageContent>>,
35    enable_thinking: Option<bool>,
36}
37
38impl From<TextMessages> for Vec<IndexMap<String, MessageContent>> {
39    fn from(value: TextMessages) -> Self {
40        value.messages
41    }
42}
43
44#[derive(Debug, Clone, PartialEq)]
45/// A chat message role.
46pub enum TextMessageRole {
47    User,
48    Assistant,
49    System,
50    Tool,
51    Custom(String),
52}
53
54impl Display for TextMessageRole {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        match self {
57            Self::User => write!(f, "user"),
58            Self::Assistant => write!(f, "assistant"),
59            Self::System => write!(f, "system"),
60            Self::Tool => write!(f, "tool"),
61            Self::Custom(c) => write!(f, "{c}"),
62        }
63    }
64}
65
66impl Default for TextMessages {
67    fn default() -> Self {
68        Self::new()
69    }
70}
71
72impl TextMessages {
73    pub fn new() -> Self {
74        Self {
75            messages: Vec::new(),
76            enable_thinking: None,
77        }
78    }
79
80    pub fn add_message(mut self, role: TextMessageRole, text: impl ToString) -> Self {
81        self.messages.push(IndexMap::from([
82            ("role".to_string(), Either::Left(role.to_string())),
83            ("content".to_string(), Either::Left(text.to_string())),
84        ]));
85        self
86    }
87
88    pub fn clear(mut self) -> Self {
89        self.messages.clear();
90        self
91    }
92
93    pub fn enable_thinking(mut self, enable_thinking: bool) -> Self {
94        self.enable_thinking = Some(enable_thinking);
95        self
96    }
97}
98
99impl RequestLike for TextMessages {
100    fn messages_ref(&self) -> &[IndexMap<String, MessageContent>] {
101        &self.messages
102    }
103    fn images_ref(&self) -> &[DynamicImage] {
104        &[]
105    }
106    fn take_messages(&mut self) -> RequestMessage {
107        let mut other = Vec::new();
108        std::mem::swap(&mut other, &mut self.messages);
109        RequestMessage::Chat {
110            messages: other,
111            enable_thinking: self.enable_thinking,
112        }
113    }
114    fn enable_search(&self) -> Option<bool> {
115        None
116    }
117    fn take_logits_processors(&mut self) -> Option<Vec<Arc<dyn CustomLogitsProcessor>>> {
118        None
119    }
120    fn take_adapters(&mut self) -> Option<Vec<String>> {
121        None
122    }
123    fn return_logprobs(&self) -> bool {
124        false
125    }
126    fn take_constraint(&mut self) -> Constraint {
127        Constraint::None
128    }
129    fn take_tools(&mut self) -> Option<(Vec<Tool>, ToolChoice)> {
130        None
131    }
132    fn take_sampling_params(&mut self) -> SamplingParams {
133        SamplingParams::deterministic()
134    }
135    fn take_web_search_options(&mut self) -> Option<WebSearchOptions> {
136        None
137    }
138}
139
140#[derive(Debug, Clone, PartialEq)]
141/// Text (chat) messages with images and/or audios.
142///
143/// No constraints, logits processors, logprobs, tools, or adapters.
144///
145/// Sampling is deterministic.
146pub struct VisionMessages {
147    messages: Vec<IndexMap<String, MessageContent>>,
148    images: Vec<DynamicImage>,
149    audios: Vec<AudioInput>,
150    enable_thinking: Option<bool>,
151}
152
153impl Default for VisionMessages {
154    fn default() -> Self {
155        Self::new()
156    }
157}
158
159impl VisionMessages {
160    pub fn new() -> Self {
161        Self {
162            images: Vec::new(),
163            messages: Vec::new(),
164            audios: Vec::new(),
165            enable_thinking: None,
166        }
167    }
168
169    pub fn add_message(mut self, role: TextMessageRole, text: impl ToString) -> Self {
170        self.messages.push(IndexMap::from([
171            ("role".to_string(), Either::Left(role.to_string())),
172            ("content".to_string(), Either::Left(text.to_string())),
173        ]));
174        self
175    }
176
177    pub fn add_image_message(
178        self,
179        role: TextMessageRole,
180        text: impl ToString,
181        images: Vec<DynamicImage>,
182        model: &Model,
183    ) -> anyhow::Result<Self> {
184        self.add_multimodal_message(role, text, images, vec![], model)
185    }
186
187    pub fn add_audio_message(
188        self,
189        role: TextMessageRole,
190        text: impl ToString,
191        audios: Vec<AudioInput>,
192        model: &Model,
193    ) -> anyhow::Result<Self> {
194        self.add_multimodal_message(role, text, vec![], audios, model)
195    }
196
197    pub fn add_multimodal_message(
198        mut self,
199        role: TextMessageRole,
200        text: impl ToString,
201        images: Vec<DynamicImage>,
202        audios: Vec<AudioInput>,
203        model: &Model,
204    ) -> anyhow::Result<Self> {
205        let config = model.config().unwrap();
206        let prefixer = match &config.category {
207            ModelCategory::Vision { prefixer } => prefixer,
208            _ => {
209                anyhow::bail!("`add_image_message` expects a vision model.")
210            }
211        };
212
213        // Images
214        let n_added_images = images.len();
215        let prefixed = prefixer.prefix_image(
216            (self.images.len()..self.images.len() + n_added_images).collect(),
217            &text.to_string(),
218        );
219        self.images.extend(images);
220
221        // Audios
222        let n_added_audios = audios.len();
223        let prefixed = prefixer.prefix_audio(
224            (self.audios.len()..self.audios.len() + n_added_audios).collect(),
225            &prefixed,
226        );
227        self.audios.extend(audios);
228
229        if n_added_images > 0 {
230            self.messages.push(IndexMap::from([
231                ("role".to_string(), Either::Left(role.to_string())),
232                (
233                    "content".to_string(),
234                    Either::Right(vec![
235                        IndexMap::from([("type".to_string(), Value::String("image".to_string()))]),
236                        IndexMap::from([
237                            ("type".to_string(), Value::String("text".to_string())),
238                            ("text".to_string(), Value::String(prefixed)),
239                        ]),
240                    ]),
241                ),
242            ]));
243        } else {
244            self.messages.push(IndexMap::from([
245                ("role".to_string(), Either::Left(role.to_string())),
246                ("content".to_string(), Either::Left(prefixed)),
247            ]));
248        }
249        Ok(self)
250    }
251
252    pub fn clear(mut self) -> Self {
253        self.messages.clear();
254        self.images.clear();
255        self.audios.clear();
256
257        self
258    }
259
260    pub fn enable_thinking(mut self, enable_thinking: bool) -> Self {
261        self.enable_thinking = Some(enable_thinking);
262        self
263    }
264}
265
266impl RequestLike for VisionMessages {
267    fn messages_ref(&self) -> &[IndexMap<String, MessageContent>] {
268        &self.messages
269    }
270    fn images_ref(&self) -> &[DynamicImage] {
271        &self.images
272    }
273    fn take_messages(&mut self) -> RequestMessage {
274        let mut other_messages = Vec::new();
275        std::mem::swap(&mut other_messages, &mut self.messages);
276        let mut other_images = Vec::new();
277        std::mem::swap(&mut other_images, &mut self.images);
278        let mut other_audios = Vec::new();
279        std::mem::swap(&mut other_audios, &mut self.audios);
280        RequestMessage::VisionChat {
281            images: other_images,
282            messages: other_messages,
283            audios: other_audios,
284            enable_thinking: self.enable_thinking,
285        }
286    }
287    fn enable_search(&self) -> Option<bool> {
288        None
289    }
290    fn take_logits_processors(&mut self) -> Option<Vec<Arc<dyn CustomLogitsProcessor>>> {
291        None
292    }
293    fn take_adapters(&mut self) -> Option<Vec<String>> {
294        None
295    }
296    fn return_logprobs(&self) -> bool {
297        false
298    }
299    fn take_constraint(&mut self) -> Constraint {
300        Constraint::None
301    }
302    fn take_tools(&mut self) -> Option<(Vec<Tool>, ToolChoice)> {
303        None
304    }
305    fn take_sampling_params(&mut self) -> SamplingParams {
306        SamplingParams::deterministic()
307    }
308    fn take_web_search_options(&mut self) -> Option<WebSearchOptions> {
309        None
310    }
311}
312
313#[derive(Clone)]
314/// A way to add messages with finer control given.
315///
316/// This includes control over:
317/// - Logits processors
318/// - Constraints
319/// - Logprobs
320/// - Tools
321/// - Sampling
322/// - Enable thinking for models that support the configuration
323pub struct RequestBuilder {
324    messages: Vec<IndexMap<String, MessageContent>>,
325    images: Vec<DynamicImage>,
326    audios: Vec<AudioInput>,
327    logits_processors: Vec<Arc<dyn CustomLogitsProcessor>>,
328    adapters: Vec<String>,
329    return_logprobs: bool,
330    constraint: Constraint,
331    tools: Vec<Tool>,
332    tool_choice: ToolChoice,
333    sampling_params: SamplingParams,
334    web_search_options: Option<WebSearchOptions>,
335    enable_thinking: Option<bool>,
336    truncate_sequence: bool,
337}
338
339impl Default for RequestBuilder {
340    fn default() -> Self {
341        Self::new()
342    }
343}
344
345impl From<TextMessages> for RequestBuilder {
346    fn from(value: TextMessages) -> Self {
347        Self {
348            messages: value.messages,
349            images: Vec::new(),
350            audios: Vec::new(),
351            logits_processors: Vec::new(),
352            adapters: Vec::new(),
353            return_logprobs: false,
354            constraint: Constraint::None,
355            tools: Vec::new(),
356            tool_choice: ToolChoice::Auto,
357            sampling_params: SamplingParams::deterministic(),
358            web_search_options: None,
359            enable_thinking: None,
360            truncate_sequence: false,
361        }
362    }
363}
364
365impl From<VisionMessages> for RequestBuilder {
366    fn from(value: VisionMessages) -> Self {
367        Self {
368            messages: value.messages,
369            images: value.images,
370            audios: value.audios,
371            logits_processors: Vec::new(),
372            adapters: Vec::new(),
373            return_logprobs: false,
374            constraint: Constraint::None,
375            tools: Vec::new(),
376            tool_choice: ToolChoice::Auto,
377            sampling_params: SamplingParams::deterministic(),
378            web_search_options: None,
379            enable_thinking: None,
380            truncate_sequence: false,
381        }
382    }
383}
384
385impl RequestBuilder {
386    pub fn new() -> Self {
387        Self {
388            messages: Vec::new(),
389            images: Vec::new(),
390            audios: Vec::new(),
391            logits_processors: Vec::new(),
392            adapters: Vec::new(),
393            return_logprobs: false,
394            constraint: Constraint::None,
395            tools: Vec::new(),
396            tool_choice: ToolChoice::Auto,
397            sampling_params: SamplingParams::deterministic(),
398            web_search_options: None,
399            enable_thinking: None,
400            truncate_sequence: false,
401        }
402    }
403
404    pub fn with_web_search_options(mut self, web_search_options: WebSearchOptions) -> Self {
405        self.web_search_options = Some(web_search_options);
406        self
407    }
408
409    /// Add a message to the request.
410    ///
411    /// For messages with tool calls, use [`Self::add_message_with_tool_call`].
412    /// For messages with tool outputs, use [`Self::add_tool_message`].
413    pub fn add_message(mut self, role: TextMessageRole, text: impl ToString) -> Self {
414        self.messages.push(IndexMap::from([
415            ("role".to_string(), Either::Left(role.to_string())),
416            ("content".to_string(), Either::Left(text.to_string())),
417        ]));
418        self
419    }
420
421    /// Add a message with the output of a tool call.
422    pub fn add_tool_message(mut self, tool_content: impl ToString, tool_id: impl ToString) -> Self {
423        self.messages.push(IndexMap::from([
424            (
425                "role".to_string(),
426                Either::Left(TextMessageRole::Tool.to_string()),
427            ),
428            (
429                "content".to_string(),
430                Either::Left(tool_content.to_string()),
431            ),
432            (
433                "tool_call_id".to_string(),
434                Either::Left(tool_id.to_string()),
435            ),
436        ]));
437        self
438    }
439
440    pub fn add_message_with_tool_call(
441        mut self,
442        role: TextMessageRole,
443        text: impl ToString,
444        tool_calls: Vec<ToolCallResponse>,
445    ) -> Self {
446        let tool_messages = tool_calls
447            .iter()
448            .map(|t| {
449                IndexMap::from([
450                    ("id".to_string(), Value::String(t.id.clone())),
451                    ("type".to_string(), Value::String(t.tp.to_string())),
452                    (
453                        "function".to_string(),
454                        json!({
455                            "name": t.function.name,
456                            "arguments": t.function.arguments,
457                        }),
458                    ),
459                ])
460            })
461            .collect();
462        self.messages.push(IndexMap::from([
463            ("role".to_string(), Either::Left(role.to_string())),
464            ("content".to_string(), Either::Left(text.to_string())),
465            ("function".to_string(), Either::Right(tool_messages)),
466        ]));
467        self
468    }
469
470    pub fn add_image_message(
471        self,
472        role: TextMessageRole,
473        text: impl ToString,
474        images: Vec<DynamicImage>,
475        model: &Model,
476    ) -> anyhow::Result<Self> {
477        self.add_multimodal_message(role, text, images, vec![], model)
478    }
479
480    pub fn add_audio_message(
481        self,
482        role: TextMessageRole,
483        text: impl ToString,
484        audios: Vec<AudioInput>,
485        model: &Model,
486    ) -> anyhow::Result<Self> {
487        self.add_multimodal_message(role, text, vec![], audios, model)
488    }
489
490    pub fn add_multimodal_message(
491        mut self,
492        role: TextMessageRole,
493        text: impl ToString,
494        images: Vec<DynamicImage>,
495        audios: Vec<AudioInput>,
496        model: &Model,
497    ) -> anyhow::Result<Self> {
498        let config = model.config().unwrap();
499        let prefixer = match &config.category {
500            ModelCategory::Vision { prefixer } => prefixer,
501            _ => {
502                anyhow::bail!("`add_image_message` expects a vision model.")
503            }
504        };
505
506        // Images
507        let n_added_images = images.len();
508        let prefixed = prefixer.prefix_image(
509            (self.images.len()..self.images.len() + n_added_images).collect(),
510            &text.to_string(),
511        );
512        self.images.extend(images);
513
514        // Audios
515        let n_added_audios = audios.len();
516        let prefixed = prefixer.prefix_audio(
517            (self.audios.len()..self.audios.len() + n_added_audios).collect(),
518            &prefixed,
519        );
520        self.audios.extend(audios);
521
522        if n_added_images > 0 {
523            self.messages.push(IndexMap::from([
524                ("role".to_string(), Either::Left(role.to_string())),
525                (
526                    "content".to_string(),
527                    Either::Right(vec![
528                        IndexMap::from([("type".to_string(), Value::String("image".to_string()))]),
529                        IndexMap::from([
530                            ("type".to_string(), Value::String("text".to_string())),
531                            ("text".to_string(), Value::String(prefixed)),
532                        ]),
533                    ]),
534                ),
535            ]));
536        } else {
537            self.messages.push(IndexMap::from([
538                ("role".to_string(), Either::Left(role.to_string())),
539                ("content".to_string(), Either::Left(prefixed)),
540            ]));
541        }
542        Ok(self)
543    }
544
545    pub fn add_logits_processor(mut self, processor: Arc<dyn CustomLogitsProcessor>) -> Self {
546        self.logits_processors.push(processor);
547        self
548    }
549
550    pub fn set_adapters(mut self, adapters: Vec<String>) -> Self {
551        self.adapters = adapters;
552        self
553    }
554
555    /// The default tool choice is auto.
556    pub fn set_tools(mut self, tools: Vec<Tool>) -> Self {
557        self.tools = tools;
558        self
559    }
560
561    pub fn set_tool_choice(mut self, tool_choice: ToolChoice) -> Self {
562        self.tool_choice = tool_choice;
563        self
564    }
565
566    pub fn return_logprobs(mut self, return_logprobs: bool) -> Self {
567        self.return_logprobs = return_logprobs;
568        self
569    }
570
571    pub fn set_constraint(mut self, constraint: Constraint) -> Self {
572        self.constraint = constraint;
573        self
574    }
575
576    /// Set the sampling parameters as given.
577    pub fn set_sampling(mut self, params: SamplingParams) -> Self {
578        self.sampling_params = params;
579        self
580    }
581
582    /// Set the sampling parameters for deterministic generation.
583    /// This sets up the parameters so that there is:
584    /// - No temperature, topk, topp, minp
585    /// - No penalties, stop tokens, or logit bias
586    /// - No maximum length
587    pub fn set_deterministic_sampler(mut self) -> Self {
588        self.sampling_params = SamplingParams::deterministic();
589        self
590    }
591
592    pub fn set_sampler_temperature(mut self, temperature: f64) -> Self {
593        self.sampling_params.temperature = Some(temperature);
594        self
595    }
596
597    pub fn set_sampler_topk(mut self, topk: usize) -> Self {
598        self.sampling_params.top_k = Some(topk);
599        self
600    }
601
602    pub fn set_sampler_topp(mut self, topp: f64) -> Self {
603        self.sampling_params.top_p = Some(topp);
604        self
605    }
606
607    pub fn set_sampler_minp(mut self, minp: f64) -> Self {
608        self.sampling_params.min_p = Some(minp);
609        self
610    }
611
612    pub fn set_sampler_topn_logprobs(mut self, top_n_logprobs: usize) -> Self {
613        self.sampling_params.top_n_logprobs = top_n_logprobs;
614        self
615    }
616
617    pub fn set_sampler_frequency_penalty(mut self, frequency_penalty: f32) -> Self {
618        self.sampling_params.frequency_penalty = Some(frequency_penalty);
619        self
620    }
621
622    pub fn set_sampler_presence_penalty(mut self, presence_penalty: f32) -> Self {
623        self.sampling_params.presence_penalty = Some(presence_penalty);
624        self
625    }
626
627    pub fn set_sampler_stop_toks(mut self, stop_toks: StopTokens) -> Self {
628        self.sampling_params.stop_toks = Some(stop_toks);
629        self
630    }
631
632    pub fn set_sampler_max_len(mut self, max_len: usize) -> Self {
633        self.sampling_params.max_len = Some(max_len);
634        self
635    }
636
637    pub fn set_sampler_logits_bias(mut self, logits_bias: HashMap<u32, f32>) -> Self {
638        self.sampling_params.logits_bias = Some(logits_bias);
639        self
640    }
641
642    pub fn set_sampler_n_choices(mut self, n_choices: usize) -> Self {
643        self.sampling_params.n_choices = n_choices;
644        self
645    }
646
647    pub fn set_sampler_dry_params(mut self, dry_params: DrySamplingParams) -> Self {
648        self.sampling_params.dry_params = Some(dry_params);
649        self
650    }
651
652    pub fn enable_thinking(mut self, enable_thinking: bool) -> Self {
653        self.enable_thinking = Some(enable_thinking);
654        self
655    }
656
657    /// Truncate prompts that exceed the model's maximum context length.
658    pub fn with_truncate_sequence(mut self, truncate_sequence: bool) -> Self {
659        self.truncate_sequence = truncate_sequence;
660        self
661    }
662}
663
664impl RequestLike for RequestBuilder {
665    fn messages_ref(&self) -> &[IndexMap<String, MessageContent>] {
666        &self.messages
667    }
668
669    fn images_ref(&self) -> &[DynamicImage] {
670        &self.images
671    }
672
673    fn take_messages(&mut self) -> RequestMessage {
674        if self.images.is_empty() && self.audios.is_empty() {
675            let mut other = Vec::new();
676            std::mem::swap(&mut other, &mut self.messages);
677            RequestMessage::Chat {
678                messages: other,
679                enable_thinking: self.enable_thinking,
680            }
681        } else {
682            let mut other_messages = Vec::new();
683            std::mem::swap(&mut other_messages, &mut self.messages);
684            let mut other_images = Vec::new();
685            std::mem::swap(&mut other_images, &mut self.images);
686            let mut other_audios = Vec::new();
687            std::mem::swap(&mut other_audios, &mut self.audios);
688            RequestMessage::VisionChat {
689                images: other_images,
690                messages: other_messages,
691                audios: other_audios,
692                enable_thinking: self.enable_thinking,
693            }
694        }
695    }
696
697    fn enable_search(&self) -> Option<bool> {
698        self.web_search_options.as_ref().map(|_| true)
699    }
700
701    fn take_logits_processors(&mut self) -> Option<Vec<Arc<dyn CustomLogitsProcessor>>> {
702        if self.logits_processors.is_empty() {
703            None
704        } else {
705            let mut other = Vec::new();
706            std::mem::swap(&mut other, &mut self.logits_processors);
707            Some(other)
708        }
709    }
710
711    fn take_adapters(&mut self) -> Option<Vec<String>> {
712        if self.adapters.is_empty() {
713            None
714        } else {
715            let mut other = Vec::new();
716            std::mem::swap(&mut other, &mut self.adapters);
717            Some(other)
718        }
719    }
720
721    fn return_logprobs(&self) -> bool {
722        self.return_logprobs
723    }
724
725    fn take_constraint(&mut self) -> Constraint {
726        let mut other = Constraint::None;
727        std::mem::swap(&mut other, &mut self.constraint);
728        other
729    }
730
731    fn take_tools(&mut self) -> Option<(Vec<Tool>, ToolChoice)> {
732        if self.tools.is_empty() {
733            None
734        } else {
735            let mut other_ts = Vec::new();
736            std::mem::swap(&mut other_ts, &mut self.tools);
737            let mut other_tc = ToolChoice::Auto;
738            std::mem::swap(&mut other_tc, &mut self.tool_choice);
739            Some((other_ts, other_tc))
740        }
741    }
742
743    fn take_sampling_params(&mut self) -> SamplingParams {
744        let mut other = SamplingParams::deterministic();
745        std::mem::swap(&mut other, &mut self.sampling_params);
746        other
747    }
748
749    fn take_web_search_options(&mut self) -> Option<WebSearchOptions> {
750        let mut other = None;
751        std::mem::swap(&mut other, &mut self.web_search_options);
752        other
753    }
754
755    fn truncate_sequence(&self) -> bool {
756        self.truncate_sequence
757    }
758}
759
760#[derive(Clone, Debug)]
761/// An individual embedding input.
762pub enum EmbeddingRequestInput {
763    /// Raw text prompt that will be tokenized.
764    Prompt(String),
765    /// Pre-tokenized input.
766    Tokens(Vec<u32>),
767}
768
769impl EmbeddingRequestInput {
770    pub fn into_request_message(self) -> RequestMessage {
771        match self {
772            Self::Prompt(prompt) => RequestMessage::Embedding { prompt },
773            Self::Tokens(prompt) => RequestMessage::EmbeddingTokens { prompt },
774        }
775    }
776}
777
778#[derive(Clone, Debug)]
779/// A validated embedding request constructed via [`EmbeddingRequestBuilder`].
780pub struct EmbeddingRequest {
781    pub inputs: Vec<EmbeddingRequestInput>,
782    pub truncate_sequence: bool,
783}
784
785impl EmbeddingRequest {
786    /// Create a new builder for an embedding request.
787    pub fn builder() -> EmbeddingRequestBuilder {
788        EmbeddingRequestBuilder::new()
789    }
790}
791
792/// Builder for configuring embedding requests.
793#[derive(Clone, Debug, Default)]
794pub struct EmbeddingRequestBuilder {
795    inputs: Vec<EmbeddingRequestInput>,
796    truncate_sequence: bool,
797}
798
799impl EmbeddingRequestBuilder {
800    /// Create an empty builder. You must add at least one input before using it.
801    pub fn new() -> Self {
802        Self::default()
803    }
804
805    /// Add a single text prompt.
806    pub fn add_prompt(mut self, prompt: impl Into<String>) -> Self {
807        self.inputs
808            .push(EmbeddingRequestInput::Prompt(prompt.into()));
809        self
810    }
811
812    /// Add multiple text prompts at once.
813    pub fn add_prompts<I, S>(mut self, prompts: I) -> Self
814    where
815        I: IntoIterator<Item = S>,
816        S: Into<String>,
817    {
818        self.inputs.extend(
819            prompts
820                .into_iter()
821                .map(|prompt| EmbeddingRequestInput::Prompt(prompt.into())),
822        );
823        self
824    }
825
826    /// Add a single pre-tokenized prompt.
827    pub fn add_tokens(mut self, tokens: impl Into<Vec<u32>>) -> Self {
828        self.inputs
829            .push(EmbeddingRequestInput::Tokens(tokens.into()));
830        self
831    }
832
833    /// Add multiple pre-tokenized prompts.
834    pub fn add_tokens_batch<I>(mut self, batches: I) -> Self
835    where
836        I: IntoIterator<Item = Vec<u32>>,
837    {
838        self.inputs
839            .extend(batches.into_iter().map(EmbeddingRequestInput::Tokens));
840        self
841    }
842
843    /// Control whether prompts longer than the model context are truncated.
844    pub fn with_truncate_sequence(mut self, truncate: bool) -> Self {
845        self.truncate_sequence = truncate;
846        self
847    }
848
849    pub fn build(self) -> anyhow::Result<EmbeddingRequest> {
850        if self.inputs.is_empty() {
851            anyhow::bail!("Embedding request must contain at least one input.");
852        }
853
854        Ok(EmbeddingRequest {
855            inputs: self.inputs,
856            truncate_sequence: self.truncate_sequence,
857        })
858    }
859}
mistralrs/messages.rs

mistralrs/
messages.rs
