openai_reassembler/

lib.rs

//! Reassemble OpenAI-compatible SSE streaming responses into non-streaming format.
//!
//! When an OpenAI-compatible API streams a response as Server-Sent Events (SSE),
//! each event contains a partial "chunk" of the final response. This crate provides
//! a function to merge those chunks into the equivalent non-streaming JSON response.
//!
//! # Supported formats
//!
//! - **Chat completions** (`/v1/chat/completions`): merges `choices[].delta` fields
//!   into `choices[].message`, concatenating string values (e.g. `content`, `refusal`)
//!   and assembling `tool_calls` by index. Other non-string delta fields use last-value-wins.
//! - **Legacy completions** (`/v1/completions`): concatenates `choices[].text`.
//! - **Responses API** (`/v1/responses`): extracts the full response from the
//!   `response.completed` event.
//! - **Multiple choices**: tracked independently by `index`.
//! - **Usage**: taken from the final chunk.
//!
//! Format detection is automatic: if any event's `event` field (from
//! `eventsource_stream::Event`) starts with `"response."`, the Responses API
//! path is used; otherwise the completions path.

use serde_json::{Map, Value};

/// Reassemble OpenAI-compatible streaming chunks into a non-streaming response.
///
/// Auto-detects the stream format and dispatches accordingly:
/// - **Responses API**: if any event's `event` field starts with `"response."`,
///   extracts the full response from the `response.completed` event.
/// - **Completions**: otherwise, merges `choices[].delta` / `choices[].text` chunks.
///
/// For the chat and legacy completions endpoints, top-level fields (`id`, `created`,
/// `model`, etc.) are taken from the first chunk. In this completions path, the
/// `object` field has the `.chunk` suffix stripped (e.g. `chat.completion.chunk`
/// → `chat.completion`). Responses API objects are left unchanged.
///
/// Events with empty data or `[DONE]` are skipped.
pub fn reassemble(events: &[eventsource_stream::Event]) -> anyhow::Result<String> {
    let is_responses_api = events.iter().any(|e| e.event.starts_with("response."));
    if is_responses_api {
        return reassemble_responses(events);
    }

    let mut base: Option<Value> = None;
    let mut choices: std::collections::BTreeMap<u64, Map<String, Value>> = Default::default();
    let mut usage = Value::Null;

    for event in events {
        if event.data.is_empty() || event.data == "[DONE]" {
            continue;
        }
        let chunk: Value = serde_json::from_str(&event.data)
            .map_err(|e| anyhow::anyhow!("Invalid chunk JSON: {}", e))?;

        if base.is_none() {
            let mut b = chunk.clone();
            if let Some(obj) = b["object"].as_str() {
                b["object"] = Value::String(obj.replace(".chunk", ""));
            }
            if let Some(m) = b.as_object_mut() {
                m.remove("choices");
                m.remove("usage");
            }
            base = Some(b);
        }

        if !chunk["usage"].is_null() {
            usage = chunk["usage"].clone();
        }

        if let Some(chunk_choices) = chunk["choices"].as_array() {
            for choice in chunk_choices {
                let index = choice["index"].as_u64().unwrap_or(0);
                let merged = choices.entry(index).or_default();

                if !choice["finish_reason"].is_null() {
                    merged.insert("finish_reason".to_string(), choice["finish_reason"].clone());
                }

                // Legacy completions: concatenate "text"
                if let Some(text) = choice["text"].as_str() {
                    let existing = merged
                        .entry("text".to_string())
                        .or_insert(Value::String(String::new()));
                    if let Value::String(s) = existing {
                        s.push_str(text);
                    }
                }

                // Chat completions: merge "delta" into "message"
                if let Some(delta) = choice["delta"].as_object() {
                    let message = merged
                        .entry("message".to_string())
                        .or_insert(Value::Object(Map::new()));
                    if let Value::Object(msg) = message {
                        for (key, value) in delta {
                            if value.is_null() {
                                continue;
                            }
                            match key.as_str() {
                                "tool_calls" => merge_tool_calls(msg, value),
                                _ => merge_delta_field(msg, key, value),
                            }
                        }
                    }
                }
            }
        }
    }

    let mut response = base.unwrap_or(Value::Object(Map::new()));
    let assembled_choices: Vec<Value> = choices
        .into_iter()
        .map(|(index, mut fields)| {
            fields.insert("index".to_string(), Value::Number(index.into()));
            if !fields.contains_key("finish_reason") {
                fields.insert("finish_reason".to_string(), Value::Null);
            }
            Value::Object(fields)
        })
        .collect();
    response["choices"] = Value::Array(assembled_choices);
    response["usage"] = usage;

    Ok(response.to_string())
}

