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    serde_json::to_vec(&JsonValue::Object(object))
182        .map_err(|err| Error::Eval(format!("failed to encode federated gateway request: {err}")))
183}
184
185fn budget_json(entries: &[(String, Expr)]) -> Map<String, JsonValue> {
186    let mut object = Map::new();
187    for (key, value) in entries {
188        object.insert(key.clone(), expr_to_json(value));
189    }
190    object
191}
192
193// Intentionally-divergent untagged Expr->JSON projection: this module's wire
194// form differs from sim_codec_json::project_expr_to_json and stays local because
195// merging it into the shared JSON projection would change gateway wire data.
196fn expr_to_json(expr: &Expr) -> JsonValue {
197    match expr {
198        Expr::Nil => JsonValue::Null,
199        Expr::Bool(value) => JsonValue::Bool(*value),
200        Expr::String(value) => JsonValue::String(value.clone()),
201        Expr::Symbol(value) if value.namespace.is_none() => {
202            JsonValue::String(value.name.as_ref().to_owned())
203        }
204        Expr::Number(value) => number_literal_json(&value.canonical),
205        Expr::List(items) => match items.as_slice() {
206            [Expr::Symbol(head), Expr::String(value)]
207                if head.namespace.is_none() && head.name.as_ref() == "plan/atom" =>
208            {
209                atom_literal_json(value)
210            }
211            _ => JsonValue::Array(items.iter().map(expr_to_json).collect()),
212        },
213        Expr::Map(entries) => JsonValue::Object(
214            entries
215                .iter()
216                .map(|(key, value)| (json_key(key), expr_to_json(value)))
217                .collect(),
218        ),
219        Expr::Bytes(bytes) => JsonValue::String(hex_encode(bytes)),
220        _ => JsonValue::String(format!("{expr:?}")),
221    }
222}
223
224fn json_key(expr: &Expr) -> String {
225    match expr {
226        Expr::String(value) => value.clone(),
227        Expr::Symbol(value) if value.namespace.is_none() => value.name.as_ref().to_owned(),
228        Expr::Symbol(value) => format!(
229            "{}/{}",
230            value.namespace.as_deref().unwrap_or(""),
231            value.name
232        ),
233        _ => format!("{expr:?}"),
234    }
235}
236
237fn atom_literal_json(value: &str) -> JsonValue {
238    match value {
239        "true" => JsonValue::Bool(true),
240        "false" => JsonValue::Bool(false),
241        _ => number_literal_json(value),
242    }
243}
244
245fn number_literal_json(value: &str) -> JsonValue {
246    if let Ok(unsigned) = value.parse::<u64>() {
247        return JsonValue::Number(Number::from(unsigned));
248    }
249    value
250        .parse::<f64>()
251        .ok()
252        .and_then(Number::from_f64)
253        .map(JsonValue::Number)
254        .unwrap_or_else(|| JsonValue::String(value.to_owned()))
255}
256
257fn responses_body_to_model_response(address: &str, body: &[u8]) -> Result<ModelResponse> {
258    let value: JsonValue = serde_json::from_slice(body)
259        .map_err(|err| Error::Eval(format!("federated gateway returned invalid json: {err}")))?;
260    let object = value
261        .as_object()
262        .ok_or_else(|| Error::Eval("federated gateway response must be an object".to_owned()))?;
263    let output_text = object
264        .get("output_text")
265        .and_then(JsonValue::as_str)
266        .ok_or_else(|| Error::Eval("federated gateway response missing output_text".to_owned()))?;
267    let mut response = ModelResponse::new(
268        Symbol::new(address.to_owned()),
269        address,
270        vec![text_part(output_text)],
271        Symbol::new("stop"),
272    );
273    response.usage = object.get("usage").map(usage_from_json).transpose()?;
274    if let Some(model) = object.get("model").and_then(JsonValue::as_str) {
275        response.extra.push((
276            Expr::Symbol(Symbol::new("federated-model")),
277            Expr::String(model.to_owned()),
278        ));
279    }
280    Ok(response)
281}
282
283fn usage_from_json(value: &JsonValue) -> Result<ModelUsage> {
284    let object = value
285        .as_object()
286        .ok_or_else(|| Error::Eval("federated gateway usage must be an object".to_owned()))?;
287    Ok(ModelUsage {
288        input_tokens: object.get("prompt_tokens").and_then(json_number_to_u64),
289        output_tokens: object.get("completion_tokens").and_then(json_number_to_u64),
290        latency_ms: None,
291        cost_usd: None,
292        extra: Vec::new(),
293    })
294}
sim_lib_openai_server/runtime/federation.rs

sim_lib_openai_server/runtime/
federation.rs
