common.rs 29.4 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Engine Protocols
//! ================
//!
//! This module contains the protocols in public API for the LLM Engine and AsyncEngine facades.
//!
//! The core components are the `CompletionRequest` and `StreamingCompletionResponse` objects.
//!
//! The `StreamingCompletionResponse` objects are the outputs of the LLM Engine; however, we
//! need some additional information to propagate intermediate results for improved observability.
//! The metadata is transferred via the other arms of the `StreamingResponse` enum.
//!

use anyhow::Result;
use derive_builder::Builder;
18
use serde::{Deserialize, Serialize};
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

use super::TokenIdType;

pub mod llm_backend;
pub mod postprocessor;
pub mod preprocessor;

/// SamplingOptionsProvider is a trait that allows the caller to extract the sampling options from
/// the object that implements it. This will mutate the object.
pub trait SamplingOptionsProvider {
    fn extract_sampling_options(&self) -> Result<SamplingOptions>;
}

pub trait StopConditionsProvider {
    fn extract_stop_conditions(&self) -> Result<StopConditions>;
}

Greg Clark's avatar
Greg Clark committed
36
37
38
39
pub trait OutputOptionsProvider {
    fn extract_output_options(&self) -> Result<OutputOptions>;
}

40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
pub enum FinishReason {
    #[serde(rename = "eos")]
    EoS,

    #[serde(rename = "length")]
    Length,

    #[serde(rename = "stop")]
    Stop,

    #[serde(rename = "error")]
    Error(String),

    #[serde(rename = "cancelled")]
    Cancelled,
56
57
58

    #[serde(rename = "content_filter")]
    ContentFilter,
59
60
}

Paul Hendricks's avatar
Paul Hendricks committed
61
62
63
64
65
66
67
68
impl std::fmt::Display for FinishReason {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            FinishReason::EoS => write!(f, "eos"),
            FinishReason::Length => write!(f, "length"),
            FinishReason::Stop => write!(f, "stop"),
            FinishReason::Error(msg) => write!(f, "error: {}", msg),
            FinishReason::Cancelled => write!(f, "cancelled"),
69
            FinishReason::ContentFilter => write!(f, "content_filter"),
Paul Hendricks's avatar
Paul Hendricks committed
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
        }
    }
}

impl std::str::FromStr for FinishReason {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "eos" => Ok(FinishReason::EoS),
            "length" => Ok(FinishReason::Length),
            "stop" => Ok(FinishReason::Stop),
            "cancelled" => Ok(FinishReason::Cancelled),
            s if s.starts_with("error: ") => Ok(FinishReason::Error(s[7..].to_string())),
            _ => Err(anyhow::anyhow!("Invalid FinishReason variant: '{}'", s)),
        }
    }
}

89
impl From<FinishReason> for dynamo_async_openai::types::CompletionFinishReason {
90
91
92
    fn from(reason: FinishReason) -> Self {
        match reason {
            FinishReason::EoS | FinishReason::Stop | FinishReason::Cancelled => {
93
                dynamo_async_openai::types::CompletionFinishReason::Stop
94
95
            }
            FinishReason::ContentFilter => {
96
                dynamo_async_openai::types::CompletionFinishReason::ContentFilter
97
            }
98
99
            FinishReason::Length => dynamo_async_openai::types::CompletionFinishReason::Length,
            FinishReason::Error(_) => dynamo_async_openai::types::CompletionFinishReason::Stop,
100
101
102
103
        }
    }
}

104
105
impl From<dynamo_async_openai::types::CompletionFinishReason> for FinishReason {
    fn from(reason: dynamo_async_openai::types::CompletionFinishReason) -> Self {
106
        match reason {
107
108
109
            dynamo_async_openai::types::CompletionFinishReason::Stop => FinishReason::Stop,
            dynamo_async_openai::types::CompletionFinishReason::Length => FinishReason::Length,
            dynamo_async_openai::types::CompletionFinishReason::ContentFilter => {
110
111
112
113
114
115
                FinishReason::ContentFilter
            }
        }
    }
}