/// Reassemble a Responses API SSE stream into a non-streaming response.
///
/// The Responses API emits typed events (`response.created`, `response.output_text.delta`,
/// etc.). The final `response.completed` event contains the full response object under
/// the `"response"` key. This function finds that event and extracts the response.
fn reassemble_responses(events: &[eventsource_stream::Event]) -> anyhow::Result<String> {
    for event in events.iter().rev() {
        if event.event == "response.completed" {
            let parsed: Value = serde_json::from_str(&event.data)
                .map_err(|e| anyhow::anyhow!("Invalid response.completed JSON: {}", e))?;
            if let Some(response) = parsed.get("response") {
                return serde_json::to_string(response).map_err(Into::into);
            }
            anyhow::bail!(
                "response.completed event JSON does not contain top-level \"response\" field"
            );
        }
    }
    anyhow::bail!("No response.completed event found in Responses API SSE stream")
}

/// Merge streamed tool_calls deltas into the accumulated message.
///
/// Tool calls arrive as an array of deltas, each with an `index` field indicating
/// which tool call slot they belong to. `id` and `type` are set once; `function.name`
/// and `function.arguments` are concatenated across chunks.
fn merge_tool_calls(msg: &mut Map<String, Value>, value: &Value) {
    let Some(arr) = value.as_array() else { return };
    let tc_list = msg
        .entry("tool_calls".to_string())
        .or_insert(Value::Array(vec![]));
    let Value::Array(existing) = tc_list else {
        return;
    };

    for tc_delta in arr {
        let idx = tc_delta["index"].as_u64().unwrap_or(0) as usize;
        while existing.len() <= idx {
            existing.push(Value::Object(Map::new()));
        }
        let slot = existing[idx].as_object_mut().unwrap();

        // Set id and type (arrive once, on the first delta for this tool call)
        for field in ["id", "type"] {
            if let Some(v) = tc_delta.get(field) {
                if !v.is_null() {
                    slot.insert(field.to_string(), v.clone());
                }
            }
        }

        // Concatenate function name and arguments
        if let Some(func) = tc_delta["function"].as_object() {
            let f = slot
                .entry("function".to_string())
                .or_insert(Value::Object(Map::new()))
                .as_object_mut()
                .unwrap();
            for field in ["name", "arguments"] {
                if let Some(s) = func.get(field).and_then(|v| v.as_str()) {
                    let existing = f
                        .entry(field.to_string())
                        .or_insert(Value::String(String::new()));
                    if let Value::String(es) = existing {
                        es.push_str(s);
                    }
                }
            }
        }
    }
}

