llmrust 0.1.1

//! Shared implementation for all OpenAI-compatible chat completion APIs.
//!
//! OpenAI, DeepSeek, Moonshot, and OpenRouter all expose the same
//! `/chat/completions` endpoint with an identical body format. This module
//! provides a single [`OpenAiCompatibleProvider`] struct that can be
//! configured for each provider via [`ProviderConfig`] and extra headers,
//! plus timeouts and connection pooling applied uniformly.

use async_trait::async_trait;
use futures::{stream::BoxStream, StreamExt};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::time::Duration;

use crate::providers::http::{build_http_client, REQUEST_TIMEOUT};
use crate::providers::stream_util::line_stream;
use crate::providers::{LlmError, Provider, ProviderConfig, Result};
use crate::types::{
    ChatRequest, ChatResponse, Content, Embedding, EmbeddingRequest, EmbeddingResponse,
    EmbeddingUsage, FinishReason, FunctionCall, LogProbs, Message, ResponseFormat, StreamChunk,
    Tool, ToolCall, ToolChoice, Usage,
};
use std::collections::HashMap;

// ── Shared request / response types ───────────────────────

#[derive(Serialize)]
struct CompChatRequest<'a> {
    model: &'a str,
    messages: &'a [CompMessage],
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    top_p: Option<f64>,
    stream: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    stream_options: Option<StreamOptions>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tools: Option<&'a [Tool]>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_choice: Option<&'a ToolChoice>,
    #[serde(skip_serializing_if = "Option::is_none")]
    response_format: Option<&'a ResponseFormat>,
    #[serde(skip_serializing_if = "Option::is_none")]
    stop: Option<&'a [String]>,
    #[serde(skip_serializing_if = "Option::is_none")]
    n: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    seed: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    presence_penalty: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    frequency_penalty: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    logprobs: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    top_logprobs: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    parallel_tool_calls: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    service_tier: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")]
    store: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    metadata: Option<&'a serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    user: Option<&'a str>,
    /// Provider-specific extra fields flattened into the request body.
    #[serde(flatten)]
    extra: &'a HashMap<String, serde_json::Value>,
}

/// Asks OpenAI-compatible servers to emit a terminal chunk carrying token
/// usage when streaming (supported by OpenAI, DeepSeek, Moonshot, OpenRouter).
#[derive(Serialize)]
struct StreamOptions {
    include_usage: bool,
}

#[derive(Serialize, Deserialize)]
struct CompMessage {
    role: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    content: Option<Content>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    tool_calls: Option<Vec<ToolCall>>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    tool_call_id: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    name: Option<String>,
}

impl From<&Message> for CompMessage {
    fn from(msg: &Message) -> Self {
        let role = match msg.role {
            crate::types::Role::System => "system",
            crate::types::Role::User => "user",
            crate::types::Role::Assistant => "assistant",
            crate::types::Role::Tool => "tool",
        }
        .to_string();
        // OpenAI expects an assistant tool-call turn to carry `content: null`
        // (omitted here) rather than an empty string alongside `tool_calls`.
        let content = if msg.content.is_empty() && msg.tool_calls.is_some() {
            None
        } else {
            Some(msg.content.clone())
        };
        Self {
            role,
            content,
            tool_calls: msg.tool_calls.clone(),
            tool_call_id: msg.tool_call_id.clone(),
            name: msg.name.clone(),
        }
    }
}

#[derive(Deserialize)]
struct CompResponse {
    choices: Vec<CompChoice>,
    model: String,
    usage: Option<CompUsage>,
}

#[derive(Deserialize)]
struct CompChoice {
    message: CompMessage,
    #[serde(default)]
    finish_reason: Option<String>,
    #[serde(default)]
    logprobs: Option<LogProbs>,
}

#[derive(Deserialize)]
struct CompUsage {
    prompt_tokens: u64,
    completion_tokens: u64,
    total_tokens: u64,
}

#[derive(Deserialize)]
struct CompStreamChunk {
    #[serde(default)]
    choices: Vec<CompStreamChoice>,
    #[serde(default)]
    usage: Option<CompUsage>,
}

#[derive(Deserialize)]
struct CompStreamChoice {
    delta: CompDelta,
    finish_reason: Option<String>,
}

#[derive(Deserialize)]
struct CompDelta {
    content: Option<String>,
    #[serde(default)]
    tool_calls: Option<Vec<CompToolCallDelta>>,
}

/// A streamed tool-call fragment. OpenAI-compatible servers stream tool calls
/// incrementally: the first fragment for a given `index` carries the call `id`
/// and function `name`, and later fragments append `arguments` text.
#[derive(Deserialize)]
struct CompToolCallDelta {
    #[serde(default)]
    index: usize,
    #[serde(default)]
    id: Option<String>,
    #[serde(default)]
    function: Option<CompFunctionDelta>,
}

