model-gateway-rs 0.2.6

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use toolcraft_request::{ByteStream, HeaderMap, Request, response::Response};

use crate::{
    error::{Error, Result},
    llm::Llm,
    model::{
        llm::{ChatMessage, ChatMessageContent, LlmInput, LlmOutput},
        role::Role,
    },
};

const DOUBAO_CHAT_ENDPOINT: &str = "api/v3/chat/completions";
const VERSIONED_DOUBAO_CHAT_ENDPOINT: &str = "chat/completions";

/// Native Volcengine Ark/Doubao chat client.
///
/// Doubao gateways commonly expose an OpenAI-compatible chat completions path,
/// while provider-specific fields such as `thinking` are passed through a
/// top-level `extra_body` object.
pub struct DoubaoLlm {
    request: Request,
    model: String,
    chat_endpoint: String,
    default_max_tokens: Option<u32>,
    default_temperature: Option<f32>,
    nested_extra_body: Map<String, Value>,
    extra_body: Map<String, Value>,
}

impl DoubaoLlm {
    pub fn new(base_url: &str, model: &str, api_key: &str) -> Result<Self> {
        let mut request = Request::new()?;
        request.set_base_url(base_url)?;

        let mut headers = HeaderMap::new();
        headers.insert("Content-Type", "application/json".to_string())?;
        headers.insert("Authorization", format!("Bearer {api_key}"))?;
        request.set_default_headers(headers);

        Ok(Self {
            request,
            model: model.to_string(),
            chat_endpoint: default_doubao_chat_endpoint(base_url).to_string(),
            default_max_tokens: None,
            default_temperature: None,
            nested_extra_body: Map::new(),
            extra_body: Map::new(),
        })
    }

    pub fn with_max_tokens(mut self, max_tokens: Option<u32>) -> Self {
        self.default_max_tokens = max_tokens;
        self
    }

    pub fn with_temperature(mut self, temperature: Option<f32>) -> Self {
        self.default_temperature = temperature;
        self
    }

    pub fn with_thinking(mut self, thinking: Option<DoubaoThinking>) -> Result<Self> {
        if let Some(thinking) = thinking {
            self.nested_extra_body
                .insert("thinking".to_string(), serde_json::to_value(thinking)?);
        } else {
            self.nested_extra_body.remove("thinking");
        }
        Ok(self)
    }

    pub fn without_thinking(self) -> Result<Self> {
        self.with_thinking(Some(DoubaoThinking::disabled()))
    }

    pub fn with_extra_body_param(
        mut self,
        key: impl Into<String>,
        value: impl Serialize,
    ) -> Result<Self> {
        self.extra_body
            .insert(key.into(), serde_json::to_value(value)?);
        Ok(self)
    }

    pub fn with_chat_endpoint(mut self, endpoint: impl Into<String>) -> Self {
        self.chat_endpoint = endpoint.into().trim_start_matches('/').to_string();
        self
    }
}

#[async_trait]
impl Llm for DoubaoLlm {
    async fn chat_once(&self, input: LlmInput) -> Result<LlmOutput> {
        let body = DoubaoChatRequest {
            model: self.model.clone(),
            messages: input.messages,
            temperature: self.default_temperature,
            max_tokens: self.default_max_tokens,
            stream: Some(false),
            nested_extra_body: non_empty_map(self.nested_extra_body.clone()),
            extra_body: self.extra_body.clone(),
        };
        let payload = serde_json::to_value(body)?;
        let response = self
            .request
            .post(&self.chat_endpoint, &payload, None)
            .await?;
        parse_chat_response(response).await
    }

    async fn chat_stream(&self, input: LlmInput) -> Result<ByteStream> {
        let body = DoubaoChatRequest {
            model: self.model.clone(),
            messages: input.messages,
            temperature: self.default_temperature,
            max_tokens: self.default_max_tokens,
            stream: Some(true),
            nested_extra_body: non_empty_map(self.nested_extra_body.clone()),
            extra_body: self.extra_body.clone(),
        };
        let payload = serde_json::to_value(body)?;
        self.request
            .post_stream(&self.chat_endpoint, &payload, None)
            .await
            .map_err(Into::into)
    }
}

async fn parse_chat_response(response: Response) -> Result<LlmOutput> {
    let status = response.status();
    let body = response.text().await?;

    if !status.is_success() {
        return Err(Error::ApiError(format_error_body(status.as_u16(), &body)));
    }

    if let Ok(error) = serde_json::from_str::<DoubaoErrorResponse>(&body) {
        return Err(Error::ApiError(error.to_string()));
    }

    let json: DoubaoChatResponse = serde_json::from_str(&body)?;
    Ok(json.into())
}

