Skip to main content

cherenkov/server/
routes.rs

1//! HTTP endpoints and admission into the generation queue.
2
3use super::http::{error, read_http, respond};
4use super::{ApiKind, Job, MODEL, failure, registry, sessions};
5use crate::control::State;
6use anyhow::Result;
7use serde_json::{Value, json};
8use std::{
9    net::TcpStream,
10    sync::{Arc, mpsc},
11    time::Duration,
12};
13
14pub(super) fn connection(
15    mut stream: TcpStream,
16    queue: &mpsc::SyncSender<Job>,
17    state: &State,
18    requests: &Arc<registry::Registry>,
19    sessions: &Arc<std::sync::Mutex<sessions::Store>>,
20) -> Result<()> {
21    stream.set_read_timeout(Some(Duration::from_secs(15)))?;
22    stream.set_write_timeout(Some(Duration::from_secs(30)))?;
23    stream.set_nodelay(true)?;
24
25    let settings = state.config();
26    let (method, path, body) = match read_http(&mut stream, settings.config.limits.request_bytes) {
27        Ok(r) => r,
28        Err(e) => return error(&mut stream, 400, &e.to_string()),
29    };
30
31    if let Some(result) =
32        session_endpoint(&method, &path, &body, &settings.config, requests, sessions)
33    {
34        return match result {
35            Ok((status, body)) => respond(&mut stream, status, &body),
36            Err(e) => request_error(&mut stream, state, &e),
37        };
38    }
39
40    match (method.as_str(), path.as_str()) {
41        ("GET", "/health") => respond(&mut stream, 200, &json!({"status":"ok"})),
42        ("GET", "/v1/models") => respond(
43            &mut stream,
44            200,
45            &json!({
46                "object":"list", "data":[{"id":MODEL,"object":"model","created":0,"owned_by":"local"}]
47            }),
48        ),
49        ("POST", "/v1/chat/completions" | "/v1/completions") => {
50            let body: Value = match serde_json::from_slice(&body) {
51                Ok(v) => v,
52                Err(e) => return error(&mut stream, 400, &format!("invalid JSON: {e}")),
53            };
54
55            if !body.is_object() {
56                return error(&mut stream, 400, "request must be an object");
57            }
58
59            if body.get("request_id").is_some_and(|v| !v.is_string()) {
60                return error(&mut stream, 400, "request_id must be a string");
61            }
62
63            let ticket = match requests.register(
64                body["request_id"].as_str(),
65                settings.config.limits.queued_requests + settings.config.limits.active_requests,
66            ) {
67                Ok(ticket) => ticket,
68                Err(e) => return request_error(&mut stream, state, &e),
69            };
70
71            if path == "/v1/completions" && !body["session_id"].is_null() {
72                return error(
73                    &mut stream,
74                    400,
75                    "retained sessions require chat completions",
76                );
77            }
78
79            let session = match sessions::Turn::begin(
80                sessions,
81                &body,
82                &ticket.id,
83                settings.config.limits.request_bytes,
84            ) {
85                Ok(turn) => turn,
86                Err(e) => return request_error(&mut stream, state, &e),
87            };
88            let job = Job {
89                stream,
90                body,
91                settings,
92                ticket,
93                session,
94                kind: if path == "/v1/chat/completions" {
95                    ApiKind::Chat
96                } else {
97                    ApiKind::Completion
98                },
99            };
100
101            state.update(|s| s.queued_requests += 1);
102
103            let (mut job, message) = match queue.try_send(job) {
104                Ok(()) => return Ok(()),
105                Err(mpsc::TrySendError::Full(job)) => (job, "generation queue is full"),
106                Err(mpsc::TrySendError::Disconnected(job)) => (job, "model unavailable"),
107            };
108
109            state.update(|s| {
110                s.queued_requests -= 1;
111                s.rejected_requests += 1;
112            });
113
114            error(&mut job.stream, 503, message)
115        }
116        _ => error(&mut stream, 404, "unknown endpoint"),
117    }
118}
119
120pub(super) fn session_endpoint(
121    method: &str,
122    path: &str,
123    body: &[u8],
124    config: &crate::config::Config,
125    requests: &registry::Registry,
126    sessions: &std::sync::Mutex<sessions::Store>,
127) -> Option<Result<(u16, Value)>> {
128    if method == "GET" && path == "/v1/requests" {
129        return Some(Ok((200, requests.list())));
130    }
131
132    if let Some(id) = path
133        .strip_prefix("/v1/requests/")
134        .and_then(|p| p.strip_suffix("/cancel"))
135    {
136        if method != "POST" {
137            return None;
138        }
139
140        return Some(Ok((
141            if requests.cancel(id) { 200 } else { 404 },
142            json!({"id":id}),
143        )));
144    }
145
146    if method == "POST" && path == "/v1/sessions" {
147        return Some(
148            serde_json::from_slice(body)
149                .map_err(Into::into)
150                .and_then(|v| {
151                    let mut store = sessions.lock().unwrap();
152                    let id = store.create(&v, config)?;
153
154                    store.show(&id)
155                })
156                .map(|v| (200, v)),
157        );
158    }
159
160    let id = path.strip_prefix("/v1/sessions/")?;
161
162    match method {
163        "GET" => Some(sessions.lock().unwrap().show(id).map(|v| (200, v))),
164        "DELETE" => Some(
165            sessions
166                .lock()
167                .unwrap()
168                .delete(id)
169                .map(|()| (200, json!({"id":id,"deleted":true}))),
170        ),
171        _ => None,
172    }
173}
174
175fn request_error(out: &mut TcpStream, state: &State, cause: &anyhow::Error) -> Result<()> {
176    let status = failure::status(cause);
177
178    state.update(|s| s.rejected_requests += u64::from(status == 503));
179
180    error(out, status, &cause.to_string())
181}