#[derive(Deserialize)]
struct CompFunctionDelta {
    #[serde(default)]
    name: Option<String>,
    #[serde(default)]
    arguments: Option<String>,
}

#[derive(Deserialize)]
struct CompErrorBody {
    error: CompErrorDetail,
}

#[derive(Deserialize)]
struct CompErrorDetail {
    message: String,
}

// ── Embedding request / response types ────────────────────────────

#[derive(Serialize)]
struct CompEmbeddingRequest<'a> {
    model: &'a str,
    input: &'a [String],
    #[serde(skip_serializing_if = "Option::is_none")]
    dimensions: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    user: Option<&'a str>,
    #[serde(flatten)]
    extra: &'a HashMap<String, serde_json::Value>,
}

#[derive(Deserialize)]
struct CompEmbeddingResponse {
    #[serde(default)]
    model: Option<String>,
    data: Vec<CompEmbedding>,
    #[serde(default)]
    usage: Option<CompEmbeddingUsage>,
}

#[derive(Deserialize)]
struct CompEmbedding {
    index: usize,
    embedding: Vec<f32>,
}

#[derive(Deserialize)]
struct CompEmbeddingUsage {
    prompt_tokens: u64,
    total_tokens: u64,
}

/// Accumulates streamed tool-call fragments across SSE lines. OpenAI-style
/// streams deliver tool calls as a series of deltas keyed by `index`: the
/// first delta for an index carries `id` and `name`, later deltas append
/// `arguments` text. [`ToolCallAccumulator::ingest`] folds each delta in and
/// [`ToolCallAccumulator::take`] drains the completed calls on the terminal
/// chunk.
#[derive(Default)]
struct ToolCallAccumulator {
    builders: Vec<ToolCallBuilder>,
}

#[derive(Default)]
struct ToolCallBuilder {
    id: String,
    name: String,
    arguments: String,
}

impl ToolCallAccumulator {
    fn ingest(&mut self, deltas: &[CompToolCallDelta]) {
        for delta in deltas {
            if self.builders.len() <= delta.index {
                self.builders
                    .resize_with(delta.index + 1, ToolCallBuilder::default);
            }
            let builder = &mut self.builders[delta.index];
            if let Some(id) = delta.id.as_deref().filter(|s| !s.is_empty()) {
                builder.id = id.to_string();
            }
            if let Some(function) = &delta.function {
                if let Some(name) = function.name.as_deref().filter(|s| !s.is_empty()) {
                    builder.name = name.to_string();
                }
                if let Some(arguments) = &function.arguments {
                    builder.arguments.push_str(arguments);
                }
            }
        }
    }

    fn take(&mut self) -> Option<Vec<ToolCall>> {
        if self.builders.is_empty() {
            return None;
        }
        let calls: Vec<ToolCall> = self
            .builders
            .drain(..)
            .filter(|b| !b.name.is_empty())
            .map(|b| ToolCall {
                id: b.id,
                call_type: "function".to_string(),
                function: FunctionCall {
                    name: b.name,
                    arguments: b.arguments,
                },
            })
            .collect();
        if calls.is_empty() {
            None
        } else {
            Some(calls)
        }
    }
}

/// Parse a single SSE line from an OpenAI-compatible stream into zero or more
/// [`StreamChunk`]s, threading a [`ToolCallAccumulator`] so streamed tool-call
/// fragments can be reassembled across lines. Lines are guaranteed complete by
/// [`line_stream`], so a failed JSON parse here indicates a genuinely
/// malformed payload rather than a chunk-boundary artifact.
///
/// A streamed response ends with a chunk carrying a `finish_reason`, optionally
/// followed by a choices-less chunk carrying only `usage` (when
/// `stream_options.include_usage` was requested), then a literal `[DONE]`.
/// Accumulated tool calls are drained onto the chunk that carries
/// `finish_reason`.
fn parse_sse_line(tools: &mut ToolCallAccumulator, line: &str) -> Vec<Result<StreamChunk>> {
    let line = line.trim();
    let Some(data) = line.strip_prefix("data: ") else {
        return Vec::new();
    };
    if data == "[DONE]" {
        return vec![Ok(StreamChunk {
            done: true,
            ..Default::default()
        })];
    }
    let parsed = match serde_json::from_str::<CompStreamChunk>(data) {
        Ok(p) => p,
        Err(e) => {
            return vec![Err(LlmError::Parse(format!(
                "failed to parse OpenAI-compatible stream chunk: {e}"
            )))];
        }
    };
    let usage = parsed.usage.map(|u| Usage {
        prompt_tokens: u.prompt_tokens,
        completion_tokens: u.completion_tokens,
        total_tokens: u.total_tokens,
    });
    match parsed.choices.first() {
        Some(choice) => {
            if let Some(deltas) = &choice.delta.tool_calls {
                tools.ingest(deltas);
            }
            let finish_reason = choice.finish_reason.clone().map(FinishReason::from);
            let tool_calls = if finish_reason.is_some() {
                tools.take()
            } else {
                None
            };
            vec![Ok(StreamChunk {
                delta: choice.delta.content.clone().unwrap_or_default(),
                done: finish_reason.is_some(),
                finish_reason,
                usage,
                tool_calls,
            })]
        }
        None => {
            if usage.is_some() {
                vec![Ok(StreamChunk {
                    usage,
                    ..Default::default()
                })]
            } else {
                Vec::new()
            }
        }
    }
}

