codecraft/core/src/provider/openai.rs

//! OpenAI-compatible chat completions (streaming). Covers Gemini (OpenAI compat),
//! opencode Go, fal/OpenRouter, Ollama, vLLM, llama.cpp, LM Studio, ...

use super::{ChatRequest, Endpoint, Provider, StreamEvent, http_error};
use crate::config::ModelConfig;
use crate::types::{Message, Part, Role, Usage};
use anyhow::{Result, anyhow};
use async_trait::async_trait;
use eventsource_stream::Eventsource;
use futures::StreamExt;
use serde_json::{Value, json};
use std::collections::BTreeMap;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

pub struct OpenAi {
    client: reqwest::Client,
    ep: Endpoint,
}

impl OpenAi {
    pub fn new(client: reqwest::Client, ep: Endpoint) -> Self {
        Self { client, ep }
    }
}

fn data_url(img: &crate::types::ImageData) -> String {
    format!("data:{};base64,{}", img.media_type, img.data)
}

pub fn to_wire_messages(system: &str, messages: &[Message]) -> Vec<Value> {
    let mut out = vec![json!({"role": "system", "content": system})];
    for m in messages {
        match m.role {
            Role::User => {
                let mut parts = Vec::new();
                let mut tool_images = Vec::new();
                for p in &m.parts {
                    match p {
                        Part::Text { text } if !text.is_empty() => {
                            parts.push(json!({"type": "text", "text": text}))
                        }
                        Part::Summary { text, .. } => parts.push(json!({
                            "type": "text",
                            "text": format!("<conversation-summary>\n{text}\n</conversation-summary>")
                        })),
                        Part::Image { image } => parts.push(json!({
                            "type": "image_url", "image_url": {"url": data_url(image)}
                        })),
                        Part::ToolResult { call_id, name, content, images, is_error, .. } => {
                            let text = if *is_error { format!("ERROR: {content}") } else { content.clone() };
                            out.push(json!({
                                "role": "tool", "tool_call_id": call_id, "name": name,
                                "content": if text.is_empty() { "(no output)".into() } else { text }
                            }));
                            tool_images.extend(images.iter().cloned());
                        }
                        _ => {}
                    }
                }
                // Tool messages cannot carry images: send them as a user message.
                if !tool_images.is_empty() {
                    let mut c = vec![json!({"type": "text", "text": "Images returned by the tool calls above:"})];
                    c.extend(tool_images.iter().map(|i| json!({"type": "image_url", "image_url": {"url": data_url(i)}})));
                    out.push(json!({"role": "user", "content": c}));
                }
                if !parts.is_empty() {
                    out.push(json!({"role": "user", "content": parts}));
                }
            }
            Role::Assistant => {
                let mut text = String::new();
                let mut reasoning = String::new();
                let mut calls = Vec::new();
                for p in &m.parts {
                    match p {
                        Part::Text { text: t } => text.push_str(t),
                        Part::Reasoning { text: t, .. } => reasoning.push_str(t),
                        Part::ToolCall { id, name, input, extra } => {
                            let mut c = json!({
                                "id": id, "type": "function",
                                "function": {"name": name, "arguments": input.to_string()}
                            });
                            if let Some(e) = extra {
                                c["extra_content"] = e.clone();
                            }
                            calls.push(c);
                        }
                        _ => {}
                    }
                }
                let mut msg = json!({"role": "assistant", "content": if text.is_empty() { Value::Null } else { Value::String(text) }});
                if !calls.is_empty() {
                    msg["tool_calls"] = Value::Array(calls);
                }
                // Some reasoning models (DeepSeek, Kimi) expect their reasoning back during tool loops.
                if !reasoning.is_empty() {
                    msg["reasoning_content"] = Value::String(reasoning);
                }
                if msg["content"].is_null() && msg.get("tool_calls").is_none() {
                    continue;
                }
                out.push(msg);
            }
        }
    }
    out
}

/// Splits streamed content into text and inline `<think>…</think>` reasoning.
/// Tags may be split across chunks, so a possible partial tag is held back.
#[derive(Default)]
pub struct ThinkSplitter {
    in_think: bool,
    pending: String,
}

const OPEN: &str = "<think>";
const CLOSE: &str = "</think>";

