flow_adapter_ai/providers/

openai_compat.rs

//! Shared OpenAI Chat Completions wire format.
//!
//! Both the cloud `openai` provider (api.openai.com) and the `local`
//! provider (Ollama / LM Studio / llama.cpp on localhost) speak the exact
//! same `/v1/chat/completions` request and response shape. This module owns
//! the body builder, non-streaming response parser, and the streaming SSE
//! reader so the two providers can't drift.

use futures_util::StreamExt;
use serde::Deserialize;
use serde_json::json;
use std::time::Instant;

use crate::error::{redact_error_body, CloudAiError};
use crate::request::{
    CloudAiRequest, CloudAiResponse, EmbeddingResponse, ToolDispatcher, ToolSpec,
};
use crate::stream::{LlmStreamEvent, LlmStreamSink};

/// The user message `content`: a plain string when there are no images, or an
/// OpenAI multimodal array (a `text` part plus one `image_url` part per image)
/// when the request carries images. The executor pre-resolves local image
/// paths to `data:` URLs, so each entry is a ready-to-send URL.
fn user_content(req: &CloudAiRequest) -> serde_json::Value {
    if req.images.is_empty() {
        return json!(req.prompt);
    }
    let mut parts = vec![json!({ "type": "text", "text": req.prompt })];
    for url in &req.images {
        parts.push(json!({ "type": "image_url", "image_url": { "url": url } }));
    }
    json!(parts)
}

/// Initial `messages` array (`[system?, user]`). Shared by the single-turn body
/// and the tool loop (which appends assistant/tool turns to it).
fn build_messages(req: &CloudAiRequest) -> Vec<serde_json::Value> {
    let mut messages = Vec::with_capacity(2);
    if let Some(s) = req.system.as_deref().filter(|s| !s.is_empty()) {
        // Chat Completions accepts a `system`-role message as the first
        // entry - strictly preferred over folding the system text into the
        // user message because the model is trained to weight system
        // instructions differently.
        messages.push(json!({ "role": "system", "content": s }));
    }
    messages.push(json!({ "role": "user", "content": user_content(req) }));
    messages
}

/// OpenAI `tools` array (`type: "function"`) from the request's `ToolSpec`s,
/// or `None` when no tools are bound.
fn tools_json(tools: &[ToolSpec]) -> Option<serde_json::Value> {
    if tools.is_empty() {
        return None;
    }
    let arr: Vec<_> = tools
        .iter()
        .map(|t| {
            json!({
                "type": "function",
                "function": {
                    "name": t.name,
                    "description": t.description,
                    "parameters": t.parameters,
                },
            })
        })
        .collect();
    Some(json!(arr))
}

/// Apply the sampling + reasoning params shared by every chat body.
fn apply_params(payload: &mut serde_json::Value, req: &CloudAiRequest) {
    if let Some(m) = req.max_tokens {
        payload["max_tokens"] = json!(m);
    }
    if let Some(t) = req.temperature {
        payload["temperature"] = json!(t);
    }
    if let Some(p) = req.top_p {
        payload["top_p"] = json!(p);
    }
    if let Some(k) = req.top_k {
        // `top_k` isn't part of upstream OpenAI Chat Completions, but
        // llama.cpp / Ollama / LM Studio accept it. Cloud OpenAI safely
        // ignores unknown fields, so emitting it unconditionally for the
        // shared body is harmless.
        payload["top_k"] = json!(k);
    }
    if let Some(stops) = req.stop.as_ref().filter(|s| !s.is_empty()) {
        payload["stop"] = json!(stops);
    }
    if let Some(think) = req.reasoning {
        // llama.cpp / vLLM honor `chat_template_kwargs.enable_thinking` to
        // toggle a model's reasoning mode per request; cloud OpenAI ignores
        // unknown fields, so emitting it for the shared body is harmless.
        payload["chat_template_kwargs"] = json!({ "enable_thinking": think });
    }
}

