use crate::{ModelBid, ModelCard, ModelEvent, ModelEventSink};
use sim_citizen::CitizenField;
use sim_citizen_derive::Citizen;
use sim_kernel::{Cx, Expr, Result, Symbol};
pub trait ModelRunner: Send + Sync {
fn card(&self) -> ModelCard;
fn infer(&self, cx: &mut Cx, request: ModelRequest) -> Result<ModelResponse>;
fn infer_stream(
&self,
cx: &mut Cx,
request: ModelRequest,
sink: &mut dyn ModelEventSink,
) -> Result<ModelResponse> {
let response = self.infer(cx, request)?;
sink.emit(ModelEvent::final_of(&response))?;
Ok(response)
}
fn bid(&self, _request: &ModelRequest) -> Result<ModelBid> {
Ok(ModelBid::unavailable("runner did not implement bidding"))
}
}
#[derive(Clone, Debug, PartialEq, Citizen)]
#[citizen(symbol = "agent-runner/ModelRequest", version = 1)]
pub struct ModelRequest {
pub task: Expr,
pub messages: Vec<Expr>,
pub extra: Vec<(Expr, Expr)>,
}
impl ModelRequest {
pub fn new(task: Expr, messages: Vec<Expr>) -> Self {
Self {
task,
messages,
extra: Vec::new(),
}
}
}
impl Default for ModelRequest {
fn default() -> Self {
Self::new(Expr::String("citizen fixture task".to_owned()), Vec::new())
}
}
#[derive(Clone, Debug, PartialEq, Citizen)]
#[citizen(symbol = "agent-runner/ModelResponse", version = 1)]
pub struct ModelResponse {
pub runner: Symbol,
pub model: String,
pub content: Vec<Expr>,
pub stop_reason: Symbol,
pub usage: Option<ModelUsage>,
pub extra: Vec<(Expr, Expr)>,
}
impl ModelResponse {
pub fn new(
runner: Symbol,
model: impl Into<String>,
content: Vec<Expr>,
stop_reason: Symbol,
) -> Self {
Self {
runner,
model: model.into(),
content,
stop_reason,
usage: None,
extra: Vec::new(),
}
}
}
impl Default for ModelResponse {
fn default() -> Self {
Self::new(
Symbol::qualified("runner", "fixture"),
"fixture-model",
vec![Expr::String("ok".to_owned())],
Symbol::new("stop"),
)
}
}
#[derive(Clone, Debug, PartialEq, Default, Citizen)]
#[citizen(symbol = "agent-runner/ModelUsage", version = 1)]
pub struct ModelUsage {
pub input_tokens: Option<u64>,
pub output_tokens: Option<u64>,
pub latency_ms: Option<u64>,
pub cost_usd: Option<f64>,
pub extra: Vec<(Expr, Expr)>,
}
impl CitizenField for ModelUsage {
fn encode_field(&self) -> Expr {
Expr::List(vec![
self.input_tokens.encode_field(),
self.output_tokens.encode_field(),
self.latency_ms.encode_field(),
self.cost_usd.encode_field(),
self.extra.encode_field(),
])
}
fn decode_field_expr(expr: &Expr, field: &'static str) -> Result<Self> {
let Expr::List(items) = expr else {
return Err(sim_citizen::field_error(field, "expected model usage list"));
};
let [input_tokens, output_tokens, latency_ms, cost_usd, extra] = items.as_slice() else {
return Err(sim_citizen::field_error(
field,
format!("expected 5 model usage field(s), found {}", items.len()),
));
};
Ok(Self {
input_tokens: Option::<u64>::decode_field_expr(input_tokens, field)?,
output_tokens: Option::<u64>::decode_field_expr(output_tokens, field)?,
latency_ms: Option::<u64>::decode_field_expr(latency_ms, field)?,
cost_usd: Option::<f64>::decode_field_expr(cost_usd, field)?,
extra: Vec::<(Expr, Expr)>::decode_field_expr(extra, field)?,
})
}
}
pub fn model_request_class_symbol() -> Symbol {
Symbol::qualified("agent-runner", "ModelRequest")
}
pub fn model_response_class_symbol() -> Symbol {
Symbol::qualified("agent-runner", "ModelResponse")
}
pub fn model_usage_class_symbol() -> Symbol {
Symbol::qualified("agent-runner", "ModelUsage")
}