llm-universal 0.1.0

use async_trait::async_trait;
use serde_json::{json, Value};

use crate::{
    send_json, ChatRequest, ChatResponse, LlmError, LlmProvider, Message, Part, Result, Role,
    StopReason, Usage,
};

/// Client for the OpenAI Chat Completions API (`POST {base_url}/chat/completions`).
///
/// Works with any OpenAI-compatible server as well: Ollama, vLLM, LM Studio, OpenRouter, Groq.
#[derive(Clone)]
pub struct OpenAiProvider {
    client: reqwest::Client,
    api_key: String,
    base_url: String,
}

impl std::fmt::Debug for OpenAiProvider {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("OpenAiProvider")
            .field("api_key", &"<redacted>")
            .field("base_url", &self.base_url)
            .finish()
    }
}

impl OpenAiProvider {
    pub fn new(api_key: impl Into<String>) -> Self {
        Self {
            client: crate::http_client(),
            api_key: api_key.into(),
            base_url: "https://api.openai.com/v1".to_string(),
        }
    }

    /// Points the client at another server. The key may be empty if the server does not need one.
    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
        self.base_url = base_url.into().trim_end_matches('/').to_string();
        self
    }

    /// Uses your own `reqwest::Client`, for example to share a connection pool or set a proxy.
    pub fn with_client(mut self, client: reqwest::Client) -> Self {
        self.client = client;
        self
    }

    fn build_body(req: &ChatRequest) -> Value {
        let mut messages = Vec::new();
        if let Some(system) = &req.system {
            messages.push(json!({ "role": "system", "content": system }));
        }
        for message in &req.messages {
            messages.extend(to_openai_messages(message));
        }

        let mut body = json!({ "model": req.model, "messages": messages });

        if !req.tools.is_empty() {
            let tools: Vec<Value> = req
                .tools
                .iter()
                .map(|t| {
                    json!({
                        "type": "function",
                        "function": {
                            "name": t.name,
                            "description": t.description,
                            "parameters": t.parameters,
                        }
                    })
                })
                .collect();
            body["tools"] = json!(tools);
        }
        if let Some(max_tokens) = req.max_tokens {
            body["max_completion_tokens"] = json!(max_tokens);
        }
        if let Some(temperature) = req.temperature {
            body["temperature"] = json!(temperature);
        }
        body
    }
}

fn to_openai_messages(message: &Message) -> Vec<Value> {
    if message.role == Role::Tool {
        // OpenAI expects one `tool` message per result.
        return message
            .parts
            .iter()
            .filter_map(|p| match p {
                Part::ToolResult {
                    call_id, content, ..
                } => Some(json!({
                    "role": "tool",
                    "tool_call_id": call_id,
                    "content": content,
                })),
                _ => None,
            })
            .collect();
    }

    let role = match message.role {
        Role::System => "system",
        Role::User => "user",
        Role::Assistant | Role::Tool => "assistant",
    };
    let text = message.text_content();
    let tool_calls: Vec<Value> = message
        .parts
        .iter()
        .filter_map(|p| match p {
            Part::ToolCall {
                id,
                name,
                arguments,
            } => Some(json!({
                "id": id,
                "type": "function",
                "function": { "name": name, "arguments": arguments.to_string() },
            })),
            _ => None,
        })
        .collect();

    let content = if text.is_empty() && !tool_calls.is_empty() {
        Value::Null
    } else {
        json!(text)
    };
    let mut msg = json!({ "role": role, "content": content });
    if !tool_calls.is_empty() {
        msg["tool_calls"] = json!(tool_calls);
    }
    vec![msg]
}