/// Build a Chat Completions request body from a `CloudAiRequest`.
pub fn build_body(req: &CloudAiRequest) -> serde_json::Value {
    let mut payload = json!({
        "model": req.model,
        "messages": build_messages(req),
    });
    apply_params(&mut payload, req);
    if let Some(tools) = tools_json(&req.tools) {
        payload["tools"] = tools;
    }
    if let Some(schema) = &req.response_schema {
        // `strict: false` keeps the server lenient about the exact JSON-Schema
        // subset (a hand-written schema needn't meet OpenAI's strict rules);
        // servers that don't support `response_format` ignore the unknown field
        // and fall back to the prompt constraint.
        payload["response_format"] = json!({
            "type": "json_schema",
            "json_schema": { "name": "output", "schema": schema, "strict": false },
        });
    }
    payload
}

/// Strip inline reasoning so only the user-facing answer is surfaced: removes
/// `<think>…</think>` spans (some local models inline their chain-of-thought
/// there). Separate `reasoning_content` fields are never read into `text`, so
/// they're already excluded. Idempotent; trims surrounding whitespace.
pub fn strip_reasoning(text: &str) -> String {
    let mut out = String::with_capacity(text.len());
    let mut rest = text;
    while let Some(start) = rest.find("<think>") {
        out.push_str(&rest[..start]);
        match rest[start..].find("</think>") {
            Some(end) => rest = &rest[start + end + "</think>".len()..],
            // Unclosed `<think>` - the model was cut off mid-thought; drop the
            // remainder so no partial reasoning leaks into the answer.
            None => {
                rest = "";
                break;
            }
        }
    }
    out.push_str(rest);
    out.trim().to_string()
}

#[derive(Debug, Deserialize)]
struct RawResponse {
    model: String,
    choices: Vec<Choice>,
    usage: Option<Usage>,
}

#[derive(Debug, Deserialize)]
struct Choice {
    message: ChoiceMessage,
    finish_reason: Option<String>,
}

#[derive(Debug, Deserialize)]
struct ChoiceMessage {
    #[serde(default)]
    content: Option<String>,
}

#[derive(Debug, Deserialize)]
struct Usage {
    #[serde(default)]
    prompt_tokens: u32,
    #[serde(default)]
    completion_tokens: u32,
}

/// Parse a Chat Completions response body. `provider` is the caller's name,
/// stamped into errors and the returned `CloudAiResponse.provider`.
pub fn parse_response(
    provider: &str,
    body: &str,
    latency_ms: u64,
) -> Result<CloudAiResponse, CloudAiError> {
    let raw: RawResponse = serde_json::from_str(body).map_err(|e| CloudAiError::Parse {
        provider: provider.into(),
        reason: e.to_string(),
    })?;

    let choice = raw.choices.into_iter().next().ok_or(CloudAiError::Shape {
        provider: provider.into(),
        detail: "no choices in response".into(),
    })?;

    let content = choice.message.content.unwrap_or_default();
    if content.is_empty() {
        // A reasoning ("thinking") model can spend its whole token budget
        // in `reasoning_content` and get cut off before emitting any
        // `content` - the server reports `finish_reason: "length"`. Surface
        // that so the user knows to raise max tokens or disable thinking,
        // rather than seeing an opaque "empty content".
        let detail = match choice.finish_reason.as_deref() {
            Some("length") => "empty assistant content (finish_reason: length - the model hit \
                 max_tokens, likely while \"thinking\"; raise max tokens or disable the model's \
                 reasoning/thinking mode)"
                .to_string(),
            other => format!(
                "empty assistant content (finish_reason: {})",
                other.unwrap_or("unknown")
            ),
        };
        return Err(CloudAiError::Shape {
            provider: provider.into(),
            detail,
        });
    }

    let usage = raw.usage.unwrap_or(Usage {
        prompt_tokens: 0,
        completion_tokens: 0,
    });
    let text = strip_reasoning(&content);

    Ok(CloudAiResponse {
        provider: provider.into(),
        model: raw.model,
        text,
        finish_reason: choice.finish_reason.unwrap_or_else(|| "unknown".into()),
        input_tokens: usage.prompt_tokens,
        output_tokens: usage.completion_tokens,
        latency_ms,
    })
}

/// Build the request body for a streaming Chat Completions call. Same shape
/// as `build_body` plus `stream: true` and a `stream_options` block that
/// asks the server to send a final `usage` chunk so the caller can report
/// token counts even on streaming responses.
pub fn build_streaming_body(req: &CloudAiRequest) -> serde_json::Value {
    let mut body = build_body(req);
    body["stream"] = json!(true);
    body["stream_options"] = json!({ "include_usage": true });
    body
}

