think.rsannotatedthink.rssource205 lines · 7.8 KB · raw
1//! The thinker: the language model that writes what Whiskers says, reached through whatever gateway
2//! the operator runs.
3//!
4//! The tablet already speaks the Messages shape to the service, and the service forwards the
5//! tablet's body byte for byte, so the request here is the tablet's body, *vetted*. What may pass is
6//! Whiskers' own policy, written once below, so no adapter can forward something looser: one model, a
7//! short reply, no streaming, no tools. The gateway may be the operator's subscription with no login
8//! of its own, and this check is its only guard.
9
10use std::fmt;
11
12use crate::error::Diagnostic;
13use crate::secrets::SecretName;
14
15/// The most tokens one reply may ask for.
16pub const REPLY_TOKENS_MAX: u64 = 1000;
17
18/// Which model Whiskers uses: the one model the service lets through.
19#[derive(Clone, Debug, PartialEq, Eq)]
20pub struct ModelName(String);
21
22#[derive(Clone, Copy, Debug, PartialEq, Eq)]
23pub struct NoModelNamed;
24
25impl ModelName {
26    pub fn new(name: &str) -> Result<Self, NoModelNamed> {
27        let name = name.trim();
28        if name.is_empty() { Err(NoModelNamed) } else { Ok(Self(name.to_owned())) }
29    }
30
31    pub fn as_str(&self) -> &str {
32        &self.0
33    }
34}
35
36impl fmt::Display for ModelName {
37    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
38        f.write_str(&self.0)
39    }
40}
41
42/// Why a request was not let through to the model.
43#[derive(Clone, Debug, PartialEq, Eq)]
44pub enum ThinkRefusal {
45    NotJson(Diagnostic),
46    /// The request names a model other than the one Whiskers uses (or none).
47    WrongModel { wanted: ModelName, asked_for: Option<String> },
48    /// `max_tokens` is missing, zero, or over [`REPLY_TOKENS_MAX`].
49    ReplyLength,
50    Streaming,
51    Tools,
52}
53
54impl fmt::Display for ThinkRefusal {
55    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
56        match self {
57            ThinkRefusal::NotJson(d) => write!(f, "bad request: {d}"),
58            ThinkRefusal::WrongModel { wanted, asked_for } => write!(f, "bad request: model must be {wanted}, not {asked_for:?}"),
59            ThinkRefusal::ReplyLength => write!(f, "bad request: max_tokens must be 1 to {REPLY_TOKENS_MAX}"),
60            ThinkRefusal::Streaming => write!(f, "bad request: streaming is not supported"),
61            ThinkRefusal::Tools => write!(f, "bad request: tools are not allowed"),
62        }
63    }
64}
65
66impl std::error::Error for ThinkRefusal {}
67
68/// A request the model may be asked: only [`ThinkRequest::vet`] makes one, so an adapter cannot be
69/// handed anything wider than the policy allows.
70#[derive(Clone, Debug, PartialEq, Eq)]
71pub struct ThinkRequest(String);
72
73impl ThinkRequest {
74    pub fn vet(body: &str, model: &ModelName) -> Result<Self, ThinkRefusal> {
75        let v: serde_json::Value = serde_json::from_str(body).map_err(|e| {
76            ::log::warn!("model request is not JSON ({} bytes): {e}", body.len());
77            ThinkRefusal::NotJson(Diagnostic::new(e.to_string()))
78        })?;
79        match v.get("model").and_then(|m| m.as_str()) {
80            Some(m) if m == model.as_str() => {}
81            other => {
82                ::log::warn!("model request names a model that is not allowed: {other:?}");
83                return Err(ThinkRefusal::WrongModel { wanted: model.clone(), asked_for: other.map(str::to_owned) });
84            }
85        }
86        match v.get("max_tokens").and_then(|m| m.as_u64()) {
87            Some(n) if (1..=REPLY_TOKENS_MAX).contains(&n) => {}
88            other => {
89                ::log::warn!("model request has max_tokens {other:?}, outside 1 to {REPLY_TOKENS_MAX}");
90                return Err(ThinkRefusal::ReplyLength);
91            }
92        }
93        if v.get("stream").and_then(|s| s.as_bool()).unwrap_or(false) {
94            ::log::warn!("model request asks for streaming");
95            return Err(ThinkRefusal::Streaming);
96        }
97        if v.get("tools").is_some() {
98            ::log::warn!("model request carries tools");
99            return Err(ThinkRefusal::Tools);
100        }
101        ::log::debug!("model request accepted");
102        Ok(Self(body.to_owned()))
103    }
104
105    /// The tablet's body, exactly as it was sent.
106    pub fn body(&self) -> &str {
107        &self.0
108    }
109}
110
111/// What the model's gateway answered. A non-success status is an *answer*, passed back to the tablet
112/// as it came (it knows what to do with a gateway's own errors); only failing to get one is an error.
113#[derive(Clone, Debug, PartialEq, Eq)]
114pub struct Reply {
115    pub status: u16,
116    pub body: String,
117}
118
119impl Reply {
120    /// Whether the model answered. Only an answer is charged for.
121    pub fn is_answer(&self) -> bool {
122        self.status == 200
123    }
124}
125
126/// Why no answer came back at all.
127#[derive(Clone, Debug, PartialEq, Eq)]
128pub enum ThinkError {
129    /// A credential the gateway needs is not set. Nothing was tried.
130    NotConfigured(SecretName),
131    Unreachable(Diagnostic),
132    /// The gateway answered but its reply could not be read.
133    Unreadable(Diagnostic),
134}
135
136impl fmt::Display for ThinkError {
137    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
138        match self {
139            ThinkError::NotConfigured(s) => write!(f, "gateway: {s} is not set"),
140            ThinkError::Unreachable(d) | ThinkError::Unreadable(d) => d.fmt(f),
141        }
142    }
143}
144
145impl std::error::Error for ThinkError {}
146
147/// The language model.
148///
149/// Contract, for every adapter:
150///
151/// - **Forwards the body unchanged** and returns what the gateway said, status and body, unchanged.
152/// - **Does not meter.** The thinking allowance is the service's, asked before and charged after.
153/// - **Says why** it got no answer, by [`ThinkError`] kind.
154#[expect(async_fn_in_trait, reason = "a Worker's futures hold JavaScript values and cannot be Send, so no Send bound may be required here")]
155pub trait Thinker {
156    /// The one model this thinker serves, which [`ThinkRequest::vet`] holds requests to.
157    fn serves(&self) -> &ModelName;
158
159    async fn think(&self, request: &ThinkRequest) -> Result<Reply, ThinkError>;
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    fn m() -> ModelName {
167        ModelName::new("m").unwrap()
168    }
169
170    const GOOD: &str = r#"{"model":"m","max_tokens":300,"system":"s","messages":[{"role":"user","content":"hi"}]}"#;
171
172    #[test]
173    fn a_whiskers_request_passes_and_is_kept_byte_for_byte() {
174        assert_eq!(ThinkRequest::vet(GOOD, &m()).unwrap().body(), GOOD);
175    }
176
177    #[test]
178    fn anything_else_is_refused_and_says_why() {
179        let cases = [
180            ("nonsense", "bad request: expected ident at line 1 column 2"),
181            (r#"{"model":"other","max_tokens":10,"messages":[]}"#, "bad request: model must be m, not Some(\"other\")"),
182            (r#"{"max_tokens":10,"messages":[]}"#, "bad request: model must be m, not None"),
183            (r#"{"model":"m","max_tokens":5000,"messages":[]}"#, "bad request: max_tokens must be 1 to 1000"),
184            (r#"{"model":"m","max_tokens":0,"messages":[]}"#, "bad request: max_tokens must be 1 to 1000"),
185            (r#"{"model":"m","messages":[]}"#, "bad request: max_tokens must be 1 to 1000"),
186            (r#"{"model":"m","max_tokens":10,"stream":true,"messages":[]}"#, "bad request: streaming is not supported"),
187            (r#"{"model":"m","max_tokens":10,"tools":[],"messages":[]}"#, "bad request: tools are not allowed"),
188        ];
189        for (bad, why) in cases {
190            assert_eq!(ThinkRequest::vet(bad, &m()).unwrap_err().to_string(), why, "{bad}");
191        }
192    }
193
194    #[test]
195    fn stream_false_is_fine_and_a_model_needs_a_name() {
196        assert!(ThinkRequest::vet(r#"{"model":"m","max_tokens":10,"stream":false,"messages":[]}"#, &m()).is_ok());
197        assert!(ModelName::new("  ").is_err());
198    }
199
200    #[test]
201    fn only_a_200_is_an_answer() {
202        assert!(Reply { status: 200, body: String::new() }.is_answer());
203        assert!(!Reply { status: 500, body: "boom".into() }.is_answer());
204    }
205}