sim_lib_openai_server/runtime/
federation.rs

1use std::{
2    collections::BTreeMap,
3    sync::{Arc, Mutex},
4};
5
6use serde_json::{Map, Number, Value as JsonValue};
7use sim_kernel::{
8    Consistency, Cx, Error, EvalFabricRef, EvalMode, EvalRequest, Expr, Result, Symbol,
9};
10use sim_lib_agent_runner_core::{ModelRequest, ModelResponse, ModelUsage};
11use sim_lib_net_core::hex_encode;
12
13use sim_codec_chat::text_part;
14use sim_codec_json::json_number_to_u64;
15
16use crate::{
17    capabilities::openai_gateway_federate_capability, objects::GatewayResponseValue,
18    plan::fixtures::request_text,
19};
20
21/// Registry of remote gateways reachable for federated inference, keyed by address.
22#[derive(Clone, Default)]
23pub struct OpenAiFederation {
24    gateways: Arc<Mutex<BTreeMap<String, OpenAiFederatedGateway>>>,
25}
26
27/// A single remote gateway: its `gateway/`-prefixed address, default model, and
28/// the eval fabric that reaches it.
29#[derive(Clone)]
30pub struct OpenAiFederatedGateway {
31    address: String,
32    default_model: String,
33    fabric: EvalFabricRef,
34}
35
36/// Policy applied to a federated request: optional privacy tier and budget entries.
37#[derive(Clone, Debug, Default, PartialEq)]
38pub struct OpenAiFederationPolicy {
39    privacy: Option<String>,
40    budget: Vec<(String, Expr)>,
41}
42
43impl OpenAiFederation {
44    /// Returns an empty federation registry.
45    pub fn new() -> Self {
46        Self::default()
47    }
48
49    /// Registers a remote gateway under its normalized address.
50    pub fn insert_gateway(
51        &self,
52        address: impl AsRef<str>,
53        default_model: impl Into<String>,
54        fabric: EvalFabricRef,
55    ) -> Result<()> {
56        let gateway = OpenAiFederatedGateway::new(address.as_ref(), default_model, fabric)?;
57        self.gateways
58            .lock()
59            .map_err(|_| Error::PoisonedLock("openai federation registry"))?
60            .insert(gateway.address.clone(), gateway);
61        Ok(())
62    }
63
64    /// Runs `request` against the federated gateway at `address` and returns its response.
65    ///
66    /// Requires the federation capability, encodes the request (with `policy`)
67    /// as a `/v1/responses` body, realizes it through the gateway's eval fabric,
68    /// and decodes the reply; errors if the gateway is unknown, denies the
69    /// capability, or returns a non-200 / non-response value.
70    pub fn infer(
71        &self,
72        cx: &mut Cx,
73        address: &str,
74        request: &ModelRequest,
75        policy: &OpenAiFederationPolicy,
76    ) -> Result<ModelResponse> {
77        cx.require(&openai_gateway_federate_capability())?;
78        let gateway = self.gateway(address)?;
79        let body = federated_request_body(&gateway.default_model, request, policy)?;
80        let reply = gateway.fabric.realize(
81            cx,
82            EvalRequest {
83                expr: Expr::Bytes(body),
84                result_shape: None,
85                required_capabilities: Vec::new(),
86                deadline: None,
87                consistency: Consistency::LocalFirst,
88                mode: EvalMode::Eval,
89                answer_limit: None,
90                stream_buffer: None,
91                stream: false,
92                trace: false,
93            },
94        )?;
95        let Some(response) = reply.value.object().downcast_ref::<GatewayResponseValue>() else {
96            return Err(Error::Eval(format!(
97                "federated gateway {address} returned a non-response value"
98            )));
99        };
100        if response.response().status() != 200 {
101            return Err(Error::Eval(format!(
102                "federated gateway {address} returned status {}",
103                response.response().status()
104            )));
105        }
106        responses_body_to_model_response(address, response.response().body())
107    }
108
109    fn gateway(&self, address: &str) -> Result<OpenAiFederatedGateway> {
110        let address = normalize_gateway_address(address)?;
111        self.gateways
112            .lock()
113            .map_err(|_| Error::PoisonedLock("openai federation registry"))?
114            .get(&address)
115            .cloned()
116            .ok_or_else(|| Error::Eval(format!("model_not_found: {address}")))
117    }
118}
119
120impl OpenAiFederatedGateway {
121    /// Returns a federated gateway after validating and normalizing `address`.
122    pub fn new(
123        address: impl AsRef<str>,
124        default_model: impl Into<String>,
125        fabric: EvalFabricRef,
126    ) -> Result<Self> {
127        Ok(Self {
128            address: normalize_gateway_address(address.as_ref())?,
129            default_model: default_model.into(),
130            fabric,
131        })
132    }
133}
134
135impl OpenAiFederationPolicy {
136    /// Returns a federation policy with the given privacy tier and budget entries.
137    pub fn new(privacy: Option<String>, budget: Vec<(String, Expr)>) -> Self {
138        Self { privacy, budget }
139    }
140
141    /// Returns a copy of the policy with `entries` appended to its budget.
142    pub fn with_budget_entries(&self, entries: Vec<(String, Expr)>) -> Self {
143        let mut budget = self.budget.clone();
144        budget.extend(entries);
145        Self {
146            privacy: self.privacy.clone(),
147            budget,
148        }
149    }
150}
151
152fn normalize_gateway_address(address: &str) -> Result<String> {
153    let trimmed = address.trim();
154    let Some(target) = trimmed.strip_prefix("gateway/") else {
155        return Err(Error::Eval(format!(
156            "federated gateway address must start with gateway/, found {address}"
157        )));
158    };
159    if target.is_empty() {
160        return Err(Error::Eval("federated gateway address is empty".to_owned()));
161    }
162    Ok(format!("gateway/{target}"))
163}
164
165fn federated_request_body(
166    model: &str,
167    request: &ModelRequest,
168    policy: &OpenAiFederationPolicy,
169) -> Result<Vec<u8>> {
170    let mut object = Map::new();
171    object.insert("model".to_owned(), JsonValue::String(model.to_owned()));
172    object.insert("input".to_owned(), JsonValue::String(request_text(request)));
173    object.insert("store".to_owned(), JsonValue::Bool(false));
174    if let Some(privacy) = &policy.privacy {
175        object.insert("privacy".to_owned(), JsonValue::String(privacy.clone()));
176    }
177    let budget = budget_json(&policy.budget);
178    if !budget.is_empty() {
179        object.insert("budget".to_owned(), JsonValue::Object(budget));
180    }
181    Ok(crate::objects::canonical_json_bytes(JsonValue::Object(
182        object,
183    )))
184}
185
186fn budget_json(entries: &[(String, Expr)]) -> Map<String, JsonValue> {
187    let mut object = Map::new();
188    for (key, value) in entries {
189        object.insert(key.clone(), expr_to_json(value));
190    }
191    object
192}
193
194// Intentionally-divergent untagged Expr->JSON projection: this module's wire
195// form differs from sim_codec_json::project_expr_to_json and stays local because
196// merging it into the shared JSON projection would change gateway wire data.
197fn expr_to_json(expr: &Expr) -> JsonValue {
198    match expr {
199        Expr::Nil => JsonValue::Null,
200        Expr::Bool(value) => JsonValue::Bool(*value),
201        Expr::String(value) => JsonValue::String(value.clone()),
202        Expr::Symbol(value) if value.namespace.is_none() => {
203            JsonValue::String(value.name.as_ref().to_owned())
204        }
205        Expr::Number(value) => number_literal_json(&value.canonical),
206        Expr::List(items) => match items.as_slice() {
207            [Expr::Symbol(head), Expr::String(value)]
208                if head.namespace.is_none() && head.name.as_ref() == "plan/atom" =>
209            {
210                atom_literal_json(value)
211            }
212            _ => JsonValue::Array(items.iter().map(expr_to_json).collect()),
213        },
214        Expr::Map(entries) => JsonValue::Object(
215            entries
216                .iter()
217                .map(|(key, value)| (json_key(key), expr_to_json(value)))
218                .collect(),
219        ),
220        Expr::Bytes(bytes) => JsonValue::String(hex_encode(bytes)),
221        _ => JsonValue::String(format!("{expr:?}")),
222    }
223}
224
225fn json_key(expr: &Expr) -> String {
226    match expr {
227        Expr::String(value) => value.clone(),
228        Expr::Symbol(value) if value.namespace.is_none() => value.name.as_ref().to_owned(),
229        Expr::Symbol(value) => format!(
230            "{}/{}",
231            value.namespace.as_deref().unwrap_or(""),
232            value.name
233        ),
234        _ => format!("{expr:?}"),
235    }
236}
237
238fn atom_literal_json(value: &str) -> JsonValue {
239    match value {
240        "true" => JsonValue::Bool(true),
241        "false" => JsonValue::Bool(false),
242        _ => number_literal_json(value),
243    }
244}
245
246fn number_literal_json(value: &str) -> JsonValue {
247    if let Ok(unsigned) = value.parse::<u64>() {
248        return JsonValue::Number(Number::from(unsigned));
249    }
250    value
251        .parse::<f64>()
252        .ok()
253        .and_then(Number::from_f64)
254        .map(JsonValue::Number)
255        .unwrap_or_else(|| JsonValue::String(value.to_owned()))
256}
257
258fn responses_body_to_model_response(address: &str, body: &[u8]) -> Result<ModelResponse> {
259    let value: JsonValue = serde_json::from_slice(body)
260        .map_err(|err| Error::Eval(format!("federated gateway returned invalid json: {err}")))?;
261    let object = value
262        .as_object()
263        .ok_or_else(|| Error::Eval("federated gateway response must be an object".to_owned()))?;
264    let output_text = object
265        .get("output_text")
266        .and_then(JsonValue::as_str)
267        .ok_or_else(|| Error::Eval("federated gateway response missing output_text".to_owned()))?;
268    let mut response = ModelResponse::new(
269        Symbol::new(address.to_owned()),
270        address,
271        vec![text_part(output_text)],
272        Symbol::new("stop"),
273    );
274    response.usage = object.get("usage").map(usage_from_json).transpose()?;
275    if let Some(model) = object.get("model").and_then(JsonValue::as_str) {
276        response.extra.push((
277            Expr::Symbol(Symbol::new("federated-model")),
278            Expr::String(model.to_owned()),
279        ));
280    }
281    Ok(response)
282}
283
284fn usage_from_json(value: &JsonValue) -> Result<ModelUsage> {
285    let object = value
286        .as_object()
287        .ok_or_else(|| Error::Eval("federated gateway usage must be an object".to_owned()))?;
288    Ok(ModelUsage {
289        input_tokens: object.get("prompt_tokens").and_then(json_number_to_u64),
290        output_tokens: object.get("completion_tokens").and_then(json_number_to_u64),
291        latency_ms: None,
292        cost_usd: None,
293        extra: Vec::new(),
294    })
295}
sim_lib_openai_server/runtime/federation.rs

sim_lib_openai_server/runtime/
federation.rs