/// Stream an OpenAI-compatible Chat Completions response, emitting each
/// content delta through `sink`. Returns the accumulated full response
/// once the stream closes so the caller still gets a `CloudAiResponse`
/// to surface in the node output.
///
/// Handles the canonical SSE shape (one `data: <json>\n` per chunk plus
/// a terminating `data: [DONE]`). Servers that don't include a `usage`
/// chunk leave token counts at zero, mirroring the non-streaming path.
pub async fn stream_response(
    builder: reqwest::RequestBuilder,
    provider: &str,
    call_id: &str,
    sink: &dyn LlmStreamSink,
) -> Result<CloudAiResponse, CloudAiError> {
    let started = Instant::now();
    let resp = builder.send().await.map_err(|e| {
        let err = CloudAiError::Http {
            provider: provider.into(),
            reason: e.to_string(),
        };
        err
    })?;

    let status = resp.status();
    if !status.is_success() {
        let raw = resp.text().await.unwrap_or_default();
        let err = CloudAiError::Status {
            provider: provider.into(),
            status: status.as_u16(),
            body: redact_error_body(&raw),
        };
        sink.emit(LlmStreamEvent::error(call_id, provider, err.to_string()))
            .await;
        return Err(err);
    }

    let mut stream = resp.bytes_stream();
    let mut buffer = String::new();
    let mut full_text = String::new();
    let mut model_name = String::new();
    let mut finish_reason = String::from("unknown");
    let mut input_tokens: u32 = 0;
    let mut output_tokens: u32 = 0;
    let mut errored: Option<String> = None;

    while let Some(chunk) = stream.next().await {
        let bytes = match chunk {
            Ok(b) => b,
            Err(e) => {
                errored = Some(e.to_string());
                break;
            }
        };
        buffer.push_str(&String::from_utf8_lossy(&bytes));
        // SSE frames are separated by `\n\n`; lines within a frame start
        // with `data: ` (other prefixes like `event:` / `id:` are ignored).
        while let Some(idx) = buffer.find("\n\n") {
            let frame = buffer[..idx].to_string();
            buffer.drain(..idx + 2);
            for raw_line in frame.split('\n') {
                let line = raw_line.trim_start();
                let Some(payload) = line.strip_prefix("data:") else {
                    continue;
                };
                let payload = payload.trim();
                if payload.is_empty() {
                    continue;
                }
                if payload == "[DONE]" {
                    sink.emit(LlmStreamEvent {
                        call_id: call_id.to_string(),
                        provider: provider.to_string(),
                        model: model_name.clone(),
                        delta: String::new(),
                        done: true,
                        error: None,
                    })
                    .await;
                    let latency_ms = started.elapsed().as_millis() as u64;
                    return Ok(CloudAiResponse {
                        provider: provider.to_string(),
                        model: if model_name.is_empty() {
                            "unknown".to_string()
                        } else {
                            model_name
                        },
                        text: full_text,
                        finish_reason,
                        input_tokens,
                        output_tokens,
                        latency_ms,
                    });
                }
                let value: serde_json::Value = match serde_json::from_str(payload) {
                    Ok(v) => v,
                    Err(_) => continue,
                };
                if let Some(name) = value.get("model").and_then(|m| m.as_str()) {
                    if model_name.is_empty() {
                        model_name = name.to_string();
                    }
                }
                if let Some(usage) = value.get("usage") {
                    input_tokens = usage
                        .get("prompt_tokens")
                        .and_then(|v| v.as_u64())
                        .map(|v| v as u32)
                        .unwrap_or(input_tokens);
                    output_tokens = usage
                        .get("completion_tokens")
                        .and_then(|v| v.as_u64())
                        .map(|v| v as u32)
                        .unwrap_or(output_tokens);
                }
                if let Some(choice) = value.get("choices").and_then(|c| c.get(0)) {
                    if let Some(reason) = choice.get("finish_reason").and_then(|r| r.as_str())
                    {
                        finish_reason = reason.to_string();
                    }
                    if let Some(delta) = choice
                        .get("delta")
                        .and_then(|d| d.get("content"))
                        .and_then(|c| c.as_str())
                    {
                        if !delta.is_empty() {
                            full_text.push_str(delta);
                            sink.emit(LlmStreamEvent {
                                call_id: call_id.to_string(),
                                provider: provider.to_string(),
                                model: model_name.clone(),
                                delta: delta.to_string(),
                                done: false,
                                error: None,
                            })
                            .await;
                        }
                    }
                }
            }
        }
    }

    if let Some(err) = errored {
        sink.emit(LlmStreamEvent::error(call_id, provider, err.clone()))
            .await;
        return Err(CloudAiError::Http {
            provider: provider.into(),
            reason: err,
        });
    }
    // Stream closed without `[DONE]` - some servers omit it. Treat the
    // accumulated text as the final response and emit a terminal event.
    sink.emit(LlmStreamEvent {
        call_id: call_id.to_string(),
        provider: provider.to_string(),
        model: model_name.clone(),
        delta: String::new(),
        done: true,
        error: None,
    })
    .await;
    let latency_ms = started.elapsed().as_millis() as u64;
    if full_text.is_empty() {
        return Err(CloudAiError::Shape {
            provider: provider.into(),
            detail: "stream closed before any content was emitted".into(),
        });
    }
    Ok(CloudAiResponse {
        provider: provider.to_string(),
        model: if model_name.is_empty() {
            "unknown".to_string()
        } else {
            model_name
        },
        text: strip_reasoning(&full_text),
        finish_reason,
        input_tokens,
        output_tokens,
        latency_ms,
    })
}