116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
/// LLM Inference Engines can accept a variety of input types. Not all Engines will support all
/// input types. For example, the trtllm::AsyncEngine only supports `PromptType::Tokens` as an
/// input type. The higher-level `Backend` class is a general wrapper around Engines that will
/// enable many of the input options that require pre/postprocessing.
#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq)]
pub enum PromptType {
    /// If allowed, this input type allowed the caller to pass a list of token_ids directly to the
    /// inference engine. This is an advanced feature that requires the caller to handle all of the
    /// necessary prompt formatting and tokenization.
    #[serde(rename = "token_ids")]
    TokenIds(Vec<TokenIdType>),

    /// If allowed, the raw text will be tokenized and converted to token_ids without any additional
    /// preprocessing. This is an advanced features that requires the caller to correctly format the
    /// prompt as defined by the model.
    #[serde(rename = "raw")]
    Raw(String),

    /// If allowed, the `CompletionContext` will be preprocessed server-side. If the `Model` trait
    /// `requires_prompt_template` returns true then the `CompletionContext` will be used to
    /// to render the formatted prompt from the template. `Completion` is the preferred `PromptType`
    /// for single turn completions.
    #[serde(rename = "completion")]
    Completion(CompletionContext),

    /// If allowed, the `ChatContext` will be preprocessed server-side. Most chat models will have
    /// a predefined prompt format/structure. If the `Model` trait `requires_prompt_template` returns
    /// true then the `ChatContext` will be used to to render the formatted prompt from the template.
    /// `ChatCompletion` is the preferred `PromptType` for multi-turn completions.
    #[serde(rename = "chat_completion")]
    ChatCompletion(ChatContext),

    /// If allowed, then `Model::requires_prompt_template()` must also return true. The `serde_json::Value`
    /// will be passed directly the prompt template. This allows for a complete generic data model and
    /// prompt template to be passed to be defined and used by the server.
    #[serde(rename = "custom_json")]
    CustomJson(serde_json::Value),
}

/// TensorRT LLM does not perform preprocessing or postprocessing. The input_ids / token_ids
/// are expected to be preprocessed by the client. The client is responsible for constructing
/// the model specific prompt template and applying the tokenizer.
///
/// TensorRT LLM will perform some server side postprocessing to ensure that generation is
/// efficiently stopped. See `StopConditions` below.
#[derive(Serialize, Deserialize, Debug, Clone, Builder)]
pub struct CompletionRequest {
    /// Type of prompt
    pub prompt: PromptType,

    /// StopConditions are conditions that the inference engine will use to stop generation.
    pub stop_conditions: StopConditions,

    /// SamplingOptions directs the inference engine to use sampling instead of greedy decoding.
    /// More documentation on how and on the order in which sampling options are applied
    /// are needed.
    pub sampling_options: SamplingOptions,

Greg Clark's avatar
Greg Clark committed
174
175
176
    #[builder(default)]
    pub output_options: OutputOptions,

177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
    /// The computed checksum of the Model Deployment Card (MDC).
    #[builder(default)]
    pub mdc_sum: Option<String>,

    /// User requested annotations for the request
    #[builder(default)]
    pub annotations: Option<Vec<String>>,
}

impl CompletionRequest {
    pub fn builder() -> CompletionRequestBuilder {
        CompletionRequestBuilder::default()
    }
}

#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq)]
/// Defines the prompt template and system prompt for a completion request.
/// If the model does not support prompt templates, the system_prompt will be ignored.
pub struct CompletionContext {
    /// Prompt sent by the user
    pub prompt: String,

    /// Optional system_prompt for models that support prompt templates with system_prompts.
    pub system_prompt: Option<String>,
}

/// ChatTurn is a struct that contains the user and assistant messages in a chat.
#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq)]
pub struct ChatTurn {
    /// The user message
    pub user: String,

    /// The assistant response
    pub assistant: String,
}

/// ChatContext is a struct that contains the role and context of a chat message
/// along with a flattened CompletionContext.
#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq)]
pub struct ChatContext {
    /// CompletionContext for this chat turn
    #[serde(flatten)]
    pub completion: CompletionContext,

    /// The history/context of the user and assistant messages in the chat context
    pub context: Vec<ChatTurn>,
}

