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

infinity_provider_models_dev/
openai.rs