// ---------- embeddings (`/v1/embeddings`) ----------

/// Build an OpenAI-compatible `/v1/embeddings` body. The text to embed is the
/// request `prompt`.
pub fn build_embeddings_body(req: &CloudAiRequest) -> serde_json::Value {
    json!({ "model": req.model, "input": req.prompt })
}

#[derive(Debug, Deserialize)]
struct RawEmbeddings {
    #[serde(default)]
    model: Option<String>,
    data: Vec<EmbeddingDatum>,
}

#[derive(Debug, Deserialize)]
struct EmbeddingDatum {
    embedding: Vec<f32>,
}

/// Parse an OpenAI-compatible embeddings response into the first vector.
pub fn parse_embeddings_response(
    provider: &str,
    model_fallback: &str,
    body: &str,
    latency_ms: u64,
) -> Result<EmbeddingResponse, CloudAiError> {
    let raw: RawEmbeddings = serde_json::from_str(body).map_err(|e| CloudAiError::Parse {
        provider: provider.into(),
        reason: e.to_string(),
    })?;
    let first = raw.data.into_iter().next().ok_or(CloudAiError::Shape {
        provider: provider.into(),
        detail: "no embedding data in response".into(),
    })?;
    let dims = first.embedding.len();
    Ok(EmbeddingResponse {
        provider: provider.into(),
        model: raw.model.unwrap_or_else(|| model_fallback.to_string()),
        embedding: first.embedding,
        dims,
        latency_ms,
    })
}

// ---------- tool-use loop ----------

#[derive(Debug, Deserialize)]
struct RawToolCall {
    #[serde(default)]
    id: String,
    #[serde(default)]
    function: RawToolFn,
}

#[derive(Debug, Default, Deserialize)]
struct RawToolFn {
    #[serde(default)]
    name: String,
    /// OpenAI passes arguments as a JSON-encoded string.
    #[serde(default)]
    arguments: String,
}

struct ParsedTurn {
    content: String,
    tool_calls: Vec<RawToolCall>,
    finish_reason: String,
    model: String,
    input_tokens: u32,
    output_tokens: u32,
}