/// Merge a single delta field into the accumulated message.
///
/// String fields (content, role, refusal, etc.) are concatenated.
/// Non-string fields use last-value-wins.
fn merge_delta_field(msg: &mut Map<String, Value>, key: &str, value: &Value) {
    if let Some(s) = value.as_str() {
        let existing = msg
            .entry(key.to_string())
            .or_insert(Value::String(String::new()));
        if let Value::String(existing_str) = existing {
            existing_str.push_str(s);
        }
    } else {
        msg.insert(key.to_string(), value.clone());
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;
    use std::sync::Once;

    static GENERATE: Once = Once::new();

    /// If BASE_URL and MODEL are set, generate all fixtures once before tests run.
    fn ensure_fixtures() {
        GENERATE.call_once(|| {
            let (Ok(base_url), Ok(model)) = (std::env::var("BASE_URL"), std::env::var("MODEL"))
            else {
                return;
            };
            let api_key = std::env::var("API_KEY").unwrap_or_else(|_| "none".to_string());
            let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
            let fixtures_dir = root.join("fixtures");
            std::fs::create_dir_all(&fixtures_dir).unwrap();

            let cases: Value = serde_json::from_str(
                &std::fs::read_to_string(root.join("test_cases.json")).unwrap(),
            )
            .unwrap();

            let rt = tokio::runtime::Runtime::new().unwrap();
            let client = reqwest::Client::new();

            for (name, case) in cases.as_object().unwrap() {
                let endpoint = case["endpoint"].as_str().unwrap();
                if endpoint.ends_with("/responses") {
                    rt.block_on(record_responses_fixture(
                        &client,
                        &base_url,
                        &api_key,
                        &model,
                        name,
                        case,
                        &fixtures_dir,
                    ));
                } else {
                    rt.block_on(record_fixture(
                        &client,
                        &base_url,
                        &api_key,
                        &model,
                        name,
                        case,
                        &fixtures_dir,
                    ));
                }
            }
        });
    }

    async fn record_fixture(
        client: &reqwest::Client,
        base_url: &str,
        api_key: &str,
        model: &str,
        name: &str,
        case: &Value,
        fixtures_dir: &PathBuf,
    ) {
        let endpoint = case["endpoint"].as_str().unwrap();
        let url = format!("{base_url}{endpoint}");
        let mut body = case["body"].as_object().unwrap().clone();
        body.insert("model".to_string(), Value::String(model.to_string()));
        body.insert("temperature".to_string(), Value::Number(0.into()));
        body.insert("seed".to_string(), Value::Number(42.into()));

        // Non-streaming
        let mut non_stream_body = body.clone();
        non_stream_body.insert("stream".to_string(), Value::Bool(false));
        eprintln!("[{name}] POST {url} (non-streaming)");
        let expected: Value = client
            .post(&url)
            .bearer_auth(api_key)
            .json(&non_stream_body)
            .send()
            .await
            .unwrap_or_else(|e| panic!("{name}: non-streaming request failed: {e}"))
            .json()
            .await
            .unwrap_or_else(|e| panic!("{name}: non-streaming parse failed: {e}"));
        eprintln!("[{name}] non-streaming response received");

        // Streaming
        let mut stream_body = body.clone();
        stream_body.insert("stream".to_string(), Value::Bool(true));
        let mut stream_opts = serde_json::Map::new();
        stream_opts.insert("include_usage".to_string(), Value::Bool(true));
        stream_body.insert("stream_options".to_string(), Value::Object(stream_opts));

        eprintln!("[{name}] POST {url} (streaming)");
        let response_text = client
            .post(&url)
            .bearer_auth(api_key)
            .json(&stream_body)
            .send()
            .await
            .unwrap_or_else(|e| panic!("{name}: streaming request failed: {e}"))
            .text()
            .await
            .unwrap_or_else(|e| panic!("{name}: streaming read failed: {e}"));

        let mut chunks: Vec<Value> = vec![];
        for line in response_text.lines() {
            if let Some(data) = line.strip_prefix("data: ") {
                if data == "[DONE]" {
                    chunks.push(Value::String("[DONE]".to_string()));
                } else if let Ok(parsed) = serde_json::from_str::<Value>(data) {
                    chunks.push(parsed);
                }
            }
        }

        eprintln!("[{name}] streaming response: {} chunks", chunks.len());

        let fixture = serde_json::json!({ "chunks": chunks, "expected": expected });
        let path = fixtures_dir.join(format!("{name}.json"));
        std::fs::write(
            &path,
            serde_json::to_string_pretty(&fixture).unwrap() + "\n",
        )
        .unwrap_or_else(|e| panic!("{name}: failed to write fixture: {e}"));
        eprintln!("[{name}] fixture written to {}", path.display());
    }

    /// Record a fixture for the Responses API.
    ///
    /// Unlike completions, responses SSE events are typed (e.g. `event: response.created`)
    /// and the non-streaming request omits `stream` entirely rather than setting it to false.
    /// Usage is always included in `response.completed` without needing `stream_options`.
    async fn record_responses_fixture(
        client: &reqwest::Client,
        base_url: &str,
        api_key: &str,
        model: &str,
        name: &str,
        case: &Value,
        fixtures_dir: &PathBuf,
    ) {
        let endpoint = case["endpoint"].as_str().unwrap();
        let url = format!("{base_url}{endpoint}");
        let mut body = case["body"].as_object().unwrap().clone();
        body.insert("model".to_string(), Value::String(model.to_string()));
        body.insert("temperature".to_string(), Value::Number(0.into()));
        body.insert("seed".to_string(), Value::Number(42.into()));

        // Non-streaming (no stream field at all for responses API)
        eprintln!("[{name}] POST {url} (non-streaming)");
        let expected: Value = client
            .post(&url)
            .bearer_auth(api_key)
            .json(&body)
            .send()
            .await
            .unwrap_or_else(|e| panic!("{name}: non-streaming request failed: {e}"))
            .json()
            .await
            .unwrap_or_else(|e| panic!("{name}: non-streaming parse failed: {e}"));
        eprintln!("[{name}] non-streaming response received");

        // Streaming
        body.insert("stream".to_string(), Value::Bool(true));

        eprintln!("[{name}] POST {url} (streaming)");
        let response_text = client
            .post(&url)
            .bearer_auth(api_key)
            .json(&body)
            .send()
            .await
            .unwrap_or_else(|e| panic!("{name}: streaming request failed: {e}"))
            .text()
            .await
            .unwrap_or_else(|e| panic!("{name}: streaming read failed: {e}"));

        // Parse SSE events preserving event types (spec-compliant: accumulate
        // data lines until a blank line delimits the event).
        let mut events: Vec<Value> = vec![];
        let mut current_event_type: Option<String> = None;
        let mut current_data_lines: Vec<String> = Vec::new();

        for raw_line in response_text.lines() {
            let line = raw_line.trim_end_matches('\r');
            if line.is_empty() {
                if !current_data_lines.is_empty() {
                    let data_str = current_data_lines.join("\n");
                    if data_str != "[DONE]" {
                        if let Ok(parsed) = serde_json::from_str::<Value>(&data_str) {
                            let event_type = current_event_type.clone().unwrap_or_default();
                            events.push(
                                serde_json::json!({ "event_type": event_type, "data": parsed }),
                            );
                        }
                    }
                }
                current_event_type = None;
                current_data_lines.clear();
            } else if let Some(event_type) = line
                .strip_prefix("event: ")
                .or_else(|| line.strip_prefix("event:"))
            {
                current_event_type = Some(event_type.to_string());
            } else if let Some(data) = line
                .strip_prefix("data: ")
                .or_else(|| line.strip_prefix("data:"))
            {
                current_data_lines.push(data.to_string());
            }
        }

        // Finalize any event not terminated by a trailing blank line
        if !current_data_lines.is_empty() {
            let data_str = current_data_lines.join("\n");
            if data_str != "[DONE]" {
                if let Ok(parsed) = serde_json::from_str::<Value>(&data_str) {
                    let event_type = current_event_type.clone().unwrap_or_default();
                    events.push(
                        serde_json::json!({ "event_type": event_type, "data": parsed }),
                    );
                }
            }
        }

        eprintln!("[{name}] streaming response: {} events", events.len());

        let fixture = serde_json::json!({ "events": events, "expected": expected });
        let path = fixtures_dir.join(format!("{name}.json"));
        std::fs::write(
            &path,
            serde_json::to_string_pretty(&fixture).unwrap() + "\n",
        )
        .unwrap_or_else(|e| panic!("{name}: failed to write fixture: {e}"));
        eprintln!("[{name}] fixture written to {}", path.display());
    }

    /// Recursively compare two JSON values, collecting mismatches.
    /// Fields in `skip` are skipped at any nesting depth.
    fn diff(
        actual: &Value,
        expected: &Value,
        path: &str,
        skip: &[String],
        errors: &mut Vec<String>,
    ) {
        match (actual, expected) {
            (Value::Object(a), Value::Object(e)) => {
                for (key, ev) in e {
                    if skip.iter().any(|s| s == key) {
                        continue;
                    }
                    let p = if path.is_empty() {
                        key.clone()
                    } else {
                        format!("{path}.{key}")
                    };
                    match a.get(key) {
                        Some(av) => diff(av, ev, &p, skip, errors),
                        None => errors.push(format!("{p}: missing from reassembled output")),
                    }
                }
                for key in a.keys() {
                    if skip.iter().any(|s| s == key) {
                        continue;
                    }
                    if !e.contains_key(key) {
                        let p = if path.is_empty() {
                            key.clone()
                        } else {
                            format!("{path}.{key}")
                        };
                        errors.push(format!("{p}: unexpected field in reassembled output"));
                    }
                }
            }
            (Value::Array(a), Value::Array(e)) => {
                if a.len() != e.len() {
                    errors.push(format!(
                        "{path}: array length {}, expected {}",
                        a.len(),
                        e.len()
                    ));
                    return;
                }
                for (i, (av, ev)) in a.iter().zip(e).enumerate() {
                    diff(av, ev, &format!("{path}[{i}]"), skip, errors);
                }
            }
            _ => {
                if actual != expected {
                    // Tool call arguments: compare as parsed JSON (whitespace may differ)
                    if path.ends_with(".arguments") {
                        if let (Some(a), Some(e)) = (actual.as_str(), expected.as_str()) {
                            let ap: Result<Value, _> = serde_json::from_str(a);
                            let ep: Result<Value, _> = serde_json::from_str(e);
                            if let (Ok(ap), Ok(ep)) = (ap, ep) {
                                if ap == ep {
                                    return;
                                }
                            }
                        }
                    }
                    errors.push(format!("{path}: got {actual}, expected {expected}"));
                }
            }
        }
    }

    fn assert_fixture(name: &str) {
        ensure_fixtures();
        let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));

        // Load allowed_mismatches from test_cases.json
        let cases: Value =
            serde_json::from_str(&std::fs::read_to_string(root.join("test_cases.json")).unwrap())
                .unwrap();
        let skip: Vec<String> = cases[name]["allowed_mismatches"]
            .as_array()
            .map(|a| a.iter().map(|v| v.as_str().unwrap().to_string()).collect())
            .unwrap_or_default();

        let path = root.join("fixtures").join(format!("{name}.json"));
        let content = std::fs::read_to_string(&path)
            .unwrap_or_else(|e| panic!("missing fixture {}: {e}", path.display()));
        let fixture: Value = serde_json::from_str(&content).unwrap();

        let events: Vec<eventsource_stream::Event> = fixture["chunks"]
            .as_array()
            .unwrap()
            .iter()
            .map(|chunk| eventsource_stream::Event {
                data: if chunk.is_string() {
                    chunk.as_str().unwrap().to_string()
                } else {
                    chunk.to_string()
                },
                ..Default::default()
            })
            .collect();

        let actual: Value = serde_json::from_str(&reassemble(&events).unwrap()).unwrap();

        let mut errors = vec![];
        diff(&actual, &fixture["expected"], "", &skip, &mut errors);
        if !errors.is_empty() {
            panic!("fixture {name}:\n{}", errors.join("\n"));
        }
    }

    /// Load a Responses API fixture and verify reassembly matches the expected response.
    ///
    /// Responses fixtures store events as `{ "event_type": ..., "data": ... }` objects
    /// under the `"events"` key (not `"chunks"`).
    fn assert_responses_fixture(name: &str) {
        ensure_fixtures();
        let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));

        let cases: Value =
            serde_json::from_str(&std::fs::read_to_string(root.join("test_cases.json")).unwrap())
                .unwrap();
        let skip: Vec<String> = cases[name]["allowed_mismatches"]
            .as_array()
            .map(|a| a.iter().map(|v| v.as_str().unwrap().to_string()).collect())
            .unwrap_or_default();

        let path = root.join("fixtures").join(format!("{name}.json"));
        let content = std::fs::read_to_string(&path)
            .unwrap_or_else(|e| panic!("missing fixture {}: {e}", path.display()));
        let fixture: Value = serde_json::from_str(&content).unwrap();

        let events: Vec<eventsource_stream::Event> = fixture["events"]
            .as_array()
            .unwrap()
            .iter()
            .map(|ev| eventsource_stream::Event {
                event: ev["event_type"]
                    .as_str()
                    .unwrap_or_default()
                    .to_string(),
                data: ev["data"].to_string(),
                ..Default::default()
            })
            .collect();

        let actual: Value = serde_json::from_str(&reassemble(&events).unwrap()).unwrap();

        let mut errors = vec![];
        diff(&actual, &fixture["expected"], "", &skip, &mut errors);
        if !errors.is_empty() {
            panic!("fixture {name}:\n{}", errors.join("\n"));
        }
    }

    macro_rules! fixture_test {
        ($name:ident) => {
            #[test]
            fn $name() {
                assert_fixture(stringify!($name));
            }
        };
    }

    macro_rules! responses_fixture_test {
        ($name:ident) => {
            #[test]
            fn $name() {
                assert_responses_fixture(stringify!($name));
            }
        };
    }

    fixture_test!(chat_simple);
    fixture_test!(chat_tools);
    fixture_test!(legacy_completion);
    responses_fixture_test!(responses_simple);
}