// ── The unified OpenAI-compatible provider ────────────

/// A generic provider for any OpenAI-compatible `/chat/completions` API.
///
/// Headers are built from:
/// 1. `Authorization: Bearer {api_key}` (always)
/// 2. Any extra headers provided at construction time (e.g. OpenRouter's
///    `HTTP-Referer` / `X-Title`)
pub struct OpenAiCompatibleProvider {
    client: Client,
    api_key: String,
    base_url: String,
    extra_headers: Vec<(String, String)>,
}

impl OpenAiCompatibleProvider {
    /// Create a new provider. Extra headers are appended after the standard
    /// `Authorization` header — they can override it if needed.
    pub fn new(
        config: ProviderConfig,
        extra_headers: impl IntoIterator<Item = (String, String)>,
    ) -> Self {
        let timeout = config
            .timeout_secs
            .map(Duration::from_secs)
            .or(Some(REQUEST_TIMEOUT));
        let base_url = config.base_url.unwrap_or_default();
        Self {
            client: build_http_client(timeout, config.custom_headers.as_ref()),
            api_key: config.api_key,
            base_url,
            extra_headers: extra_headers.into_iter().collect(),
        }
    }

    /// Send a fully-built request body and return the raw HTTP response.
    async fn send(&self, body: &CompChatRequest<'_>) -> Result<reqwest::Response> {
        let mut rb = self
            .client
            .post(format!("{}/chat/completions", self.base_url))
            .header("Authorization", format!("Bearer {}", self.api_key));

        for (k, v) in &self.extra_headers {
            rb = rb.header(k.as_str(), v.as_str());
        }

        let resp = rb.json(body).send().await?;
        Ok(resp)
    }

    /// Common HTTP error -> LlmError conversion.
    async fn parse_error(resp: reqwest::Response) -> LlmError {
        let status = resp.status().as_u16();
        let text = resp.text().await.unwrap_or_default();
        let msg = serde_json::from_str::<CompErrorBody>(&text)
            .map(|e| e.error.message)
            .unwrap_or(text);
        LlmError::Api {
            status,
            message: msg,
        }
    }

    /// Send an embeddings request and return the raw HTTP response.
    async fn send_embed(&self, body: &CompEmbeddingRequest<'_>) -> Result<reqwest::Response> {
        let mut rb = self
            .client
            .post(format!("{}/embeddings", self.base_url))
            .header("Authorization", format!("Bearer {}", self.api_key));

        for (k, v) in &self.extra_headers {
            rb = rb.header(k.as_str(), v.as_str());
        }

        let resp = rb.json(body).send().await?;
        Ok(resp)
    }

    /// Parse a non-streaming response into [`ChatResponse`].
    fn parse_response(parsed: CompResponse) -> ChatResponse {
        let usage = parsed.usage.map(|u| Usage {
            prompt_tokens: u.prompt_tokens,
            completion_tokens: u.completion_tokens,
            total_tokens: u.total_tokens,
        });

        let (content, tool_calls, finish_reason, logprobs) = match parsed.choices.into_iter().next()
        {
            Some(choice) => {
                let content = match choice.message.content {
                    Some(c) => c.as_text(),
                    None => String::new(),
                };
                (
                    content,
                    choice.message.tool_calls,
                    choice.finish_reason.map(FinishReason::from),
                    choice.logprobs,
                )
            }
            None => (String::new(), None, None, None),
        };

        ChatResponse {
            content,
            model: parsed.model,
            usage,
            tool_calls,
            finish_reason,
            logprobs,
        }
    }
}

#[async_trait]
impl Provider for OpenAiCompatibleProvider {
    async fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
        crate::providers::warn_if_unsupported_n("openai-compatible", req.n);
        let messages: Vec<CompMessage> = req.messages.iter().map(CompMessage::from).collect();

