1use 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
37const BODY_LOG_CHARS: usize = 200;
39
40fn 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
55fn 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#[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 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 pub fn model_id(&self) -> &str {
104 &self.model_id
105 }
106
107 pub fn endpoint(&self) -> &str {
109 &self.endpoint
110 }
111
112 pub fn endpoint_host(&self) -> Option<String> {
114 super::endpoint_host(&self.endpoint)
115 }
116
117 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 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 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 fn log_excerpt(&self, text: &str) -> String {
148 self.key
149 .excerpt(&self.redact_endpoint(text), BODY_LOG_CHARS)
150 }
151
152 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 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 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 _ => 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 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 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}