Skip to main content

rlx_runtime/
mock_requests.rs

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