cratefield_core/ports/text_model.rs
1//! The `TextModel` port (issue #429): a text completion asked for by
2//! **tier**, never by vendor. A module says [`ModelTier::Fast`] for its
3//! drafting and [`ModelTier::Strong`] for its judging; which provider
4//! answers each tier is the venture's wiring, decided once in `src/lib.rs`
5//! and invisible to the module (ADR 0002). [`RoutingTextModel`] is the seam,
6//! the way [`RoutingPush`](crate::RoutingPush) is for transports.
7//!
8//! **Tools arrived in issue #665; streaming stays out of scope.** A
9//! [`Prompt`] may carry [`ToolSpec`]s and a [`ToolChoice`]; a
10//! [`Completion`] may carry the [`ToolCall`]s the model asked for; and a
11//! [`Turn`] can quote an assistant's requested calls and the
12//! [`ToolResult`]s a user turn carries back. An adapter that can carry
13//! tools says so through [`TextModel::supports`]; one that cannot is never
14//! asked — [`RoutingTextModel`] and [`run_tool_loop`](crate::run_tool_loop)
15//! refuse a tools-bearing prompt up front with
16//! [`TextModelError::Unsupported`]. Embeddings live on the separate
17//! [`Embedder`](crate::Embedder) port (issue #561), never here.
18//!
19//! **Structured output builds on the schema a prompt already carries.**
20//! [`TextModelExt`](crate::TextModelExt) wraps `complete` with
21//! `complete_json`/`complete_as`: it checks the schema on the tool loop's
22//! subset, asks for it through [`Prompt::json_schema`], validates the
23//! answer, and allows one repair retry — so a caller holds a value that
24//! conforms or a [`TextModelError::SchemaViolation`], never a partial one.
25//!
26//! Streaming is the deliberate omission, not a gap:
27//! `response_to_worker` buffers a whole harness response to
28//! `MAX_RESPONSE_BUFFER` (1 MiB, `crates/runtime-cloudflare/src/lib.rs`),
29//! so a streamed completion has nowhere to arrive on this runtime — a port
30//! that promised deltas would be a port the Workers twin could not keep.
31//! A tool loop still is one request and one buffered answer per step.
32//!
33//! **Image input arrived in issue #628.** A [`Turn`] may carry [`Part`]s —
34//! text and images in order — built with [`Turn::user_parts`]; a text-only
35//! turn is unchanged. An adapter that can carry images says so through
36//! [`Capability::Images`], and the router, the tool loop and every adapter
37//! call [`Prompt::check_images`] **before** any network call, so an
38//! over-limit prompt is refused locally rather than paid for. Adapters
39//! encode with [`encode_image`], needing no base64 dependency of their own.
40//!
41//! There is no outcome enum on this port, unlike [`Mailer`](crate::Mailer)
42//! and [`Push`](crate::Push), and that is deliberate: a completion has no
43//! "delivered but not configured" middle state — either text came back or
44//! nothing did. The unwired answer is therefore
45//! [`TextModelError::NotConfigured`], an error variant the caller can match,
46//! so a module that cannot degrade without its model fails loudly instead of
47//! silently producing nothing.
48
49use std::sync::Arc;
50use std::time::Duration;
51
52use async_trait::async_trait;
53use base64::Engine as _;
54use serde::{Deserialize, Serialize};
55use serde_json::Value;
56
57use super::http::MAX_RESPONSE_BYTES;
58
59/// The token ceiling a [`Prompt`] starts with: enough for a drafted reply,
60/// small enough that a forgotten `.max_tokens(..)` cannot turn into a run
61/// away bill. Explicit in the type so "how long can the answer get" is a
62/// field a caller can read, not an adapter's private default.
63pub const DEFAULT_MAX_TOKENS: u32 = 1024;
64
65/// The most images one [`Prompt`] may carry (issue #628): four, inside the
66/// vendors' own per-request range. [`Prompt::check_images`] refuses more.
67pub const MAX_PROMPT_IMAGES: usize = 4;
68
69/// The most base64 bytes one image may encode to (issue #628): 3.75 MiB,
70/// the vendors' per-image ceiling. [`encoded_image_len`] gives a raw
71/// image's encoded length; [`Prompt::check_images`] refuses any over it.
72pub const MAX_IMAGE_ENCODED_BYTES: usize = 3_932_160;
73
74/// The most base64 bytes **all** of a prompt's images may encode to
75/// together (issue #628): 4 MiB, the same [`MAX_RESPONSE_BYTES`] ceiling
76/// the [`HttpClient`](crate::HttpClient) port puts on a response body.
77/// [`Prompt::check_images`] refuses a total over it.
78pub const MAX_PROMPT_IMAGE_ENCODED_BYTES: usize = MAX_RESPONSE_BYTES;
79
80/// Which class of model a completion asks for — a **tier**, never a vendor
81/// or a model name. A venture maps each tier onto a provider in its own
82/// wiring, and can move drafting from one vendor to another without
83/// touching a module.
84#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
85#[serde(rename_all = "snake_case")]
86pub enum ModelTier {
87 /// The cheap, quick class: drafting, summarising, classifying.
88 Fast,
89 /// The best-available class: judging, long synthesis, the one call
90 /// where quality is the point.
91 Strong,
92}
93
94impl ModelTier {
95 /// The name used in errors and logs.
96 pub fn name(&self) -> &'static str {
97 match self {
98 ModelTier::Fast => "fast",
99 ModelTier::Strong => "strong",
100 }
101 }
102}
103
104impl std::fmt::Display for ModelTier {
105 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
106 f.write_str(self.name())
107 }
108}
109
110/// Who is speaking in a [`Turn`]. A prompt is a conversation, and a
111/// provider needs to know which side each part came from.
112#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
113#[serde(rename_all = "snake_case")]
114pub enum Role {
115 /// The human (or module) side of the conversation.
116 User,
117 /// The model's side, as a previous completion was recorded.
118 Assistant,
119}
120
121impl Role {
122 /// The name used in errors and logs.
123 pub fn name(&self) -> &'static str {
124 match self {
125 Role::User => "user",
126 Role::Assistant => "assistant",
127 }
128 }
129}
130
131impl std::fmt::Display for Role {
132 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
133 f.write_str(self.name())
134 }
135}
136
137/// Something a [`TextModel`] adapter may or may not be able to do beyond a
138/// plain completion. A port whose shape grows over time needs a way for the
139/// router to ask "can the model behind this tier actually do this" before a
140/// caller wastes a request on one that cannot.
141///
142/// `#[non_exhaustive]`: today the capabilities are [`Self::Tools`] and
143/// [`Self::Images`], and the next one — a provider-native structured
144/// output — should not be a breaking change for every `match` a caller
145/// writes.
146#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
147#[non_exhaustive]
148pub enum Capability {
149 /// The adapter carries tools: a prompt's [`Prompt::tools`] reach the
150 /// provider, and the answer may carry [`Completion::tool_calls`]. An
151 /// adapter that does not implement this is never asked to.
152 Tools,
153 /// The adapter carries images: a prompt's [`Part::Image`]s reach the
154 /// provider. An adapter without vision is refused with
155 /// [`TextModelError::Unsupported`], never sent an image to drop.
156 Images,
157}
158
159impl Capability {
160 /// The name used in errors and logs.
161 pub fn name(&self) -> &'static str {
162 match self {
163 Capability::Tools => "tools",
164 Capability::Images => "images",
165 }
166 }
167}
168
169impl std::fmt::Display for Capability {
170 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
171 f.write_str(self.name())
172 }
173}
174
175/// A tool the model is offered: a name, the instruction a model reads when
176/// deciding whether to call it, and the JSON Schema of the arguments object
177/// the call must carry (issue #665).
178///
179/// `#[non_exhaustive]`: build one with [`ToolSpec::new`] rather than a
180/// struct literal, the same rule [`Prompt`] and [`Completion`] follow —
181/// what a tool declares grows, and it should not break every caller.
182#[derive(Debug, Clone, PartialEq)]
183#[non_exhaustive]
184pub struct ToolSpec {
185 pub name: String,
186 pub description: String,
187 /// A JSON Schema (draft 2020-12) of the arguments object a call must
188 /// provide. [`run_tool_loop`](crate::run_tool_loop) validates it up
189 /// front and then validates every call's arguments against it, and
190 /// refuses a schema using a keyword it cannot honour rather than
191 /// silently skipping the check.
192 pub parameters: Value,
193}
194
195impl ToolSpec {
196 /// A tool named `name`, described by `description`, taking arguments
197 /// conforming to `parameters`.
198 #[must_use]
199 pub fn new(name: impl Into<String>, description: impl Into<String>, parameters: Value) -> Self {
200 Self {
201 name: name.into(),
202 description: description.into(),
203 parameters,
204 }
205 }
206}
207
208/// How the model should choose among a prompt's [`Prompt::tools`]: freely,
209/// not at all, mandatorily, or pinned to one named tool (issue #665).
210///
211/// `#[non_exhaustive]`: the wire protocols have more ways to steer this
212/// ("any tool but this one", a per-tool probability), and adding one must
213/// not break every caller.
214#[derive(Debug, Clone, PartialEq, Eq)]
215#[non_exhaustive]
216pub enum ToolChoice {
217 /// The model decides whether to call a tool or answer directly.
218 Auto,
219 /// The model must answer without calling a tool.
220 None,
221 /// The model must call some tool, whichever it picks.
222 Required,
223 /// The model must call the tool with this name.
224 Tool(String),
225}
226
227/// One tool call the model asked for: which call it is, which tool, and the
228/// arguments the model produced (issue #665). The arguments are exactly
229/// what the model sent — [`run_tool_loop`](crate::run_tool_loop) validates
230/// them against the tool's schema before any executor sees them.
231#[derive(Debug, Clone, PartialEq)]
232pub struct ToolCall {
233 pub id: String,
234 pub name: String,
235 pub arguments: Value,
236}
237
238impl ToolCall {
239 /// A call identified by `id` asking for the tool named `name` with
240 /// `arguments`.
241 #[must_use]
242 pub fn new(id: impl Into<String>, name: impl Into<String>, arguments: Value) -> Self {
243 Self {
244 id: id.into(),
245 name: name.into(),
246 arguments,
247 }
248 }
249}
250
251/// The result of running one [`ToolCall`], fed back to the model as
252/// context (issue #665). A tool-level failure is a [`ToolResult::error`],
253/// not a failed loop: the model is given the error text and may try again.
254#[derive(Debug, Clone, PartialEq, Eq)]
255pub struct ToolResult {
256 /// The [`ToolCall::id`] this answers.
257 pub tool_call_id: String,
258 pub content: String,
259 /// Whether `content` is a tool error the model should see rather than a
260 /// successful result.
261 pub is_error: bool,
262}
263
264impl ToolResult {
265 /// A successful result for the call `tool_call_id`.
266 #[must_use]
267 pub fn ok(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
268 Self {
269 tool_call_id: tool_call_id.into(),
270 content: content.into(),
271 is_error: false,
272 }
273 }
274
275 /// A failed result for the call `tool_call_id`: `content` is the error
276 /// the model is shown so it can recover.
277 #[must_use]
278 pub fn error(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
279 Self {
280 tool_call_id: tool_call_id.into(),
281 content: content.into(),
282 is_error: true,
283 }
284 }
285}
286
287/// The media type of an image a [`Part`] carries (issue #628): the closed
288/// set the providers accept, so a caller picks one rather than handing the
289/// provider a bare string to validate. `#[non_exhaustive]`, so a `match`
290/// carries a wildcard.
291#[derive(Debug, Clone, Copy, PartialEq, Eq)]
292#[non_exhaustive]
293pub enum ImageMediaType {
294 /// `image/png`.
295 Png,
296 /// `image/jpeg`.
297 Jpeg,
298 /// `image/gif`.
299 Gif,
300 /// `image/webp`.
301 Webp,
302}
303
304impl ImageMediaType {
305 /// The MIME type a provider expects on the wire.
306 #[must_use]
307 pub const fn as_str(self) -> &'static str {
308 match self {
309 ImageMediaType::Png => "image/png",
310 ImageMediaType::Jpeg => "image/jpeg",
311 ImageMediaType::Gif => "image/gif",
312 ImageMediaType::Webp => "image/webp",
313 }
314 }
315}
316
317/// One piece of a [`Turn`]: text, or an image with its media type and raw
318/// bytes (issue #628). A parts-bearing turn carries them **in order**;
319/// `#[non_exhaustive]`, so a part is built with [`Part::text`] and
320/// [`Part::image`] rather than a struct literal.
321#[derive(Debug, Clone, PartialEq, Eq)]
322#[non_exhaustive]
323pub enum Part {
324 /// A run of text.
325 Text(String),
326 /// An image: its media type and its raw (undecoded) bytes. Adapters
327 /// encode with [`encode_image`] when they put it on the wire.
328 Image {
329 media_type: ImageMediaType,
330 bytes: Vec<u8>,
331 },
332}
333
334impl Part {
335 /// A text part.
336 #[must_use]
337 pub fn text(text: impl Into<String>) -> Self {
338 Part::Text(text.into())
339 }
340
341 /// An image part of `media_type` carrying the raw `bytes`.
342 #[must_use]
343 pub fn image(media_type: ImageMediaType, bytes: impl Into<Vec<u8>>) -> Self {
344 Part::Image {
345 media_type,
346 bytes: bytes.into(),
347 }
348 }
349}
350
351/// The standard, padded base64 encoding of an image's raw bytes (issue
352/// #628) — the form every provider expects for inline image data. Here so
353/// an adapter needs no base64 dependency of its own.
354#[must_use]
355pub fn encode_image(bytes: &[u8]) -> String {
356 base64::engine::general_purpose::STANDARD.encode(bytes)
357}
358
359/// The length of [`encode_image`]'s output for `raw_len` raw bytes (issue
360/// #628): `4 * ceil(raw_len / 3)`, padding included. Pure arithmetic, so
361/// [`Prompt::check_images`] can size an image without encoding it — to
362/// refuse an over-limit image before the allocation, not after.
363#[must_use]
364pub const fn encoded_image_len(raw_len: usize) -> usize {
365 raw_len.div_ceil(3) * 4
366}
367
368/// One message of the conversation a [`Prompt`] carries.
369///
370/// `#[non_exhaustive]`: a turn grew its tool fields in issue #665 and its
371/// image [`Part`]s in issue #628 and will grow again, so it is built with
372/// [`Turn::user`], [`Turn::user_parts`], [`Turn::assistant`],
373/// [`Turn::assistant_tool_calls`] or [`Turn::tool_results`] rather than a
374/// struct literal.
375#[derive(Debug, Clone, PartialEq)]
376#[non_exhaustive]
377pub struct Turn {
378 pub role: Role,
379 pub content: String,
380 /// The turn's content as ordered [`Part`]s, where it carries any
381 /// (issue #628). When non-empty it **is** the whole content, in order,
382 /// and [`content`](Self::content) is empty; a text-only turn keeps it
383 /// empty so it serialises exactly as before images existed.
384 pub parts: Vec<Part>,
385 /// The calls an assistant turn asked for. Carried on a
386 /// [`Role::Assistant`] turn; empty on every other turn.
387 pub tool_calls: Vec<ToolCall>,
388 /// The results a user turn carries back for the previous assistant
389 /// turn's calls. Carried on a [`Role::User`] turn; empty otherwise.
390 pub tool_results: Vec<ToolResult>,
391}
392
393impl Turn {
394 /// A turn spoken by the [`Role::User`] side.
395 #[must_use]
396 pub fn user(content: impl Into<String>) -> Self {
397 Turn {
398 role: Role::User,
399 content: content.into(),
400 parts: Vec::new(),
401 tool_calls: Vec::new(),
402 tool_results: Vec::new(),
403 }
404 }
405
406 /// A user turn carrying `parts` — text and images in order — with an
407 /// empty [`content`](Self::content) (issue #628). The way a prompt
408 /// sends an image: providers expect text and image in one message, so
409 /// the parts travel together rather than across two turns.
410 #[must_use]
411 pub fn user_parts(parts: impl IntoIterator<Item = Part>) -> Self {
412 Turn {
413 role: Role::User,
414 content: String::new(),
415 parts: parts.into_iter().collect(),
416 tool_calls: Vec::new(),
417 tool_results: Vec::new(),
418 }
419 }
420
421 /// A turn spoken by the [`Role::Assistant`] side — a previous
422 /// completion, quoted back as context.
423 #[must_use]
424 pub fn assistant(content: impl Into<String>) -> Self {
425 Turn {
426 role: Role::Assistant,
427 content: content.into(),
428 parts: Vec::new(),
429 tool_calls: Vec::new(),
430 tool_results: Vec::new(),
431 }
432 }
433
434 /// An assistant turn that asked to call `calls`. `content` is the text
435 /// the model wrote alongside the calls (often empty).
436 #[must_use]
437 pub fn assistant_tool_calls(content: impl Into<String>, calls: Vec<ToolCall>) -> Self {
438 Turn {
439 role: Role::Assistant,
440 content: content.into(),
441 parts: Vec::new(),
442 tool_calls: calls,
443 tool_results: Vec::new(),
444 }
445 }
446
447 /// A user turn carrying the `results` of the previous assistant turn's
448 /// tool calls — the shape a provider expects the tool answers back in.
449 #[must_use]
450 pub fn tool_results(results: Vec<ToolResult>) -> Self {
451 Turn {
452 role: Role::User,
453 content: String::new(),
454 parts: Vec::new(),
455 tool_calls: Vec::new(),
456 tool_results: results,
457 }
458 }
459}
460
461/// One completion request: the tier asked for, the conversation so far, and
462/// what the caller will accept back.
463///
464/// `#[non_exhaustive]`: build one with [`Prompt::new`] and the builder
465/// methods rather than a struct literal. What a v2 port has to carry —
466/// tools, temperature, a stop sequence — should not be a breaking change
467/// for every caller, the same way [`Message`](crate::Message) is built.
468#[derive(Debug, Clone, PartialEq)]
469#[non_exhaustive]
470pub struct Prompt {
471 pub tier: ModelTier,
472 pub system: Option<String>,
473 pub messages: Vec<Turn>,
474 /// A JSON Schema (draft 2020-12) the answer should conform to. When set,
475 /// the adapter asks its provider for structured output; a provider that
476 /// honours it fills [`Completion::json`], and one that cannot leaves the
477 /// raw answer in [`Completion::text`] with `json` `None`, rather than
478 /// failing the call.
479 pub json_schema: Option<Value>,
480 pub max_tokens: u32,
481 /// The tools the model may call. Empty by default: a prompt that sends
482 /// no tools behaves exactly as it did before tools existed, and an
483 /// adapter is never asked to carry tools it was not given.
484 pub tools: Vec<ToolSpec>,
485 /// How the model should choose among [`Prompt::tools`]. `None` leaves
486 /// the choice to the provider's own default; the builder sets it.
487 pub tool_choice: Option<ToolChoice>,
488}
489
490impl Prompt {
491 /// An empty prompt for `tier`: no system prompt, no messages, no
492 /// schema, no tools, and [`DEFAULT_MAX_TOKENS`] as the ceiling.
493 /// Everything else is a builder method.
494 #[must_use]
495 pub fn new(tier: ModelTier) -> Self {
496 Self {
497 tier,
498 system: None,
499 messages: Vec::new(),
500 json_schema: None,
501 max_tokens: DEFAULT_MAX_TOKENS,
502 tools: Vec::new(),
503 tool_choice: None,
504 }
505 }
506
507 /// Offers the model one more tool it may call.
508 #[must_use]
509 pub fn tool(mut self, tool: ToolSpec) -> Self {
510 self.tools.push(tool);
511 self
512 }
513
514 /// Offers the model every tool in `tools`.
515 #[must_use]
516 pub fn tools(mut self, tools: impl IntoIterator<Item = ToolSpec>) -> Self {
517 self.tools.extend(tools);
518 self
519 }
520
521 /// Steers how the model chooses among the offered tools.
522 #[must_use]
523 pub fn tool_choice(mut self, tool_choice: ToolChoice) -> Self {
524 self.tool_choice = Some(tool_choice);
525 self
526 }
527
528 /// The standing instruction the model answers under.
529 #[must_use]
530 pub fn system(mut self, system: impl Into<String>) -> Self {
531 self.system = Some(system.into());
532 self
533 }
534
535 /// Appends a [`Turn::user`] — the common case, a one-message prompt.
536 #[must_use]
537 pub fn user(mut self, content: impl Into<String>) -> Self {
538 self.messages.push(Turn::user(content));
539 self
540 }
541
542 /// Appends a [`Turn::user_parts`] — a user message carrying text and
543 /// images in order (issue #628).
544 #[must_use]
545 pub fn user_parts(mut self, parts: impl IntoIterator<Item = Part>) -> Self {
546 self.messages.push(Turn::user_parts(parts));
547 self
548 }
549
550 /// Appends a [`Turn::assistant`].
551 #[must_use]
552 pub fn assistant(mut self, content: impl Into<String>) -> Self {
553 self.messages.push(Turn::assistant(content));
554 self
555 }
556
557 /// Appends a whole turn, for a conversation built elsewhere.
558 #[must_use]
559 pub fn turn(mut self, turn: Turn) -> Self {
560 self.messages.push(turn);
561 self
562 }
563
564 /// Asks for structured output conforming to `schema`; a successful
565 /// completion then carries the parsed value in [`Completion::json`].
566 #[must_use]
567 pub fn json_schema(mut self, schema: Value) -> Self {
568 self.json_schema = Some(schema);
569 self
570 }
571
572 /// The token ceiling for the answer.
573 #[must_use]
574 pub fn max_tokens(mut self, max_tokens: u32) -> Self {
575 self.max_tokens = max_tokens;
576 self
577 }
578
579 /// Whether any turn carries an image [`Part`] (issue #628). The router
580 /// and the tool loop gate [`Capability::Images`] on this, exactly as
581 /// they gate [`Capability::Tools`] on [`Prompt::tools`].
582 #[must_use]
583 pub fn has_images(&self) -> bool {
584 self.messages.iter().any(|turn| {
585 turn.parts
586 .iter()
587 .any(|part| matches!(part, Part::Image { .. }))
588 })
589 }
590
591 /// Refuses a prompt whose images exceed a bound (issue #628), without
592 /// encoding anything: the count against [`MAX_PROMPT_IMAGES`], each
593 /// image's [`encoded_image_len`] against [`MAX_IMAGE_ENCODED_BYTES`],
594 /// and their sum against [`MAX_PROMPT_IMAGE_ENCODED_BYTES`].
595 ///
596 /// # Errors
597 ///
598 /// [`TextModelError::ImageLimit`], naming which bound was crossed:
599 /// [`ImageLimit::TooMany`], [`ImageLimit::ImageTooLarge`] (with the
600 /// 0-based index among the prompt's images) or
601 /// [`ImageLimit::TotalTooLarge`].
602 pub fn check_images(&self) -> Result<(), TextModelError> {
603 let count = self
604 .messages
605 .iter()
606 .flat_map(|turn| &turn.parts)
607 .filter(|part| matches!(part, Part::Image { .. }))
608 .count();
609 if count > MAX_PROMPT_IMAGES {
610 return Err(TextModelError::ImageLimit(ImageLimit::TooMany { count }));
611 }
612
613 let mut total: usize = 0;
614 let mut index: usize = 0;
615 for turn in &self.messages {
616 for part in &turn.parts {
617 let Part::Image { bytes, .. } = part else {
618 continue;
619 };
620 let encoded_bytes = encoded_image_len(bytes.len());
621 if encoded_bytes > MAX_IMAGE_ENCODED_BYTES {
622 return Err(TextModelError::ImageLimit(ImageLimit::ImageTooLarge {
623 index,
624 encoded_bytes,
625 }));
626 }
627 total = total.saturating_add(encoded_bytes);
628 index += 1;
629 }
630 }
631 if total > MAX_PROMPT_IMAGE_ENCODED_BYTES {
632 return Err(TextModelError::ImageLimit(ImageLimit::TotalTooLarge {
633 encoded_bytes: total,
634 }));
635 }
636 Ok(())
637 }
638}
639
640/// A completed completion: the text, the model that answered (a provider
641/// identifier, for the log line), and the token usage.
642///
643/// The two token counts have a fixed relationship: [`Completion::input_tokens`]
644/// is the **total** prompt tokens the provider processed, and
645/// [`Completion::cached_input_tokens`] — where the provider reports one —
646/// is the subset of that total served from the provider's own prompt
647/// cache. A caller costing a completion never subtracts; the total already
648/// includes the cached part.
649///
650/// `#[non_exhaustive]` for the same reason [`Prompt`] is: what a provider
651/// reports back grows, and it should not break every caller.
652#[derive(Debug, Clone, PartialEq)]
653#[non_exhaustive]
654pub struct Completion {
655 pub text: String,
656 /// The parsed answer, when the prompt carried a
657 /// [`Prompt::json_schema`]. `None` when it did not, or the provider's
658 /// answer could not be parsed — in which case `text` still holds what
659 /// came back.
660 pub json: Option<Value>,
661 pub model: String,
662 /// The total prompt tokens the provider processed, cached tokens
663 /// included — a provider that reports caching reports them alongside
664 /// the uncached remainder, not inside it. The discount question is
665 /// [`Completion::cached_input_tokens`], never a smaller total here.
666 pub input_tokens: u64,
667 pub output_tokens: u64,
668 /// The subset of [`Completion::input_tokens`] the provider served from
669 /// its own prompt cache, where the provider reports one. `None` when it
670 /// does not: an absent report and a reported zero mean different things
671 /// (a vendor that does not do caching versus a cache miss), so the
672 /// option is never collapsed.
673 pub cached_input_tokens: Option<u64>,
674 /// The tools the model asked to call. Empty for a plain answer, and
675 /// empty for a prompt that offered no tools.
676 pub tool_calls: Vec<ToolCall>,
677}
678
679impl Completion {
680 /// A completion with just the text and the model that wrote it; the
681 /// usage, parsed JSON and tool calls are builder methods.
682 #[must_use]
683 pub fn new(text: impl Into<String>, model: impl Into<String>) -> Self {
684 Self {
685 text: text.into(),
686 json: None,
687 model: model.into(),
688 input_tokens: 0,
689 output_tokens: 0,
690 cached_input_tokens: None,
691 tool_calls: Vec::new(),
692 }
693 }
694
695 /// The tools the model asked to call, in the order it named them.
696 #[must_use]
697 pub fn tool_calls(mut self, tool_calls: Vec<ToolCall>) -> Self {
698 self.tool_calls = tool_calls;
699 self
700 }
701
702 /// The parsed structured answer, for a prompt that asked for one.
703 #[must_use]
704 pub fn json(mut self, json: Value) -> Self {
705 self.json = Some(json);
706 self
707 }
708
709 /// The token counts the provider reported.
710 #[must_use]
711 pub fn usage(mut self, input_tokens: u64, output_tokens: u64) -> Self {
712 self.input_tokens = input_tokens;
713 self.output_tokens = output_tokens;
714 self
715 }
716
717 /// The subset of the input tokens the provider served from its own
718 /// prompt cache, where the provider reports one.
719 /// [`Completion::input_tokens`] stays the total, cached included.
720 #[must_use]
721 pub fn cached_input_tokens(mut self, cached_input_tokens: u64) -> Self {
722 self.cached_input_tokens = Some(cached_input_tokens);
723 self
724 }
725}
726
727/// Which image bound a [`Prompt`] crossed (issue #628): the count, one
728/// image's encoded size, or the whole prompt's. Carries only sizes and
729/// indices — never user content — so it is safe in a log line as-is.
730///
731/// `#[non_exhaustive]`: another bound should not break every `match`.
732#[derive(Debug, Clone, Copy, PartialEq, Eq)]
733#[non_exhaustive]
734pub enum ImageLimit {
735 /// The prompt carries more than [`MAX_PROMPT_IMAGES`] images.
736 TooMany { count: usize },
737 /// The image at 0-based `index` among the prompt's images encodes to
738 /// more than [`MAX_IMAGE_ENCODED_BYTES`] base64 bytes.
739 ImageTooLarge { index: usize, encoded_bytes: usize },
740 /// The prompt's images encode to more than
741 /// [`MAX_PROMPT_IMAGE_ENCODED_BYTES`] base64 bytes together.
742 TotalTooLarge { encoded_bytes: usize },
743}
744
745impl std::fmt::Display for ImageLimit {
746 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
747 match self {
748 Self::TooMany { count } => write!(
749 f,
750 "a prompt may carry at most {MAX_PROMPT_IMAGES} images, but this one carries {count}"
751 ),
752 Self::ImageTooLarge {
753 index,
754 encoded_bytes,
755 } => write!(
756 f,
757 "image {index} encodes to {encoded_bytes} bytes, over the \
758 {MAX_IMAGE_ENCODED_BYTES}-byte per-image ceiling"
759 ),
760 Self::TotalTooLarge { encoded_bytes } => write!(
761 f,
762 "the prompt's images encode to {encoded_bytes} bytes, over the \
763 {MAX_PROMPT_IMAGE_ENCODED_BYTES}-byte total"
764 ),
765 }
766 }
767}
768
769/// Completion failures.
770///
771/// [`NotConfigured`](Self::NotConfigured) sits on the **error** enum here,
772/// unlike [`SendOutcome::NotConfigured`](crate::SendOutcome) and
773/// [`PushOutcome::NotConfigured`](crate::PushOutcome): a completion has no
774/// "delivered but not configured" middle state, so an unwired tier is an
775/// error the caller matches, not an outcome it inspects.
776///
777/// `Transient` deliberately carries **only** `retry_after`, where
778/// [`PushError::Transient`](crate::PushError) also carries a message: with
779/// no provider text of its own there is nothing to scrub, and provider text
780/// belongs on [`Rejected`](Self::Rejected) and [`Transport`](Self::Transport).
781///
782/// The two variants that carry provider text are sanitized in `Display`,
783/// the same way [`PushError`](crate::PushError)'s and
784/// [`MailError`](crate::MailError)'s are (issue #235). What an adapter wraps
785/// is the provider's own words, and a prompt is exactly the kind of value
786/// that rides back in them — a drafting module quotes a customer's note, and
787/// the provider's `4xx` quotes it straight back. `Display` therefore runs it
788/// through [`crate::logging::scrub_text`]; `Debug` still shows the raw
789/// string for tests.
790///
791/// `#[non_exhaustive]`: a port's failure set grows with its shape, and
792/// [`Unsupported`](Self::Unsupported) is exactly such a growth — a caller
793/// that matches the variants it knows and treats anything else as a
794/// transport failure should not have to recompile when one is added.
795#[derive(Debug, Clone, PartialEq, Eq)]
796#[non_exhaustive]
797pub enum TextModelError {
798 /// The tier asked for has no adapter — the venture did not wire it.
799 /// Nothing is wrong with the prompt: a module may degrade, the way it
800 /// degrades on a `NotConfigured` mailer, and
801 /// [`RoutingTextModel`] answers this for every tier it has no
802 /// adapter for.
803 NotConfigured,
804 /// The provider refused the request (a `4xx`), or refused the prompt's
805 /// content or schema; not retryable without a change.
806 Rejected(String),
807 /// A transient failure (a `5xx`, a `429`, a transport error): retry
808 /// later, and not before `retry_after` when the provider named one.
809 /// Carries no message — provider text belongs on
810 /// [`Rejected`](Self::Rejected) and [`Transport`](Self::Transport).
811 Transient { retry_after: Option<Duration> },
812 /// The request never completed as a conversation — the adapter could
813 /// not reach the provider, or the answer did not survive the hop.
814 Transport(String),
815 /// The model behind this tier cannot do what the prompt asked for —
816 /// a prompt carrying [`Prompt::tools`] or image [`Part`]s reached an
817 /// adapter whose [`TextModel::supports`] is false for
818 /// [`Capability::Tools`] or [`Capability::Images`], so the adapter was
819 /// never called. A fixed sentence, never provider text, so there is
820 /// nothing to scrub.
821 Unsupported(Capability),
822 /// The schema itself cannot be honoured: it uses a JSON Schema keyword
823 /// the built-in validator does not implement (a `$ref`, a `pattern`,
824 /// an `if`), or is malformed. This is a **caller bug** — the schema was
825 /// written by the caller, not the model — so the model was never
826 /// called. Raised by
827 /// [`TextModelExt::complete_json`](crate::TextModelExt::complete_json)
828 /// before the first request.
829 InvalidSchema(String),
830 /// The model's answer did not conform to the required schema, after the
831 /// one repair retry
832 /// [`TextModelExt::complete_json`](crate::TextModelExt::complete_json)
833 /// allows. Carries only the schema path and reason the answer failed
834 /// (for example `the property "price" is invalid: …`) — **never the
835 /// model's own output**, which may hold the personal data the schema
836 /// was there to structure.
837 SchemaViolation(String),
838 /// The prompt's images crossed a bound [`Prompt::check_images`]
839 /// enforces: too many, one too large, or too large together.
840 ImageLimit(ImageLimit),
841}
842
843impl std::fmt::Display for TextModelError {
844 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
845 let scrub = crate::logging::scrub_text;
846 match self {
847 Self::NotConfigured => f.write_str("no text model is wired for this tier"),
848 Self::Rejected(message) => write!(f, "completion rejected: {}", scrub(message)),
849 Self::Transient { .. } => f.write_str("completion failed, retryable"),
850 Self::Transport(message) => {
851 write!(f, "completion transport failed: {}", scrub(message))
852 }
853 Self::Unsupported(capability) => {
854 write!(f, "the text model does not support {capability}")
855 }
856 Self::InvalidSchema(reason) => {
857 write!(f, "the JSON schema cannot be honoured: {}", scrub(reason))
858 }
859 Self::SchemaViolation(reason) => write!(
860 f,
861 "the model's answer did not match the schema: {}",
862 scrub(reason)
863 ),
864 Self::ImageLimit(limit) => write!(f, "the prompt's images exceed a limit: {limit}"),
865 }
866 }
867}
868
869impl std::error::Error for TextModelError {}
870
871impl TextModelError {
872 /// How long the provider asked the caller to wait, where it said. An
873 /// [`Unsupported`](Self::Unsupported) capability is not a back-off: no
874 /// wait makes a model grow a capability it lacks.
875 pub fn retry_after(&self) -> Option<Duration> {
876 match self {
877 TextModelError::Transient { retry_after } => *retry_after,
878 _ => None,
879 }
880 }
881}
882
883/// Completes a prompt, over whichever provider the venture wired for the
884/// [`Prompt::tier`] it was asked for. An adapter serves any tier it is
885/// configured for (typically one); [`RoutingTextModel`] dispatches between
886/// the adapters a venture configured.
887#[async_trait]
888pub trait TextModel: Send + Sync {
889 /// Completes `prompt`.
890 ///
891 /// [`Prompt::json_schema`] is a request: an adapter whose provider
892 /// cannot honour structured output answers with plain text in
893 /// [`Completion::text`] and `None` in [`Completion::json`], rather than
894 /// failing the call.
895 async fn complete(&self, prompt: &Prompt) -> Result<Completion, TextModelError>;
896
897 /// Whether this adapter can do `capability` for a prompt routed to
898 /// `tier`. Defaults to `false`: a model that has not opted in is
899 /// refused a tools-bearing prompt by [`RoutingTextModel`] and
900 /// [`run_tool_loop`](crate::run_tool_loop) before it is ever called,
901 /// rather than sent a request it would silently drop the tools from.
902 /// An adapter that carries [`Capability::Tools`] or
903 /// [`Capability::Images`] overrides this to answer `true` for it.
904 ///
905 /// `tier` is the router's own routing key, not the adapter's: an
906 /// adapter serves whatever tier it was wired for and is free to ignore
907 /// it — an adapter only honours `tier` if the same deployment could be
908 /// wired for two tiers with different capabilities. The parameter is
909 /// on the method so a caller never has to ask the router which model a
910 /// tier holds just to ask what it can do.
911 fn supports(&self, _tier: ModelTier, _capability: Capability) -> bool {
912 false
913 }
914}
915
916/// Dispatches by [`Prompt::tier`] to the adapter a venture configured for
917/// that tier, so venture code holds one `Arc<dyn TextModel>` and never
918/// matches on the tier or the vendor itself (ADR 0002). The router lives
919/// above the adapters, in core — the same place
920/// [`RoutingPush`](crate::RoutingPush) does, and for the same reason
921/// (ADR 0015).
922///
923/// This is the seam that lets a venture put drafting on one vendor and an
924/// independent judge on another without either module knowing.
925///
926/// A tier with no adapter is [`TextModelError::NotConfigured`] —
927/// deliberately not [`TextModelError::Rejected`]: nothing is wrong with the
928/// prompt, the venture simply did not wire that tier.
929///
930/// The router also enforces capability agreement (issue #665): a prompt
931/// that carries [`Prompt::tools`] is refused with
932/// [`TextModelError::Unsupported`] before a tier that does not report
933/// [`Capability::Tools`] for that tier is called, so a venture can assert
934/// once at compose time — `router.supports(ModelTier::Strong,
935/// Capability::Tools)` — that the tier it routed tools to can carry them.
936/// [`RoutingTextModel::supports`] answers for the single routed tier: the
937/// adapter wired for `tier`, or `false` when that tier is unwired. It is
938/// deliberately **not** an aggregate over every wired tier: a prompt goes
939/// to exactly one tier, so what matters is that tier's answer, and a
940/// strong-tier judge that cannot carry tools must not make a fast-tier
941/// drafting model look incapable.
942///
943/// The same agreement covers images (issue #628): a prompt carrying a
944/// [`Part::Image`] is refused with [`TextModelError::Unsupported`] before a
945/// tier that does not report [`Capability::Images`] is called, and an
946/// over-limit prompt by [`Prompt::check_images`] before the tier is looked
947/// up at all.
948#[derive(Default, Clone)]
949pub struct RoutingTextModel {
950 fast: Option<Arc<dyn TextModel>>,
951 strong: Option<Arc<dyn TextModel>>,
952}
953
954impl RoutingTextModel {
955 /// A router with no adapters: every tier is `NotConfigured`.
956 #[must_use]
957 pub fn new() -> Self {
958 Self::default()
959 }
960
961 /// The adapter for [`ModelTier::Fast`].
962 #[must_use]
963 pub fn fast(mut self, model: Arc<dyn TextModel>) -> Self {
964 self.fast = Some(model);
965 self
966 }
967
968 /// The adapter for [`ModelTier::Strong`].
969 #[must_use]
970 pub fn strong(mut self, model: Arc<dyn TextModel>) -> Self {
971 self.strong = Some(model);
972 self
973 }
974
975 /// The adapter that serves `tier`, if one is configured.
976 #[must_use]
977 pub fn route_for(&self, tier: ModelTier) -> Option<&Arc<dyn TextModel>> {
978 match tier {
979 ModelTier::Fast => self.fast.as_ref(),
980 ModelTier::Strong => self.strong.as_ref(),
981 }
982 }
983}
984
985impl std::fmt::Debug for RoutingTextModel {
986 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
987 f.debug_struct("RoutingTextModel")
988 .field("fast", &self.fast.is_some())
989 .field("strong", &self.strong.is_some())
990 .finish()
991 }
992}
993
994#[async_trait]
995impl TextModel for RoutingTextModel {
996 async fn complete(&self, prompt: &Prompt) -> Result<Completion, TextModelError> {
997 // An over-limit image prompt is refused before anything is routed:
998 // the limits are about the request a provider would be handed, so
999 // they hold whatever tier answers.
1000 prompt.check_images()?;
1001 match self.route_for(prompt.tier) {
1002 Some(model) => {
1003 // A tools-bearing prompt to a model that cannot carry tools
1004 // is refused here, not silently sent with the tools
1005 // dropped: the caller asked for something the wiring cannot
1006 // honour, and a fabricated plain answer would hide that.
1007 if !prompt.tools.is_empty() && !model.supports(prompt.tier, Capability::Tools) {
1008 return Err(TextModelError::Unsupported(Capability::Tools));
1009 }
1010 // The same rule for images: a prompt carrying one to a
1011 // model without vision is refused, never sent with the image
1012 // silently dropped.
1013 if prompt.has_images() && !model.supports(prompt.tier, Capability::Images) {
1014 return Err(TextModelError::Unsupported(Capability::Images));
1015 }
1016 model.complete(prompt).await
1017 }
1018 None => Err(TextModelError::NotConfigured),
1019 }
1020 }
1021
1022 fn supports(&self, tier: ModelTier, capability: Capability) -> bool {
1023 // The routed tier's own answer, and `false` for an unwired tier:
1024 // nothing is wired to do anything. A prompt goes to one tier, so
1025 // this is the tier that must be able to carry the capability — not
1026 // an aggregate over the tiers the router happens to hold.
1027 self.route_for(tier)
1028 .is_some_and(|model| model.supports(tier, capability))
1029 }
1030}
1031
1032#[cfg(test)]
1033mod tests {
1034 use super::*;
1035
1036 // -----------------------------------------------------------------
1037 // ModelTier, Role, and the wire form
1038
1039 #[test]
1040 fn a_tier_and_a_role_round_trip_through_json() {
1041 for tier in [ModelTier::Fast, ModelTier::Strong] {
1042 let json = serde_json::to_string(&tier).expect("serialises");
1043 let back: ModelTier = serde_json::from_str(&json).expect("deserialises");
1044 assert_eq!(tier, back);
1045 }
1046 for role in [Role::User, Role::Assistant] {
1047 let json = serde_json::to_string(&role).expect("serialises");
1048 let back: Role = serde_json::from_str(&json).expect("deserialises");
1049 assert_eq!(role, back);
1050 }
1051 }
1052
1053 #[test]
1054 fn the_wire_form_is_snake_case() {
1055 // The point of `rename_all`: the persisted name is the prose name,
1056 // so a config file reads `"strong"`, not `"Strong"`.
1057 assert_eq!(serde_json::to_value(ModelTier::Strong).unwrap(), "strong");
1058 assert_eq!(serde_json::to_value(Role::User).unwrap(), "user");
1059 }
1060
1061 #[test]
1062 fn the_tier_names_itself_for_logs_and_errors() {
1063 assert_eq!(ModelTier::Fast.name(), "fast");
1064 assert_eq!(ModelTier::Strong.to_string(), "strong");
1065 assert_eq!(Role::Assistant.name(), "assistant");
1066 assert_eq!(Role::User.to_string(), "user");
1067 }
1068
1069 // -----------------------------------------------------------------
1070 // Prompt and Completion builders
1071
1072 #[test]
1073 fn a_prompt_starts_empty_and_builds_up() {
1074 let prompt = Prompt::new(ModelTier::Fast);
1075 assert_eq!(prompt.tier, ModelTier::Fast);
1076 assert_eq!(prompt.system, None);
1077 assert!(prompt.messages.is_empty());
1078 assert_eq!(prompt.json_schema, None);
1079 assert_eq!(prompt.max_tokens, DEFAULT_MAX_TOKENS);
1080 assert_eq!(DEFAULT_MAX_TOKENS, 1024);
1081
1082 let prompt = prompt
1083 .system("Draft the reply.")
1084 .user("Hello")
1085 .assistant("Hi there")
1086 .turn(Turn::user("And again"))
1087 .json_schema(serde_json::json!({ "type": "object" }))
1088 .max_tokens(256);
1089
1090 assert_eq!(prompt.system.as_deref(), Some("Draft the reply."));
1091 assert_eq!(
1092 prompt.messages,
1093 vec![
1094 Turn::user("Hello"),
1095 Turn::assistant("Hi there"),
1096 Turn::user("And again"),
1097 ]
1098 );
1099 assert_eq!(
1100 prompt.json_schema,
1101 Some(serde_json::json!({ "type": "object" }))
1102 );
1103 assert_eq!(prompt.max_tokens, 256);
1104 }
1105
1106 #[test]
1107 fn a_completion_starts_with_text_and_model_and_builds_up() {
1108 let completion = Completion::new("the answer", "vendor-1");
1109 assert_eq!(completion.text, "the answer");
1110 assert_eq!(completion.model, "vendor-1");
1111 assert_eq!(completion.json, None);
1112 assert_eq!(completion.input_tokens, 0);
1113 assert_eq!(completion.output_tokens, 0);
1114
1115 let completion = completion
1116 .json(serde_json::json!({ "reply": "the answer" }))
1117 .usage(12, 34);
1118 assert_eq!(
1119 completion.json,
1120 Some(serde_json::json!({ "reply": "the answer" }))
1121 );
1122 assert_eq!(completion.input_tokens, 12);
1123 assert_eq!(completion.output_tokens, 34);
1124 }
1125
1126 #[test]
1127 fn cached_input_tokens_defaults_to_absent_and_never_shrinks_the_total() {
1128 // A vendor that says nothing about caching stays `None` — a
1129 // reported zero (a cache miss) and an absent report (a vendor
1130 // without caching) are different facts.
1131 let completion = Completion::new("the answer", "vendor-1");
1132 assert_eq!(completion.cached_input_tokens, None);
1133
1134 // The builder records the cache read without shrinking the total:
1135 // `input_tokens` is what the vendor processed, cached included.
1136 let completion = Completion::new("the answer", "vendor-1")
1137 .usage(17, 5)
1138 .cached_input_tokens(9);
1139 assert_eq!(completion.input_tokens, 17);
1140 assert_eq!(completion.cached_input_tokens, Some(9));
1141 }
1142
1143 // -----------------------------------------------------------------
1144 // TextModelError
1145
1146 #[test]
1147 fn display_sanitizes_the_provider_text() {
1148 // Issue #235. The text an adapter wraps is the provider's own
1149 // words, and a drafting prompt quotes whatever a user wrote: what
1150 // the provider echoes back must not survive into a log line or a
1151 // dead-letter row carrying its URLs, tokens or addresses.
1152 let error = TextModelError::Rejected(
1153 "provider 400 for https://api.example.test/v1/complete?token=secret-abcdef".to_owned(),
1154 );
1155 let text = error.to_string();
1156 assert!(text.contains("completion rejected"), "{text}");
1157 assert!(!text.contains("secret-abcdef"), "{text}");
1158 assert!(text.contains("?[redacted]"), "{text}");
1159
1160 let error = TextModelError::Transport("timeout quoting alice@example.test".to_owned());
1161 let text = error.to_string();
1162 assert!(text.contains("completion transport failed"), "{text}");
1163 assert!(!text.contains('@'), "{text}");
1164
1165 // `Debug` still shows the raw string for a failing test to read.
1166 assert!(format!("{error:?}").contains("alice@example.test"));
1167 }
1168
1169 #[test]
1170 fn transient_says_retryable_and_carries_no_text_to_scrub() {
1171 // `Transient` has no message field on purpose — the fixed sentence
1172 // is the whole `Display`, and a provider's words would be an
1173 // unsanitised leak by construction.
1174 let error = TextModelError::Transient {
1175 retry_after: Some(Duration::from_secs(30)),
1176 };
1177 assert_eq!(error.to_string(), "completion failed, retryable");
1178 }
1179
1180 #[test]
1181 fn transient_carries_an_optional_retry_after() {
1182 let throttled = TextModelError::Transient {
1183 retry_after: Some(Duration::from_secs(30)),
1184 };
1185 assert_eq!(throttled.retry_after(), Some(Duration::from_secs(30)));
1186 let plain = TextModelError::Transient { retry_after: None };
1187 assert_eq!(plain.retry_after(), None);
1188 assert_eq!(TextModelError::NotConfigured.retry_after(), None);
1189 assert_eq!(
1190 TextModelError::Rejected("422".to_owned()).retry_after(),
1191 None,
1192 "a rejection is not a back-off"
1193 );
1194 assert_eq!(
1195 TextModelError::Transport("timed out".to_owned()).retry_after(),
1196 None
1197 );
1198 }
1199
1200 #[test]
1201 fn not_configured_names_the_missing_wiring() {
1202 assert_eq!(
1203 TextModelError::NotConfigured.to_string(),
1204 "no text model is wired for this tier"
1205 );
1206 }
1207
1208 #[test]
1209 fn the_structured_output_failures_name_themselves_and_never_back_off() {
1210 let invalid =
1211 TextModelError::InvalidSchema("the `pattern` keyword is not supported".to_owned());
1212 assert_eq!(
1213 invalid.to_string(),
1214 "the JSON schema cannot be honoured: the `pattern` keyword is not supported"
1215 );
1216 let violation =
1217 TextModelError::SchemaViolation("the property \"price\" is invalid".to_owned());
1218 assert!(
1219 violation.to_string().contains("did not match the schema"),
1220 "{violation}"
1221 );
1222 assert_eq!(invalid.retry_after(), None);
1223 assert_eq!(violation.retry_after(), None);
1224 }
1225
1226 // -----------------------------------------------------------------
1227 // RoutingTextModel
1228
1229 struct Recording {
1230 label: &'static str,
1231 tools: bool,
1232 images: bool,
1233 seen: std::sync::atomic::AtomicUsize,
1234 }
1235
1236 impl Recording {
1237 /// A model that reports no capability — the default for an adapter
1238 /// that has not opted into tools.
1239 fn new(label: &'static str) -> Arc<Self> {
1240 Arc::new(Self {
1241 label,
1242 tools: false,
1243 images: false,
1244 seen: std::sync::atomic::AtomicUsize::new(0),
1245 })
1246 }
1247
1248 /// A model that reports [`Capability::Tools`].
1249 fn tool_capable(label: &'static str) -> Arc<Self> {
1250 Arc::new(Self {
1251 label,
1252 tools: true,
1253 images: false,
1254 seen: std::sync::atomic::AtomicUsize::new(0),
1255 })
1256 }
1257
1258 /// A model that reports [`Capability::Images`].
1259 fn image_capable(label: &'static str) -> Arc<Self> {
1260 Arc::new(Self {
1261 label,
1262 tools: false,
1263 images: true,
1264 seen: std::sync::atomic::AtomicUsize::new(0),
1265 })
1266 }
1267
1268 fn count(&self) -> usize {
1269 self.seen.load(std::sync::atomic::Ordering::Relaxed)
1270 }
1271 }
1272
1273 #[async_trait]
1274 impl TextModel for Recording {
1275 async fn complete(&self, _prompt: &Prompt) -> Result<Completion, TextModelError> {
1276 self.seen.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1277 Ok(Completion::new("recorded", self.label))
1278 }
1279
1280 fn supports(&self, _tier: ModelTier, capability: Capability) -> bool {
1281 (capability == Capability::Tools && self.tools)
1282 || (capability == Capability::Images && self.images)
1283 }
1284 }
1285
1286 #[test]
1287 fn the_router_dispatches_by_tier() {
1288 let fast = Recording::new("fast-vendor");
1289 let strong = Recording::new("strong-vendor");
1290 let router = RoutingTextModel::new()
1291 .fast(fast.clone())
1292 .strong(strong.clone());
1293
1294 let completion =
1295 pollster::block_on(router.complete(&Prompt::new(ModelTier::Fast).user("hi"))).unwrap();
1296 assert_eq!(
1297 completion.model, "fast-vendor",
1298 "the fast tier reached its adapter"
1299 );
1300 let completion =
1301 pollster::block_on(router.complete(&Prompt::new(ModelTier::Strong).user("hi")))
1302 .unwrap();
1303 assert_eq!(
1304 completion.model, "strong-vendor",
1305 "the strong tier reached its own, different adapter"
1306 );
1307 assert_eq!(fast.count(), 1);
1308 assert_eq!(strong.count(), 1);
1309 }
1310
1311 #[test]
1312 fn a_tier_with_no_adapter_is_not_configured_not_rejected() {
1313 let router = RoutingTextModel::new().fast(Recording::new("fast-vendor"));
1314 let error = pollster::block_on(router.complete(&Prompt::new(ModelTier::Strong).user("hi")))
1315 .unwrap_err();
1316 assert_eq!(error, TextModelError::NotConfigured);
1317 assert!(router.route_for(ModelTier::Strong).is_none());
1318 }
1319
1320 #[test]
1321 fn an_empty_router_is_not_configured_for_every_tier() {
1322 let router = RoutingTextModel::new();
1323 for tier in [ModelTier::Fast, ModelTier::Strong] {
1324 let error =
1325 pollster::block_on(router.complete(&Prompt::new(tier).user("hi"))).unwrap_err();
1326 assert_eq!(error, TextModelError::NotConfigured);
1327 assert!(router.route_for(tier).is_none());
1328 }
1329 }
1330
1331 #[test]
1332 fn debug_prints_which_tiers_are_wired_and_nothing_else() {
1333 // The adapters themselves are `Arc<dyn TextModel>` and have no
1334 // meaningful `Debug`; printing their presence is the whole report.
1335 let router = RoutingTextModel::new().fast(Recording::new("fast-vendor"));
1336 let printed = format!("{router:?}");
1337 assert!(printed.contains("RoutingTextModel"), "{printed}");
1338 assert!(printed.contains("fast: true"), "{printed}");
1339 assert!(printed.contains("strong: false"), "{printed}");
1340 }
1341
1342 // -----------------------------------------------------------------
1343 // Tools (issue #665)
1344
1345 #[test]
1346 fn a_capability_names_itself_for_logs_and_errors() {
1347 assert_eq!(Capability::Tools.name(), "tools");
1348 assert_eq!(Capability::Tools.to_string(), "tools");
1349 assert_eq!(Capability::Images.name(), "images");
1350 assert_eq!(Capability::Images.to_string(), "images");
1351 }
1352
1353 #[test]
1354 fn a_tool_spec_carries_its_name_description_and_schema() {
1355 let tool = ToolSpec::new(
1356 "lookup",
1357 "Look a thing up.",
1358 serde_json::json!({ "type": "object" }),
1359 );
1360 assert_eq!(tool.name, "lookup");
1361 assert_eq!(tool.description, "Look a thing up.");
1362 assert_eq!(tool.parameters, serde_json::json!({ "type": "object" }));
1363 }
1364
1365 #[test]
1366 fn a_tool_call_and_result_round_trip_their_fields() {
1367 let call = ToolCall::new("call-1", "lookup", serde_json::json!({ "q": "x" }));
1368 assert_eq!(call.id, "call-1");
1369 assert_eq!(call.name, "lookup");
1370 assert_eq!(call.arguments, serde_json::json!({ "q": "x" }));
1371
1372 let ok = ToolResult::ok("call-1", "the answer");
1373 assert_eq!(ok.tool_call_id, "call-1");
1374 assert_eq!(ok.content, "the answer");
1375 assert!(!ok.is_error);
1376
1377 let error = ToolResult::error("call-1", "it broke");
1378 assert!(error.is_error);
1379 assert_eq!(error.content, "it broke");
1380 }
1381
1382 #[test]
1383 fn a_turn_can_quote_tool_calls_and_carry_results_back() {
1384 let call = ToolCall::new("call-1", "lookup", serde_json::json!({}));
1385 let assistant = Turn::assistant_tool_calls("thinking", vec![call.clone()]);
1386 assert_eq!(assistant.role, Role::Assistant);
1387 assert_eq!(assistant.content, "thinking");
1388 assert_eq!(assistant.tool_calls, vec![call]);
1389 assert!(assistant.tool_results.is_empty());
1390
1391 let results = Turn::tool_results(vec![ToolResult::ok("call-1", "fine")]);
1392 assert_eq!(results.role, Role::User);
1393 assert_eq!(results.content, "");
1394 assert!(results.tool_calls.is_empty());
1395 assert_eq!(results.tool_results.len(), 1);
1396
1397 // The plain constructors still carry neither list, so an existing
1398 // caller that builds a prompt with them is unaffected.
1399 assert!(Turn::user("hi").tool_calls.is_empty());
1400 assert!(Turn::assistant("hi").tool_results.is_empty());
1401 }
1402
1403 #[test]
1404 fn a_prompt_starts_without_tools_and_builds_them_up() {
1405 let prompt = Prompt::new(ModelTier::Fast);
1406 assert!(prompt.tools.is_empty());
1407 assert_eq!(prompt.tool_choice, None);
1408
1409 let prompt = prompt
1410 .tool(ToolSpec::new("a", "A", serde_json::json!({})))
1411 .tools(vec![
1412 ToolSpec::new("b", "B", serde_json::json!({})),
1413 ToolSpec::new("c", "C", serde_json::json!({})),
1414 ])
1415 .tool_choice(ToolChoice::Required);
1416 assert_eq!(prompt.tools.len(), 3);
1417 assert_eq!(
1418 prompt
1419 .tools
1420 .iter()
1421 .map(|tool| tool.name.as_str())
1422 .collect::<Vec<_>>(),
1423 vec!["a", "b", "c"]
1424 );
1425 assert_eq!(prompt.tool_choice, Some(ToolChoice::Required));
1426 }
1427
1428 #[test]
1429 fn a_completion_starts_without_tool_calls_and_builds_them_up() {
1430 let completion = Completion::new("the answer", "vendor-1");
1431 assert!(completion.tool_calls.is_empty());
1432 let completion = completion.tool_calls(vec![ToolCall::new(
1433 "call-1",
1434 "lookup",
1435 serde_json::json!({}),
1436 )]);
1437 assert_eq!(completion.tool_calls.len(), 1);
1438 }
1439
1440 #[test]
1441 fn an_unsupported_capability_names_what_is_missing_and_never_backs_off() {
1442 let error = TextModelError::Unsupported(Capability::Tools);
1443 assert_eq!(error.to_string(), "the text model does not support tools");
1444 assert_eq!(error.retry_after(), None);
1445 }
1446
1447 #[test]
1448 fn the_default_capability_answer_is_no() {
1449 // An adapter that has not opted into tools reports none, so the
1450 // router and the tool loop refuse a tools-bearing prompt before it
1451 // is called. An adapter ignores the tier it is asked about.
1452 assert!(!Recording::new("plain").supports(ModelTier::Fast, Capability::Tools));
1453 assert!(Recording::tool_capable("capable").supports(ModelTier::Strong, Capability::Tools));
1454 }
1455
1456 #[test]
1457 fn tools_to_a_tier_that_cannot_carry_them_are_refused_before_the_call() {
1458 let fast = Recording::new("fast-vendor");
1459 let router = RoutingTextModel::new().fast(fast.clone());
1460 let prompt = Prompt::new(ModelTier::Fast).user("hi").tool(ToolSpec::new(
1461 "lookup",
1462 "Look up.",
1463 serde_json::json!({}),
1464 ));
1465
1466 let error = pollster::block_on(router.complete(&prompt)).unwrap_err();
1467 assert_eq!(error, TextModelError::Unsupported(Capability::Tools));
1468 assert_eq!(
1469 fast.count(),
1470 0,
1471 "the adapter was never called, tools and all"
1472 );
1473 }
1474
1475 #[test]
1476 fn tools_to_a_capable_tier_pass_through_untouched() {
1477 let fast = Recording::tool_capable("fast-vendor");
1478 let router = RoutingTextModel::new().fast(fast.clone());
1479 let prompt = Prompt::new(ModelTier::Fast).tool(ToolSpec::new(
1480 "lookup",
1481 "Look up.",
1482 serde_json::json!({}),
1483 ));
1484
1485 let completion = pollster::block_on(router.complete(&prompt)).unwrap();
1486 assert_eq!(completion.model, "fast-vendor");
1487 assert_eq!(fast.count(), 1);
1488 }
1489
1490 #[test]
1491 fn a_prompt_without_tools_reaches_a_plain_model_as_before() {
1492 // The capability gate is only for tools: a prompt carrying none is
1493 // served by a model that reports no capabilities, unchanged.
1494 let fast = Recording::new("fast-vendor");
1495 let router = RoutingTextModel::new().fast(fast.clone());
1496 let completion =
1497 pollster::block_on(router.complete(&Prompt::new(ModelTier::Fast).user("hi"))).unwrap();
1498 assert_eq!(completion.model, "fast-vendor");
1499 assert_eq!(fast.count(), 1);
1500 }
1501
1502 #[test]
1503 fn the_router_answers_a_capability_per_tier_and_false_when_unwired() {
1504 let capable = Recording::tool_capable("capable");
1505 let plain = Recording::new("plain");
1506 let router = RoutingTextModel::new()
1507 .fast(capable.clone())
1508 .strong(plain.clone());
1509
1510 // The routed tier's own answer: a prompt goes to one tier, so a
1511 // capable fast tier reads capable even when the strong tier is not,
1512 // and vice versa.
1513 assert!(router.supports(ModelTier::Fast, Capability::Tools));
1514 assert!(!router.supports(ModelTier::Strong, Capability::Tools));
1515
1516 // Nothing wired for a tier: no capability, whatever the other tier
1517 // holds.
1518 let unwired = RoutingTextModel::new();
1519 assert!(!unwired.supports(ModelTier::Fast, Capability::Tools));
1520 assert!(!unwired.supports(ModelTier::Strong, Capability::Tools));
1521 }
1522
1523 // -----------------------------------------------------------------
1524 // Images (issue #628)
1525
1526 /// A raw byte length whose standard, padded base64 encoding is exactly
1527 /// `encoded`, for `encoded` a multiple of four.
1528 fn raw_len_for_encoded(encoded: usize) -> usize {
1529 (encoded / 4) * 3
1530 }
1531
1532 #[test]
1533 fn an_image_media_type_names_its_wire_mime() {
1534 assert_eq!(ImageMediaType::Png.as_str(), "image/png");
1535 assert_eq!(ImageMediaType::Jpeg.as_str(), "image/jpeg");
1536 assert_eq!(ImageMediaType::Gif.as_str(), "image/gif");
1537 assert_eq!(ImageMediaType::Webp.as_str(), "image/webp");
1538 }
1539
1540 #[test]
1541 fn encode_image_is_standard_padded_base64_and_encoded_image_len_predicts_it() {
1542 // The canonical RFC 4648 test vectors, padded.
1543 assert_eq!(encode_image(b""), "");
1544 assert_eq!(encode_image(b"f"), "Zg==");
1545 assert_eq!(encode_image(b"fo"), "Zm8=");
1546 assert_eq!(encode_image(b"foo"), "Zm9v");
1547 assert_eq!(encode_image(b"foob"), "Zm9vYg==");
1548 assert_eq!(encode_image(b"fooba"), "Zm9vYmE=");
1549 assert_eq!(encode_image(b"foobar"), "Zm9vYmFy");
1550
1551 // `encoded_image_len` is pure arithmetic and agrees with the
1552 // encoding for every remainder class.
1553 for raw in 0..=8usize {
1554 assert_eq!(
1555 encoded_image_len(raw),
1556 encode_image(&vec![0u8; raw]).len(),
1557 "raw length {raw}"
1558 );
1559 }
1560 assert_eq!(encoded_image_len(0), 0);
1561 assert_eq!(encoded_image_len(1), 4);
1562 assert_eq!(encoded_image_len(3), 4);
1563 assert_eq!(encoded_image_len(4), 8);
1564 }
1565
1566 #[test]
1567 fn a_parts_turn_carries_its_parts_in_order_and_a_text_turn_carries_none() {
1568 let with_image = Prompt::new(ModelTier::Fast).user_parts([
1569 Part::text("look"),
1570 Part::image(ImageMediaType::Jpeg, vec![1, 2, 3]),
1571 ]);
1572 let turn = &with_image.messages[0];
1573 assert_eq!(turn.role, Role::User);
1574 assert_eq!(turn.content, "", "the parts are the whole content");
1575 assert_eq!(
1576 turn.parts,
1577 vec![
1578 Part::Text("look".to_owned()),
1579 Part::Image {
1580 media_type: ImageMediaType::Jpeg,
1581 bytes: vec![1, 2, 3],
1582 },
1583 ]
1584 );
1585 assert!(with_image.has_images());
1586 assert!(
1587 !Prompt::new(ModelTier::Fast)
1588 .user_parts([Part::text("no image here")])
1589 .has_images()
1590 );
1591
1592 // The compatibility promise: a turn built the old way carries no
1593 // parts, so an adapter serialises it exactly as before images.
1594 let text_only = Prompt::new(ModelTier::Fast).user("hi");
1595 assert!(text_only.messages[0].parts.is_empty());
1596 assert!(!text_only.has_images());
1597 assert_eq!(text_only.check_images(), Ok(()));
1598 }
1599
1600 #[test]
1601 fn check_images_holds_each_bound_as_a_ceiling() {
1602 // At the bound the prompt passes; one step over, the matching error.
1603 let four = Prompt::new(ModelTier::Fast).user_parts(
1604 (0..MAX_PROMPT_IMAGES).map(|_| Part::image(ImageMediaType::Png, vec![0u8])),
1605 );
1606 assert_eq!(four.check_images(), Ok(()));
1607 let each_raw = raw_len_for_encoded(MAX_IMAGE_ENCODED_BYTES);
1608 assert_eq!(encoded_image_len(each_raw), MAX_IMAGE_ENCODED_BYTES);
1609 let at_ceiling = Prompt::new(ModelTier::Fast)
1610 .user_parts([Part::image(ImageMediaType::Png, vec![0u8; each_raw])]);
1611 assert_eq!(at_ceiling.check_images(), Ok(()));
1612
1613 // One image too many.
1614 let five = Prompt::new(ModelTier::Fast).user_parts(
1615 (0..=MAX_PROMPT_IMAGES).map(|_| Part::image(ImageMediaType::Png, vec![0u8])),
1616 );
1617 assert_eq!(
1618 five.check_images(),
1619 Err(TextModelError::ImageLimit(ImageLimit::TooMany {
1620 count: MAX_PROMPT_IMAGES + 1,
1621 }))
1622 );
1623
1624 // A raw length encoding to one base64 group over the per-image cap.
1625 let over = MAX_IMAGE_ENCODED_BYTES + 4;
1626 let over_raw = raw_len_for_encoded(over);
1627 assert_eq!(encoded_image_len(over_raw), over);
1628 let too_large = Prompt::new(ModelTier::Fast)
1629 .user_parts([Part::image(ImageMediaType::Png, vec![0u8; over_raw])]);
1630 assert_eq!(
1631 too_large.check_images(),
1632 Err(TextModelError::ImageLimit(ImageLimit::ImageTooLarge {
1633 index: 0,
1634 encoded_bytes: over,
1635 }))
1636 );
1637
1638 // Two images just under the per-image cap, together over the total:
1639 // the count and each image are legal, the sum is not.
1640 let each = MAX_IMAGE_ENCODED_BYTES - 4;
1641 let raw = raw_len_for_encoded(each);
1642 assert!(2 * each > MAX_PROMPT_IMAGE_ENCODED_BYTES);
1643 let too_total = Prompt::new(ModelTier::Fast).user_parts([
1644 Part::image(ImageMediaType::Png, vec![0u8; raw]),
1645 Part::image(ImageMediaType::Jpeg, vec![0u8; raw]),
1646 ]);
1647 assert_eq!(
1648 too_total.check_images(),
1649 Err(TextModelError::ImageLimit(ImageLimit::TotalTooLarge {
1650 encoded_bytes: 2 * each,
1651 }))
1652 );
1653 }
1654
1655 #[test]
1656 fn the_image_limit_error_is_displayed_without_any_user_content() {
1657 let error = TextModelError::ImageLimit(ImageLimit::TooMany { count: 9 });
1658 assert_eq!(
1659 error.to_string(),
1660 "the prompt's images exceed a limit: a prompt may carry at most 4 images, but this \
1661 one carries 9"
1662 );
1663 assert_eq!(error.retry_after(), None);
1664 }
1665
1666 #[test]
1667 fn images_to_a_tier_that_cannot_carry_them_are_refused_before_the_call() {
1668 let fast = Recording::new("fast-vendor");
1669 let router = RoutingTextModel::new().fast(fast.clone());
1670 let prompt = Prompt::new(ModelTier::Fast)
1671 .user_parts([Part::image(ImageMediaType::Png, vec![0u8; 4])]);
1672
1673 let error = pollster::block_on(router.complete(&prompt)).unwrap_err();
1674 assert_eq!(error, TextModelError::Unsupported(Capability::Images));
1675 assert_eq!(fast.count(), 0, "the adapter was never called");
1676
1677 // The capable mirror: the same prompt passes through untouched, but
1678 // an over-limit one is still refused by the router before the tier,
1679 // whatever the deployment claims.
1680 let capable = Recording::image_capable("fast-vendor");
1681 let router = RoutingTextModel::new().fast(capable.clone());
1682 let prompt = Prompt::new(ModelTier::Fast)
1683 .user_parts([Part::image(ImageMediaType::Webp, vec![0u8; 4])]);
1684 assert_eq!(
1685 pollster::block_on(router.complete(&prompt)).unwrap().model,
1686 "fast-vendor"
1687 );
1688 assert_eq!(capable.count(), 1);
1689
1690 let over = Prompt::new(ModelTier::Fast).user_parts(
1691 (0..=MAX_PROMPT_IMAGES).map(|_| Part::image(ImageMediaType::Png, vec![0u8])),
1692 );
1693 assert_eq!(
1694 pollster::block_on(router.complete(&over)).unwrap_err(),
1695 TextModelError::ImageLimit(ImageLimit::TooMany {
1696 count: MAX_PROMPT_IMAGES + 1,
1697 })
1698 );
1699 assert_eq!(capable.count(), 1, "the over-limit prompt never reached it");
1700 }
1701}