/// TensorRT LLM server-side stop conditions. These options allow for the server to evaluate
/// the generated sequence and stop generation if the sequence meets a stop condition.
#[derive(Serialize, Deserialize, Debug, Clone, Default)]
pub struct StopConditions {
    /// The maximum number of tokens to generate
    pub max_tokens: Option<u32>,

    /// List of strings that stop the generation when they are generated.
    /// The returned output will not contain the stop strings.
    pub stop: Option<Vec<String>>,

    /// List of tokens that stop the generation when they are
    /// generated. The returned output will NOT contain the stop tokens.
    pub stop_token_ids_hidden: Option<Vec<TokenIdType>>,

    /// The minimum number of tokens to generate
    /// To ignore_eos, set min_tokens to max_tokens
    pub min_tokens: Option<u32>,

    /// Whether to ignore the EOS token and continue generating
    /// tokens after the EOS token is generated.
    // TODO(ignore_eos) - improve this my masking the EOS token with logit bias
    pub ignore_eos: Option<bool>,
248
249
250
251

    /// Maximum number of thinking tokens allowed
    /// NOTE: Currently a passthrough - no enforcement logic implemented
    pub max_thinking_tokens: Option<u32>,
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
}

impl StopConditions {
    pub fn apply_ignore_eos(&mut self) {
        if self.ignore_eos.unwrap_or(false) {
            self.min_tokens = self.max_tokens;
            self.stop = None;
            self.stop_token_ids_hidden = None;
        }
    }
}

/// Temperature range for sampling.
pub const TEMPERATURE_RANGE: (f32, f32) = (0.0, 1.0);

/// Top P range for sampling.
pub const TOP_P_RANGE: (f32, f32) = (0.0, 1.0);

/// Frequency Penalty range for sampling.
pub const FREQUENCY_PENALTY_RANGE: (f32, f32) = (-1.0, 1.0);

/// Collection of options that control the sampling behavior of the inference engine.
#[derive(Serialize, Deserialize, Debug, Clone, Default)]
pub struct SamplingOptions {
    /// Number of output sequences to return for the given prompt
277
    pub n: Option<u8>,
278
279
280
281
282
283

    /// Number of output sequences that are generated from the prompt.
    /// From these `best_of` sequences, the top `n` sequences are returned.
    /// `best_of` must be greater than or equal to `n`. This is treated as
    /// the beam width when `use_beam_search` is True. By default, `best_of`
    /// is set to `n`.
284
    pub best_of: Option<u8>,
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330

    /// Float that penalizes new tokens based on whether they
    /// appear in the generated text so far. Values > 0 encourage the model
    /// to use new tokens, while values < 0 encourage the model to repeat
    /// tokens.
    pub presence_penalty: Option<f32>,

    /// Float that penalizes new tokens based on their
    /// frequency in the generated text so far. Values > 0 encourage the
    /// model to use new tokens, while values < 0 encourage the model to
    /// repeat tokens.
    pub frequency_penalty: Option<f32>,

    /// Float that penalizes new tokens based on whether
    /// they appear in the prompt and the generated text so far. Values > 1
    /// encourage the model to use new tokens, while values < 1 encourage
    /// the model to repeat tokens.
    pub repetition_penalty: Option<f32>,

    /// Float that controls the randomness of the sampling. Lower
    /// values make the model more deterministic, while higher values make
    /// the model more random. Zero means greedy sampling.
    pub temperature: Option<f32>,

    /// Float that controls the cumulative probability of the top tokens
    /// to consider. Must be in (0, 1]. Set to 1 to consider all tokens.
    pub top_p: Option<f32>,

    /// Integer that controls the number of top tokens to consider. Set
    /// to -1 to consider all tokens.
    pub top_k: Option<i32>,

    /// Float that represents the minimum probability for a token to be
    /// considered, relative to the probability of the most likely token.
    /// Must be in [0, 1]. Set to 0 to disable this.
    pub min_p: Option<f32>,

    /// Whether to use beam search instead of sampling.
    pub use_beam_search: Option<bool>,

    /// Float that penalizes sequences based on their length.
    /// Used in beam search.
    pub length_penalty: Option<f32>,

