use std::{fmt, sync::Arc};
use sim_kernel::{Cx, Error, Result};
use sim_lib_agent_runner_core::{ModelCard, ModelRequest, ModelResponse, ModelRunner};
#[derive(Clone, Default)]
pub struct OpenAiRunnerRegistry {
entries: Vec<OpenAiRunnerEntry>,
}
#[derive(Clone)]
struct OpenAiRunnerEntry {
model: String,
runner: Arc<dyn ModelRunner>,
}
impl fmt::Debug for OpenAiRunnerRegistry {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("OpenAiRunnerRegistry")
.field(
"models",
&self
.entries
.iter()
.map(|entry| entry.model.as_str())
.collect::<Vec<_>>(),
)
.finish()
}
}
impl OpenAiRunnerRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn with_runner(mut self, model: impl Into<String>, runner: Arc<dyn ModelRunner>) -> Self {
self.register(model, runner);
self
}
pub fn register(&mut self, model: impl Into<String>, runner: Arc<dyn ModelRunner>) {
self.entries.push(OpenAiRunnerEntry {
model: model.into(),
runner,
});
}
pub fn cards(&self) -> Vec<ModelCard> {
self.entries
.iter()
.map(|entry| {
let mut card = entry.runner.card();
if card.model != entry.model {
card.extra.push((
sim_kernel::Expr::Symbol(sim_kernel::Symbol::new("runner-model")),
sim_kernel::Expr::String(card.model.clone()),
));
card.model = entry.model.clone();
}
card
})
.collect()
}
pub fn infer(&self, cx: &mut Cx, model: &str, request: ModelRequest) -> Result<ModelResponse> {
let runner = self
.entries
.iter()
.find(|entry| entry.model == model)
.map(|entry| entry.runner.clone())
.ok_or_else(|| Error::Eval(format!("model_not_found: {model}")))?;
runner.infer(cx, request)
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
}