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

sim_lib_openai_server/runtime/
federation.rs
