Skip to main content

cpm_planner/llm/
jev.rs

1//! [`JudgmentModel`] over Jev, via `rig-typesafeai`, pointed at OpenRouter.
2//!
3//! The request and reply are `rig-typesafeai`'s (`JevConfig` encodes the
4//! request and `DynamicQuery` validates questions and answers), sent by
5//! rig's HTTP driver through [`Evaluate::evaluate`]. rig-typesafeai 0.45
6//! reads Jev's `usage` leniently: token counters and OpenRouter's numeric
7//! `usage.cost` are kept, and a counter or cost of another shape is
8//! unreported rather than failing the judgment, so this module carries no
9//! reply-decoding workaround.
10//!
11//! Error messages are tool-facing: scrubbed of the key, capped, and limited
12//! to the class plus an HTTP status or short reason. Upstream bodies are
13//! logged at `debug` only (scrubbed, then truncated, and Debug-escaped so a
14//! body cannot forge log lines).
15//!
16//! Follow-up (not implemented): the reply body size is not capped. rig's
17//! driver reads replies through `HttpClientExt::send_streaming`, so a cap
18//! needs a wrapper implementing all three `HttpClientExt` methods over a
19//! boxed byte stream (plus direct `rig-reqwest`/`bytes`/`futures` deps), and
20//! its overflow error would surface as a transport error rather than
21//! `Decode`. The per-call timeout bounds how long a large reply can stream.
22
23use std::collections::BTreeMap;
24use std::fmt;
25use std::time::Duration;
26
27use async_trait::async_trait;
28use rig_core::driver::Model;
29use rig_core::error::ProviderError;
30use rig_typesafeai::{DynamicQuery, Evaluate, JevConfig};
31use serde_json::Value;
32
33use super::{
34    ApiKey, Decisions, JudgmentError, JudgmentErrorKind, JudgmentModel, LlmConfig, Question,
35};
36
37/// Longest slice of a provider error body written to the debug log.
38const BODY_LOG_CHARS: usize = 200;
39
40/// rig-typesafeai 0.45 reads Jev's input and output tokens but not its
41/// `total_tokens`; report their sum so callers keep a total.
42fn with_total_tokens(usage: rig_core::completion::Usage) -> rig_core::completion::Usage {
43    if usage.total_tokens.is_some()
44        || (usage.input_tokens.is_none() && usage.output_tokens.is_none())
45    {
46        return usage;
47    }
48    let total = usage
49        .input_tokens
50        .unwrap_or(0)
51        .saturating_add(usage.output_tokens.unwrap_or(0));
52    usage.total_tokens(total)
53}
54
55/// A reqwest client for one judge, configured like rig's shared client.
56///
57/// rig's `.client()` sends through one process-wide client whose pool hands
58/// keep-alive connections to every caller, while each connection's driver
59/// task runs on the tokio runtime that opened it. A judge on another runtime
60/// could then check out a connection whose runtime was idle (the request
61/// stalled until the timeout) or shutting down (an immediate transport error,
62/// "runtime dropped the dispatch task"). A client per judge keeps a judge's
63/// connections to itself. If reqwest cannot build a client (TLS setup), the
64/// shared one is used: it reports that build error on every send.
65fn own_http_client() -> rig_reqwest::ReqwestClient {
66    match rig_reqwest::reqwest::Client::builder().build() {
67        Ok(client) => rig_reqwest::ReqwestClient::from(client),
68        Err(_) => rig_reqwest::ReqwestClient::default(),
69    }
70}
71
72/// Jev on OpenRouter's System One endpoint. Every call is bounded by the
73/// configured timeout. Each judge owns its HTTP client and connection pool;
74/// pooled connections are driven by the runtime that opened them, so use a
75/// judge from one tokio runtime (clones share the pool).
76#[derive(Clone)]
77pub struct JevJudge {
78    model: Model<JevConfig>,
79    model_id: String,
80    endpoint: String,
81    timeout: Duration,
82    key: ApiKey,
83}
84
85impl JevJudge {
86    /// Build the client from `config`. No network traffic happens here.
87    pub fn new(config: &LlmConfig) -> Self {
88        let model = JevConfig::new(config.api_key().expose())
89            .model(config.jev_model())
90            .with_endpoint(config.jev_endpoint())
91            .connect(own_http_client())
92            .evaluation();
93        Self {
94            model,
95            model_id: config.jev_model().to_string(),
96            endpoint: config.jev_endpoint().to_string(),
97            timeout: config.timeout(),
98            key: config.api_key().clone(),
99        }
100    }
101
102    /// The configured Jev model id.
103    pub fn model_id(&self) -> &str {
104        &self.model_id
105    }
106
107    /// The configured endpoint URL.
108    pub fn endpoint(&self) -> &str {
109        &self.endpoint
110    }
111
112    /// Host of the configured endpoint, from the parsed URL.
113    pub fn endpoint_host(&self) -> Option<String> {
114        super::endpoint_host(&self.endpoint)
115    }
116
117    /// A tool-facing error: endpoint query/fragment removed, scrubbed of the
118    /// key, then capped at
119    /// [`MAX_ERROR_MESSAGE_CHARS`](super::MAX_ERROR_MESSAGE_CHARS).
120    fn error(&self, kind: JudgmentErrorKind, message: impl AsRef<str>) -> JudgmentError {
121        JudgmentError::scrubbed(kind, self.redact_endpoint(message.as_ref()), &self.key)
122    }
123
124    /// `text` with the configured endpoint's query string and fragment
125    /// replaced by [`REDACTED`](super::REDACTED): transport errors quote the
126    /// request URL, and a query may carry a token.
127    fn redact_endpoint(&self, text: &str) -> String {
128        let Ok(url) = url::Url::parse(&self.endpoint) else {
129            return text.to_string();
130        };
131        let mut out = text.to_string();
132        for part in [url.query(), url.fragment()].into_iter().flatten() {
133            if !part.is_empty() {
134                out = out.replace(part, super::REDACTED);
135            }
136        }
137        // The raw endpoint may differ from the parsed form (encoding).
138        if let Some((_, rest)) = self.endpoint.split_once(['?', '#']) {
139            for part in rest.split(['?', '#']).filter(|p| !p.is_empty()) {
140                out = out.replace(part, super::REDACTED);
141            }
142        }
143        out
144    }
145
146    /// [`ApiKey::excerpt`] of `text` with the endpoint redacted, for logs.
147    fn log_excerpt(&self, text: &str) -> String {
148        self.key
149            .excerpt(&self.redact_endpoint(text), BODY_LOG_CHARS)
150    }
151
152    /// The upstream body goes to the debug log only (scrubbed, then cut);
153    /// the tool-facing message carries the status alone.
154    fn reply_error(
155        &self,
156        kind: JudgmentErrorKind,
157        reply: &rig_core::ProviderResponseError,
158    ) -> JudgmentError {
159        let status = reply.status.map(|s| s.as_u16());
160        tracing::debug!(
161            ?status,
162            body = ?self.log_excerpt(reply.body.trim()),
163            "jev upstream error reply"
164        );
165        match status {
166            Some(code) => self.error(kind, format!("status {code}")),
167            None => self.error(kind, "provider error reply without an HTTP status"),
168        }
169    }
170
171    /// A failure whose text may be provider-authored: logged at debug,
172    /// replaced by `summary` in the tool-facing message.
173    fn opaque_error(
174        &self,
175        kind: JudgmentErrorKind,
176        summary: &str,
177        error: &ProviderError,
178    ) -> JudgmentError {
179        tracing::debug!(
180            detail = ?self.log_excerpt(&error.to_string()),
181            "jev provider failure"
182        );
183        self.error(kind, summary)
184    }
185
186    fn classify(&self, error: &ProviderError) -> JudgmentError {
187        use JudgmentErrorKind as K;
188        match error {
189            ProviderError::InvalidAuthentication(reply) => self.reply_error(K::Unauthorized, reply),
190            ProviderError::ProviderResponse(reply) => {
191                let kind = match reply.status.map(|s| s.as_u16()) {
192                    Some(401 | 403) => K::Unauthorized,
193                    Some(429) => K::RateLimited,
194                    _ => K::Upstream,
195                };
196                self.reply_error(kind, reply)
197            }
198            ProviderError::CacheExpired { response, .. } => self.reply_error(K::Upstream, response),
199            // rig-authored, short reasons (serde position, id mismatch).
200            ProviderError::Json(_)
201            | ProviderError::Response(_)
202            | ProviderError::Truncated
203            | ProviderError::MismatchedDimensions { .. } => {
204                self.error(K::Decode, error.to_string())
205            }
206            ProviderError::Http(_) => self.error(K::Transport, error.to_string()),
207            ProviderError::Url(_)
208            | ProviderError::Request(_)
209            | ProviderError::UnsupportedOption(_) => {
210                self.error(K::InvalidRequest, error.to_string())
211            }
212            ProviderError::Provider(_) | ProviderError::Relayed(_) => {
213                self.opaque_error(K::Upstream, "provider reported a failure", error)
214            }
215            // `ProviderError` is `#[non_exhaustive]`.
216            _ => self.opaque_error(K::Upstream, "unclassified provider failure", error),
217        }
218    }
219}
220
221impl fmt::Debug for JevJudge {
222    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
223        f.debug_struct("JevJudge")
224            .field("model", &self.model_id)
225            .field("endpoint", &self.endpoint)
226            .field("timeout", &self.timeout)
227            .field("key", &self.key)
228            .finish_non_exhaustive()
229    }
230}
231
232#[async_trait]
233impl JudgmentModel for JevJudge {
234    async fn decide(
235        &self,
236        state: Value,
237        questions: BTreeMap<String, Question>,
238    ) -> Result<Decisions, JudgmentError> {
239        let query = DynamicQuery::new(questions).map_err(|e| self.classify(&e))?;
240        // rig validates the state and the question IDs before sending, then
241        // validates and decodes the answers, so failures keep their kind.
242        let result = tokio::time::timeout(self.timeout, self.model.evaluate(&state, query))
243            .await
244            .map_err(|_| {
245                self.error(
246                    JudgmentErrorKind::Timeout,
247                    format!("no reply within {} ms", self.timeout.as_millis()),
248                )
249            })?
250            .map_err(|e| self.classify(&e))?;
251        Ok(Decisions {
252            answers: result.answers,
253            // Provider-authored: scrubbed of the key, then cut.
254            model: self
255                .key
256                .excerpt(&result.model, super::MAX_REPORTED_MODEL_CHARS),
257            usage: result.usage.map(with_total_tokens),
258        })
259    }
260
261    fn endpoint_host(&self) -> Option<String> {
262        JevJudge::endpoint_host(self)
263    }
264
265    fn model_id(&self) -> Option<String> {
266        Some(self.model_id.clone())
267    }
268}