impl ThinkSplitter {
    /// Returns (is_reasoning, text) pieces that are safe to emit.
    pub fn push(&mut self, chunk: &str) -> Vec<(bool, String)> {
        self.pending.push_str(chunk);
        let mut out = Vec::new();
        loop {
            let tag = if self.in_think { CLOSE } else { OPEN };
            if let Some(pos) = self.pending.find(tag) {
                let before: String = self.pending[..pos].to_string();
                if !before.is_empty() {
                    out.push((self.in_think, before));
                }
                self.pending.drain(..pos + tag.len());
                self.in_think = !self.in_think;
                continue;
            }
            // Keep a trailing prefix of the tag (e.g. "<thi") until the next chunk.
            let keep = (1..tag.len()).rev().find(|n| self.pending.ends_with(&tag[..*n])).unwrap_or(0);
            let emit_len = self.pending.len() - keep;
            if emit_len > 0 {
                let text: String = self.pending.drain(..emit_len).collect();
                out.push((self.in_think, text));
            }
            break;
        }
        // Leading whitespace after </think> is noise.
        out.retain(|(_, t)| !t.is_empty());
        out
    }

    pub fn finish(&mut self) -> Vec<(bool, String)> {
        if self.pending.is_empty() {
            return vec![];
        }
        vec![(self.in_think, std::mem::take(&mut self.pending))]
    }
}

#[derive(Default)]
struct PendingCall {
    id: String,
    name: String,
    args: String,
    extra: Option<Value>,
}

#[async_trait]
impl Provider for OpenAi {
    async fn stream(
        &self,
        req: ChatRequest,
        tx: mpsc::Sender<StreamEvent>,
        cancel: CancellationToken,
    ) -> Result<()> {
        let tools: Vec<Value> = req
            .tools
            .iter()
            .map(|t| json!({
                "type": "function",
                "function": {"name": t.name, "description": t.description, "parameters": t.input_schema}
            }))
            .collect();
        let mut body = json!({
            "model": req.model.id,
            "messages": to_wire_messages(&req.system, &req.messages),
            "stream": true,
            "stream_options": {"include_usage": true},
            "max_tokens": req.max_tokens,
        });
        if !tools.is_empty() {
            body["tools"] = Value::Array(tools);
        }
        if let Some(effort) = &req.model.effort {
            body["reasoning_effort"] = json!(effort);
        }

        for (k, v) in &req.model.options {
            body[k] = v.clone();
        }
        let rb = self.client.post(format!("{}/chat/completions", self.ep.base_url)).json(&body);
        let resp = self.ep.with_session(self.ep.apply(rb), req.session_id.as_deref()).send().await?;
        if !resp.status().is_success() {
            return Err(http_error(resp).await);
        }

        let mut events = resp.bytes_stream().eventsource();
        let mut calls: BTreeMap<u64, PendingCall> = BTreeMap::new();
        let mut think = ThinkSplitter::default();
        let mut usage = Usage::default();
        let mut stop_reason = String::from("stop");

        loop {
            let ev = tokio::select! {
                _ = cancel.cancelled() => return Ok(()),
                ev = events.next() => ev,
            };
            let Some(ev) = ev else { break };
            let ev = ev.map_err(|e| anyhow!("stream error: {e}"))?;
            if ev.data.trim() == "[DONE]" {
                break;
            }
            if ev.data.trim().is_empty() {
                continue;
            }
            let data: Value = serde_json::from_str(&ev.data)?;
            if let Some(err) = data.get("error") {
                return Err(anyhow!("provider error: {}", err["message"].as_str().unwrap_or(&err.to_string())));
            }
            if let Some(u) = data.get("usage").filter(|u| u.is_object()) {
                let cached = u["prompt_tokens_details"]["cached_tokens"].as_u64().unwrap_or(0);
                usage.input_tokens = u["prompt_tokens"].as_u64().unwrap_or(0).saturating_sub(cached);
                usage.cache_read_tokens = cached;
                usage.output_tokens = u["completion_tokens"].as_u64().unwrap_or(0);
                usage.read_reported_cost(u);
            }
            let Some(choice) = data["choices"].get(0) else { continue };
            let d = &choice["delta"];
            if let Some(t) = d["content"].as_str().filter(|t| !t.is_empty()) {
                // Models served by llama.cpp & co. may inline their reasoning as <think>…</think>.
                for (is_reasoning, chunk) in think.push(t) {
                    let ev = if is_reasoning { StreamEvent::ReasoningDelta(chunk) } else { StreamEvent::TextDelta(chunk) };
                    tx.send(ev).await.ok();
                }
            }
            let r = d["reasoning_content"].as_str().or_else(|| d["reasoning"].as_str());
            if let Some(r) = r.filter(|r| !r.is_empty()) {
                tx.send(StreamEvent::ReasoningDelta(r.into())).await.ok();
            }
            if let Some(tcs) = d["tool_calls"].as_array() {
                for (pos, tc) in tcs.iter().enumerate() {
                    let idx = tc["index"].as_u64().unwrap_or(pos as u64);
                    let pc = calls.entry(idx).or_default();
                    if let Some(id) = tc["id"].as_str().filter(|s| !s.is_empty()) {
                        pc.id = id.into();
                    }
                    if let Some(n) = tc["function"]["name"].as_str() {
                        pc.name.push_str(n);
                    }
                    if let Some(a) = tc["function"]["arguments"].as_str() {
                        pc.args.push_str(a);
                    }
                    if let Some(e) = tc.get("extra_content") {
                        pc.extra = Some(e.clone());
                    }
                }
            }
            if let Some(f) = choice["finish_reason"].as_str() {
                stop_reason = f.to_string();
            }
        }

        for (is_reasoning, chunk) in think.finish() {
            let ev = if is_reasoning { StreamEvent::ReasoningDelta(chunk) } else { StreamEvent::TextDelta(chunk) };
            tx.send(ev).await.ok();
        }
        for (i, pc) in calls {
            let input = if pc.args.trim().is_empty() {
                json!({})
            } else {
                serde_json::from_str(&pc.args).unwrap_or_else(|_| json!({"__invalid_json": pc.args}))
            };
            let id = if pc.id.is_empty() { format!("call_{}_{i}", crate::types::new_id()) } else { pc.id };
            tx.send(StreamEvent::ToolCall { id, name: pc.name, input, extra: pc.extra }).await.ok();
        }
        tx.send(StreamEvent::Usage(usage)).await.ok();
        tx.send(StreamEvent::Done { stop_reason }).await.ok();
        Ok(())
    }

