rlx_runtime/
mock_requests.rs1use serde::{Deserialize, Serialize};
27
28#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
30pub struct ChatMessage {
31 pub role: String, pub content: String,
33}
34
35#[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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
47pub struct EmbeddingRequest {
48 pub model: String,
49 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
62pub 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 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}