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, PartialEq, Eq)]
pub struct Turn {
pub role: Role,
pub content: String,
}
impl Turn {
#[must_use]
pub fn user(content: impl Into<String>) -> Self {
Turn {
role: Role::User,
content: content.into(),
}
}
#[must_use]
pub fn assistant(content: impl Into<String>) -> Self {
Turn {
role: Role::Assistant,
content: content.into(),
}
}
}
#[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,
}
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,
}
}
#[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,
}
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,
}
}
#[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
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TextModelError {
NotConfigured,
Rejected(String),
Transient { retry_after: Option<Duration> },
Transport(String),
}
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))
}
}
}
}
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>;
}
#[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) => model.complete(prompt).await,
None => Err(TextModelError::NotConfigured),
}
}
}
#[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 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,
seen: std::sync::atomic::AtomicUsize,
}
impl Recording {
fn new(label: &'static str) -> Arc<Self> {
Arc::new(Self {
label,
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))
}
}
#[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}");
}
}