use cuttlefish_core::spec::ModelRef;
use cuttlefish_host::{
backend::{BackendFactory, Registry},
infer::{InferBackend, InferRequest, InferResult},
};
use std::sync::Arc;
struct FakeFactory;
struct FakeBackend(String);
#[async_trait::async_trait]
impl InferBackend for FakeBackend {
async fn infer(
&self,
_req: InferRequest<'_>,
_on_token: &mut (dyn for<'t> FnMut(&'t str) -> bool + Send),
) -> anyhow::Result<InferResult> {
Ok(InferResult {
text: String::new(),
tokens_in: 0,
tokens_out: 0,
})
}
fn model_name(&self) -> String {
self.0.clone()
}
}
impl BackendFactory for FakeFactory {
fn provider(&self) -> &'static str {
"invented"
}
fn describe(&self) -> &'static str {
"a provider that did not exist when the registry was written"
}
fn build(&self, target: &str) -> anyhow::Result<Arc<dyn InferBackend>> {
Ok(Arc::new(FakeBackend(target.to_string())))
}
}
#[test]
fn builtins_are_registered() {
let providers = Registry::with_builtins().providers();
assert!(providers.contains(&"stub"), "got {providers:?}");
assert!(providers.contains(&"ollama"), "got {providers:?}");
}
#[test]
fn a_new_provider_needs_no_changes_elsewhere() {
let mut registry = Registry::with_builtins();
registry.register(Box::new(FakeFactory));
let backend = match registry.resolve(&ModelRef::new("invented", "some-target")) {
Ok(b) => b,
Err(e) => panic!("a registered provider must resolve: {e}"),
};
assert_eq!(backend.model_name(), "some-target");
}
#[test]
fn an_unknown_provider_lists_what_is_available() {
let err = Registry::with_builtins()
.resolve(&ModelRef::new("gpt9", "x"))
.err()
.expect("an unknown provider must not resolve");
let msg = err.to_string();
assert!(msg.contains("gpt9"), "the bad name must appear: {msg}");
assert!(msg.contains("ollama"), "alternatives must be listed: {msg}");
assert!(msg.contains("stub"), "alternatives must be listed: {msg}");
}
#[test]
fn registering_twice_replaces_rather_than_duplicates() {
let mut registry = Registry::new();
registry.register(Box::new(FakeFactory));
registry.register(Box::new(FakeFactory));
assert_eq!(registry.providers(), vec!["invented"]);
}
#[test]
fn an_empty_registry_resolves_nothing() {
assert!(Registry::new()
.resolve(&ModelRef::new("stub", "x"))
.is_err());
}
#[test]
fn the_stub_provider_takes_its_reply_from_the_spec() {
let backend = match Registry::with_builtins().resolve(&ModelRef::new("stub", "canned reply")) {
Ok(b) => b,
Err(e) => panic!("stub must resolve: {e}"),
};
assert_eq!(backend.model_name(), "stub");
}
#[test]
fn ollama_requires_a_model_name() {
let err = Registry::with_builtins()
.resolve(&ModelRef::new("ollama", ""))
.err()
.expect("an empty model name is unusable");
assert!(err.to_string().contains("model name"), "{err}");
}
#[test]
fn resolution_errors_name_the_provider_and_target() {
let err = Registry::with_builtins()
.resolve(&ModelRef::new("ollama", ""))
.err()
.unwrap();
let msg = err.to_string();
assert!(msg.contains("ollama"), "{msg}");
}