sim-lib-openai-server 0.1.2

OpenAI-compatible gateway skeleton for SIM.
Documentation
use sim_kernel::Expr;

use crate::{
    plan::eval_helpers::{keyword_value, request_field, request_field_text},
    runtime::OpenAiFederationPolicy,
};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum PlanPrivacy {
    LocalOnly,
    AllowRemote,
}

#[derive(Clone, Debug, PartialEq)]
pub(crate) struct EvalContext {
    pub(crate) privacy: PlanPrivacy,
    privacy_label: Option<String>,
    budget: Vec<(String, Expr)>,
    trace: bool,
}

impl EvalContext {
    pub(crate) fn from_request(request: &Expr) -> Self {
        let privacy_label = request_field_text(request, "privacy");
        Self {
            privacy: if privacy_label.as_deref() == Some("local-only") {
                PlanPrivacy::LocalOnly
            } else {
                PlanPrivacy::AllowRemote
            },
            privacy_label,
            budget: request_budget_entries(request),
            trace: false,
        }
    }

    pub(crate) fn local(&self) -> Self {
        Self {
            privacy: PlanPrivacy::LocalOnly,
            privacy_label: Some("local-only".to_owned()),
            budget: self.budget.clone(),
            trace: self.trace,
        }
    }

    pub(crate) fn trace(&self) -> Self {
        Self {
            trace: true,
            privacy: self.privacy,
            privacy_label: self.privacy_label.clone(),
            budget: self.budget.clone(),
        }
    }

    pub(crate) fn budgeted(&self, args: &[Expr]) -> Self {
        Self {
            privacy: self.privacy,
            privacy_label: self.privacy_label.clone(),
            budget: self
                .budget
                .iter()
                .cloned()
                .chain(budget_keyword_entries(args))
                .collect(),
            trace: self.trace,
        }
    }

    pub(crate) fn federation_policy(&self) -> OpenAiFederationPolicy {
        OpenAiFederationPolicy::new(self.privacy_label.clone(), self.budget.clone())
    }
}

fn request_budget_entries(request: &Expr) -> Vec<(String, Expr)> {
    let mut entries = Vec::new();
    if let Some(Expr::Map(fields)) = request_field(request, "budget") {
        entries.extend(fields.iter().filter_map(|(key, value)| {
            field_name(key).map(|name| (name.to_owned(), value.clone()))
        }));
    }
    for name in [
        "max-tokens",
        "max-output-tokens",
        "max-cost",
        "max-latency-ms",
    ] {
        if let Some(value) = request_field(request, name) {
            entries.push((name.to_owned(), value.clone()));
        }
    }
    entries
}

fn budget_keyword_entries(args: &[Expr]) -> Vec<(String, Expr)> {
    ["max-tokens", "max-cost", "max-latency-ms"]
        .into_iter()
        .filter_map(|name| keyword_value(args, name).map(|value| (name.to_owned(), value.clone())))
        .collect()
}

fn field_name(expr: &Expr) -> Option<&str> {
    match expr {
        Expr::Symbol(symbol) if symbol.namespace.is_none() => Some(symbol.name.as_ref()),
        Expr::String(value) => Some(value.as_str()),
        _ => None,
    }
}