use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub const DEFAULT_MAX_TOKENS: u32 = 1024;
#[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,
}
impl Capability {
pub fn name(&self) -> &'static str {
match self {
Capability::Tools => "tools",
}
}
}
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)]
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, PartialEq)]
#[non_exhaustive]
pub struct Turn {
pub role: Role,
pub content: String,
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(),
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(),
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(),
tool_calls: calls,
tool_results: Vec::new(),
}
}
#[must_use]
pub fn tool_results(results: Vec<ToolResult>) -> Self {
Turn {
role: Role::User,
content: String::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 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
}
}
#[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)]
#[non_exhaustive]
pub enum TextModelError {
NotConfigured,
Rejected(String),
Transient { retry_after: Option<Duration> },
Transport(String),
Unsupported(Capability),
}
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}")
}
}
}
}
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 supports(&self, _tier: ModelTier, _capability: Capability) -> bool {
false
}
}
#[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(),
}
}
}
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> {
match self.route_for(prompt.tier) {
Some(model) => {
if !prompt.tools.is_empty() && !model.supports(prompt.tier, Capability::Tools) {
return Err(TextModelError::Unsupported(Capability::Tools));
}
model.complete(prompt).await
}
None => Err(TextModelError::NotConfigured),
}
}
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::*;
#[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"
);
}
struct Recording {
label: &'static str,
tools: bool,
seen: std::sync::atomic::AtomicUsize,
}
impl Recording {
fn new(label: &'static str) -> Arc<Self> {
Arc::new(Self {
label,
tools: false,
seen: std::sync::atomic::AtomicUsize::new(0),
})
}
fn tool_capable(label: &'static str) -> Arc<Self> {
Arc::new(Self {
label,
tools: 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 {
self.tools && capability == Capability::Tools
}
}
#[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");
}
#[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));
}
}