use std::collections::BTreeSet;
use serde_json::{Value, json};
use sim_kernel::{Args, Cx, Error, Expr, Result, Symbol, Value as SimValue};
use sim_lib_agent_runner_core::ModelCard;
use crate::server::{GatewayRequest, GatewayResponse, GatewayRouteState};
pub const MODELS_PATH: &str = "/v1/models";
const FIXTURE_ECHO_MODEL: &str = "fixture/echo";
const SIM_BROWSE_PLAN_MODEL: &str = "sim/browse/plan";
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OpenAiModel {
id: String,
owned_by: String,
}
impl OpenAiModel {
pub fn new(id: impl Into<String>, owned_by: impl Into<String>) -> Self {
Self {
id: id.into(),
owned_by: owned_by.into(),
}
}
pub fn fixture_echo() -> Self {
Self::new(FIXTURE_ECHO_MODEL, "sim")
}
pub fn sim_native(id: impl Into<String>) -> Self {
Self::new(id, "sim")
}
pub fn from_model_card(card: ModelCard) -> Self {
Self::new(card.model, card.provider.to_string())
}
pub fn id(&self) -> &str {
&self.id
}
fn to_json(&self) -> Value {
json!({
"id": self.id,
"object": "model",
"created": 0,
"owned_by": self.owned_by,
})
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ModelCatalog {
models: Vec<OpenAiModel>,
}
impl ModelCatalog {
pub fn default_fixture() -> Self {
Self::from_parts(Vec::new())
}
pub fn from_runner_card_exprs(cards: impl IntoIterator<Item = Expr>) -> Result<Self> {
let cards = cards
.into_iter()
.map(ModelCard::try_from)
.collect::<Result<Vec<_>>>()?;
Ok(Self::from_model_cards(cards))
}
pub fn from_model_cards(cards: impl IntoIterator<Item = ModelCard>) -> Self {
let runner_models = cards
.into_iter()
.map(OpenAiModel::from_model_card)
.collect::<Vec<_>>();
Self::from_parts(runner_models)
}
pub fn from_runner_args(cx: &mut Cx, args: Vec<SimValue>) -> Result<Self> {
if args.is_empty() {
return Ok(Self::default_fixture());
}
let cards = cx.call_function(&runner_cards_symbol(), Args::new(args))?;
let expr = cards.object().as_expr(cx)?;
let Expr::List(items) = expr else {
return Err(Error::Eval(
"runner/cards must return a list of model-card transcripts".to_owned(),
));
};
Self::from_runner_card_exprs(items)
}
pub fn models(&self) -> &[OpenAiModel] {
&self.models
}
fn from_parts(runner_models: Vec<OpenAiModel>) -> Self {
let mut seen = BTreeSet::new();
let mut models = Vec::new();
push_unique(&mut models, &mut seen, OpenAiModel::fixture_echo());
for model in runner_models {
push_unique(&mut models, &mut seen, model);
}
push_unique(
&mut models,
&mut seen,
OpenAiModel::sim_native(SIM_BROWSE_PLAN_MODEL),
);
Self { models }
}
fn to_json(&self) -> Value {
json!({
"object": "list",
"data": self.models.iter().map(OpenAiModel::to_json).collect::<Vec<_>>(),
})
}
}
pub fn handle_models(_request: &GatewayRequest, state: &GatewayRouteState) -> GatewayResponse {
models_response_for_catalog(&ModelCatalog::from_model_cards(state.runners().cards()))
}
pub fn models_response() -> GatewayResponse {
models_response_for_catalog(&ModelCatalog::default_fixture())
}
pub fn models_response_for_runner_args(
cx: &mut Cx,
args: Vec<SimValue>,
) -> Result<GatewayResponse> {
let catalog = ModelCatalog::from_runner_args(cx, args)?;
Ok(models_response_for_catalog(&catalog))
}
pub fn models_response_for_catalog(catalog: &ModelCatalog) -> GatewayResponse {
GatewayResponse::json(200, catalog.to_json().to_string().into_bytes())
}
pub fn runner_cards_symbol() -> Symbol {
Symbol::qualified("runner", "cards")
}
fn push_unique(models: &mut Vec<OpenAiModel>, seen: &mut BTreeSet<String>, model: OpenAiModel) {
if seen.insert(model.id.clone()) {
models.push(model);
}
}