/// Parse one assistant turn (content + any tool calls + usage) from a
/// non-streaming Chat Completions body.
fn parse_turn(provider: &str, body: &str) -> Result<ParsedTurn, CloudAiError> {
    #[derive(Deserialize)]
    struct R {
        #[serde(default)]
        model: String,
        choices: Vec<C>,
        usage: Option<Usage>,
    }
    #[derive(Deserialize)]
    struct C {
        message: M,
        #[serde(default)]
        finish_reason: Option<String>,
    }
    #[derive(Deserialize)]
    struct M {
        #[serde(default)]
        content: Option<String>,
        #[serde(default)]
        tool_calls: Vec<RawToolCall>,
    }
    let r: R = serde_json::from_str(body).map_err(|e| CloudAiError::Parse {
        provider: provider.into(),
        reason: e.to_string(),
    })?;
    let c = r.choices.into_iter().next().ok_or(CloudAiError::Shape {
        provider: provider.into(),
        detail: "no choices in response".into(),
    })?;
    let usage = r.usage.unwrap_or(Usage {
        prompt_tokens: 0,
        completion_tokens: 0,
    });
    Ok(ParsedTurn {
        content: c.message.content.unwrap_or_default(),
        tool_calls: c.message.tool_calls,
        finish_reason: c.finish_reason.unwrap_or_else(|| "unknown".into()),
        model: r.model,
        input_tokens: usage.prompt_tokens,
        output_tokens: usage.completion_tokens,
    })
}

/// Run the OpenAI-compatible tool-use loop: expose `tools`, and on each turn,
/// if the model requests tool calls, run them through `dispatcher`, append the
/// results, and re-call - until the model answers without tool calls (returned
/// as the `CloudAiResponse`) or `max_iters` is exhausted. `api_key` empty means
/// no bearer header (local servers). Non-streaming: tool-call deltas don't
/// reassemble cleanly over SSE.
#[allow(clippy::too_many_arguments)]
pub async fn run_tools_loop(
    client: &reqwest::Client,
    url: &str,
    provider: &str,
    api_key: &str,
    req: &CloudAiRequest,
    tools: &[ToolSpec],
    dispatcher: &dyn ToolDispatcher,
    max_iters: usize,
) -> Result<CloudAiResponse, CloudAiError> {
    let started = Instant::now();
    let mut messages = build_messages(req);
    let tools_body = tools_json(tools).unwrap_or_else(|| json!([]));
    let mut input_tokens: u32 = 0;
    let mut output_tokens: u32 = 0;

    for _ in 0..max_iters.max(1) {
        let mut body = json!({
            "model": req.model,
            "messages": messages,
            "tools": tools_body,
        });
        apply_params(&mut body, req);

        let mut builder = client.post(url).json(&body);
        if !api_key.is_empty() {
            builder = builder.bearer_auth(api_key);
        }
        let resp = builder.send().await.map_err(|e| CloudAiError::Http {
            provider: provider.into(),
            reason: e.to_string(),
        })?;
        let status = resp.status();
        let raw = resp.text().await.unwrap_or_default();
        if !status.is_success() {
            return Err(CloudAiError::Status {
                provider: provider.into(),
                status: status.as_u16(),
                body: redact_error_body(&raw),
            });
        }

        let turn = parse_turn(provider, &raw)?;
        input_tokens = input_tokens.saturating_add(turn.input_tokens);
        output_tokens = output_tokens.saturating_add(turn.output_tokens);

        if turn.tool_calls.is_empty() {
            return Ok(CloudAiResponse {
                provider: provider.into(),
                model: if turn.model.is_empty() {
                    req.model.clone()
                } else {
                    turn.model
                },
                text: strip_reasoning(&turn.content),
                finish_reason: turn.finish_reason,
                input_tokens,
                output_tokens,
                latency_ms: started.elapsed().as_millis() as u64,
            });
        }

        // Echo the assistant's tool-call turn back into the history…
        messages.push(json!({
            "role": "assistant",
            "content": turn.content,
            "tool_calls": turn.tool_calls.iter().map(|tc| json!({
                "id": tc.id,
                "type": "function",
                "function": { "name": tc.function.name, "arguments": tc.function.arguments },
            })).collect::<Vec<_>>(),
        }));

        // …then run each tool and append its result (errors fed back as the
        // result so the model can recover rather than the run hard-failing).
        for tc in &turn.tool_calls {
            let args: serde_json::Value =
                serde_json::from_str(&tc.function.arguments).unwrap_or_else(|_| json!({}));
            let content = match dispatcher.call(&tc.function.name, &args).await {
                Ok(v) => v.to_string(),
                Err(e) => json!({ "error": e }).to_string(),
            };
            messages.push(json!({
                "role": "tool",
                "tool_call_id": tc.id,
                "content": content,
            }));
        }
    }

    Err(CloudAiError::Shape {
        provider: provider.into(),
        detail: format!("tool loop did not converge within {max_iters} iterations"),
    })
}