lib.rs

crates/standin/src/lib.rs

1. //! A stand-in for DeepInfra (or any OpenAI-compatible server).
2. //!
3. //! The model id picks the behavior, so one server covers every path:
4. //!
5. //! | model | behavior |
6. //! |---------------------------|-------------------------------------------------------------|
7. //! | `standin/echo` | "You said: ..." word by word; usage in a final empty chunk |
8. //! | `standin/cached` | like echo, with 40 cached prompt tokens and a cache price |
9. //! | `standin/inline-usage` | like echo, usage on the last content chunk |
10. //! | `standin/slow` | 200 words, one every 50 ms (for cancelling) |
11. //! | `standin/no-usage` | like echo, no usage at all (as most local servers) |
12. //! | `standin/reasoning` | reasoning deltas, then the answer |
13. //! | `standin/vision` | "I see N image(s)." |
14. //! | `standin/midstream-error` | three words, then an error event |
15. //! | `standin/402` `/429` `/500` | that HTTP status (429 with Retry-After: 7) |
16. //!
17. //! The key `bad-key` gets a 401 in DeepInfra's shape. Every request is recorded,
18. //! and a stream the client abandons is counted in `disconnects`.
20. use std::convert::Infallible;
21. use std::net::SocketAddr;
22. use std::sync::atomic::{AtomicUsize, Ordering};
23. use std::sync::{Arc, Mutex};
24. use std::time::Duration;
26. use axum::body::Body;
27. use axum::extract::State;
28. use axum::http::{HeaderMap, StatusCode};
29. use axum::response::{IntoResponse, Response};
30. use axum::routing::{get, post};
31. use axum::{Json, Router};
32. use bytes::Bytes;
33. use serde_json::{json, Value};
34. use tokio::sync::mpsc;
35. use tokio_stream::wrappers::ReceiverStream;
37. #[derive(Debug, Clone)]
38. pub struct Recorded {
39. pub authorization: Option<String>,
40. pub body: Value,
41. }
43. #[derive(Default)]
44. pub struct Shared {
45. pub requests: Mutex<Vec<Recorded>>,
46. pub disconnects: AtomicUsize,
47. }
49. pub struct Standin {
50. pub addr: SocketAddr,
51. pub shared: Arc<Shared>,
52. }
54. impl Standin {
55. /// The base URL a client should use, e.g. `http://127.0.0.1:PORT/v1`.
56. pub fn base_url(&self) -> String {
57. format!("http://{}/v1", self.addr)
58. }
59. pub fn requests(&self) -> Vec<Recorded> {
60. self.shared.requests.lock().unwrap().clone()
61. }
62. pub fn disconnects(&self) -> usize {
63. self.shared.disconnects.load(Ordering::SeqCst)
64. }
65. }
67. pub fn models() -> Value {
68. let m = |id: &str, i: f64, o: f64, cache: Option<f64>, tags: &[&str]| {
69. let mut pricing = json!({"input_tokens": i, "output_tokens": o});
70. if let Some(c) = cache {
71. pricing["cache_read_tokens"] = json!(c);
72. }
73. json!({"id": id, "object": "model", "owned_by": "standin",
74. "metadata": {"context_length": 32768, "max_tokens": 4096, "pricing": pricing, "tags": tags}})
75. };
76. json!({"object": "list", "data": [
77. m("standin/echo", 0.5, 1.5, None, &["chat"]),
78. m("standin/cached", 0.5, 1.5, Some(0.1), &["chat", "prompt_cache"]),
79. m("standin/inline-usage", 0.5, 1.5, None, &["chat"]),
80. m("standin/slow", 0.5, 1.5, None, &["chat"]),
81. m("standin/no-usage", 0.5, 1.5, None, &["chat"]),
82. m("standin/reasoning", 0.5, 1.5, None, &["chat", "reasoning"]),
83. m("standin/vision", 0.5, 1.5, None, &["chat", "vision"]),
84. m("standin/midstream-error", 0.5, 1.5, None, &["chat"]),
85. m("standin/402", 0.5, 1.5, None, &["chat"]),
86. m("standin/429", 0.5, 1.5, None, &["chat"]),
87. m("standin/500", 0.5, 1.5, None, &["chat"]),
88. json!({"id": "standin/embed", "object": "model", "metadata": {"pricing": {"input_tokens": 0.01}, "tags": ["embed"]}}),
89. json!({"id": "standin/image", "object": "model", "metadata": {"pricing": {"per_image_unit": 0.005}, "tags": ["image-gen"]}}),
90. ]})
91. }
93. pub fn router(shared: Arc<Shared>) -> Router {
94. Router::new().route("/v1/models", get(list_models)).route("/v1/chat/completions", post(chat)).with_state(shared)
95. }
97. pub async fn spawn() -> Standin {
98. spawn_on("127.0.0.1:0").await
99. }
101. pub async fn spawn_on(bind: &str) -> Standin {
102. let shared = Arc::new(Shared::default());
103. let listener = tokio::net::TcpListener::bind(bind).await.expect("bind");
104. let addr = listener.local_addr().unwrap();
105. let app = router(shared.clone());
106. tokio::spawn(async move {
107. axum::serve(listener, app).await.ok();
108. });
109. Standin { addr, shared }
110. }
112. fn text_of(content: &Value) -> (String, usize) {
113. match content {
114. Value::String(s) => (s.clone(), 0),
115. Value::Array(parts) => {
116. let mut text = String::new();
117. let mut images = 0;
118. for p in parts {
119. match p.get("type").and_then(Value::as_str) {
120. Some("text") => text.push_str(p.get("text").and_then(Value::as_str).unwrap_or("")),
121. Some("image_url") => images += 1,
122. _ => {}
123. }
124. }
125. (text, images)
126. }
127. _ => (String::new(), 0),
128. }
129. }
131. fn chunk(model: &str, delta: Value, finish: Option<&str>) -> String {
132. json!({"id": "standin-1", "object": "chat.completion.chunk", "model": model,
133. "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]})
134. .to_string()
135. }
137. /// DeepInfra's list needs no key, but a bad one gets a 401.
138. async fn list_models(headers: HeaderMap) -> Response {
139. if headers.get("authorization").and_then(|v| v.to_str().ok()) == Some("Bearer bad-key") {
140. let e = json!({"detail": "User is not authorized to access this resource"});
141. return (StatusCode::UNAUTHORIZED, Json(e)).into_response();
142. }
143. Json(models()).into_response()
144. }
146. async fn chat(State(shared): State<Arc<Shared>>, headers: HeaderMap, Json(body): Json<Value>) -> Response {
147. let authorization = headers.get("authorization").and_then(|v| v.to_str().ok()).map(str::to_string);
148. shared.requests.lock().unwrap().push(Recorded { authorization: authorization.clone(), body: body.clone() });
150. if authorization.as_deref() == Some("Bearer bad-key") {
151. let e = json!({"error": {"message": "Invalid API key", "type": "invalid_request_error", "code": "invalid_api_key"}});
152. return (StatusCode::UNAUTHORIZED, Json(e)).into_response();
153. }
154. let model = body.get("model").and_then(Value::as_str).unwrap_or("").to_string();
155. match model.as_str() {
156. "standin/402" => {
157. return (StatusCode::PAYMENT_REQUIRED, Json(json!({"error": {"message": "Insufficient balance"}})))
158. .into_response()
159. }
160. "standin/429" => {
161. return (
162. StatusCode::TOO_MANY_REQUESTS,
163. [("retry-after", "7")],
164. Json(json!({"error": {"message": "Rate limited"}})),
165. )
166. .into_response()
167. }
168. "standin/500" => {
169. return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"detail": "upstream exploded"}))).into_response()
170. }
171. _ => {}
172. }
174. let messages = body.get("messages").and_then(Value::as_array).cloned().unwrap_or_default();
175. let mut prompt_words = 0usize;
176. let mut images = 0usize;
177. let mut last_user = String::new();
178. for m in &messages {
179. let (t, n) = text_of(m.get("content").unwrap_or(&Value::Null));
180. prompt_words += t.split_whitespace().count();
181. images += n;
182. if m.get("role").and_then(Value::as_str) == Some("user") {
183. last_user = t;
184. }
185. }
186. let prompt_tokens = (prompt_words + 85 * images + 4 * messages.len()) as i64;
188. let reply: String = match model.as_str() {
189. "standin/vision" => format!("I see {images} image{}.", if images == 1 { "" } else { "s" }),
190. "standin/slow" => (1..=200).map(|i| format!("word{i}")).collect::<Vec<_>>().join(" "),
191. _ => format!("You said: {last_user}"),
192. };
193. let words: Vec<String> = {
194. let parts: Vec<&str> = reply.split(' ').collect();
195. parts.iter().enumerate().map(|(i, w)| if i == 0 { w.to_string() } else { format!(" {w}") }).collect()
196. };
197. let completion_tokens = words.len() as i64;
198. let want_usage = body.pointer("/stream_options/include_usage").and_then(Value::as_bool).unwrap_or(false);
200. let (tx, rx) = mpsc::channel::<Result<Bytes, Infallible>>(8);
201. let shared2 = shared.clone();
202. tokio::spawn(async move {
203. let send = |s: String| {
204. let tx = tx.clone();
205. async move { tx.send(Ok(Bytes::from(format!("data: {s}\n\n")))).await.is_ok() }
206. };
207. let usage = |cached: i64| {
208. json!({"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens,
209. "total_tokens": prompt_tokens + completion_tokens,
210. "prompt_tokens_details": {"cached_tokens": cached}})
211. };
212. let mut ok = send(chunk(&model, json!({"role": "assistant", "content": ""}), None)).await;
213. if model == "standin/reasoning" && ok {
214. for r in ["Checking", " the", " question."] {
215. ok = ok && send(chunk(&model, json!({"reasoning_content": r}), None)).await;
216. }
217. }
218. let n = words.len();
219. for (i, w) in words.into_iter().enumerate() {
220. if !ok {
221. break;
222. }
223. if model == "standin/midstream-error" && i == 3 {
224. let _ = send(json!({"error": {"message": "model overloaded"}}).to_string()).await;
225. return;
226. }
227. if model == "standin/slow" {
228. tokio::time::sleep(Duration::from_millis(50)).await;
229. }
230. let last = i + 1 == n;
231. if last && model == "standin/inline-usage" {
232. let mut c: Value = serde_json::from_str(&chunk(&model, json!({"content": w}), Some("stop"))).unwrap();
233. c["usage"] = usage(0);
234. ok = send(c.to_string()).await;
235. } else {
236. ok = send(chunk(&model, json!({"content": w}), if last { Some("stop") } else { None })).await;
237. }
238. }
239. if ok && want_usage && model != "standin/no-usage" && model != "standin/inline-usage" {
240. let cached = if model == "standin/cached" { 40.min(prompt_tokens) } else { 0 };
241. let mut c =
242. json!({"id": "standin-1", "object": "chat.completion.chunk", "choices": [], "usage": usage(cached)});
243. if model == "standin/cached" {
244. c["usage"]["estimated_cost"] = json!(0.0000123);
245. }
246. ok = send(c.to_string()).await;
247. }
248. if ok {
249. ok = send("[DONE]".into()).await;
250. }
251. if !ok {
252. shared2.disconnects.fetch_add(1, Ordering::SeqCst);
253. }
254. });
255. Response::builder()
256. .header("content-type", "text/event-stream")
257. .header("cache-control", "no-cache")
258. .body(Body::from_stream(ReceiverStream::new(rx)))
259. .unwrap()
260. }