fn parse_response(v: &Value) -> Result<ChatResponse> {
    let choice = v["choices"]
        .get(0)
        .ok_or_else(|| LlmError::InvalidResponse("missing field `choices[0]`".into()))?;
    let message = &choice["message"];

    let mut parts = Vec::new();
    if let Some(text) = message["content"].as_str() {
        if !text.is_empty() {
            parts.push(Part::Text {
                text: text.to_string(),
            });
        }
    }
    // A refusal arrives in its own `refusal` field while `content` is null.
    let refusal = message["refusal"].as_str();
    if let Some(text) = refusal {
        parts.push(Part::Text {
            text: text.to_string(),
        });
    }
    if let Some(calls) = message["tool_calls"].as_array() {
        for call in calls {
            // Arguments are a JSON string. An empty string means a tool without arguments;
            // invalid JSON (e.g. cut off by max tokens) is kept raw instead of failing the call.
            let raw = call["function"]["arguments"].as_str().unwrap_or_default();
            let arguments = if raw.trim().is_empty() {
                json!({})
            } else {
                serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.to_string()))
            };
            parts.push(Part::ToolCall {
                id: call["id"].as_str().unwrap_or_default().to_string(),
                name: call["function"]["name"]
                    .as_str()
                    .unwrap_or_default()
                    .to_string(),
                arguments,
            });
        }
    }

    let stop_reason = match choice["finish_reason"].as_str() {
        _ if refusal.is_some() => StopReason::Refusal,
        Some("stop") => StopReason::EndTurn,
        Some("tool_calls") | Some("function_call") => StopReason::ToolUse,
        Some("length") => StopReason::MaxTokens,
        Some("content_filter") => StopReason::Refusal,
        other => StopReason::Other(other.unwrap_or("unknown").to_string()),
    };

    let tokens =
        |field: &str| u32::try_from(v["usage"][field].as_u64().unwrap_or(0)).unwrap_or(u32::MAX);
    let usage = Usage {
        input_tokens: tokens("prompt_tokens"),
        output_tokens: tokens("completion_tokens"),
    };

    Ok(ChatResponse {
        model: v["model"].as_str().unwrap_or_default().to_string(),
        message: Message::new(Role::Assistant, parts),
        stop_reason,
        usage,
    })
}

#[async_trait]
impl LlmProvider for OpenAiProvider {
    async fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
        crate::validate("OpenAiProvider", &self.api_key, &self.base_url, req)?;
        let url = format!("{}/chat/completions", self.base_url);
        let mut builder = self.client.post(url).json(&Self::build_body(req));
        if !self.api_key.is_empty() {
            builder = builder.bearer_auth(&self.api_key);
        }
        let body = send_json(builder).await?;
        parse_response(&body)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Tool;

    #[test]
    fn builds_tool_request() {
        let req = ChatRequest::new("gpt-4o")
            .system("sys")
            .max_tokens(100)
            .tool(Tool::new("f", "d", json!({"type": "object"})))
            .message(Message::user("hi"));
        let body = OpenAiProvider::build_body(&req);
        assert_eq!(body["messages"][0]["role"], "system");
        assert_eq!(body["messages"][1]["content"], "hi");
        assert_eq!(body["tools"][0]["function"]["name"], "f");
        assert_eq!(body["max_completion_tokens"], 100);
    }

    #[test]
    fn parses_tool_call_response() {
        let v = json!({
            "model": "gpt-4o",
            "choices": [{
                "finish_reason": "tool_calls",
                "message": {
                    "content": null,
                    "tool_calls": [{
                        "id": "call_1",
                        "type": "function",
                        "function": { "name": "f", "arguments": "{\"x\":1}" }
                    }]
                }
            }],
            "usage": { "prompt_tokens": 5, "completion_tokens": 7 }
        });
        let resp = parse_response(&v).unwrap();
        assert_eq!(resp.stop_reason, StopReason::ToolUse);
        assert_eq!(resp.tool_calls()[0].2["x"], 1);
        assert_eq!(resp.usage.output_tokens, 7);
    }

    #[test]
    fn tolerates_empty_and_truncated_arguments() {
        let v = json!({
            "choices": [{
                "finish_reason": "length",
                "message": {
                    "content": null,
                    "tool_calls": [
                        { "id": "a", "function": { "name": "f", "arguments": "" } },
                        { "id": "b", "function": { "name": "f", "arguments": "{\"x\":" } }
                    ]
                }
            }]
        });
        let resp = parse_response(&v).unwrap();
        let calls = resp.tool_calls();
        assert_eq!(calls[0].2, json!({}));
        assert_eq!(calls[1].2, json!("{\"x\":"));
        assert_eq!(resp.stop_reason, StopReason::MaxTokens);
    }

    #[test]
    fn maps_refusal_field() {
        let v = json!({
            "choices": [{
                "finish_reason": "stop",
                "message": { "content": null, "refusal": "I can't help with that." }
            }]
        });
        let resp = parse_response(&v).unwrap();
        assert_eq!(resp.stop_reason, StopReason::Refusal);
        assert_eq!(resp.text(), "I can't help with that.");
    }

    #[test]
    fn debug_redacts_api_key() {
        let debug = format!("{:?}", OpenAiProvider::new("sk-secret"));
        assert!(!debug.contains("sk-secret"));
    }
}