certifai/src/infrastructure/provider_client.rs

//! Unified LLM provider dispatch.
//!
//! Routes chat completion requests to LiteLLM, OpenAI, Anthropic, or
//! HuggingFace based on the session's provider setting. All providers
//! except Anthropic use the OpenAI-compatible chat completions format.

use reqwest::Client;
use serde::{Deserialize, Serialize};

use super::server_state::ServerState;

/// OpenAI-compatible chat message used for request bodies.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderMessage {
    pub role: String,
    pub content: String,
}

/// Send a chat completion request to the configured provider.
///
/// # Arguments
///
/// * `state` - Server state (for default LiteLLM URL/model)
/// * `provider` - Provider name (`"litellm"`, `"openai"`, `"anthropic"`, `"huggingface"`)
/// * `model` - Model ID
/// * `messages` - Conversation history
/// * `api_key` - API key (required for non-LiteLLM providers; LiteLLM uses server config)
/// * `stream` - Whether to request streaming
///
/// # Returns
///
/// The raw `reqwest::Response` for the caller to consume (streaming or not).
///
/// # Errors
///
/// Returns an error if the HTTP request fails.
pub async fn send_chat_request(
    state: &ServerState,
    provider: &str,
    model: &str,
    messages: &[ProviderMessage],
    api_key: Option<&str>,
    stream: bool,
) -> Result<reqwest::Response, reqwest::Error> {
    let client = Client::new();

    match provider {
        "openai" => {
            let body = serde_json::json!({
                "model": model,
                "messages": messages,
                "stream": stream,
            });
            client
                .post("https://api.openai.com/v1/chat/completions")
                .header("content-type", "application/json")
                .header(
                    "Authorization",
                    format!("Bearer {}", api_key.unwrap_or_default()),
                )
                .json(&body)
                .send()
                .await
        }
        "anthropic" => {
            // Anthropic uses a different API format -- translate.
            // Extract system message separately, convert roles.
            let system_msg: String = messages
                .iter()
                .filter(|m| m.role == "system")
                .map(|m| m.content.clone())
                .collect::<Vec<_>>()
                .join("\n");

            let anthropic_msgs: Vec<serde_json::Value> = messages
                .iter()
                .filter(|m| m.role != "system")
                .map(|m| {
                    serde_json::json!({
                        "role": m.role,
                        "content": m.content,
                    })
                })
                .collect();

            let mut body = serde_json::json!({
                "model": model,
                "messages": anthropic_msgs,
                "max_tokens": 4096,
                "stream": stream,
            });
            if !system_msg.is_empty() {
                body["system"] = serde_json::Value::String(system_msg);
            }

            client
                .post("https://api.anthropic.com/v1/messages")
                .header("content-type", "application/json")
                .header("x-api-key", api_key.unwrap_or_default())
                .header("anthropic-version", "2023-06-01")
                .json(&body)
                .send()
                .await
        }
        "huggingface" => {
            let url = format!(
                "https://api-inference.huggingface.co/models/{}/v1/chat/completions",
                model
            );
            let body = serde_json::json!({
                "model": model,
                "messages": messages,
                "stream": stream,
            });
            client
                .post(&url)
                .header("content-type", "application/json")
                .header(
                    "Authorization",
                    format!("Bearer {}", api_key.unwrap_or_default()),
                )
                .json(&body)
                .send()
                .await
        }
        // Default: LiteLLM proxy (OpenAI-compatible endpoint)
        _ => {
            let base_url = &state.services.litellm_url;
            let resolved_model = if model.is_empty() {
                &state.services.litellm_model
            } else {
                model
            };
            let url = format!("{}/v1/chat/completions", base_url.trim_end_matches('/'));
            let body = serde_json::json!({
                "model": resolved_model,
                "messages": messages,
                "stream": stream,
            });
            let litellm_key = &state.services.litellm_api_key;
            let mut request = client
                .post(&url)
                .header("content-type", "application/json")
                .json(&body);
            if !litellm_key.is_empty() {
                request = request.header("Authorization", format!("Bearer {litellm_key}"));
            }
            request.send().await
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use pretty_assertions::assert_eq;

    #[test]
    fn provider_message_serde_round_trip() {
        let msg = ProviderMessage {
            role: "assistant".into(),
            content: "Hello, world!".into(),
        };
        let json = serde_json::to_string(&msg).expect("serialize ProviderMessage");
        let back: ProviderMessage =
            serde_json::from_str(&json).expect("deserialize ProviderMessage");
        assert_eq!(msg.role, back.role);
        assert_eq!(msg.content, back.content);
    }

    #[test]
    fn provider_message_deserialize_from_json_literal() {
        let json = r#"{"role":"user","content":"What is Rust?"}"#;
        let msg: ProviderMessage = serde_json::from_str(json).expect("deserialize from literal");
        assert_eq!(msg.role, "user");
        assert_eq!(msg.content, "What is Rust?");
    }
}