use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use base64::Engine as _;
use futures_util::SinkExt as _;
use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::http::MAX_RESPONSE_BYTES;
use crate::stream::{BoxStream, pump};
pub const DEFAULT_MAX_TOKENS: u32 = 1024;
pub const MAX_PROMPT_IMAGES: usize = 4;
pub const MAX_IMAGE_ENCODED_BYTES: usize = 3_932_160;
pub const MAX_PROMPT_IMAGE_ENCODED_BYTES: usize = MAX_RESPONSE_BYTES;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelTier {
Fast,
Strong,
}
impl ModelTier {
pub fn name(&self) -> &'static str {
match self {
ModelTier::Fast => "fast",
ModelTier::Strong => "strong",
}
}
}
impl std::fmt::Display for ModelTier {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.name())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Role {
User,
Assistant,
}
impl Role {
pub fn name(&self) -> &'static str {
match self {
Role::User => "user",
Role::Assistant => "assistant",
}
}
}
impl std::fmt::Display for Role {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.name())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum Capability {
Tools,
Images,
}
impl Capability {
pub fn name(&self) -> &'static str {
match self {
Capability::Tools => "tools",
Capability::Images => "images",
}
}
}
impl std::fmt::Display for Capability {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.name())
}
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub struct ToolSpec {
pub name: String,
pub description: String,
pub parameters: Value,
}
impl ToolSpec {
#[must_use]
pub fn new(name: impl Into<String>, description: impl Into<String>, parameters: Value) -> Self {
Self {
name: name.into(),
description: description.into(),
parameters,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum ToolChoice {
Auto,
None,
Required,
Tool(String),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToolCall {
pub id: String,
pub name: String,
pub arguments: Value,
}
impl ToolCall {
#[must_use]
pub fn new(id: impl Into<String>, name: impl Into<String>, arguments: Value) -> Self {
Self {
id: id.into(),
name: name.into(),
arguments,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolResult {
pub tool_call_id: String,
pub content: String,
pub is_error: bool,
}
impl ToolResult {
#[must_use]
pub fn ok(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
Self {
tool_call_id: tool_call_id.into(),
content: content.into(),
is_error: false,
}
}
#[must_use]
pub fn error(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
Self {
tool_call_id: tool_call_id.into(),
content: content.into(),
is_error: true,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum ImageMediaType {
Png,
Jpeg,
Gif,
Webp,
}
impl ImageMediaType {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
ImageMediaType::Png => "image/png",
ImageMediaType::Jpeg => "image/jpeg",
ImageMediaType::Gif => "image/gif",
ImageMediaType::Webp => "image/webp",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum Part {
Text(String),
Image {
media_type: ImageMediaType,
bytes: Vec<u8>,
},
}
impl Part {
#[must_use]
pub fn text(text: impl Into<String>) -> Self {
Part::Text(text.into())
}
#[must_use]
pub fn image(media_type: ImageMediaType, bytes: impl Into<Vec<u8>>) -> Self {
Part::Image {
media_type,
bytes: bytes.into(),
}
}
}
#[must_use]
pub fn encode_image(bytes: &[u8]) -> String {
base64::engine::general_purpose::STANDARD.encode(bytes)
}
#[must_use]
pub const fn encoded_image_len(raw_len: usize) -> usize {
raw_len.div_ceil(3) * 4
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub struct Turn {
pub role: Role,
pub content: String,
pub parts: Vec<Part>,
pub tool_calls: Vec<ToolCall>,
pub tool_results: Vec<ToolResult>,
}
impl Turn {
#[must_use]
pub fn user(content: impl Into<String>) -> Self {
Turn {
role: Role::User,
content: content.into(),
parts: Vec::new(),
tool_calls: Vec::new(),
tool_results: Vec::new(),
}
}
#[must_use]
pub fn user_parts(parts: impl IntoIterator<Item = Part>) -> Self {
Turn {
role: Role::User,
content: String::new(),
parts: parts.into_iter().collect(),
tool_calls: Vec::new(),
tool_results: Vec::new(),
}
}
#[must_use]
pub fn assistant(content: impl Into<String>) -> Self {
Turn {
role: Role::Assistant,
content: content.into(),
parts: Vec::new(),
tool_calls: Vec::new(),
tool_results: Vec::new(),
}
}
#[must_use]
pub fn assistant_tool_calls(content: impl Into<String>, calls: Vec<ToolCall>) -> Self {
Turn {
role: Role::Assistant,
content: content.into(),
parts: Vec::new(),
tool_calls: calls,
tool_results: Vec::new(),
}
}
#[must_use]
pub fn tool_results(results: Vec<ToolResult>) -> Self {
Turn {
role: Role::User,
content: String::new(),
parts: Vec::new(),
tool_calls: Vec::new(),
tool_results: results,
}
}
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub struct Prompt {
pub tier: ModelTier,
pub system: Option<String>,
pub messages: Vec<Turn>,
pub json_schema: Option<Value>,
pub max_tokens: u32,
pub tools: Vec<ToolSpec>,
pub tool_choice: Option<ToolChoice>,
}
impl Prompt {
#[must_use]
pub fn new(tier: ModelTier) -> Self {
Self {
tier,
system: None,
messages: Vec::new(),
json_schema: None,
max_tokens: DEFAULT_MAX_TOKENS,
tools: Vec::new(),
tool_choice: None,
}
}
#[must_use]
pub fn tool(mut self, tool: ToolSpec) -> Self {
self.tools.push(tool);
self
}
#[must_use]
pub fn tools(mut self, tools: impl IntoIterator<Item = ToolSpec>) -> Self {
self.tools.extend(tools);
self
}
#[must_use]
pub fn tool_choice(mut self, tool_choice: ToolChoice) -> Self {
self.tool_choice = Some(tool_choice);
self
}
#[must_use]
pub fn system(mut self, system: impl Into<String>) -> Self {
self.system = Some(system.into());
self
}
#[must_use]
pub fn user(mut self, content: impl Into<String>) -> Self {
self.messages.push(Turn::user(content));
self
}
#[must_use]
pub fn user_parts(mut self, parts: impl IntoIterator<Item = Part>) -> Self {
self.messages.push(Turn::user_parts(parts));
self
}
#[must_use]
pub fn assistant(mut self, content: impl Into<String>) -> Self {
self.messages.push(Turn::assistant(content));
self
}
#[must_use]
pub fn turn(mut self, turn: Turn) -> Self {
self.messages.push(turn);
self
}
#[must_use]
pub fn json_schema(mut self, schema: Value) -> Self {
self.json_schema = Some(schema);
self
}
#[must_use]
pub fn max_tokens(mut self, max_tokens: u32) -> Self {
self.max_tokens = max_tokens;
self
}
#[must_use]
pub fn has_images(&self) -> bool {
self.messages.iter().any(|turn| {
turn.parts
.iter()
.any(|part| matches!(part, Part::Image { .. }))
})
}
pub fn check_images(&self) -> Result<(), TextModelError> {
let count = self
.messages
.iter()
.flat_map(|turn| &turn.parts)
.filter(|part| matches!(part, Part::Image { .. }))
.count();
if count > MAX_PROMPT_IMAGES {
return Err(TextModelError::ImageLimit(ImageLimit::TooMany { count }));
}
let mut total: usize = 0;
let mut index: usize = 0;
for turn in &self.messages {
for part in &turn.parts {
let Part::Image { bytes, .. } = part else {
continue;
};
let encoded_bytes = encoded_image_len(bytes.len());
if encoded_bytes > MAX_IMAGE_ENCODED_BYTES {
return Err(TextModelError::ImageLimit(ImageLimit::ImageTooLarge {
index,
encoded_bytes,
}));
}
total = total.saturating_add(encoded_bytes);
index += 1;
}
}
if total > MAX_PROMPT_IMAGE_ENCODED_BYTES {
return Err(TextModelError::ImageLimit(ImageLimit::TotalTooLarge {
encoded_bytes: total,
}));
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub struct Completion {
pub text: String,
pub json: Option<Value>,
pub model: String,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_input_tokens: Option<u64>,
pub tool_calls: Vec<ToolCall>,
}
impl Completion {
#[must_use]
pub fn new(text: impl Into<String>, model: impl Into<String>) -> Self {
Self {
text: text.into(),
json: None,
model: model.into(),
input_tokens: 0,
output_tokens: 0,
cached_input_tokens: None,
tool_calls: Vec::new(),
}
}
#[must_use]
pub fn tool_calls(mut self, tool_calls: Vec<ToolCall>) -> Self {
self.tool_calls = tool_calls;
self
}
#[must_use]
pub fn json(mut self, json: Value) -> Self {
self.json = Some(json);
self
}
#[must_use]
pub fn usage(mut self, input_tokens: u64, output_tokens: u64) -> Self {
self.input_tokens = input_tokens;
self.output_tokens = output_tokens;
self
}
#[must_use]
pub fn cached_input_tokens(mut self, cached_input_tokens: u64) -> Self {
self.cached_input_tokens = Some(cached_input_tokens);
self
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum FinishReason {
Stop,
Length,
ToolUse,
ContentFilter,
Other(String),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", content = "data", rename_all = "snake_case")]
#[non_exhaustive]
pub enum TextDelta {
Text(String),
Reasoning(String),
ToolCallStarted {
index: usize,
id: String,
name: String,
},
ToolCallArguments {
index: usize,
chunk: String,
},
ToolCallFinished {
index: usize,
call: ToolCall,
},
Usage {
input_tokens: u64,
output_tokens: u64,
cached_input_tokens: Option<u64>,
},
Finish {
reason: FinishReason,
model: String,
},
}
#[must_use]
pub fn completion_deltas(completion: &Completion) -> Vec<TextDelta> {
let mut deltas = Vec::new();
if !completion.text.is_empty() {
deltas.push(TextDelta::Text(completion.text.clone()));
}
for (index, call) in completion.tool_calls.iter().enumerate() {
deltas.push(TextDelta::ToolCallStarted {
index,
id: call.id.clone(),
name: call.name.clone(),
});
deltas.push(TextDelta::ToolCallArguments {
index,
chunk: call.arguments.to_string(),
});
deltas.push(TextDelta::ToolCallFinished {
index,
call: call.clone(),
});
}
deltas.push(TextDelta::Usage {
input_tokens: completion.input_tokens,
output_tokens: completion.output_tokens,
cached_input_tokens: completion.cached_input_tokens,
});
let reason = if completion.tool_calls.is_empty() {
FinishReason::Stop
} else {
FinishReason::ToolUse
};
deltas.push(TextDelta::Finish {
reason,
model: completion.model.clone(),
});
deltas
}
#[derive(Debug, Default)]
pub struct CompletionBuilder {
text: String,
tool_calls: std::collections::BTreeMap<usize, ToolCall>,
usage: Option<(u64, u64, Option<u64>)>,
model: Option<String>,
finish_reason: Option<FinishReason>,
}
impl CompletionBuilder {
pub fn push(&mut self, delta: &TextDelta) {
match delta {
TextDelta::Text(text) => self.text.push_str(text),
TextDelta::Reasoning(_)
| TextDelta::ToolCallStarted { .. }
| TextDelta::ToolCallArguments { .. } => {}
TextDelta::ToolCallFinished { index, call } => {
self.tool_calls.insert(*index, call.clone());
}
TextDelta::Usage {
input_tokens,
output_tokens,
cached_input_tokens,
} => {
self.usage = Some((*input_tokens, *output_tokens, *cached_input_tokens));
}
TextDelta::Finish { reason, model } => {
self.finish_reason = Some(reason.clone());
self.model = Some(model.clone());
}
}
}
#[must_use]
pub fn finish_reason(&self) -> Option<&FinishReason> {
self.finish_reason.as_ref()
}
#[must_use]
pub fn finish(self) -> Completion {
let mut completion = Completion::new(self.text, self.model.unwrap_or_default());
completion.tool_calls = self.tool_calls.into_values().collect();
if let Some((input_tokens, output_tokens, cached_input_tokens)) = self.usage {
completion.input_tokens = input_tokens;
completion.output_tokens = output_tokens;
completion.cached_input_tokens = cached_input_tokens;
}
completion
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum ImageLimit {
TooMany { count: usize },
ImageTooLarge { index: usize, encoded_bytes: usize },
TotalTooLarge { encoded_bytes: usize },
}
impl std::fmt::Display for ImageLimit {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::TooMany { count } => write!(
f,
"a prompt may carry at most {MAX_PROMPT_IMAGES} images, but this one carries {count}"
),
Self::ImageTooLarge {
index,
encoded_bytes,
} => write!(
f,
"image {index} encodes to {encoded_bytes} bytes, over the \
{MAX_IMAGE_ENCODED_BYTES}-byte per-image ceiling"
),
Self::TotalTooLarge { encoded_bytes } => write!(
f,
"the prompt's images encode to {encoded_bytes} bytes, over the \
{MAX_PROMPT_IMAGE_ENCODED_BYTES}-byte total"
),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum TextModelError {
NotConfigured,
Rejected(String),
Transient { retry_after: Option<Duration> },
Transport(String),
Unsupported(Capability),
InvalidSchema(String),
SchemaViolation(String),
ImageLimit(ImageLimit),
}
impl std::fmt::Display for TextModelError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let scrub = crate::logging::scrub_text;
match self {
Self::NotConfigured => f.write_str("no text model is wired for this tier"),
Self::Rejected(message) => write!(f, "completion rejected: {}", scrub(message)),
Self::Transient { .. } => f.write_str("completion failed, retryable"),
Self::Transport(message) => {
write!(f, "completion transport failed: {}", scrub(message))
}
Self::Unsupported(capability) => {
write!(f, "the text model does not support {capability}")
}
Self::InvalidSchema(reason) => {
write!(f, "the JSON schema cannot be honoured: {}", scrub(reason))
}
Self::SchemaViolation(reason) => write!(
f,
"the model's answer did not match the schema: {}",
scrub(reason)
),
Self::ImageLimit(limit) => write!(f, "the prompt's images exceed a limit: {limit}"),
}
}
}
impl std::error::Error for TextModelError {}
impl TextModelError {
pub fn retry_after(&self) -> Option<Duration> {
match self {
TextModelError::Transient { retry_after } => *retry_after,
_ => None,
}
}
}
#[async_trait]
pub trait TextModel: Send + Sync {
async fn complete(&self, prompt: &Prompt) -> Result<Completion, TextModelError>;
fn stream<'a>(
&'a self,
prompt: &'a Prompt,
) -> BoxStream<'a, Result<TextDelta, TextModelError>> {
Box::pin(
futures_util::stream::once(self.complete(prompt)).flat_map(|outcome| {
futures_util::stream::iter(match outcome {
Ok(completion) => completion_deltas(&completion)
.into_iter()
.map(Ok)
.collect::<Vec<_>>(),
Err(err) => vec![Err(err)],
})
}),
)
}
fn supports(&self, _tier: ModelTier, _capability: Capability) -> bool {
false
}
}
pub fn stream_owned(
model: Arc<dyn TextModel>,
prompt: Prompt,
) -> BoxStream<'static, Result<TextDelta, TextModelError>> {
pump(move |mut tx| async move {
let mut stream = model.stream(&prompt);
while let Some(delta) = stream.next().await {
if tx.send(delta).await.is_err() {
break;
}
}
})
}
#[derive(Default, Clone)]
pub struct RoutingTextModel {
fast: Option<Arc<dyn TextModel>>,
strong: Option<Arc<dyn TextModel>>,
}
impl RoutingTextModel {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn fast(mut self, model: Arc<dyn TextModel>) -> Self {
self.fast = Some(model);
self
}
#[must_use]
pub fn strong(mut self, model: Arc<dyn TextModel>) -> Self {
self.strong = Some(model);
self
}
#[must_use]
pub fn route_for(&self, tier: ModelTier) -> Option<&Arc<dyn TextModel>> {
match tier {
ModelTier::Fast => self.fast.as_ref(),
ModelTier::Strong => self.strong.as_ref(),
}
}
fn routed(&self, prompt: &Prompt) -> Result<&Arc<dyn TextModel>, TextModelError> {
prompt.check_images()?;
let model = self
.route_for(prompt.tier)
.ok_or(TextModelError::NotConfigured)?;
if !prompt.tools.is_empty() && !model.supports(prompt.tier, Capability::Tools) {
return Err(TextModelError::Unsupported(Capability::Tools));
}
if prompt.has_images() && !model.supports(prompt.tier, Capability::Images) {
return Err(TextModelError::Unsupported(Capability::Images));
}
Ok(model)
}
}
impl std::fmt::Debug for RoutingTextModel {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RoutingTextModel")
.field("fast", &self.fast.is_some())
.field("strong", &self.strong.is_some())
.finish()
}
}
#[async_trait]
impl TextModel for RoutingTextModel {
async fn complete(&self, prompt: &Prompt) -> Result<Completion, TextModelError> {
let model = self.routed(prompt)?;
model.complete(prompt).await
}
fn stream<'a>(
&'a self,
prompt: &'a Prompt,
) -> BoxStream<'a, Result<TextDelta, TextModelError>> {
match self.routed(prompt) {
Ok(model) => model.stream(prompt),
Err(err) => Box::pin(futures_util::stream::iter([Err(err)])),
}
}
fn supports(&self, tier: ModelTier, capability: Capability) -> bool {
self.route_for(tier)
.is_some_and(|model| model.supports(tier, capability))
}
}
#[cfg(test)]
mod tests {
use super::*;
use futures_core::Stream;
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::task::{Context, Poll};
#[test]
fn a_tier_and_a_role_round_trip_through_json() {
for tier in [ModelTier::Fast, ModelTier::Strong] {
let json = serde_json::to_string(&tier).expect("serialises");
let back: ModelTier = serde_json::from_str(&json).expect("deserialises");
assert_eq!(tier, back);
}
for role in [Role::User, Role::Assistant] {
let json = serde_json::to_string(&role).expect("serialises");
let back: Role = serde_json::from_str(&json).expect("deserialises");
assert_eq!(role, back);
}
}
#[test]
fn the_wire_form_is_snake_case() {
assert_eq!(serde_json::to_value(ModelTier::Strong).unwrap(), "strong");
assert_eq!(serde_json::to_value(Role::User).unwrap(), "user");
}
#[test]
fn the_tier_names_itself_for_logs_and_errors() {
assert_eq!(ModelTier::Fast.name(), "fast");
assert_eq!(ModelTier::Strong.to_string(), "strong");
assert_eq!(Role::Assistant.name(), "assistant");
assert_eq!(Role::User.to_string(), "user");
}
#[test]
fn a_prompt_starts_empty_and_builds_up() {
let prompt = Prompt::new(ModelTier::Fast);
assert_eq!(prompt.tier, ModelTier::Fast);
assert_eq!(prompt.system, None);
assert!(prompt.messages.is_empty());
assert_eq!(prompt.json_schema, None);
assert_eq!(prompt.max_tokens, DEFAULT_MAX_TOKENS);
assert_eq!(DEFAULT_MAX_TOKENS, 1024);
let prompt = prompt
.system("Draft the reply.")
.user("Hello")
.assistant("Hi there")
.turn(Turn::user("And again"))
.json_schema(serde_json::json!({ "type": "object" }))
.max_tokens(256);
assert_eq!(prompt.system.as_deref(), Some("Draft the reply."));
assert_eq!(
prompt.messages,
vec![
Turn::user("Hello"),
Turn::assistant("Hi there"),
Turn::user("And again"),
]
);
assert_eq!(
prompt.json_schema,
Some(serde_json::json!({ "type": "object" }))
);
assert_eq!(prompt.max_tokens, 256);
}
#[test]
fn a_completion_starts_with_text_and_model_and_builds_up() {
let completion = Completion::new("the answer", "vendor-1");
assert_eq!(completion.text, "the answer");
assert_eq!(completion.model, "vendor-1");
assert_eq!(completion.json, None);
assert_eq!(completion.input_tokens, 0);
assert_eq!(completion.output_tokens, 0);
let completion = completion
.json(serde_json::json!({ "reply": "the answer" }))
.usage(12, 34);
assert_eq!(
completion.json,
Some(serde_json::json!({ "reply": "the answer" }))
);
assert_eq!(completion.input_tokens, 12);
assert_eq!(completion.output_tokens, 34);
}
#[test]
fn cached_input_tokens_defaults_to_absent_and_never_shrinks_the_total() {
let completion = Completion::new("the answer", "vendor-1");
assert_eq!(completion.cached_input_tokens, None);
let completion = Completion::new("the answer", "vendor-1")
.usage(17, 5)
.cached_input_tokens(9);
assert_eq!(completion.input_tokens, 17);
assert_eq!(completion.cached_input_tokens, Some(9));
}
#[test]
fn display_sanitizes_the_provider_text() {
let error = TextModelError::Rejected(
"provider 400 for https://api.example.test/v1/complete?token=secret-abcdef".to_owned(),
);
let text = error.to_string();
assert!(text.contains("completion rejected"), "{text}");
assert!(!text.contains("secret-abcdef"), "{text}");
assert!(text.contains("?[redacted]"), "{text}");
let error = TextModelError::Transport("timeout quoting alice@example.test".to_owned());
let text = error.to_string();
assert!(text.contains("completion transport failed"), "{text}");
assert!(!text.contains('@'), "{text}");
assert!(format!("{error:?}").contains("alice@example.test"));
}
#[test]
fn transient_says_retryable_and_carries_no_text_to_scrub() {
let error = TextModelError::Transient {
retry_after: Some(Duration::from_secs(30)),
};
assert_eq!(error.to_string(), "completion failed, retryable");
}
#[test]
fn transient_carries_an_optional_retry_after() {
let throttled = TextModelError::Transient {
retry_after: Some(Duration::from_secs(30)),
};
assert_eq!(throttled.retry_after(), Some(Duration::from_secs(30)));
let plain = TextModelError::Transient { retry_after: None };
assert_eq!(plain.retry_after(), None);
assert_eq!(TextModelError::NotConfigured.retry_after(), None);
assert_eq!(
TextModelError::Rejected("422".to_owned()).retry_after(),
None,
"a rejection is not a back-off"
);
assert_eq!(
TextModelError::Transport("timed out".to_owned()).retry_after(),
None
);
}
#[test]
fn not_configured_names_the_missing_wiring() {
assert_eq!(
TextModelError::NotConfigured.to_string(),
"no text model is wired for this tier"
);
}
#[test]
fn the_structured_output_failures_name_themselves_and_never_back_off() {
let invalid =
TextModelError::InvalidSchema("the `pattern` keyword is not supported".to_owned());
assert_eq!(
invalid.to_string(),
"the JSON schema cannot be honoured: the `pattern` keyword is not supported"
);
let violation =
TextModelError::SchemaViolation("the property \"price\" is invalid".to_owned());
assert!(
violation.to_string().contains("did not match the schema"),
"{violation}"
);
assert_eq!(invalid.retry_after(), None);
assert_eq!(violation.retry_after(), None);
}
struct Recording {
label: &'static str,
tools: bool,
images: bool,
seen: std::sync::atomic::AtomicUsize,
}
impl Recording {
fn new(label: &'static str) -> Arc<Self> {
Arc::new(Self {
label,
tools: false,
images: false,
seen: std::sync::atomic::AtomicUsize::new(0),
})
}
fn tool_capable(label: &'static str) -> Arc<Self> {
Arc::new(Self {
label,
tools: true,
images: false,
seen: std::sync::atomic::AtomicUsize::new(0),
})
}
fn image_capable(label: &'static str) -> Arc<Self> {
Arc::new(Self {
label,
tools: false,
images: true,
seen: std::sync::atomic::AtomicUsize::new(0),
})
}
fn count(&self) -> usize {
self.seen.load(std::sync::atomic::Ordering::Relaxed)
}
}
#[async_trait]
impl TextModel for Recording {
async fn complete(&self, _prompt: &Prompt) -> Result<Completion, TextModelError> {
self.seen.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
Ok(Completion::new("recorded", self.label))
}
fn supports(&self, _tier: ModelTier, capability: Capability) -> bool {
(capability == Capability::Tools && self.tools)
|| (capability == Capability::Images && self.images)
}
}
#[test]
fn the_router_dispatches_by_tier() {
let fast = Recording::new("fast-vendor");
let strong = Recording::new("strong-vendor");
let router = RoutingTextModel::new()
.fast(fast.clone())
.strong(strong.clone());
let completion =
pollster::block_on(router.complete(&Prompt::new(ModelTier::Fast).user("hi"))).unwrap();
assert_eq!(
completion.model, "fast-vendor",
"the fast tier reached its adapter"
);
let completion =
pollster::block_on(router.complete(&Prompt::new(ModelTier::Strong).user("hi")))
.unwrap();
assert_eq!(
completion.model, "strong-vendor",
"the strong tier reached its own, different adapter"
);
assert_eq!(fast.count(), 1);
assert_eq!(strong.count(), 1);
}
#[test]
fn a_tier_with_no_adapter_is_not_configured_not_rejected() {
let router = RoutingTextModel::new().fast(Recording::new("fast-vendor"));
let error = pollster::block_on(router.complete(&Prompt::new(ModelTier::Strong).user("hi")))
.unwrap_err();
assert_eq!(error, TextModelError::NotConfigured);
assert!(router.route_for(ModelTier::Strong).is_none());
}
#[test]
fn an_empty_router_is_not_configured_for_every_tier() {
let router = RoutingTextModel::new();
for tier in [ModelTier::Fast, ModelTier::Strong] {
let error =
pollster::block_on(router.complete(&Prompt::new(tier).user("hi"))).unwrap_err();
assert_eq!(error, TextModelError::NotConfigured);
assert!(router.route_for(tier).is_none());
}
}
#[test]
fn debug_prints_which_tiers_are_wired_and_nothing_else() {
let router = RoutingTextModel::new().fast(Recording::new("fast-vendor"));
let printed = format!("{router:?}");
assert!(printed.contains("RoutingTextModel"), "{printed}");
assert!(printed.contains("fast: true"), "{printed}");
assert!(printed.contains("strong: false"), "{printed}");
}
#[test]
fn a_capability_names_itself_for_logs_and_errors() {
assert_eq!(Capability::Tools.name(), "tools");
assert_eq!(Capability::Tools.to_string(), "tools");
assert_eq!(Capability::Images.name(), "images");
assert_eq!(Capability::Images.to_string(), "images");
}
#[test]
fn a_tool_spec_carries_its_name_description_and_schema() {
let tool = ToolSpec::new(
"lookup",
"Look a thing up.",
serde_json::json!({ "type": "object" }),
);
assert_eq!(tool.name, "lookup");
assert_eq!(tool.description, "Look a thing up.");
assert_eq!(tool.parameters, serde_json::json!({ "type": "object" }));
}
#[test]
fn a_tool_call_and_result_round_trip_their_fields() {
let call = ToolCall::new("call-1", "lookup", serde_json::json!({ "q": "x" }));
assert_eq!(call.id, "call-1");
assert_eq!(call.name, "lookup");
assert_eq!(call.arguments, serde_json::json!({ "q": "x" }));
let ok = ToolResult::ok("call-1", "the answer");
assert_eq!(ok.tool_call_id, "call-1");
assert_eq!(ok.content, "the answer");
assert!(!ok.is_error);
let error = ToolResult::error("call-1", "it broke");
assert!(error.is_error);
assert_eq!(error.content, "it broke");
}
#[test]
fn a_turn_can_quote_tool_calls_and_carry_results_back() {
let call = ToolCall::new("call-1", "lookup", serde_json::json!({}));
let assistant = Turn::assistant_tool_calls("thinking", vec![call.clone()]);
assert_eq!(assistant.role, Role::Assistant);
assert_eq!(assistant.content, "thinking");
assert_eq!(assistant.tool_calls, vec![call]);
assert!(assistant.tool_results.is_empty());
let results = Turn::tool_results(vec![ToolResult::ok("call-1", "fine")]);
assert_eq!(results.role, Role::User);
assert_eq!(results.content, "");
assert!(results.tool_calls.is_empty());
assert_eq!(results.tool_results.len(), 1);
assert!(Turn::user("hi").tool_calls.is_empty());
assert!(Turn::assistant("hi").tool_results.is_empty());
}
#[test]
fn a_prompt_starts_without_tools_and_builds_them_up() {
let prompt = Prompt::new(ModelTier::Fast);
assert!(prompt.tools.is_empty());
assert_eq!(prompt.tool_choice, None);
let prompt = prompt
.tool(ToolSpec::new("a", "A", serde_json::json!({})))
.tools(vec![
ToolSpec::new("b", "B", serde_json::json!({})),
ToolSpec::new("c", "C", serde_json::json!({})),
])
.tool_choice(ToolChoice::Required);
assert_eq!(prompt.tools.len(), 3);
assert_eq!(
prompt
.tools
.iter()
.map(|tool| tool.name.as_str())
.collect::<Vec<_>>(),
vec!["a", "b", "c"]
);
assert_eq!(prompt.tool_choice, Some(ToolChoice::Required));
}
#[test]
fn a_completion_starts_without_tool_calls_and_builds_them_up() {
let completion = Completion::new("the answer", "vendor-1");
assert!(completion.tool_calls.is_empty());
let completion = completion.tool_calls(vec![ToolCall::new(
"call-1",
"lookup",
serde_json::json!({}),
)]);
assert_eq!(completion.tool_calls.len(), 1);
}
#[test]
fn an_unsupported_capability_names_what_is_missing_and_never_backs_off() {
let error = TextModelError::Unsupported(Capability::Tools);
assert_eq!(error.to_string(), "the text model does not support tools");
assert_eq!(error.retry_after(), None);
}
#[test]
fn the_default_capability_answer_is_no() {
assert!(!Recording::new("plain").supports(ModelTier::Fast, Capability::Tools));
assert!(Recording::tool_capable("capable").supports(ModelTier::Strong, Capability::Tools));
}
#[test]
fn tools_to_a_tier_that_cannot_carry_them_are_refused_before_the_call() {
let fast = Recording::new("fast-vendor");
let router = RoutingTextModel::new().fast(fast.clone());
let prompt = Prompt::new(ModelTier::Fast).user("hi").tool(ToolSpec::new(
"lookup",
"Look up.",
serde_json::json!({}),
));
let error = pollster::block_on(router.complete(&prompt)).unwrap_err();
assert_eq!(error, TextModelError::Unsupported(Capability::Tools));
assert_eq!(
fast.count(),
0,
"the adapter was never called, tools and all"
);
}
#[test]
fn tools_to_a_capable_tier_pass_through_untouched() {
let fast = Recording::tool_capable("fast-vendor");
let router = RoutingTextModel::new().fast(fast.clone());
let prompt = Prompt::new(ModelTier::Fast).tool(ToolSpec::new(
"lookup",
"Look up.",
serde_json::json!({}),
));
let completion = pollster::block_on(router.complete(&prompt)).unwrap();
assert_eq!(completion.model, "fast-vendor");
assert_eq!(fast.count(), 1);
}
#[test]
fn a_prompt_without_tools_reaches_a_plain_model_as_before() {
let fast = Recording::new("fast-vendor");
let router = RoutingTextModel::new().fast(fast.clone());
let completion =
pollster::block_on(router.complete(&Prompt::new(ModelTier::Fast).user("hi"))).unwrap();
assert_eq!(completion.model, "fast-vendor");
assert_eq!(fast.count(), 1);
}
#[test]
fn the_router_answers_a_capability_per_tier_and_false_when_unwired() {
let capable = Recording::tool_capable("capable");
let plain = Recording::new("plain");
let router = RoutingTextModel::new()
.fast(capable.clone())
.strong(plain.clone());
assert!(router.supports(ModelTier::Fast, Capability::Tools));
assert!(!router.supports(ModelTier::Strong, Capability::Tools));
let unwired = RoutingTextModel::new();
assert!(!unwired.supports(ModelTier::Fast, Capability::Tools));
assert!(!unwired.supports(ModelTier::Strong, Capability::Tools));
}
fn raw_len_for_encoded(encoded: usize) -> usize {
(encoded / 4) * 3
}
#[test]
fn an_image_media_type_names_its_wire_mime() {
assert_eq!(ImageMediaType::Png.as_str(), "image/png");
assert_eq!(ImageMediaType::Jpeg.as_str(), "image/jpeg");
assert_eq!(ImageMediaType::Gif.as_str(), "image/gif");
assert_eq!(ImageMediaType::Webp.as_str(), "image/webp");
}
#[test]
fn encode_image_is_standard_padded_base64_and_encoded_image_len_predicts_it() {
assert_eq!(encode_image(b""), "");
assert_eq!(encode_image(b"f"), "Zg==");
assert_eq!(encode_image(b"fo"), "Zm8=");
assert_eq!(encode_image(b"foo"), "Zm9v");
assert_eq!(encode_image(b"foob"), "Zm9vYg==");
assert_eq!(encode_image(b"fooba"), "Zm9vYmE=");
assert_eq!(encode_image(b"foobar"), "Zm9vYmFy");
for raw in 0..=8usize {
assert_eq!(
encoded_image_len(raw),
encode_image(&vec![0u8; raw]).len(),
"raw length {raw}"
);
}
assert_eq!(encoded_image_len(0), 0);
assert_eq!(encoded_image_len(1), 4);
assert_eq!(encoded_image_len(3), 4);
assert_eq!(encoded_image_len(4), 8);
}
#[test]
fn a_parts_turn_carries_its_parts_in_order_and_a_text_turn_carries_none() {
let with_image = Prompt::new(ModelTier::Fast).user_parts([
Part::text("look"),
Part::image(ImageMediaType::Jpeg, vec![1, 2, 3]),
]);
let turn = &with_image.messages[0];
assert_eq!(turn.role, Role::User);
assert_eq!(turn.content, "", "the parts are the whole content");
assert_eq!(
turn.parts,
vec![
Part::Text("look".to_owned()),
Part::Image {
media_type: ImageMediaType::Jpeg,
bytes: vec![1, 2, 3],
},
]
);
assert!(with_image.has_images());
assert!(
!Prompt::new(ModelTier::Fast)
.user_parts([Part::text("no image here")])
.has_images()
);
let text_only = Prompt::new(ModelTier::Fast).user("hi");
assert!(text_only.messages[0].parts.is_empty());
assert!(!text_only.has_images());
assert_eq!(text_only.check_images(), Ok(()));
}
#[test]
fn check_images_holds_each_bound_as_a_ceiling() {
let four = Prompt::new(ModelTier::Fast).user_parts(
(0..MAX_PROMPT_IMAGES).map(|_| Part::image(ImageMediaType::Png, vec![0u8])),
);
assert_eq!(four.check_images(), Ok(()));
let each_raw = raw_len_for_encoded(MAX_IMAGE_ENCODED_BYTES);
assert_eq!(encoded_image_len(each_raw), MAX_IMAGE_ENCODED_BYTES);
let at_ceiling = Prompt::new(ModelTier::Fast)
.user_parts([Part::image(ImageMediaType::Png, vec![0u8; each_raw])]);
assert_eq!(at_ceiling.check_images(), Ok(()));
let five = Prompt::new(ModelTier::Fast).user_parts(
(0..=MAX_PROMPT_IMAGES).map(|_| Part::image(ImageMediaType::Png, vec![0u8])),
);
assert_eq!(
five.check_images(),
Err(TextModelError::ImageLimit(ImageLimit::TooMany {
count: MAX_PROMPT_IMAGES + 1,
}))
);
let over = MAX_IMAGE_ENCODED_BYTES + 4;
let over_raw = raw_len_for_encoded(over);
assert_eq!(encoded_image_len(over_raw), over);
let too_large = Prompt::new(ModelTier::Fast)
.user_parts([Part::image(ImageMediaType::Png, vec![0u8; over_raw])]);
assert_eq!(
too_large.check_images(),
Err(TextModelError::ImageLimit(ImageLimit::ImageTooLarge {
index: 0,
encoded_bytes: over,
}))
);
let each = MAX_IMAGE_ENCODED_BYTES - 4;
let raw = raw_len_for_encoded(each);
assert!(2 * each > MAX_PROMPT_IMAGE_ENCODED_BYTES);
let too_total = Prompt::new(ModelTier::Fast).user_parts([
Part::image(ImageMediaType::Png, vec![0u8; raw]),
Part::image(ImageMediaType::Jpeg, vec![0u8; raw]),
]);
assert_eq!(
too_total.check_images(),
Err(TextModelError::ImageLimit(ImageLimit::TotalTooLarge {
encoded_bytes: 2 * each,
}))
);
}
#[test]
fn the_image_limit_error_is_displayed_without_any_user_content() {
let error = TextModelError::ImageLimit(ImageLimit::TooMany { count: 9 });
assert_eq!(
error.to_string(),
"the prompt's images exceed a limit: a prompt may carry at most 4 images, but this \
one carries 9"
);
assert_eq!(error.retry_after(), None);
}
#[test]
fn images_to_a_tier_that_cannot_carry_them_are_refused_before_the_call() {
let fast = Recording::new("fast-vendor");
let router = RoutingTextModel::new().fast(fast.clone());
let prompt = Prompt::new(ModelTier::Fast)
.user_parts([Part::image(ImageMediaType::Png, vec![0u8; 4])]);
let error = pollster::block_on(router.complete(&prompt)).unwrap_err();
assert_eq!(error, TextModelError::Unsupported(Capability::Images));
assert_eq!(fast.count(), 0, "the adapter was never called");
let capable = Recording::image_capable("fast-vendor");
let router = RoutingTextModel::new().fast(capable.clone());
let prompt = Prompt::new(ModelTier::Fast)
.user_parts([Part::image(ImageMediaType::Webp, vec![0u8; 4])]);
assert_eq!(
pollster::block_on(router.complete(&prompt)).unwrap().model,
"fast-vendor"
);
assert_eq!(capable.count(), 1);
let over = Prompt::new(ModelTier::Fast).user_parts(
(0..=MAX_PROMPT_IMAGES).map(|_| Part::image(ImageMediaType::Png, vec![0u8])),
);
assert_eq!(
pollster::block_on(router.complete(&over)).unwrap_err(),
TextModelError::ImageLimit(ImageLimit::TooMany {
count: MAX_PROMPT_IMAGES + 1,
})
);
assert_eq!(capable.count(), 1, "the over-limit prompt never reached it");
}
struct Native {
seen: AtomicUsize,
cancelled: Arc<AtomicUsize>,
}
impl Native {
fn new() -> Arc<Self> {
Arc::new(Self {
seen: AtomicUsize::new(0),
cancelled: Arc::new(AtomicUsize::new(0)),
})
}
fn completions(&self) -> usize {
self.seen.load(Ordering::Relaxed)
}
fn cancellations(&self) -> usize {
self.cancelled.load(Ordering::Relaxed)
}
}
struct DropGuard {
counter: Arc<AtomicUsize>,
armed: bool,
}
impl DropGuard {
fn new(counter: Arc<AtomicUsize>) -> Self {
Self {
counter,
armed: true,
}
}
}
impl Drop for DropGuard {
fn drop(&mut self) {
if self.armed {
self.counter.fetch_add(1, Ordering::Relaxed);
}
}
}
struct NativeStream {
guard: Option<DropGuard>,
items: u8,
exhausted: bool,
}
impl Stream for NativeStream {
type Item = Result<TextDelta, TextModelError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let _ = &self.guard;
self.items += 1;
match self.items {
1 => Poll::Ready(Some(Ok(TextDelta::Text("native".to_owned())))),
2 => Poll::Ready(Some(Ok(TextDelta::Finish {
reason: FinishReason::Stop,
model: "native-vendor".to_owned(),
}))),
3 => {
cx.waker().wake_by_ref();
Poll::Pending
}
_ => {
self.exhausted = true;
Poll::Ready(None)
}
}
}
}
impl Drop for NativeStream {
fn drop(&mut self) {
if self.exhausted
&& let Some(guard) = self.guard.as_mut()
{
guard.armed = false;
}
}
}
#[async_trait]
impl TextModel for Native {
async fn complete(&self, _prompt: &Prompt) -> Result<Completion, TextModelError> {
self.seen.fetch_add(1, Ordering::Relaxed);
Ok(Completion::new("buffered", "native-vendor"))
}
fn stream<'a>(
&'a self,
_prompt: &'a Prompt,
) -> BoxStream<'a, Result<TextDelta, TextModelError>> {
Box::pin(NativeStream {
guard: Some(DropGuard::new(Arc::clone(&self.cancelled))),
items: 0,
exhausted: false,
})
}
}
#[test]
fn completion_deltas_decomposes_text_tools_usage_and_finish() {
let call = ToolCall::new("call-1", "lookup", serde_json::json!({ "q": "x" }));
let completion = Completion::new("hello", "vendor-1")
.usage(3, 5)
.cached_input_tokens(2)
.tool_calls(vec![call.clone()]);
assert_eq!(
completion_deltas(&completion),
vec![
TextDelta::Text("hello".to_owned()),
TextDelta::ToolCallStarted {
index: 0,
id: "call-1".to_owned(),
name: "lookup".to_owned(),
},
TextDelta::ToolCallArguments {
index: 0,
chunk: r#"{"q":"x"}"#.to_owned(),
},
TextDelta::ToolCallFinished { index: 0, call },
TextDelta::Usage {
input_tokens: 3,
output_tokens: 5,
cached_input_tokens: Some(2),
},
TextDelta::Finish {
reason: FinishReason::ToolUse,
model: "vendor-1".to_owned(),
},
]
);
}
#[test]
fn a_text_only_completion_finishes_stop_and_an_empty_one_carries_no_text_delta() {
let completion = Completion::new("", "vendor-1");
let deltas = completion_deltas(&completion);
assert_eq!(deltas.len(), 2, "usage and finish only: {deltas:?}");
assert!(matches!(
deltas[1],
TextDelta::Finish {
reason: FinishReason::Stop,
..
}
));
}
#[pollster::test]
async fn the_default_stream_is_one_round_trip_decomposed() {
let model = Recording::new("fast-vendor");
let prompt = Prompt::new(ModelTier::Fast).user("hi");
let deltas: Vec<_> = model.stream(&prompt).collect().await;
let completion = model.complete(&prompt).await.unwrap();
assert_eq!(
deltas,
completion_deltas(&completion)
.into_iter()
.map(Ok)
.collect::<Vec<_>>(),
"the default stream is the completion, as deltas"
);
assert_eq!(model.count(), 2, "the stream made its one round trip");
}
#[pollster::test]
async fn a_builder_round_trips_deltas_back_to_the_completion_minus_json() {
let call = ToolCall::new("call-1", "lookup", serde_json::json!({ "q": "x" }));
let completion = Completion::new("the answer", "vendor-1")
.json(serde_json::json!({ "reply": "the answer" }))
.usage(12, 34)
.cached_input_tokens(9)
.tool_calls(vec![call.clone()]);
let mut builder = CompletionBuilder::default();
for delta in completion_deltas(&completion) {
builder.push(&delta);
}
let rebuilt = builder.finish();
let expected = Completion::new("the answer", "vendor-1")
.usage(12, 34)
.cached_input_tokens(9)
.tool_calls(vec![call]);
assert_eq!(rebuilt, expected);
assert_eq!(
rebuilt.json, None,
"structured output is a `complete` feature"
);
}
#[test]
fn the_builder_takes_usage_last_and_needs_only_finished_calls() {
let mut builder = CompletionBuilder::default();
builder.push(&TextDelta::Text("a".to_owned()));
builder.push(&TextDelta::Reasoning(
"never reaches a completion".to_owned(),
));
builder.push(&TextDelta::ToolCallArguments {
index: 0,
chunk: "{\"x\":".to_owned(),
});
builder.push(&TextDelta::Usage {
input_tokens: 1,
output_tokens: 1,
cached_input_tokens: None,
});
builder.push(&TextDelta::Usage {
input_tokens: 7,
output_tokens: 9,
cached_input_tokens: Some(4),
});
builder.push(&TextDelta::Finish {
reason: FinishReason::Length,
model: "vendor-9".to_owned(),
});
assert_eq!(builder.finish_reason(), Some(&FinishReason::Length));
let completion = builder.finish();
assert_eq!(completion.text, "a");
assert!(
completion.tool_calls.is_empty(),
"an unfinished call is not a call the tool loop may run"
);
assert_eq!(completion.input_tokens, 7);
assert_eq!(completion.output_tokens, 9);
assert_eq!(completion.cached_input_tokens, Some(4));
assert_eq!(completion.model, "vendor-9");
}
#[test]
fn a_delta_serialises_with_a_type_tag_and_round_trips() {
let deltas = vec![
TextDelta::Text("hi".to_owned()),
TextDelta::Reasoning("hmm".to_owned()),
TextDelta::ToolCallStarted {
index: 0,
id: "c1".to_owned(),
name: "lookup".to_owned(),
},
TextDelta::ToolCallArguments {
index: 0,
chunk: "{}".to_owned(),
},
TextDelta::ToolCallFinished {
index: 0,
call: ToolCall::new("c1", "lookup", serde_json::json!({})),
},
TextDelta::Usage {
input_tokens: 1,
output_tokens: 2,
cached_input_tokens: None,
},
TextDelta::Finish {
reason: FinishReason::ToolUse,
model: "vendor-1".to_owned(),
},
];
for delta in &deltas {
let json = serde_json::to_value(delta).expect("serialises");
assert!(json.get("type").is_some(), "tagged: {json}");
let back: TextDelta = serde_json::from_value(json).expect("deserialises");
assert_eq!(&back, delta);
}
assert_eq!(
serde_json::to_value(TextDelta::Text("hi".to_owned())).unwrap(),
serde_json::json!({ "type": "text", "data": "hi" })
);
}
#[pollster::test]
async fn the_router_forwards_stream_to_the_routed_tier() {
let native = Native::new();
let router = RoutingTextModel::new().fast(Arc::clone(&native) as Arc<dyn TextModel>);
let prompt = Prompt::new(ModelTier::Fast).user("hi");
let deltas: Vec<_> = router.stream(&prompt).collect().await;
assert_eq!(
deltas,
vec![
Ok(TextDelta::Text("native".to_owned())),
Ok(TextDelta::Finish {
reason: FinishReason::Stop,
model: "native-vendor".to_owned(),
}),
],
"the inner model's native stream, not a flattened complete"
);
assert_eq!(
native.completions(),
0,
"forwarding never fell back to the buffered round trip"
);
}
#[pollster::test]
async fn the_router_refuses_a_stream_like_it_refuses_a_complete() {
let native = Native::new();
let router = RoutingTextModel::new().fast(Arc::clone(&native) as Arc<dyn TextModel>);
let deltas: Vec<_> = router
.stream(&Prompt::new(ModelTier::Strong).user("hi"))
.collect()
.await;
assert_eq!(deltas, vec![Err(TextModelError::NotConfigured)]);
let incapable = Recording::new("plain");
let router = RoutingTextModel::new().fast(Arc::clone(&incapable) as Arc<dyn TextModel>);
let prompt = Prompt::new(ModelTier::Fast).user("hi").tool(ToolSpec::new(
"lookup",
"Look up.",
serde_json::json!({}),
));
let deltas: Vec<_> = router.stream(&prompt).collect().await;
assert_eq!(
deltas,
vec![Err(TextModelError::Unsupported(Capability::Tools))]
);
assert_eq!(incapable.count(), 0, "the adapter was never called");
}
#[pollster::test]
async fn dropping_a_forwarded_stream_cancels_the_inner() {
let native = Native::new();
let router = RoutingTextModel::new().fast(Arc::clone(&native) as Arc<dyn TextModel>);
{
let prompt = Prompt::new(ModelTier::Fast).user("hi");
let mut stream = router.stream(&prompt);
let first = stream.next().await;
assert_eq!(first, Some(Ok(TextDelta::Text("native".to_owned()))));
}
assert_eq!(
native.cancellations(),
1,
"the drop reached the inner stream"
);
}
#[pollster::test]
async fn stream_owned_yields_a_static_stream_of_the_models_deltas() {
let native = Native::new();
let stream: BoxStream<'static, Result<TextDelta, TextModelError>> =
stream_owned(native.clone(), Prompt::new(ModelTier::Fast).user("hi"));
let deltas: Vec<_> = stream.collect().await;
assert_eq!(
deltas,
vec![
Ok(TextDelta::Text("native".to_owned())),
Ok(TextDelta::Finish {
reason: FinishReason::Stop,
model: "native-vendor".to_owned(),
}),
]
);
assert_eq!(native.completions(), 0);
assert_eq!(
native.cancellations(),
0,
"a drained stream is not a cancelled one"
);
}
#[pollster::test]
async fn dropping_a_stream_owned_stream_cancels_the_model() {
let native = Native::new();
{
let mut stream = stream_owned(native.clone(), Prompt::new(ModelTier::Fast).user("hi"));
let first = stream.next().await;
assert_eq!(first, Some(Ok(TextDelta::Text("native".to_owned()))));
}
assert_eq!(
native.cancellations(),
1,
"the pump dropped the inner stream"
);
}
}