10use std::fmt;

12use crate::error::Diagnostic;
13use crate::secrets::SecretName;

16pub const REPLY_TOKENS_MAX: u64 = 1000;

19#[derive(Clone, Debug, PartialEq, Eq)]
20pub struct ModelName(String);

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}

43#[derive(Clone, Debug, PartialEq, Eq)]
44pub enum ThinkRefusal {
45    NotJson(Diagnostic),

47    WrongModel { wanted: ModelName, asked_for: Option<String> },

49    ReplyLength,
50    Streaming,
51    Tools,
52}

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 {}

70#[derive(Clone, Debug, PartialEq, Eq)]
71pub struct ThinkRequest(String);

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    }

106    pub fn body(&self) -> &str {
107        &self.0
108    }
109}

113#[derive(Clone, Debug, PartialEq, Eq)]
114pub struct Reply {
115    pub status: u16,
116    pub body: String,
117}

119impl Reply {

121    pub fn is_answer(&self) -> bool {
122        self.status == 200
123    }
124}

127#[derive(Clone, Debug, PartialEq, Eq)]
128pub enum ThinkError {

130    NotConfigured(SecretName),
131    Unreachable(Diagnostic),

133    Unreadable(Diagnostic),
134}

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 {}

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 {

157    fn serves(&self) -> &ModelName;

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}