    /// The seed to use when sampling
    pub seed: Option<i64>,
331

332
333
334
    /// Whether to include the stop string in the output.
    pub include_stop_str_in_output: Option<bool>,

335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
    /// Guided Decoding Options
    pub guided_decoding: Option<GuidedDecodingOptions>,
}

/// Guided Decoding Options
///
/// Only one of `json`, `regex`, `choice`, or `grammar` should be set.
#[derive(Serialize, Deserialize, Debug, Clone, Default)]
pub struct GuidedDecodingOptions {
    /// If specified, the output will follow the JSON schema. Can be a string, an object, or null.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub json: Option<serde_json::Value>,

    /// If specified, the output will follow the regex pattern. Can be a string or null.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub regex: Option<String>,

    /// If specified, the output will be exactly one of the choices.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub choice: Option<Vec<String>>,

    /// If specified, the output will follow the context-free grammar. Can be a string or null.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub grammar: Option<String>,

    /// If specified, the backend to use for guided decoding, can be backends like xgrammar or custom guided decoding backend
    #[serde(skip_serializing_if = "Option::is_none")]
    pub backend: Option<String>,
363
364
365
366

    /// If specified, whitespace pattern to use for guided decoding. Can be a string or null.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub whitespace_pattern: Option<String>,
367
368
369
370
371
372
373
374
375
376
}

impl GuidedDecodingOptions {
    /// Construct without validation
    pub fn new(
        json: Option<serde_json::Value>,
        regex: Option<String>,
        choice: Option<Vec<String>>,
        grammar: Option<String>,
        backend: Option<String>,
377
        whitespace_pattern: Option<String>,
378
379
380
381
382
383
384
    ) -> Self {
        Self {
            json,
            regex,
            choice,
            grammar,
            backend,
385
            whitespace_pattern,
386
387
388
389
390
391
392
393
394
395
        }
    }

    /// Construct and validate (fallible)
    pub fn validated(
        json: Option<serde_json::Value>,
        regex: Option<String>,
        choice: Option<Vec<String>>,
        grammar: Option<String>,
        backend: Option<String>,
396
        whitespace_pattern: Option<String>,
397
    ) -> Result<Self> {
398
        let instance = Self::new(json, regex, choice, grammar, backend, whitespace_pattern);
399
400
401
402
403
404
405
406
407
408
409
        instance.validate()?;
        Ok(instance)
    }

    /// Construct only if one field is Some (fallible)
    pub fn from_optional(
        json: Option<serde_json::Value>,
        regex: Option<String>,
        choice: Option<Vec<String>>,
        grammar: Option<String>,
        backend: Option<String>,
410
        whitespace_pattern: Option<String>,
411
412
    ) -> Result<Option<Self>> {
        let is_empty_choice = choice.as_ref().is_none_or(|v| v.is_empty());
413
414
415
416
417
418
        if json.is_none()
            && regex.is_none()
            && is_empty_choice
            && grammar.is_none()
            && whitespace_pattern.is_none()
        {
419
420
            return Ok(None);
        }
421
        let instance = Self::validated(json, regex, choice, grammar, backend, whitespace_pattern)?;
422
423
424
425
426
427
428
429
430
431
        Ok(Some(instance))
    }

    /// Validate that only one guided decoding option is set
    pub fn validate(&self) -> Result<()> {
        let count = [
            self.json.is_some(),
            self.regex.is_some(),
            self.choice.as_ref().is_some_and(|v| !v.is_empty()),
            self.grammar.is_some(),
432
            self.whitespace_pattern.is_some(),
433
434
435
436
437
438
439
440
441
442
443
444
445
446
        ]
        .iter()
        .filter(|&&v| v)
        .count();

        if count > 1 {
            Err(anyhow::anyhow!(
                "Only one of json, regex, choice, or grammar can be set, but multiple are specified: {:?}",
                self
            ))
        } else {
            Ok(())
        }
    }
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
}

impl SamplingOptions {
    pub fn force_greedy(&mut self) {
        self.presence_penalty = None;
        self.frequency_penalty = None;
        self.repetition_penalty = None;
        self.temperature = None;
        self.top_p = None;
        self.top_k = None;
        self.min_p = None;
    }
}