        let body = CompChatRequest {
            model: &req.model,
            messages: &messages,
            temperature: req.temperature,
            max_tokens: req.max_tokens,
            top_p: req.top_p,
            stream: false,
            stream_options: None,
            tools: req.tools.as_deref(),
            tool_choice: req.tool_choice.as_ref(),
            response_format: req.response_format.as_ref(),
            stop: req.stop.as_deref(),
            n: req.n,
            seed: req.seed,
            presence_penalty: req.presence_penalty,
            frequency_penalty: req.frequency_penalty,
            logprobs: req.logprobs,
            top_logprobs: req.top_logprobs,
            parallel_tool_calls: req.parallel_tool_calls,
            service_tier: req.service_tier.as_deref(),
            store: req.store,
            metadata: req.metadata.as_ref(),
            user: req.user.as_deref(),
            extra: &req.extra,
        };

        tracing::debug!(
            provider = "openai-compatible",
            model = &req.model,
            "sending chat request"
        );
        let resp = self.send(&body).await?;

        if !resp.status().is_success() {
            let status = resp.status().as_u16();
            let err = Self::parse_error(resp).await;
            tracing::error!(
                provider = "openai-compatible",
                status,
                error_kind = "api_error",
                "API error"
            );
            return Err(err);
        }

        let parsed: CompResponse = resp
            .json()
            .await
            .map_err(|e| LlmError::Parse(format!("OpenAI-compatible parse: {e}")))?;