    async fn list_models(&self) -> Result<Vec<ModelConfig>> {
        let rb = self.client.get(format!("{}/models", self.ep.base_url));
        let resp = self.ep.apply(rb).send().await?;
        if !resp.status().is_success() {
            return Err(http_error(resp).await);
        }
        let v: Value = resp.json().await?;
        let mut models: Vec<ModelConfig> = v["data"]
            .as_array()
            .map(|a| {
                a.iter()
                    .filter_map(|m| m["id"].as_str())
                    // Gemini lists "models/<id>"; the chat endpoint wants the bare id.
                    .map(|id| ModelConfig {
                        id: id.trim_start_matches("models/").to_string(),
                        ..Default::default()
                    })
                    .collect()
            })
            .unwrap_or_default();
        models.sort_by(|a, b| a.id.cmp(&b.id));
        Ok(models)
    }
}

#[cfg(test)]
mod tests {
    use super::ThinkSplitter;

    fn run(chunks: &[&str]) -> (String, String) {
        let mut s = ThinkSplitter::default();
        let mut text = String::new();
        let mut reasoning = String::new();
        let mut parts: Vec<(bool, String)> = chunks.iter().flat_map(|c| s.push(c)).collect();
        parts.extend(s.finish());
        for (r, t) in parts {
            if r { reasoning.push_str(&t) } else { text.push_str(&t) }
        }
        (text, reasoning)
    }

    #[test]
    fn splits_think_tags_across_chunks() {
        let (text, reasoning) = run(&["<thi", "nk>plan ", "the fix</th", "ink>Here is ", "the answer."]);
        assert_eq!(reasoning, "plan the fix");
        assert_eq!(text, "Here is the answer.");
    }

    #[test]
    fn plain_text_and_lookalikes_pass_through() {
        let (text, reasoning) = run(&["a < b and <thin", "g> stays"]);
        assert_eq!(text, "a < b and <thing> stays");
        assert!(reasoning.is_empty());
    }
}