cpm_planner/llm/
jev.rs

1//! [`JudgmentModel`] over Jev, via `rig-typesafeai`, pointed at OpenRouter.
2//!
3//! The request is `rig-typesafeai`'s (its [`JevConfig`] wire encodes it and
4//! [`DynamicQuery`] validates questions and answers), sent by rig's HTTP
5//! driver. Only the reply type is ours (`Reply`): rig-typesafeai 0.44
6//! decodes `usage` as rig-core's `Usage`, whose `cost` must be an object
7//! with a `total`, while OpenRouter reports `usage.cost` as a bare number,
8//! so every OpenRouter reply failed to decode. `Reply` reads `usage`
9//! leniently: `cost` may be a number or an object, unknown fields are
10//! ignored, and malformed accounting is dropped rather than failing the
11//! judgment.
12//!
13//! Error messages are tool-facing: scrubbed of the key, capped, and limited
14//! to the class plus an HTTP status or short reason. Upstream bodies are
15//! logged at `debug` only (scrubbed, then truncated, and Debug-escaped so a
16//! body cannot forge log lines).
17//!
18//! Follow-up (not implemented): the reply body size is not capped. rig's
19//! driver reads replies through `HttpClientExt::send_streaming`, so a cap
20//! needs a wrapper implementing all three `HttpClientExt` methods over a
21//! boxed byte stream (plus direct `rig-reqwest`/`bytes`/`futures` deps), and
22//! its overflow error would surface as a transport error rather than
23//! `Decode`. The per-call timeout bounds how long a large reply can stream.
24
25use std::collections::BTreeMap;
26use std::fmt;
27use std::time::Duration;
28
29use async_trait::async_trait;
30use rig_core::completion::request::Cost;
31use rig_core::driver::Model;
32use rig_core::error::{EncodeError, ProviderError};
33use rig_core::operation::Whole;
34use rig_core::wire::{
35    Call, Descriptor, Encoded, Free, Json, Mode, Operation, Wire, WireFrame, document,
36};
37use rig_typesafeai::{DynamicQuery, JevConfig, Query};
38use serde::Deserialize;
39
40use super::{
41    Answer, ApiKey, Decisions, JudgmentError, JudgmentErrorKind, JudgmentModel, LlmConfig,
42    Question, Usage,
43};
44
45/// Longest slice of a provider error body written to the debug log.
46const BODY_LOG_CHARS: usize = 200;
47
48/// One Jev evaluation whose reply decodes as [`Reply`]; otherwise
49/// rig-typesafeai's `Evaluation` (same request, one whole JSON reply),
50/// except that the reply's `provider_request_id` and `raw` document are not
51/// captured: nothing here reads them.
52struct Evaluation;
53
54impl Operation for Evaluation {
55    type Request = rig_typesafeai::types::Request;
56    type Event = std::convert::Infallible;
57    type End = Reply;
58    type Response = Reply;
59    type Fold = Whole<Self>;
60    type Emit = Free;
61
62    fn fold(_request: &Self::Request, _call: &mut Call<'_>) -> Whole<Self> {
63        Whole::<Self>::new()
64    }
65}
66
67/// rig-typesafeai's [`JevConfig`] wire (endpoint, model, bearer key, request
68/// body), answering with [`Reply`].
69#[derive(Clone)]
70struct JevWire(JevConfig);
71
72impl Wire for JevWire {
73    type Op = Evaluation;
74    type Payload = Encoded;
75    type Frame = WireFrame;
76    type Decoder<'id> = Json;
77    type Reassembler = document::Unreassembled;
78
79    fn describe(&self) -> Descriptor<'_> {
80        self.0.describe()
81    }
82
83    fn encode(
84        &self,
85        request: rig_typesafeai::types::Request,
86        mode: Mode,
87    ) -> Result<Encoded, EncodeError> {
88        self.0.encode(request, mode)
89    }
90
91    fn decoder<'id>(&self) -> Self::Decoder<'id> {
92        Json
93    }
94}
95
96/// A Jev reply. Fields other than these are ignored.
97#[derive(Debug, Deserialize)]
98struct Reply {
99    model: String,
100    answers: UniqueAnswers,
101    /// Read by [`lenient_usage`] after decoding, so bad accounting never
102    /// fails the reply.
103    #[serde(default)]
104    usage: Option<serde_json::Value>,
105}
106
107/// The `answers` object; like rig-typesafeai, an empty or repeated question
108/// id is a decode error rather than a silently merged map.
109#[derive(Debug)]
110struct UniqueAnswers(BTreeMap<String, Answer>);
111
112impl<'de> Deserialize<'de> for UniqueAnswers {
113    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
114        struct Visit;
115        impl<'de> serde::de::Visitor<'de> for Visit {
116            type Value = UniqueAnswers;
117            fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
118                f.write_str("an object keyed by unique question IDs")
119            }
120            fn visit_map<M: serde::de::MapAccess<'de>>(
121                self,
122                mut map: M,
123            ) -> Result<UniqueAnswers, M::Error> {
124                let mut answers = BTreeMap::new();
125                while let Some((id, answer)) = map.next_entry::<String, Answer>()? {
126                    if id.is_empty() || answers.insert(id, answer).is_some() {
127                        return Err(serde::de::Error::custom(
128                            "question IDs must be nonempty and unique",
129                        ));
130                    }
131                }
132                Ok(UniqueAnswers(answers))
133            }
134        }
135        deserializer.deserialize_map(Visit)
136    }
137}
138
139/// The `usage` counters rig-core's [`Usage`] reads, each a `u64`.
140const USAGE_COUNTERS: [&str; 7] = [
141    "input_tokens",
142    "output_tokens",
143    "total_tokens",
144    "cached_input_tokens",
145    "cache_creation_input_tokens",
146    "tool_use_prompt_tokens",
147    "reasoning_tokens",
148];
149
150/// `usage` as rig-core's [`Usage`], tolerating provider dialects, and
151/// whether part of it was not understood and dropped. A bare number `cost`
152/// (OpenRouter) is its total and an object `cost` is read as [`Cost`]. Each
153/// part is kept or dropped on its own: a bad `cost` or a counter that is not
154/// a `u64` is dropped alone, unknown fields are ignored, and a `usage` that
155/// is not an object is `None`. Token accounting never fails a judgment.
156fn lenient_usage(raw: Option<&serde_json::Value>) -> (Option<Usage>, bool) {
157    use serde_json::Value;
158    let mut usage = match raw {
159        None | Some(Value::Null) => return (None, false),
160        Some(Value::Object(usage)) => usage.clone(),
161        Some(_) => return (None, true),
162    };
163    let mut dropped = false;
164    let cost = match usage.remove("cost") {
165        None | Some(Value::Null) => None,
166        Some(Value::Number(total)) => total.as_f64().map(Cost::from_total),
167        Some(other) => serde_json::from_value::<Cost>(other).ok(),
168    };
169    dropped |= cost.is_none() && raw.is_some_and(|u| !u["cost"].is_null());
170    usage.retain(|key, value| {
171        let bad = USAGE_COUNTERS.contains(&key.as_str()) && !value.is_u64() && !value.is_null();
172        dropped |= bad;
173        !bad
174    });
175    match serde_json::from_value::<Usage>(Value::Object(usage)) {
176        Ok(usage) => (Some(usage.cost(cost)), dropped),
177        // Unreachable after the filter above; kept total for safety.
178        Err(_) => (cost.map(|cost| Usage::new().cost(cost)), true),
179    }
180}
181
182/// Why `state` breaks rig-typesafeai's state rule (a string, object or
183/// array), if it does.
184fn state_problem(state: &serde_json::Value) -> Option<&'static str> {
185    match state {
186        serde_json::Value::String(_)
187        | serde_json::Value::Object(_)
188        | serde_json::Value::Array(_) => None,
189        serde_json::Value::Null => Some("state cannot be null"),
190        _ => Some("question content must be a string, object, array, or null"),
191    }
192}
193
194/// A reqwest client for one judge, configured like rig's shared client.
195///
196/// rig's `.client()` sends through one process-wide client whose pool hands
197/// keep-alive connections to every caller, while each connection's driver
198/// task runs on the tokio runtime that opened it. A judge on another runtime
199/// could then check out a connection whose runtime was idle (the request
200/// stalled until the timeout) or shutting down (an immediate transport error,
201/// "runtime dropped the dispatch task"). A client per judge keeps a judge's
202/// connections to itself. If reqwest cannot build a client (TLS setup), the
203/// shared one is used: it reports that build error on every send.
204fn own_http_client() -> rig_reqwest::ReqwestClient {
205    match rig_reqwest::reqwest::Client::builder().build() {
206        Ok(client) => rig_reqwest::ReqwestClient::from(client),
207        Err(_) => rig_reqwest::ReqwestClient::default(),
208    }
209}
210
211/// Jev on OpenRouter's System One endpoint. Every call is bounded by the
212/// configured timeout. Each judge owns its HTTP client and connection pool;
213/// pooled connections are driven by the runtime that opened them, so use a
214/// judge from one tokio runtime (clones share the pool).
215#[derive(Clone)]
216pub struct JevJudge {
217    model: Model<JevWire, rig_reqwest::ReqwestClient>,
218    model_id: String,
219    endpoint: String,
220    timeout: Duration,
221    key: ApiKey,
222}
223
224impl JevJudge {
225    /// Build the client from `config`. No network traffic happens here.
226    pub fn new(config: &LlmConfig) -> Self {
227        let wire = JevConfig::new(config.api_key().expose())
228            .model(config.jev_model())
229            .with_endpoint(config.jev_endpoint());
230        Self {
231            model: Model::new(JevWire(wire), own_http_client()),
232            model_id: config.jev_model().to_string(),
233            endpoint: config.jev_endpoint().to_string(),
234            timeout: config.timeout(),
235            key: config.api_key().clone(),
236        }
237    }
238
239    /// The configured Jev model id.
240    pub fn model_id(&self) -> &str {
241        &self.model_id
242    }
243
244    /// The configured endpoint URL.
245    pub fn endpoint(&self) -> &str {
246        &self.endpoint
247    }
248
249    /// Host of the configured endpoint, from the parsed URL.
250    pub fn endpoint_host(&self) -> Option<String> {
251        super::endpoint_host(&self.endpoint)
252    }
253
254    /// A tool-facing error: endpoint query/fragment removed, scrubbed of the
255    /// key, then capped at
256    /// [`MAX_ERROR_MESSAGE_CHARS`](super::MAX_ERROR_MESSAGE_CHARS).
257    fn error(&self, kind: JudgmentErrorKind, message: impl AsRef<str>) -> JudgmentError {
258        JudgmentError::scrubbed(kind, self.redact_endpoint(message.as_ref()), &self.key)
259    }
260
261    /// `text` with the configured endpoint's query string and fragment
262    /// replaced by [`REDACTED`](super::REDACTED): transport errors quote the
263    /// request URL, and a query may carry a token.
264    fn redact_endpoint(&self, text: &str) -> String {
265        let Ok(url) = url::Url::parse(&self.endpoint) else {
266            return text.to_string();
267        };
268        let mut out = text.to_string();
269        for part in [url.query(), url.fragment()].into_iter().flatten() {
270            if !part.is_empty() {
271                out = out.replace(part, super::REDACTED);
272            }
273        }
274        // The raw endpoint may differ from the parsed form (encoding).
275        if let Some((_, rest)) = self.endpoint.split_once(['?', '#']) {
276            for part in rest.split(['?', '#']).filter(|p| !p.is_empty()) {
277                out = out.replace(part, super::REDACTED);
278            }
279        }
280        out
281    }
282
283    /// [`ApiKey::excerpt`] of `text` with the endpoint redacted, for logs.
284    fn log_excerpt(&self, text: &str) -> String {
285        self.key
286            .excerpt(&self.redact_endpoint(text), BODY_LOG_CHARS)
287    }
288
289    /// The upstream body goes to the debug log only (scrubbed, then cut);
290    /// the tool-facing message carries the status alone.
291    fn reply_error(
292        &self,
293        kind: JudgmentErrorKind,
294        reply: &rig_core::ProviderResponseError,
295    ) -> JudgmentError {
296        let status = reply.status.map(|s| s.as_u16());
297        tracing::debug!(
298            ?status,
299            body = ?self.log_excerpt(reply.body.trim()),
300            "jev upstream error reply"
301        );
302        match status {
303            Some(code) => self.error(kind, format!("status {code}")),
304            None => self.error(kind, "provider error reply without an HTTP status"),
305        }
306    }
307
308    /// A failure whose text may be provider-authored: logged at debug,
309    /// replaced by `summary` in the tool-facing message.
310    fn opaque_error(
311        &self,
312        kind: JudgmentErrorKind,
313        summary: &str,
314        error: &ProviderError,
315    ) -> JudgmentError {
316        tracing::debug!(
317            detail = ?self.log_excerpt(&error.to_string()),
318            "jev provider failure"
319        );
320        self.error(kind, summary)
321    }
322
323    fn classify(&self, error: &ProviderError) -> JudgmentError {
324        use JudgmentErrorKind as K;
325        match error {
326            ProviderError::InvalidAuthentication(reply) => self.reply_error(K::Unauthorized, reply),
327            ProviderError::ProviderResponse(reply) => {
328                let kind = match reply.status.map(|s| s.as_u16()) {
329                    Some(401 | 403) => K::Unauthorized,
330                    Some(429) => K::RateLimited,
331                    _ => K::Upstream,
332                };
333                self.reply_error(kind, reply)
334            }
335            ProviderError::CacheExpired { response, .. } => self.reply_error(K::Upstream, response),
336            // rig-authored, short reasons (serde position, id mismatch).
337            ProviderError::Json(_)
338            | ProviderError::Response(_)
339            | ProviderError::Truncated
340            | ProviderError::MismatchedDimensions { .. } => {
341                self.error(K::Decode, error.to_string())
342            }
343            ProviderError::Http(_) => self.error(K::Transport, error.to_string()),
344            ProviderError::Url(_)
345            | ProviderError::Request(_)
346            | ProviderError::UnsupportedOption(_) => {
347                self.error(K::InvalidRequest, error.to_string())
348            }
349            ProviderError::Provider(_) | ProviderError::Relayed(_) => {
350                self.opaque_error(K::Upstream, "provider reported a failure", error)
351            }
352            // `ProviderError` is `#[non_exhaustive]`.
353            _ => self.opaque_error(K::Upstream, "unclassified provider failure", error),
354        }
355    }
356}
357
358impl fmt::Debug for JevJudge {
359    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
360        f.debug_struct("JevJudge")
361            .field("model", &self.model_id)
362            .field("endpoint", &self.endpoint)
363            .field("timeout", &self.timeout)
364            .field("key", &self.key)
365            .finish_non_exhaustive()
366    }
367}
368
369#[async_trait]
370impl JudgmentModel for JevJudge {
371    async fn decide(
372        &self,
373        state: serde_json::Value,
374        questions: BTreeMap<String, Question>,
375    ) -> Result<Decisions, JudgmentError> {
376        let query = DynamicQuery::new(questions).map_err(|e| self.classify(&e))?;
377        if let Some(reason) = state_problem(&state) {
378            return Err(self.classify(&ProviderError::request(reason)));
379        }
380        let request = rig_typesafeai::types::Request {
381            state,
382            questions: serde_json::value::to_raw_value(&query)
383                .map_err(|e| self.classify(&e.into()))?,
384        };
385        let reply = tokio::time::timeout(self.timeout, self.model.call(request))
386            .await
387            .map_err(|_| {
388                self.error(
389                    JudgmentErrorKind::Timeout,
390                    format!("no reply within {} ms", self.timeout.as_millis()),
391                )
392            })?
393            .map_err(|e| self.classify(&e))?;
394        let (usage, dropped) = lenient_usage(reply.usage.as_ref());
395        if dropped && let Some(raw) = &reply.usage {
396            tracing::debug!(
397                usage = ?self.log_excerpt(&raw.to_string()),
398                "jev usage not understood; dropped"
399            );
400        }
401        // rig's own answer validation: ids match, kinds and distributions fit.
402        let answers = query
403            .decode(reply.answers.0)
404            .map_err(|e| self.classify(&e))?;
405        Ok(Decisions {
406            answers,
407            // Provider-authored: scrubbed of the key, then cut.
408            model: self
409                .key
410                .excerpt(&reply.model, super::MAX_REPORTED_MODEL_CHARS),
411            usage,
412        })
413    }
414
415    fn endpoint_host(&self) -> Option<String> {
416        JevJudge::endpoint_host(self)
417    }
418
419    fn model_id(&self) -> Option<String> {
420        Some(self.model_id.clone())
421    }
422}
cpm_planner/llm/jev.rs

cpm_planner/llm/
jev.rs
