llm-universal 0.2.0

use async_trait::async_trait;
use serde_json::Value;

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

/// Client for the Cohere Chat API v2 (`POST {base_url}/chat`).
///
/// Requests use the same message and tool shapes as OpenAI; responses differ.
#[derive(Clone)]
pub struct CohereProvider {
    client: reqwest::Client,
    api_key: String,
    base_url: String,
}

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

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

    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 body = OpenAiProvider::build_body(req);
        // Cohere names the limit `max_tokens`.
        if let Some(max_tokens) = body
            .as_object_mut()
            .and_then(|o| o.remove("max_completion_tokens"))
        {
            body["max_tokens"] = max_tokens;
        }
        body
    }
}

fn parse_response(v: &Value, requested_model: &str) -> Result<ChatResponse> {
    let message = &v["message"];
    if !message.is_object() {
        return Err(LlmError::InvalidResponse("missing field `message`".into()));
    }

    let mut parts: Vec<Part> = message["content"]
        .as_array()
        .into_iter()
        .flatten()
        .filter(|block| block["type"] == "text")
        .filter_map(|block| block["text"].as_str())
        .filter(|text| !text.is_empty())
        .map(|text| Part::Text {
            text: text.to_string(),
        })
        .collect();
    parts.extend(parse_tool_calls(&message["tool_calls"]));

    let stop_reason = match v["finish_reason"].as_str() {
        Some("COMPLETE" | "STOP_SEQUENCE") => StopReason::EndTurn,
        Some("TOOL_CALL") => StopReason::ToolUse,
        Some("MAX_TOKENS") => StopReason::MaxTokens,
        other => StopReason::Other(other.unwrap_or("unknown").to_string()),
    };

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

    Ok(ChatResponse {
        // The response does not echo the model.
        model: requested_model.to_string(),
        message: Message::new(Role::Assistant, parts),
        stop_reason,
        usage,
    })
}

#[async_trait]
impl LlmProvider for CohereProvider {
    async fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
        crate::validate(
            "CohereProvider",
            &self.api_key,
            true,
            &self.base_url,
            64,
            req,
        )?;
        let url = format!("{}/chat", self.base_url);
        let builder = self
            .client
            .post(url)
            .header(
                reqwest::header::AUTHORIZATION,
                crate::secret_header("CohereProvider", &format!("Bearer {}", self.api_key))?,
            )
            .json(&Self::build_body(req));
        let body = send_json(builder).await?;
        parse_response(&body, &req.model)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn renames_max_tokens() {
        let req = ChatRequest::new("command-a")
            .max_tokens(50)
            .message(Message::user("hi"));
        let body = CohereProvider::build_body(&req);
        assert_eq!(body["max_tokens"], 50);
        assert!(body.get("max_completion_tokens").is_none());
        assert_eq!(body["messages"][0]["content"], "hi");
    }

    #[test]
    fn parses_tool_call_response() {
        let v = json!({
            "finish_reason": "TOOL_CALL",
            "message": {
                "role": "assistant",
                "tool_plan": "I will look up the weather.",
                "tool_calls": [{
                    "id": "get_weather_abc",
                    "type": "function",
                    "function": { "name": "get_weather", "arguments": "{\"city\":\"Berlin\"}" }
                }]
            },
            "usage": { "tokens": { "input_tokens": 12, "output_tokens": 8 } }
        });
        let resp = parse_response(&v, "command-a").unwrap();
        assert_eq!(resp.stop_reason, StopReason::ToolUse);
        assert_eq!(resp.tool_calls()[0].2["city"], "Berlin");
        assert_eq!(resp.usage.input_tokens, 12);
        assert_eq!(resp.model, "command-a");
    }

    #[test]
    fn parses_text_response() {
        let v = json!({
            "finish_reason": "COMPLETE",
            "message": { "role": "assistant", "content": [{ "type": "text", "text": "Hello" }] }
        });
        let resp = parse_response(&v, "command-a").unwrap();
        assert_eq!(resp.text(), "Hello");
        assert_eq!(resp.stop_reason, StopReason::EndTurn);
    }
}