infinity_provider_models_dev/
openai.rs

1//! `OpenAI` Chat Completions backend with SSE streaming.
2//!
3//! Covers every `OpenAI`-compatible endpoint: `OpenAI` itself plus `DeepSeek`,
4//! `Groq`, `Cerebras`, `xAI`, `Mistral`, `Gemini` (`OpenAI`-compat path), `Together`,
5//! `Fireworks`, `OpenRouter`, `NVIDIA NIM`, `Moonshot`/`Kimi`, `ZAI`, `MiniMax`, `Qwen`
6//! (`DashScope` compat), `HuggingFace` router, `Vercel` AI Gateway, `OpenCode` Zen,
7//! `Azure OpenAI`, `Ollama` / `LM Studio` / `vLLM`, and any custom
8//! `INFINITY_OPENAI_BASE_URL` proxy.
9
10use async_stream::stream;
11use futures_util::StreamExt;
12use infinity_provider_protocol::{
13    CompletionError, CompletionRequest, ErrorClass, FinalResponse, ModelStream, StreamChunk,
14    ToolCall, ToolCallDeltaContent, ToolDefinition, Usage,
15};
16use reqwest::Client;
17use serde_json::{Value, json};
18
19use crate::convert::{
20    classify_http_status, error_message_from_body, to_openai_messages, to_openai_tools,
21};
22
23const REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60);
24/// Upper bound for one model-list discovery probe.
25const DISCOVERY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
26
27#[derive(Debug)]
28struct AccumulatedToolCall {
29    id: String,
30    name: String,
31    args: String,
32    name_sent: bool,
33}
34
35/// Stream one completion from an `OpenAI`-compatible endpoint.
36///
37/// # Errors
38///
39/// Returns [`CompletionError::provider`] when the endpoint rejects the
40/// request (with the provider's retry classification) and
41/// [`CompletionError::ResponseError`] when a live stream breaks mid-flight.
42/// Individual stream items may also yield errors.
43pub async fn invoke(
44    client: &Client,
45    base_url: &str,
46    api_key: &str,
47    extra_headers: &[(String, String)],
48    upstream_model: &str,
49    default_max_tokens: Option<u64>,
50    request: CompletionRequest,
51) -> Result<ModelStream, CompletionError> {
52    let body = request_body(upstream_model, default_max_tokens, &request);
53
54    let url = format!("{}/chat/completions", base_url.trim_end_matches('/'));
55    let mut req = client
56        .post(&url)
57        .timeout(REQUEST_TIMEOUT)
58        .header("Content-Type", "application/json")
59        .header("Accept", "text/event-stream");
60    if !api_key.is_empty() {
61        req = req.bearer_auth(api_key);
62    }
63    for (k, v) in extra_headers {
64        req = req.header(k.as_str(), v.as_str());
65    }
66    let resp = req.json(&body).send().await.map_err(transport_error)?;
67
68    if !resp.status().is_success() {
69        let status = resp.status().as_u16();
70        let text = resp.text().await.unwrap_or_default();
71        let class = classify_http_status(status, &text);
72        return Err(CompletionError::provider(
73            class,
74            format!("{base_url} -> {status}: {}", error_message_from_body(&text)),
75        ));
76    }
77
78    let mut byte_stream = resp.bytes_stream();
79    let s = stream! {
80        let mut buf = String::new();
81        let mut tools: Vec<AccumulatedToolCall> = Vec::new();
82        let mut usage: Option<Usage> = None;
83
84        while let Some(chunk) = byte_stream.next().await {
85            let bytes = match chunk {
86                Ok(b) => b,
87                Err(e) => {
88                    yield Err(CompletionError::ResponseError(format!("stream read failed: {e}")));
89                    break;
90                }
91            };
92            buf.push_str(&String::from_utf8_lossy(&bytes));
93            // Process complete `data:` frames; keep the tail.
94            while let Some(end) = buf.find("\n\n").or_else(|| buf.find("\r\n\r\n")) {
95                let frame_end = end.saturating_add(2);
96                let frame: String = buf.drain(..frame_end).collect();
97                buf = buf.trim_start_matches(['\n', '\r']).to_owned();
98                for line in frame.lines() {
99                    let line = line.trim();
100                    let Some(data) = line.strip_prefix("data:") else { continue };
101                    let data = data.trim();
102                    if data.is_empty() || data == "[DONE]" {
103                        continue;
104                    }
105                    let Ok(v) = serde_json::from_str::<Value>(data) else { continue };
106                    if let Some(u) = parse_usage(&v) {
107                        usage = Some(u);
108                    }
109                    for out in parse_choices(&v, &mut tools) {
110                        yield out;
111                    }
112                }
113            }
114        }
115        for tc in std::mem::take(&mut tools) {
116            yield Ok(finish_tool_call(tc));
117        }
118        yield Ok(StreamChunk::Final(FinalResponse { usage }));
119    };
120    Ok(Box::pin(s))
121}
122
123fn request_body(
124    upstream_model: &str,
125    default_max_tokens: Option<u64>,
126    request: &CompletionRequest,
127) -> Value {
128    let (messages, system) =
129        to_openai_messages(request.preamble.as_deref(), &request.chat_history, true);
130    let mut map = serde_json::Map::new();
131    map.insert("model".to_owned(), json!(upstream_model));
132    map.insert(
133        "messages".to_owned(),
134        Value::Array(with_system(messages, system)),
135    );
136    map.insert("stream".to_owned(), Value::Bool(true));
137    map.insert("stream_options".to_owned(), json!({"include_usage": true}));
138    if let Some(max) = request.max_tokens.or(default_max_tokens) {
139        map.insert("max_tokens".to_owned(), json!(max));
140    }
141    if !request.tools.is_empty() {
142        map.insert(
143            "tools".to_owned(),
144            Value::Array(to_openai_tools(&request.tools)),
145        );
146    }
147    let mut body = Value::Object(map);
148    if let Some(extra) = request.additional_params.as_ref() {
149        merge_value(&mut body, extra);
150    }
151    body
152}
153
154fn with_system(mut messages: Vec<Value>, system: Option<String>) -> Vec<Value> {
155    if let Some(s) = system {
156        messages.insert(0, json!({"role": "system", "content": s}));
157    }
158    messages
159}
160
161fn merge_value(dst: &mut Value, src: &Value) {
162    match (dst, src) {
163        (Value::Object(d), Value::Object(s)) => {
164            for (k, v) in s {
165                merge_value(d.entry(k.clone()).or_insert(Value::Null), v);
166            }
167        }
168        (d, s) => *d = s.clone(),
169    }
170}
171
172fn parse_usage(v: &Value) -> Option<Usage> {
173    let u = v.get("usage")?;
174    Some(Usage {
175        input_tokens: u.get("prompt_tokens").and_then(Value::as_u64).unwrap_or(0),
176        output_tokens: u
177            .get("completion_tokens")
178            .and_then(Value::as_u64)
179            .unwrap_or(0),
180        total_tokens: u.get("total_tokens").and_then(Value::as_u64).unwrap_or(0),
181        cached_input_tokens: u
182            .pointer("/prompt_tokens_details/cached_tokens")
183            .and_then(Value::as_u64)
184            .unwrap_or(0),
185    })
186}
187
188/// Grow `tools` to cover `index` and return the slot for it.
189fn slot_for(
190    tools: &mut Vec<AccumulatedToolCall>,
191    index: usize,
192) -> Option<&mut AccumulatedToolCall> {
193    while tools.len() <= index {
194        tools.push(AccumulatedToolCall {
195            id: String::new(),
196            name: String::new(),
197            args: String::new(),
198            name_sent: false,
199        });
200    }
201    tools.get_mut(index)
202}
203
204fn parse_choices(
205    v: &Value,
206    tools: &mut Vec<AccumulatedToolCall>,
207) -> Vec<Result<StreamChunk, CompletionError>> {
208    let choices = v
209        .get("choices")
210        .and_then(Value::as_array)
211        .cloned()
212        .unwrap_or_default();
213    let mut out = Vec::new();
214    for choice in &choices {
215        out.extend(parse_choice(choice, tools));
216    }
217    out
218}
219
220fn parse_choice(
221    choice: &Value,
222    tools: &mut Vec<AccumulatedToolCall>,
223) -> Vec<Result<StreamChunk, CompletionError>> {
224    if let Some(err) = choice.get("error") {
225        let msg = err
226            .get("message")
227            .and_then(Value::as_str)
228            .unwrap_or("stream error");
229        let code = err.get("code").and_then(Value::as_str).unwrap_or("");
230        return vec![Err(CompletionError::provider(
231            classify_http_status(400, &format!("{code} {msg}")),
232            msg.to_owned(),
233        ))];
234    }
235    let mut out = Vec::new();
236    let delta = choice.pointer("/delta").unwrap_or(&Value::Null).clone();
237    if let Some(text) = delta.get("content").and_then(Value::as_str) {
238        if !text.is_empty() {
239            out.push(Ok(StreamChunk::Text(text.to_owned())));
240        }
241    }
242    if let Some(reasoning) = delta.get("reasoning_content").and_then(Value::as_str) {
243        if !reasoning.is_empty() {
244            out.push(Ok(StreamChunk::ReasoningDelta {
245                id: None,
246                text: reasoning.to_owned(),
247            }));
248        }
249    }
250    if let Some(calls) = delta.get("tool_calls").and_then(Value::as_array) {
251        accumulate_delta_tool_calls(calls, tools, &mut out);
252    }
253    // Some servers send the full tool call only in the final message.
254    if let Some(msg) = choice.get("message").and_then(Value::as_object) {
255        apply_full_message(msg, tools, &mut out);
256    }
257    out
258}
259
260fn accumulate_delta_tool_calls(
261    calls: &[Value],
262    tools: &mut Vec<AccumulatedToolCall>,
263    out: &mut Vec<Result<StreamChunk, CompletionError>>,
264) {
265    for call in calls {
266        let raw_index = call.get("index").and_then(Value::as_u64).unwrap_or(0);
267        let Ok(index) = usize::try_from(raw_index) else {
268            continue;
269        };
270        let Some(slot) = slot_for(tools, index) else {
271            continue;
272        };
273        if let Some(id) = call.get("id").and_then(Value::as_str) {
274            if !id.is_empty() {
275                id.clone_into(&mut slot.id);
276            }
277        }
278        let func = call.get("function").unwrap_or(&Value::Null);
279        if let Some(name) = func.get("name").and_then(Value::as_str) {
280            if !name.is_empty() {
281                name.clone_into(&mut slot.name);
282            }
283        }
284        if let Some(frag) = func.get("arguments").and_then(Value::as_str) {
285            slot.args.push_str(frag);
286            let id = if slot.id.is_empty() {
287                format!("tool-{index}")
288            } else {
289                slot.id.clone()
290            };
291            if !slot.name_sent && !slot.name.is_empty() {
292                slot.name_sent = true;
293                out.push(Ok(StreamChunk::ToolCallDelta {
294                    id: id.clone(),
295                    content: ToolCallDeltaContent::Name(slot.name.clone()),
296                }));
297            }
298            if !frag.is_empty() {
299                out.push(Ok(StreamChunk::ToolCallDelta {
300                    id,
301                    content: ToolCallDeltaContent::Delta(frag.to_owned()),
302                }));
303            }
304        }
305        if slot.id.is_empty() {
306            slot.id = format!("tool-{index}");
307        }
308    }
309}
310
311fn apply_full_message(
312    msg: &serde_json::Map<String, Value>,
313    tools: &mut Vec<AccumulatedToolCall>,
314    out: &mut Vec<Result<StreamChunk, CompletionError>>,
315) {
316    let empty = Vec::new();
317    let calls = msg
318        .get("tool_calls")
319        .and_then(Value::as_array)
320        .unwrap_or(&empty);
321    for (index, call) in calls.iter().enumerate() {
322        let id = call
323            .get("id")
324            .and_then(Value::as_str)
325            .unwrap_or("")
326            .to_owned();
327        let func = call.get("function").unwrap_or(&Value::Null);
328        let name = func
329            .get("name")
330            .and_then(Value::as_str)
331            .unwrap_or("")
332            .to_owned();
333        let args = func
334            .get("arguments")
335            .and_then(Value::as_str)
336            .unwrap_or("")
337            .to_owned();
338        if name.is_empty() {
339            continue;
340        }
341        let Some(slot) = slot_for(tools, index) else {
342            continue;
343        };
344        slot.id = if id.is_empty() {
345            format!("tool-{index}")
346        } else {
347            id
348        };
349        slot.name = name;
350        slot.args = args;
351        slot.name_sent = true;
352    }
353    if let Some(text) = msg.get("content").and_then(Value::as_str) {
354        if !text.is_empty() {
355            out.push(Ok(StreamChunk::Text(text.to_owned())));
356        }
357    }
358}
359
360fn finish_tool_call(tc: AccumulatedToolCall) -> StreamChunk {
361    if tc.name.is_empty() && tc.args.is_empty() {
362        return StreamChunk::Text(String::new());
363    }
364    let args: Value = if tc.args.trim().is_empty() {
365        json!({})
366    } else {
367        serde_json::from_str(&tc.args).unwrap_or_else(|_| json!({"_raw": tc.args}))
368    };
369    StreamChunk::ToolCall(ToolCall::new(tc.id, tc.name, args))
370}
371
372fn transport_error(e: reqwest::Error) -> CompletionError {
373    if e.is_timeout() || e.is_connect() || e.is_body() || e.is_decode() {
374        CompletionError::provider(ErrorClass::Transient, format!("transport error: {e}"))
375    } else if e.is_request() {
376        CompletionError::RequestError(Box::new(e))
377    } else {
378        CompletionError::provider(ErrorClass::Transient, format!("transport error: {e}"))
379    }
380}
381
382/// Extract model ids from an OpenAI-style `GET /models` body.
383///
384/// Accepts `{data: [{id, ...}]}` and bare `[{id, ...}]` shapes.
385#[must_use]
386pub fn discovered_ids(body: &Value) -> Vec<String> {
387    let items = body
388        .get("data")
389        .and_then(Value::as_array)
390        .or_else(|| body.as_array());
391    let Some(items) = items else {
392        return Vec::new();
393    };
394    items
395        .iter()
396        .filter_map(|m| m.get("id").and_then(Value::as_str))
397        .map(str::to_owned)
398        .filter(|id| !id.is_empty())
399        .collect()
400}
401
402/// Probe an OpenAI-compatible `GET {base}/models` endpoint.
403///
404/// Returns the advertised ids, or an empty list when the endpoint is down,
405/// refuses, or answers in an unknown shape. Never fails: discovery is a
406/// best-effort enrichment of the curated catalog.
407#[must_use]
408pub async fn discover_models(client: &Client, base_url: &str, api_key: &str) -> Vec<String> {
409    let url = format!("{}/models", base_url.trim_end_matches('/'));
410    let mut req = client
411        .get(&url)
412        .timeout(DISCOVERY_TIMEOUT)
413        .header("Accept", "application/json");
414    if !api_key.is_empty() {
415        req = req.bearer_auth(api_key);
416    }
417    let resp = match req.send().await {
418        Ok(resp) => resp,
419        Err(e) => {
420            tracing::debug!("discovery: {url} unreachable: {e}");
421            return Vec::new();
422        }
423    };
424    if !resp.status().is_success() {
425        tracing::debug!("discovery: {url} -> {}", resp.status());
426        return Vec::new();
427    }
428    match resp.json::<Value>().await {
429        Ok(body) => discovered_ids(&body),
430        Err(e) => {
431            tracing::debug!("discovery: {url} returned invalid JSON: {e}");
432            Vec::new()
433        }
434    }
435}
436
437/// Build a non-streaming request body (used by tests).
438#[must_use]
439pub fn build_body(
440    upstream_model: &str,
441    max_tokens: Option<u64>,
442    tools: &[ToolDefinition],
443    messages: Vec<Value>,
444    system: Option<String>,
445) -> Value {
446    let mut map = serde_json::Map::new();
447    map.insert("model".to_owned(), json!(upstream_model));
448    map.insert(
449        "messages".to_owned(),
450        Value::Array(with_system(messages, system)),
451    );
452    map.insert("stream".to_owned(), Value::Bool(true));
453    if let Some(max) = max_tokens {
454        map.insert("max_tokens".to_owned(), json!(max));
455    }
456    if !tools.is_empty() {
457        map.insert("tools".to_owned(), Value::Array(to_openai_tools(tools)));
458    }
459    Value::Object(map)
460}
461
462#[cfg(test)]
463mod tests {
464    use super::*;
465    use serde_json::json;
466
467    #[test]
468    fn parses_openai_list_shape() {
469        let body = json!({"data": [{"id": "a"}, {"id": "b"}]});
470        assert_eq!(discovered_ids(&body), vec!["a".to_owned(), "b".to_owned()]);
471    }
472
473    #[test]
474    fn parses_bare_array_and_skips_bad_entries() {
475        let body = json!([{"id": "a"}, {"other": 1}, {"id": ""}]);
476        assert_eq!(discovered_ids(&body), vec!["a".to_owned()]);
477    }
478
479    #[test]
480    fn rejects_unknown_shapes() {
481        assert_eq!(discovered_ids(&json!({"models": []})), Vec::<String>::new());
482        assert_eq!(discovered_ids(&Value::Null), Vec::<String>::new());
483    }
484}
infinity_provider_models_dev/openai.rs

infinity_provider_models_dev/
openai.rs
