rlx_runtime/
mock_requests.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Mock request payloads for tests (plan #64).
6//!
7//! Borrowed from MAX's `serve/mocks/mock_api_requests.py`. One
8//! source of truth for "what does an OpenAI / OpenAI-style request
9//! actually look like" so tests don't redefine payloads each time.
10//!
11//! These are *test fixtures*, not a serving impl — the structs are
12//! intentionally small and serde-friendly. When RLX grows a serving
13//! crate it should consume these as a starting point.
14
15use serde::{Deserialize, Serialize};
16
17/// A single message in a chat completion request.
18#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
19pub struct ChatMessage {
20    pub role: String, // "system" | "user" | "assistant"
21    pub content: String,
22}
23
24/// OpenAI-shaped chat completion request.
25#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
26pub struct ChatCompletionRequest {
27    pub model: String,
28    pub messages: Vec<ChatMessage>,
29    pub max_tokens: Option<u32>,
30    pub temperature: Option<f32>,
31    pub stream: Option<bool>,
32}
33
34/// OpenAI-shaped embedding request.
35#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
36pub struct EmbeddingRequest {
37    pub model: String,
38    /// Either a single string or a list — accept both via the
39    /// untagged `Input` enum below.
40    pub input: Input,
41    pub encoding_format: Option<String>,
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
45#[serde(untagged)]
46pub enum Input {
47    Single(String),
48    Batch(Vec<String>),
49}
50
51// ── Canned fixtures ───────────────────────────────────────────────
52
53pub fn chat_simple() -> ChatCompletionRequest {
54    ChatCompletionRequest {
55        model: "gpt-4o-mini".to_string(),
56        messages: vec![ChatMessage {
57            role: "user".into(),
58            content: "What is the capital of France?".into(),
59        }],
60        max_tokens: Some(64),
61        temperature: Some(0.0),
62        stream: Some(false),
63    }
64}
65
66pub fn chat_system_user() -> ChatCompletionRequest {
67    ChatCompletionRequest {
68        model: "gpt-4o-mini".to_string(),
69        messages: vec![
70            ChatMessage {
71                role: "system".into(),
72                content: "You are a terse oracle.".into(),
73            },
74            ChatMessage {
75                role: "user".into(),
76                content: "Color of the sky?".into(),
77            },
78        ],
79        max_tokens: Some(8),
80        temperature: Some(0.7),
81        stream: Some(false),
82    }
83}
84
85pub fn chat_streaming() -> ChatCompletionRequest {
86    let mut r = chat_simple();
87    r.stream = Some(true);
88    r
89}
90
91pub fn embed_single() -> EmbeddingRequest {
92    EmbeddingRequest {
93        model: "text-embedding-3-small".to_string(),
94        input: Input::Single("Hello, World!".to_string()),
95        encoding_format: Some("float".to_string()),
96    }
97}
98
99pub fn embed_batch() -> EmbeddingRequest {
100    EmbeddingRequest {
101        model: "text-embedding-3-small".to_string(),
102        input: Input::Batch(vec![
103            "Hello, World!".into(),
104            "fastembed-rs is licensed under Apache-2.0".into(),
105            "Some other short text here".into(),
106        ]),
107        encoding_format: Some("float".to_string()),
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn fixtures_round_trip_through_json() {
117        // Serialize + deserialize each canned fixture; the round
118        // trip catches any field rename / type drift early.
119        let cases: Vec<serde_json::Value> = vec![
120            serde_json::to_value(chat_simple()).unwrap(),
121            serde_json::to_value(chat_system_user()).unwrap(),
122            serde_json::to_value(chat_streaming()).unwrap(),
123            serde_json::to_value(embed_single()).unwrap(),
124            serde_json::to_value(embed_batch()).unwrap(),
125        ];
126        for v in cases {
127            let s = serde_json::to_string(&v).unwrap();
128            let _: serde_json::Value = serde_json::from_str(&s).unwrap();
129        }
130    }
131
132    #[test]
133    fn embed_input_accepts_both_string_and_array() {
134        let single =
135            serde_json::from_str::<EmbeddingRequest>(r#"{"model":"x","input":"hi"}"#).unwrap();
136        let batch =
137            serde_json::from_str::<EmbeddingRequest>(r#"{"model":"x","input":["a","b"]}"#).unwrap();
138        assert!(matches!(single.input, Input::Single(_)));
139        assert!(matches!(batch.input, Input::Batch(_)));
140    }
141}
rlx_runtime/mock_requests.rs

rlx_runtime/
mock_requests.rs