        let result = Self::parse_response(parsed);
        tracing::debug!(provider = "openai-compatible", model = &result.model, finish_reason = ?result.finish_reason, "chat response received");
        Ok(result)
    }

    async fn stream(&self, req: &ChatRequest) -> Result<BoxStream<'static, Result<StreamChunk>>> {
        crate::providers::warn_if_unsupported_n("openai-compatible", req.n);
        let messages: Vec<CompMessage> = req.messages.iter().map(CompMessage::from).collect();

        let body = CompChatRequest {
            model: &req.model,
            messages: &messages,
            temperature: req.temperature,
            max_tokens: req.max_tokens,
            top_p: req.top_p,
            stream: true,
            stream_options: Some(StreamOptions {
                include_usage: true,
            }),
            tools: req.tools.as_deref(),
            tool_choice: req.tool_choice.as_ref(),
            response_format: req.response_format.as_ref(),
            stop: req.stop.as_deref(),
            n: req.n,
            seed: req.seed,
            presence_penalty: req.presence_penalty,
            frequency_penalty: req.frequency_penalty,
            logprobs: req.logprobs,
            top_logprobs: req.top_logprobs,
            parallel_tool_calls: req.parallel_tool_calls,
            service_tier: req.service_tier.as_deref(),
            store: req.store,
            metadata: req.metadata.as_ref(),
            user: req.user.as_deref(),
            extra: &req.extra,
        };

        tracing::debug!(
            provider = "openai-compatible",
            model = &req.model,
            "sending stream request"
        );
        let resp = self.send(&body).await?;

        if !resp.status().is_success() {
            let status = resp.status().as_u16();
            let err = Self::parse_error(resp).await;
            tracing::error!(
                provider = "openai-compatible",
                status,
                error_kind = "api_error",
                "API error"
            );
            return Err(err);
        }

        let byte_stream = resp
            .bytes_stream()
            .map(|r| r.map_err(|e| LlmError::Stream(e.to_string())));

        let stream = line_stream(byte_stream)
            .scan(ToolCallAccumulator::default(), |tools, line_result| {
                let chunks = match line_result {
                    Ok(line) => parse_sse_line(tools, &line),
                    Err(e) => vec![Err(e)],
                };
                futures::future::ready(Some(futures::stream::iter(chunks)))
            })
            .flatten();

        Ok(stream.boxed())
    }

    async fn embed(&self, req: &EmbeddingRequest) -> Result<EmbeddingResponse> {
        tracing::debug!(
            provider = "openai-compatible",
            model = %req.model,
            input_count = req.input.len(),
            "sending embedding request"
        );

        let body = CompEmbeddingRequest {
            model: &req.model,
            input: &req.input,
            dimensions: req.dimensions,
            user: req.user.as_deref(),
            extra: &req.extra,
        };

        let resp = self.send_embed(&body).await?;

        if !resp.status().is_success() {
            let status = resp.status().as_u16();
            let err = Self::parse_error(resp).await;
            tracing::error!(
                provider = "openai-compatible",
                status,
                error_kind = "api_error",
                "embedding API error"
            );
            return Err(err);
        }

        let parsed: CompEmbeddingResponse = resp
            .json()
            .await
            .map_err(|e| LlmError::Parse(format!("OpenAI-compatible embeddings parse: {e}")))?;

        let data: Vec<Embedding> = parsed
            .data
            .into_iter()
            .map(|d| Embedding {
                index: d.index,
                embedding: d.embedding,
            })
            .collect();

        let usage = parsed.usage.map(|u| EmbeddingUsage {
            prompt_tokens: u.prompt_tokens,
            total_tokens: u.total_tokens,
        });

        let model = parsed.model.unwrap_or_else(|| req.model.clone());

        let result = EmbeddingResponse { model, data, usage };

        tracing::debug!(
            provider = "openai-compatible",
            model = %result.model,
            embedding_count = result.data.len(),
            "embedding response received"
        );
        Ok(result)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{ContentPart, FunctionCall};

    #[test]
    fn serializes_tools_and_tool_choice() {
        let tools = vec![Tool::function(
            "get_weather",
            Some("Get weather".to_string()),
            serde_json::json!({"type": "object"}),
        )];
        let choice = ToolChoice::auto();
        let messages = vec![CompMessage::from(&Message::user("hi"))];
        let body = CompChatRequest {
            model: "gpt-4o",
            messages: &messages,
            temperature: None,
            max_tokens: None,
            top_p: None,
            stream: false,
            stream_options: None,
            tools: Some(tools.as_slice()),
            tool_choice: Some(&choice),
            response_format: None,
            stop: None,
            n: None,
            seed: None,
            presence_penalty: None,
            frequency_penalty: None,
            logprobs: None,
            top_logprobs: None,
            parallel_tool_calls: None,
            service_tier: None,
            store: None,
            metadata: None,
            user: None,
            extra: &HashMap::new(),
        };

        let v = serde_json::to_value(&body).unwrap();
        assert_eq!(v["tools"][0]["type"], "function");
        assert_eq!(v["tools"][0]["function"]["name"], "get_weather");
        assert_eq!(v["tool_choice"], "auto");
    }

    #[test]
    fn parses_tool_calls_from_response() {
        let raw = serde_json::json!({
            "model": "gpt-4o",
            "choices": [{
                "message": {
                    "role": "assistant",
                    "content": null,
                    "tool_calls": [{
                        "id": "call_1",
                        "type": "function",
                        "function": {"name": "get_weather", "arguments": "{}"}
                    }]
                },
                "finish_reason": "tool_calls"
            }]
        })
        .to_string();

        let parsed: CompResponse = serde_json::from_str(&raw).unwrap();
        let resp = OpenAiCompatibleProvider::parse_response(parsed);
        assert_eq!(resp.content, "");
        assert_eq!(resp.finish_reason, Some(FinishReason::ToolCalls));
        let calls = resp.tool_calls.expect("tool_calls present");
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].function.name, "get_weather");
    }

    #[test]
    fn parses_logprobs_from_response() {
        let raw = serde_json::json!({
            "model": "gpt-4o",
            "choices": [{
                "message": {
                    "role": "assistant",
                    "content": "Hi"
                },
                "finish_reason": "stop",
                "logprobs": {
                    "content": [{
                        "token": "Hi",
                        "logprob": -0.25,
                        "bytes": [72, 105],
                        "top_logprobs": [{
                            "token": "Hello",
                            "logprob": -0.5,
                            "bytes": [72, 101, 108, 108, 111]
                        }]
                    }]
                }
            }]
        })
        .to_string();

        let parsed: CompResponse = serde_json::from_str(&raw).unwrap();
        let resp = OpenAiCompatibleProvider::parse_response(parsed);
        assert_eq!(resp.content, "Hi");
        assert_eq!(resp.finish_reason, Some(FinishReason::Stop));

        let logprobs = resp.logprobs.expect("logprobs present");
        assert_eq!(logprobs.content.len(), 1);
        assert_eq!(logprobs.content[0].token, "Hi");
        assert_eq!(logprobs.content[0].logprob, -0.25);
        assert_eq!(logprobs.content[0].bytes, Some(vec![72, 105]));
        assert_eq!(logprobs.content[0].top_logprobs.len(), 1);
        assert_eq!(logprobs.content[0].top_logprobs[0].token, "Hello");
        assert_eq!(logprobs.content[0].top_logprobs[0].logprob, -0.5);
    }

    #[test]
    fn tool_message_serializes_with_id() {
        let comp = CompMessage::from(&Message::tool("call_1", "result"));
        let v = serde_json::to_value(&comp).unwrap();
        assert_eq!(v["role"], "tool");
        assert_eq!(v["content"], "result");
        assert_eq!(v["tool_call_id"], "call_1");
    }

    #[test]
    fn assistant_tool_call_message_omits_empty_content() {
        let msg = Message::assistant_tool_calls(vec![ToolCall {
            id: "call_1".to_string(),
            call_type: "function".to_string(),
            function: FunctionCall {
                name: "f".to_string(),
                arguments: "{}".to_string(),
            },
        }]);
        let comp = CompMessage::from(&msg);
        let v = serde_json::to_value(&comp).unwrap();
        assert!(v.get("content").is_none(), "content should be omitted");
        assert_eq!(v["tool_calls"][0]["id"], "call_1");
    }

    #[test]
    fn user_image_message_serializes_as_content_parts() {
        let msg = Message::user_with_parts(vec![
            ContentPart::text("what is this?"),
            ContentPart::image_url("https://example.com/cat.png"),
        ]);
        let comp = CompMessage::from(&msg);
        let v = serde_json::to_value(&comp).unwrap();
        assert_eq!(v["role"], "user");
        assert_eq!(v["content"][0]["type"], "text");
        assert_eq!(v["content"][0]["text"], "what is this?");
        assert_eq!(v["content"][1]["type"], "image_url");
        assert_eq!(
            v["content"][1]["image_url"]["url"],
            "https://example.com/cat.png"
        );
    }

    #[test]
    fn serializes_sampling_params_and_response_format() {
        let messages = vec![CompMessage::from(&Message::user("hi"))];
        let rf = ResponseFormat::json_object();
        let stop = vec!["STOP".to_string()];
        let metadata = serde_json::json!({"trace_id": "abc"});
        let body = CompChatRequest {
            model: "gpt-4o",
            messages: &messages,
            temperature: Some(0.2),
            max_tokens: Some(64),
            top_p: Some(0.9),
            stream: false,
            stream_options: None,
            tools: None,
            tool_choice: None,
            response_format: Some(&rf),
            stop: Some(stop.as_slice()),
            n: Some(2),
            seed: Some(7),
            presence_penalty: Some(0.5),
            frequency_penalty: Some(-0.25),
            logprobs: Some(true),
            top_logprobs: Some(3),
            parallel_tool_calls: Some(false),
            service_tier: Some("flex"),
            store: Some(true),
            metadata: Some(&metadata),
            user: Some("user-123"),
            extra: &HashMap::new(),
        };

        let v = serde_json::to_value(&body).unwrap();
        assert_eq!(v["response_format"]["type"], "json_object");
        assert_eq!(v["stop"][0], "STOP");
        assert_eq!(v["n"], 2);
        assert_eq!(v["seed"], 7);
        assert_eq!(v["presence_penalty"], 0.5);
        assert_eq!(v["frequency_penalty"], -0.25);
        assert_eq!(v["logprobs"], true);
        assert_eq!(v["top_logprobs"], 3);
        assert_eq!(v["parallel_tool_calls"], false);
        assert_eq!(v["service_tier"], "flex");
        assert_eq!(v["store"], true);
        assert_eq!(v["metadata"]["trace_id"], "abc");
        assert_eq!(v["user"], "user-123");
    }

    #[test]
    fn omits_unset_sampling_params() {
        let messages = vec![CompMessage::from(&Message::user("hi"))];
        let body = CompChatRequest {
            model: "gpt-4o",
            messages: &messages,
            temperature: None,
            max_tokens: None,
            top_p: None,
            stream: false,
            stream_options: None,
            tools: None,
            tool_choice: None,
            response_format: None,
            stop: None,
            n: None,
            seed: None,
            presence_penalty: None,
            frequency_penalty: None,
            logprobs: None,
            top_logprobs: None,
            parallel_tool_calls: None,
            service_tier: None,
            store: None,
            metadata: None,
            user: None,
            extra: &HashMap::new(),
        };

        let v = serde_json::to_value(&body).unwrap();
        assert!(v.get("response_format").is_none());
        assert!(v.get("stop").is_none());
        assert!(v.get("seed").is_none());
        assert!(v.get("n").is_none());
        assert!(v.get("presence_penalty").is_none());
        assert!(v.get("frequency_penalty").is_none());
        assert!(v.get("logprobs").is_none());
        assert!(v.get("top_logprobs").is_none());
        assert!(v.get("parallel_tool_calls").is_none());
        assert!(v.get("service_tier").is_none());
        assert!(v.get("store").is_none());
        assert!(v.get("metadata").is_none());
        assert!(v.get("user").is_none());
    }

    #[test]
    fn stream_accumulates_tool_call_fragments() {
        let mut tools = ToolCallAccumulator::default();
        let lines = [
            r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"get_weather","arguments":"{\"ci"}}]},"finish_reason":null}]}"#,
            r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"ty\":\"SF\"}"}}]},"finish_reason":null}]}"#,
            r#"data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}"#,
        ];
        let mut final_chunk = None;
        for line in lines {
            for chunk in parse_sse_line(&mut tools, line) {
                final_chunk = Some(chunk.unwrap());
            }
        }
        let chunk = final_chunk.expect("a terminal chunk");
        assert!(chunk.done);
        assert_eq!(chunk.finish_reason, Some(FinishReason::ToolCalls));
        let calls = chunk.tool_calls.expect("tool calls present");
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].id, "call_1");
        assert_eq!(calls[0].function.name, "get_weather");
        assert_eq!(calls[0].function.arguments, "{\"city\":\"SF\"}");
    }

    #[test]
    fn stream_without_tool_calls_has_none() {
        let mut tools = ToolCallAccumulator::default();
        let chunks = parse_sse_line(
            &mut tools,
            r#"data: {"choices":[{"delta":{"content":"hi"},"finish_reason":"stop"}]}"#,
        );
        let chunk = chunks.into_iter().next().unwrap().unwrap();
        assert_eq!(chunk.delta, "hi");
        assert!(chunk.done);
        assert!(chunk.tool_calls.is_none());
    }

    #[test]
    fn compat_stream_malformed_data_returns_parse_error() {
        let mut tools = ToolCallAccumulator::default();
        let chunks = parse_sse_line(&mut tools, "data: {not valid json}");
        let err = chunks.into_iter().next().unwrap().unwrap_err();
        assert!(matches!(err, LlmError::Parse(_)));
        assert!(err.to_string().contains("stream chunk"));
    }

    #[test]
    fn compat_stream_ignores_non_data_lines() {
        let mut tools = ToolCallAccumulator::default();
        let chunks = parse_sse_line(&mut tools, ": keep-alive");
        assert!(chunks.is_empty());
    }

    #[test]
    fn compat_stream_done_still_returns_done_chunk() {
        let mut tools = ToolCallAccumulator::default();
        let chunks = parse_sse_line(&mut tools, "data: [DONE]");
        let chunk = chunks.into_iter().next().unwrap().unwrap();
        assert!(chunk.done);
    }

    // ── embedding contract tests ──────────────────────────────────

    use std::sync::{Arc, Mutex};

    #[test]
    fn embed_request_serializes_correctly() {
        let req = EmbeddingRequest::new("model", "hi")
            .with_dimensions(1024)
            .with_user("user-1")
            .with_extra("task", "search");
        let body = CompEmbeddingRequest {
            model: &req.model,
            input: &req.input,
            dimensions: req.dimensions,
            user: req.user.as_deref(),
            extra: &req.extra,
        };
        let v = serde_json::to_value(&body).expect("embedding request should serialize");
        assert_eq!(v["model"], "model");
        assert_eq!(v["input"][0], "hi");
        assert_eq!(v["dimensions"], 1024);
        assert_eq!(v["user"], "user-1");
        assert_eq!(v["task"], "search");
    }

    #[test]
    fn embed_response_parses_correctly() {
        let json = r#"{"model":"text-embedding-3-small","data":[{"index":0,"embedding":[0.1,0.2]},{"index":1,"embedding":[0.3,0.4]}],"usage":{"prompt_tokens":5,"total_tokens":5}}"#;
        let parsed: CompEmbeddingResponse =
            serde_json::from_str(json).expect("embedding response should parse");
        assert_eq!(parsed.model.as_deref(), Some("text-embedding-3-small"));
        assert_eq!(parsed.data.len(), 2);
        assert_eq!(parsed.data[0].index, 0);
        assert_eq!(parsed.data[0].embedding, vec![0.1_f32, 0.2]);
        assert_eq!(parsed.data[1].index, 1);
        assert_eq!(parsed.data[1].embedding, vec![0.3_f32, 0.4]);
        let usage = parsed.usage.expect("usage present");
        assert_eq!(usage.prompt_tokens, 5);
        assert_eq!(usage.total_tokens, 5);
    }

    #[test]
    fn embed_response_falls_back_to_request_model() {
        let json = r#"{"data":[]}"#;
        let parsed: CompEmbeddingResponse =
            serde_json::from_str(json).expect("should parse without model");
        assert!(parsed.model.is_none());
        let resp = EmbeddingResponse {
            model: parsed.model.unwrap_or_else(|| "fallback".to_string()),
            data: vec![],
            usage: None,
        };
        assert_eq!(resp.model, "fallback");
    }

    #[test]
    fn embed_error_response_parses_message() {
        let json = r#"{"error":{"message":"bad input"}}"#;
        let parsed: CompErrorBody = serde_json::from_str(json).expect("error body should parse");
        assert!(parsed.error.message.contains("bad input"));
    }

    #[tokio::test]
    async fn openai_compatible_embed_via_lmrs_client_strips_prefix() {
        use crate::LmrsClient;
        let llm = LmrsClient::new();
        let seen: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
        let seen_clone = Arc::clone(&seen);

        struct CapturingEmbedProvider {
            seen: Arc<Mutex<Option<String>>>,
        }

        #[async_trait]
        impl Provider for CapturingEmbedProvider {
            async fn chat(&self, _: &ChatRequest) -> Result<ChatResponse> {
                unimplemented!()
            }
            async fn stream(
                &self,
                _: &ChatRequest,
            ) -> Result<BoxStream<'static, Result<StreamChunk>>> {
                unimplemented!()
            }
            async fn embed(&self, req: &EmbeddingRequest) -> Result<EmbeddingResponse> {
                *self.seen.lock().unwrap() = Some(req.model.clone());
                Ok(EmbeddingResponse {
                    model: req.model.clone(),
                    data: vec![],
                    usage: None,
                })
            }
        }

        llm.set_custom(
            "openai",
            Arc::new(CapturingEmbedProvider { seen: seen_clone }),
        )
        .await;

        llm.embed("openai/text-embedding-3-small", "hello")
            .await
            .expect("embed should succeed via LmrsClient");

        assert_eq!(
            seen.lock().unwrap().as_deref(),
            Some("text-embedding-3-small")
        );
    }

    // ── HTTP integration tests (axum, requires proxy feature) ────

    #[cfg(feature = "proxy")]
    #[tokio::test]
    async fn openai_compatible_embed_posts_to_embeddings_endpoint() {
        let embed_ok = axum::Router::new().route(
            "/embeddings",
            axum::routing::post(
                |headers: axum::http::HeaderMap,

                 axum::Json(body): axum::Json<serde_json::Value>| async move {
                    assert!(headers.get("authorization").is_some());
                    assert_eq!(body["model"], "text-embedding-3-small");
                    assert_eq!(body["input"][0], "hello");
                    assert_eq!(body["input"][1], "world");
                    axum::Json(serde_json::json!({"data":[],"model":"m"}))
                },
            ),
        );
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let url = format!("http://{}", addr);
        tokio::spawn(async move { axum::serve(listener, embed_ok).await.unwrap() });
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;

        let config = ProviderConfig::new("sk-test").with_base_url(&url);
        let provider = OpenAiCompatibleProvider::new(config, []);
        let req = EmbeddingRequest::batch("text-embedding-3-small", ["hello", "world"]);
        provider.embed(&req).await.expect("embed should succeed");
    }

    #[cfg(feature = "proxy")]
    #[tokio::test]
    async fn openai_compatible_embed_sends_dimensions_user_and_extra() {
        let embed_ok = axum::Router::new().route(
            "/embeddings",
            axum::routing::post(
                |axum::Json(body): axum::Json<serde_json::Value>| async move {
                    assert_eq!(body["dimensions"], 1024);
                    assert_eq!(body["user"], "u");
                    assert_eq!(body["task"], "search");
                    axum::Json(serde_json::json!({"data":[],"model":"m"}))
                },
            ),
        );
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let url = format!("http://{}", addr);
        tokio::spawn(async move { axum::serve(listener, embed_ok).await.unwrap() });
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;

        let config = ProviderConfig::new("sk-test").with_base_url(&url);
        let provider = OpenAiCompatibleProvider::new(config, []);
        let req = EmbeddingRequest::new("m", "hi")
            .with_dimensions(1024)
            .with_user("u")
            .with_extra("task", "search");
        provider.embed(&req).await.expect("embed should succeed");
    }

    #[cfg(feature = "proxy")]
    #[tokio::test]
    async fn openai_compatible_embed_maps_api_error() {
        let embed_400 = axum::Router::new().route(
            "/embeddings",
            axum::routing::post(|| async {
                (
                    axum::http::StatusCode::BAD_REQUEST,
                    axum::Json(serde_json::json!({"error":{"message":"bad input"}})),
                )
            }),
        );
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let url = format!("http://{}", addr);
        tokio::spawn(async move { axum::serve(listener, embed_400).await.unwrap() });
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;

        let config = ProviderConfig::new("sk-test").with_base_url(&url);
        let provider = OpenAiCompatibleProvider::new(config, []);
        let req = EmbeddingRequest::new("m", "hi");
        let err = provider.embed(&req).await.unwrap_err();
        assert!(
            matches!(&err, LlmError::Api { status: 400, message } if message.contains("bad input")),
            "expected Api 400, got: {err:?}"
        );
    }

    #[cfg(feature = "proxy")]
    #[tokio::test]
    async fn openai_provider_embed_strips_prefix_with_real_wrapper() {
        use crate::providers::openai::OpenAIProvider;

        let embed_ok = axum::Router::new().route(
            "/embeddings",
            axum::routing::post(
                |axum::Json(body): axum::Json<serde_json::Value>| async move {
                    assert_eq!(body["model"], "text-embedding-3-small");
                    axum::Json(serde_json::json!({"data":[],"model":"text-embedding-3-small"}))
                },
            ),
        );
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let url = format!("http://{}", addr);
        tokio::spawn(async move { axum::serve(listener, embed_ok).await.unwrap() });
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;

        let config = ProviderConfig::new("sk-test").with_base_url(&url);
        let provider: Arc<dyn Provider> = Arc::new(OpenAIProvider::new(config));

        use crate::LmrsClient;
        let llm = LmrsClient::new();
        llm.set_custom("openai", provider).await;
        llm.embed("openai/text-embedding-3-small", "hello")
            .await
            .expect("embed should succeed via real OpenAIProvider");
    }
}