fn format_error_body(status: u16, body: &str) -> String {
    match serde_json::from_str::<DoubaoErrorResponse>(body) {
        Ok(error) => format!("status={status}, {error}"),
        Err(_) => format!("status={status}, body={body}"),
    }
}

fn default_doubao_chat_endpoint(base_url: &str) -> &'static str {
    let base_url = base_url.trim_end_matches('/');

    if base_url.ends_with("/api/v3") || base_url.ends_with("/v1") {
        VERSIONED_DOUBAO_CHAT_ENDPOINT
    } else {
        DOUBAO_CHAT_ENDPOINT
    }
}

fn non_empty_map(map: Map<String, Value>) -> Option<Map<String, Value>> {
    if map.is_empty() { None } else { Some(map) }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct DoubaoThinking {
    #[serde(rename = "type")]
    kind: DoubaoThinkingType,
}

impl DoubaoThinking {
    pub fn enabled() -> Self {
        Self {
            kind: DoubaoThinkingType::Enabled,
        }
    }

    pub fn disabled() -> Self {
        Self {
            kind: DoubaoThinkingType::Disabled,
        }
    }

    pub fn auto() -> Self {
        Self {
            kind: DoubaoThinkingType::Auto,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum DoubaoThinkingType {
    Enabled,
    Disabled,
    Auto,
}

#[derive(Debug, Clone, Serialize)]
struct DoubaoChatRequest {
    model: String,
    messages: Vec<ChatMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    stream: Option<bool>,
    #[serde(rename = "extra_body")]
    #[serde(skip_serializing_if = "Option::is_none")]
    nested_extra_body: Option<Map<String, Value>>,
    #[serde(flatten)]
    extra_body: Map<String, Value>,
}

#[derive(Debug, Deserialize)]
struct DoubaoChatResponse {
    choices: Vec<DoubaoChoice>,
    usage: Option<DoubaoUsage>,
}

#[derive(Debug, Deserialize)]
struct DoubaoChoice {
    message: DoubaoMessage,
}

#[derive(Debug, Deserialize)]
struct DoubaoMessage {
    role: Role,
    content: Option<ChatMessageContent>,
}

#[derive(Debug, Deserialize)]
struct DoubaoUsage {
    total_tokens: Option<u32>,
}

#[derive(Debug, Deserialize)]
struct DoubaoErrorResponse {
    error: DoubaoError,
}

impl std::fmt::Display for DoubaoErrorResponse {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.error)
    }
}

#[derive(Debug, Deserialize)]
struct DoubaoError {
    message: String,
    #[serde(rename = "type")]
    kind: Option<String>,
    param: Option<String>,
    code: Option<String>,
}

impl std::fmt::Display for DoubaoError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "message={}", self.message)?;
        if let Some(kind) = &self.kind {
            write!(f, ", type={kind}")?;
        }
        if let Some(param) = &self.param {
            write!(f, ", param={param}")?;
        }
        if let Some(code) = &self.code {
            write!(f, ", code={code}")?;
        }
        Ok(())
    }
}

impl From<DoubaoChatResponse> for LlmOutput {
    fn from(response: DoubaoChatResponse) -> Self {
        let message = response
            .choices
            .into_iter()
            .next()
            .map(|choice| ChatMessage {
                role: choice.message.role,
                content: choice.message.content.unwrap_or_default(),
            });

        LlmOutput {
            message,
            usage: response.usage.and_then(|usage| usage.total_tokens),
        }
    }
}

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

    use super::*;

    #[test]
    fn root_base_url_uses_doubao_api_v3_chat_completions_path() {
        assert_eq!(
            default_doubao_chat_endpoint("https://ark.cn-beijing.volces.com"),
            "api/v3/chat/completions"
        );
    }

    #[test]
    fn versioned_base_url_uses_chat_completions_path() {
        assert_eq!(
            default_doubao_chat_endpoint("https://ark.cn-beijing.volces.com/api/v3"),
            "chat/completions"
        );
        assert_eq!(
            default_doubao_chat_endpoint("http://10.100.11.51:4000/v1"),
            "chat/completions"
        );
    }

    #[test]
    fn chat_request_serializes_thinking_under_extra_body() {
        let mut nested_extra_body = Map::new();
        nested_extra_body.insert("thinking".to_string(), json!({"type": "disabled"}));

        let request = DoubaoChatRequest {
            model: "doubao-seed-2-0-pro-260215".to_string(),
            messages: vec![ChatMessage::user("你好，测试一下")],
            temperature: None,
            max_tokens: None,
            stream: Some(false),
            nested_extra_body: Some(nested_extra_body),
            extra_body: Map::new(),
        };

        let value = serde_json::to_value(request).unwrap();

        assert_eq!(value["extra_body"]["thinking"], json!({"type": "disabled"}));
        assert!(value.get("thinking").is_none());
    }
}