/// Collection of options that control what information the inference engine returns in the response.
#[derive(Serialize, Deserialize, Debug, Clone, Default)]
pub struct OutputOptions {
    /// Number of log probabilities to return per output token.
    /// Note that the implementation follows the OpenAI API: The return
    /// result includes the log probabilities on the `logprobs` most likely
    /// tokens, as well the chosen tokens. The API will always return the
    /// log probability of the sampled token, so there  may be up to
    /// `logprobs+1` elements in the response
    pub logprobs: Option<u32>,

    /// Number of log probabilities to return per prompt token.
    pub prompt_logprobs: Option<u32>,

    /// Whether to skip special tokens in the output.
    pub skip_special_tokens: Option<bool>,

    /// If true, the Context object will contain the prompt that was pass to
    /// the tokenizer. This is useful for inspecting the behavior of prompt
    /// templates that are applied during the backend preprocessing.
    pub formatted_prompt: Option<bool>,
}

// Struct for log probability information
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ChatCompletionLogprobs {
    /// A list of message content tokens with log probability information.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<Vec<ChatCompletionTokenLogprob>>,

    /// A list of message refusal tokens with log probability information.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub refusal: Option<Vec<ChatCompletionTokenLogprob>>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ChatCompletionTokenLogprob {
    /// The token.
    pub token: String,

    /// The log probability of this token, if it is within the top 20 most likely tokens.
    /// Otherwise, the value `-9999.0` signifies that the token is very unlikely.
    pub logprob: f64,

    /// A list of integers representing the UTF-8 bytes representation of the token.
    /// Useful in instances where characters are represented by multiple tokens and their
    /// byte representations must be combined to generate the correct text representation.
    /// Can be `None` if there is no bytes representation for the token.
    pub bytes: Option<Vec<u8>>,

    /// List of the most likely tokens and their log probability, at this token position.
    /// In rare cases, there may be fewer than the requested number of `top_logprobs` returned.
    pub top_logprobs: Vec<TopLogprob>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct TopLogprob {
    /// The token.
    pub token: String,

    /// The log probability of this token.
    pub logprob: f64,

    /// A list of integers representing the UTF-8 bytes representation of the token.
    /// Can be `None` if there is no bytes representation for the token.
    pub bytes: Option<Vec<u8>>,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum StreamState {
    Active,
    Finished(FinishReason),
}

#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(rename_all = "snake_case")]
pub enum Logits {
    All(Vec<f32>),
    Sparse(Vec<(u32, f32)>),
}

#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(rename_all = "snake_case")]
pub enum LogProbs {
    Normalized(Logits),
    Raw(Logits),
}

/// At each SequencePosition we hold position specific data
pub struct SequencePositionData {
    pub token_id: TokenIdType,

    /// The log probability of the token
    pub logprobs: Option<LogProbs>,
}

557
558
559
560
561
562
#[derive(Debug)]
pub struct StreamingCompletionResponse {
    pub delta: Delta,
    pub logprobs: Option<ChatCompletionLogprobs>,
}

563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
// todo(ryan) - we need to create a DeltaBuilder which is a mutable object that can be passed
// around from the low-level compute engine to the high-level api. The DeltaBuilder will allow
// us to construct the Delta object at multiple layers in the streaming response path.
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct Delta {
    pub is_complete: bool,

    pub finish_reason: Option<FinishReason>,

    // new token_ids
    pub token_ids: Option<Vec<u32>>,

    // tokens
    pub tokens: Option<Vec<String>>,

    // decoded text
    pub text: Option<String>,

    // current sequence length
    // when stream, we expect this to increase by 1 on each response
    pub sequence_length: Option<usize>,

    // if the number of slots for a given request is greater than 1
    // this indicates the index of the slot for the response
    pub index: Option<usize>,

    /// cumulative log probabilities
    pub cum_log_probs: Option<f64>,

    /// error message from engine
    /// if this is set, is_complete should also be true
    pub err_msg: Option<String>,

    /// usage info
    pub usage: Option<Usage>,
}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct Usage {
    pub input_tokens_count: usize,
    pub output_tokens_count: usize,
}

impl CompletionContext {
    /// Create a new CompletionContext
    pub fn new(prompt: String, system_prompt: Option<String>) -> Self {
        Self {
            prompt,
            system_prompt,
        }
    }

    /// Create a new CompletionContext with only a prompt
    pub fn from_prompt(prompt: String) -> Self {
        Self {
            prompt,
            system_prompt: None,
        }
    }

    /// Create a new CompletionContext with a prompt and system prompt
    pub fn with_system_prompt(prompt: String, system_prompt: String) -> Self {
        Self {
            prompt,
            system_prompt: Some(system_prompt),
        }
    }
}

// todo(ryan) - create a builder for chat context
impl From<CompletionContext> for PromptType {
    fn from(context: CompletionContext) -> Self {
        PromptType::Completion(context)
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_completion_context_new() {
        let prompt = "Hello, world!".to_string();
        let system_prompt = Some("This is a system prompt.".to_string());
        let context = CompletionContext::new(prompt.clone(), system_prompt.clone());

        assert_eq!(context.prompt, prompt);
        assert_eq!(context.system_prompt, system_prompt);
    }

    #[test]
    fn test_completion_context_from_prompt() {
        let prompt = "Hello, world!".to_string();
        let context = CompletionContext::from_prompt(prompt.clone());

        assert_eq!(context.prompt, prompt);
        assert_eq!(context.system_prompt, None);
    }

    #[test]
    fn test_completion_context_with_system_prompt() {
        let prompt = "Hello, world!".to_string();
        let system_prompt = "This is a system prompt.".to_string();
        let context = CompletionContext::with_system_prompt(prompt.clone(), system_prompt.clone());

        assert_eq!(context.prompt, prompt);
        assert_eq!(context.system_prompt, Some(system_prompt));
    }

    #[test]
    fn test_completion_context_into_prompt_type() {
        let prompt = "Hello, world!".to_string();
        let system_prompt = "This is a system prompt.".to_string();
        let context = CompletionContext::with_system_prompt(prompt.clone(), system_prompt.clone());
        let prompt_type: PromptType = context.into();

        if let PromptType::Completion(completion_context) = prompt_type {
            assert_eq!(completion_context.prompt, prompt);
            assert_eq!(completion_context.system_prompt, Some(system_prompt));
        } else {
            panic!("Expected a Completion variant");
        }
    }
687
688
689
690
691
692
693
694
695
696
697
698

    #[test]
    fn test_guided_decoding_options_new_and_exclusive() {
        // Only JSON set
        let json_val = serde_json::json!({"type": "object"});
        let backend = Some("xgrammar".to_string());
        let opts = GuidedDecodingOptions::validated(
            Some(json_val.clone()),
            None,
            None,
            None,
            backend.clone(),
699
            None,
700
701
702
703
704
705
706
707
        );
        assert!(opts.is_ok());
        let opts = opts.unwrap();
        assert_eq!(opts.json, Some(json_val));
        assert!(opts.regex.is_none());
        assert!(opts.choice.is_none());
        assert!(opts.grammar.is_none());
        assert_eq!(opts.backend, backend);
708
        assert!(opts.whitespace_pattern.is_none());
709
710
711

        // Only regex set
        let regex = Some(r"\d+".to_string());
712
        let opts = GuidedDecodingOptions::validated(None, regex.clone(), None, None, None, None);
713
714
715
716
717
718
        assert!(opts.is_ok());
        let opts = opts.unwrap();
        assert_eq!(opts.regex, regex);
        assert!(opts.json.is_none());
        assert!(opts.choice.is_none());
        assert!(opts.grammar.is_none());
719
        assert!(opts.whitespace_pattern.is_none());
720
721
722

        // Only choice set
        let choice = Some(vec!["A".to_string(), "B".to_string()]);
723
        let opts = GuidedDecodingOptions::validated(None, None, choice.clone(), None, None, None);
724
725
726
727
728
729
        assert!(opts.is_ok());
        let opts = opts.unwrap();
        assert_eq!(opts.choice, choice);
        assert!(opts.json.is_none());
        assert!(opts.regex.is_none());
        assert!(opts.grammar.is_none());
730
        assert!(opts.whitespace_pattern.is_none());
731
732
733

        // Only grammar set
        let grammar = Some("root ::= 'yes' | 'no'".to_string());
734
        let opts = GuidedDecodingOptions::validated(None, None, None, grammar.clone(), None, None);
735
736
737
738
739
740
        assert!(opts.is_ok());
        let opts = opts.unwrap();
        assert_eq!(opts.grammar, grammar);
        assert!(opts.json.is_none());
        assert!(opts.regex.is_none());
        assert!(opts.choice.is_none());
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
        assert!(opts.whitespace_pattern.is_none());

        // Only whitespace_pattern set
        let whitespace_pattern = Some(r"\s+".to_string());
        let opts = GuidedDecodingOptions::validated(
            None,
            None,
            None,
            None,
            None,
            whitespace_pattern.clone(),
        );
        assert!(opts.is_ok());
        let opts = opts.unwrap();
        assert_eq!(opts.whitespace_pattern, whitespace_pattern);
        assert!(opts.json.is_none());
        assert!(opts.regex.is_none());
        assert!(opts.choice.is_none());
        assert!(opts.grammar.is_none());
760
761
762
763
764
765
766
767

        // Multiple fields set (should error)
        let opts = GuidedDecodingOptions::validated(
            Some(serde_json::json!({})),
            Some(r"\d+".to_string()),
            None,
            None,
            None,
768
            None,
769
770
771
772
773
774
775
776
777
        );
        assert!(opts.is_err());

        let opts = GuidedDecodingOptions::validated(
            None,
            Some(r"\d+".to_string()),
            Some(vec!["A".to_string()]),
            None,
            None,
778
            None,
779
780
781
782
783
784
785
786
787
        );
        assert!(opts.is_err());

        let opts = GuidedDecodingOptions::validated(
            Some(serde_json::json!({})),
            None,
            Some(vec!["A".to_string()]),
            Some("root ::= 'yes'".to_string()),
            None,
788
            None,
789
790
791
792
        );
        assert!(opts.is_err());

        // All fields None (should be ok, but not useful)
793
        let opts = GuidedDecodingOptions::validated(None, None, None, None, None, None);
794
795
796
797
798
799
        assert!(opts.is_ok());
    }

    #[test]
    fn test_guided_decoding_options_from_optional() {
        // All None returns Ok(None)
800
        let opts = GuidedDecodingOptions::from_optional(None, None, None, None, None, None);
801
802
803
804
805
        assert!(opts.is_ok());
        assert!(opts.unwrap().is_none());

        // Only one set returns Ok(Some)
        let regex = Some(r"\w+".to_string());
806
807
        let opts =
            GuidedDecodingOptions::from_optional(None, regex.clone(), None, None, None, None);
808
809
810
811
812
813
814
815
816
817
818
819
820
        assert!(opts.is_ok());
        let val = opts.unwrap();
        assert!(val.is_some());
        let val = val.unwrap();
        assert_eq!(val.regex, regex);

        // Multiple set returns Err
        let opts = GuidedDecodingOptions::from_optional(
            Some(serde_json::json!({})),
            Some(r"\d+".to_string()),
            None,
            None,
            None,
821
            None,
822
823
824
825
        );
        assert!(opts.is_err());

        // Choice set but empty vector should not count as set
826
        let opts = GuidedDecodingOptions::from_optional(None, None, Some(vec![]), None, None, None);
827
828
829
830
831
832
833
834
835
836
837
        assert!(opts.is_ok());
        let val = opts.unwrap();
        assert!(val.is_none());

        // Choice set with non-empty vector
        let opts = GuidedDecodingOptions::from_optional(
            None,
            None,
            Some(vec!["A".to_string()]),
            None,
            None,
838
            None,
839
840
841
842
843
844
845
        );
        assert!(opts.is_ok());
        let val = opts.unwrap();
        assert!(val.is_some());
        let val = val.unwrap();
        assert_eq!(val.choice, Some(vec!["A".to_string()]));
    }
846
}