anda_engine 0.11.17

//! DeepSeek API client implementation for Anda Engine
//!
//! This module provides integration with DeepSeek's API, including:
//! - Client configuration and management
//! - Completion model handling
//! - Response parsing and conversion to Anda's internal formats

use anda_core::{
    AgentOutput, BoxError, BoxPinFut, CompletionFeatures, CompletionRequest, ContentPart,
    FunctionDefinition, Json, Message, Resource, Usage as ModelUsage,
};
use log::{Level::Debug, log_enabled};
use serde::{Deserialize, Serialize};
use serde_json::json;

use super::{CompletionFeaturesDyn, pruned_placeholder, request_client_builder};
use crate::{rfc3339_datetime, unix_ms};

// ================================================================
// Main DeepSeek Client
// ================================================================
const API_BASE_URL: &str = "https://api.deepseek.com";
pub static DEFAULT_COMPLETION_MODEL: &str = "deepseek-chat";

/// DeepSeek API client configuration and HTTP client
#[derive(Clone)]
pub struct Client {
    endpoint: String,
    api_key: String,
    http: reqwest::Client,
}

impl Client {
    /// Creates a new DeepSeek client instance with the provided API key
    ///
    /// # Arguments
    /// * `api_key` - DeepSeek API key for authentication
    ///
    /// # Returns
    /// Configured DeepSeek client instance
    pub fn new(api_key: &str, endpoint: Option<String>) -> Self {
        let endpoint = endpoint.unwrap_or_else(|| API_BASE_URL.to_string());
        let endpoint = if endpoint.is_empty() {
            API_BASE_URL.to_string()
        } else {
            endpoint
        };
        Self {
            endpoint,
            api_key: api_key.to_string(),
            http: request_client_builder()
                .build()
                .expect("DeepSeek reqwest client should build"),
        }
    }

    /// Sets a custom HTTP client for the client
    pub fn with_client(self, http: reqwest::Client) -> Self {
        Self {
            endpoint: self.endpoint,
            api_key: self.api_key,
            http,
        }
    }

    /// Creates a POST request builder for the specified API path
    fn post(&self, path: &str) -> reqwest::RequestBuilder {
        let url = format!("{}{}", self.endpoint, path);
        self.http.post(url).bearer_auth(&self.api_key)
    }

    /// Creates a new completion model instance using the default DeepSeek model
    pub fn completion_model(&self, model: &str) -> CompletionModel {
        CompletionModel::new(
            self.clone(),
            if model.is_empty() {
                DEFAULT_COMPLETION_MODEL
            } else {
                model
            },
        )
    }
}

/// Token usage statistics from DeepSeek API responses
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Usage {
    /// Number of tokens used in the prompt
    pub prompt_tokens: usize,
    /// Number of tokens used in the completion
    pub completion_tokens: usize,
    /// Number of tokens retrieved from the cache
    #[serde(default)]
    pub prompt_cache_hit_tokens: usize,
}

impl std::fmt::Display for Usage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Prompt tokens: {} completion tokens: {}",
            self.prompt_tokens, self.completion_tokens
        )
    }
}

/// Completion response from DeepSeek API
#[derive(Debug, Deserialize, Serialize)]
pub struct CompletionResponse {
    /// Unique identifier for the completion
    pub id: String,
    /// Object type (typically "chat.completion")
    pub object: String,
    /// Creation timestamp
    pub created: u64,
    /// Model used for the completion
    pub model: String,
    /// List of completion choices
    pub choices: Vec<Choice>,
    /// Token usage statistics
    pub usage: Option<Usage>,
}

impl CompletionResponse {
    fn try_into(
        mut self,
        raw_history: Vec<Json>,
        chat_history: Vec<Message>,
    ) -> Result<AgentOutput, BoxError> {
        let timestamp = unix_ms();
        let mut output = AgentOutput {
            raw_history,
            chat_history,
            usage: self
                .usage
                .as_ref()
                .map(|u| ModelUsage {
                    input_tokens: u.prompt_tokens as u64,
                    output_tokens: u.completion_tokens as u64,
                    cached_tokens: u.prompt_cache_hit_tokens as u64,
                    requests: 1,
                })
                .unwrap_or_default(),
            ..Default::default()
        };

        let choice = self.choices.pop().ok_or("No completion choice")?;
        if !matches!(choice.finish_reason.as_str(), "stop" | "tool_calls") {
            output.failed_reason = Some(choice.finish_reason);
        } else {
            output.raw_history.push(json!(&choice.message));
            let mut msg: Message = choice.message.into();
            msg.name = Some(self.model);
            msg.timestamp = Some(timestamp);
            output.content = msg.text().unwrap_or_default();
            output.tool_calls = msg.tool_calls();
            output.chat_history.push(msg);
        }

        Ok(output)
    }

    fn maybe_failed(&self) -> bool {
        !self.choices.iter().any(|choice| {
            matches!(choice.finish_reason.as_str(), "stop" | "tool_calls")
                && (!choice.message.content.is_empty() || choice.message.tool_calls.is_some())
        })
    }
}

#[derive(Debug, Deserialize, Serialize)]
pub struct MessageInput {
    pub role: String,

    #[serde(default)]
    pub content: String,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
}

fn to_message_input(msg: &Message) -> Vec<MessageInput> {
    let mut res = Vec::new();
    for content in msg.content.iter() {
        match content {
            ContentPart::Text { text } => res.push(MessageInput {
                role: msg.role.clone(),
                content: text.clone(),
                tool_call_id: None,
            }),
            ContentPart::ToolOutput {
                output, call_id, ..
            } => res.push(MessageInput {
                role: msg.role.clone(),
                content: serde_json::to_string(output).unwrap_or_default(),
                tool_call_id: call_id.clone(),
            }),
            v => res.push(MessageInput {
                role: msg.role.clone(),
                content: serde_json::to_string(v).unwrap_or_default(),
                tool_call_id: None,
            }),
        }
    }
    res
}

/// Individual completion choice from DeepSeek API
#[derive(Debug, Deserialize, Serialize)]
pub struct Choice {
    pub index: usize,
    pub message: MessageOutput,
    pub finish_reason: String,
}

/// Output message structure from DeepSeek API
#[derive(Debug, Deserialize, Serialize)]
pub struct MessageOutput {
    pub role: String,
    #[serde(default)]
    pub content: String,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_content: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCallOutput>>,
}

impl From<MessageOutput> for Message {
    fn from(msg: MessageOutput) -> Self {
        let mut content = Vec::new();
        if !msg.content.is_empty() {
            content.push(ContentPart::Text { text: msg.content });
        }
        if let Some(text) = msg.reasoning_content {
            content.push(ContentPart::Reasoning { text });
        }
        if let Some(tool_calls) = msg.tool_calls {
            for tc in tool_calls {
                content.push(ContentPart::ToolCall {
                    name: tc.function.name,
                    args: serde_json::from_str(&tc.function.arguments).unwrap_or_default(),
                    call_id: Some(tc.id),
                });
            }
        }
        Self {
            role: msg.role,
            content,
            ..Default::default()
        }
    }
}

/// Tool call output structure from DeepSeek API
#[derive(Debug, Deserialize, Serialize)]
pub struct ToolCallOutput {
    pub id: String,
    pub r#type: String,
    pub function: Function,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ToolDefinition {
    pub r#type: String,
    pub function: FunctionDefinition,
}

impl From<FunctionDefinition> for ToolDefinition {
    fn from(f: FunctionDefinition) -> Self {
        Self {
            r#type: "function".into(),
            function: f,
        }
    }
}

#[derive(Debug, Deserialize, Serialize)]
pub struct Function {
    pub name: String,
    pub arguments: String,
}

/// Completion model wrapper for DeepSeek API
#[derive(Clone)]
pub struct CompletionModel {
    /// DeepSeek client instance
    client: Client,
    /// Model identifier
    pub model: String,
}

impl CompletionModel {
    /// Creates a new completion model instance
    ///
    /// # Arguments
    /// * `client` - DeepSeek client instance
    /// * `model` - Model identifier string
    pub fn new(client: Client, model: &str) -> Self {
        Self {
            client,
            model: model.to_string(),
        }
    }
}

impl CompletionFeatures for CompletionModel {
    fn model_name(&self) -> String {
        self.model.clone()
    }

    async fn completion(
        &self,
        req: CompletionRequest,
        _resources: Vec<Resource>,
    ) -> Result<AgentOutput, BoxError> {
        CompletionFeaturesDyn::completion(self, req).await
    }
}

impl CompletionFeaturesDyn for CompletionModel {
    fn model_name(&self) -> String {
        self.model.clone()
    }

    /// Prune a DeepSeek chat-completions message JSON in-place.
    ///
    /// DeepSeek uses the same flat shape as legacy OpenAI chat-completions:
    /// `{role, content: String, tool_call_id?}`. Tool outputs carry
    /// `tool_call_id`; we truncate their `content` to a short placeholder.
    /// Other messages (user/assistant text) are left alone since `content` is
    /// already a plain string.
    fn prune_raw_message(&self, value: &mut Json) -> usize {
        let Some(obj) = value.as_object_mut() else {
            return 0;
        };
        let is_tool_output = obj.get("role").and_then(|r| r.as_str()) == Some("tool")
            || obj.contains_key("tool_call_id");
        if !is_tool_output {
            return 0;
        }
        let Some(c) = obj.get_mut("content") else {
            return 0;
        };
        let placeholder = pruned_placeholder(1);
        if c.as_str() == Some(placeholder.as_str()) {
            return 0;
        }
        *c = Json::String(placeholder);
        1
    }

    fn completion(&self, mut req: CompletionRequest) -> BoxPinFut<Result<AgentOutput, BoxError>> {
        let model = self.model.clone();
        let client = self.client.clone();

        Box::pin(async move {
            let timestamp = unix_ms();
            let mut raw_history: Vec<Json> = Vec::new();
            let mut chat_history: Vec<Message> = Vec::new();

            if !req.instructions.is_empty() {
                raw_history.push(json!(MessageInput {
                    role: "system".into(),
                    content: req.instructions.clone(),
                    tool_call_id: None,
                }));
            };

            raw_history.append(&mut req.raw_history);
            let skip_raw = raw_history.len();

            for msg in req.chat_history {
                let val = to_message_input(&msg);
                for v in val {
                    raw_history.push(serde_json::to_value(&v)?);
                }
            }

            if let Some(mut msg) = req
                .documents
                .to_message(&rfc3339_datetime(timestamp).unwrap())
            {
                msg.timestamp = Some(timestamp);
                let val = to_message_input(&msg);
                for v in val {
                    raw_history.push(serde_json::to_value(&v)?);
                }
                chat_history.push(msg);
            }

            let mut content = req.content;
            if !req.prompt.is_empty() {
                content.push(req.prompt.into());
            }
            if !content.is_empty() {
                let msg = Message {
                    role: req.role.unwrap_or_else(|| "user".to_string()),
                    content,
                    timestamp: Some(timestamp),
                    ..Default::default()
                };

                let val = to_message_input(&msg);
                for v in val {
                    raw_history.push(serde_json::to_value(&v)?);
                }
                chat_history.push(msg);
            }

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

            let body = body.as_object_mut().unwrap();
            if let Some(temperature) = req.temperature {
                body.insert("temperature".to_string(), Json::from(temperature));
            }

            if let Some(max_tokens) = req.max_output_tokens {
                body.insert("max_tokens".to_string(), Json::from(max_tokens));
            }

            if req.output_schema.is_some() {
                // DeepSeek only supports `{"type": "json_object"}`
                body.insert(
                    "response_format".to_string(),
                    json!({"type": "json_object"}),
                );
            }

            if let Some(stop) = req.stop {
                body.insert("stop".to_string(), Json::from(stop));
            }

            if !req.tools.is_empty() {
                body.insert(
                    "tools".to_string(),
                    json!(
                        req.tools
                            .into_iter()
                            .map(ToolDefinition::from)
                            .collect::<Vec<_>>()
                    ),
                );
                body.insert("tool_choice".to_string(), Json::from("auto"));
            };

            if log_enabled!(Debug)
                && let Ok(val) = serde_json::to_string(&body)
            {
                log::debug!(request = val; "Completion request");
            }

            let response = client.post("/chat/completions").json(body).send().await?;
            if response.status().is_success() {
                let text = response.text().await?;
                match serde_json::from_str::<CompletionResponse>(&text) {
                    Ok(res) => {
                        if log_enabled!(Debug) {
                            log::debug!(
                                model = model,
                                request:serde = body,
                                response:serde = res;
                                "Completion response");
                        } else if res.maybe_failed() {
                            log::warn!(
                                model = model,
                                request:serde = body,
                                response:serde = res;
                                "Completion maybe failed");
                        }
                        if skip_raw > 0 {
                            raw_history.drain(0..skip_raw);
                        }
                        res.try_into(raw_history, chat_history)
                    }
                    Err(err) => Err(format!(
                        "Invalid completion response, model: {}, error: {}, body: {}",
                        model, err, text
                    )
                    .into()),
                }
            } else {
                let status = response.status();
                let msg = response.text().await?;
                log::error!(
                    model = model,
                    request:serde = body;
                    "Completion request failed: {status}, body: {msg}",
                );
                Err(format!("Completion failed, model: {}, error: {}", model, msg).into())
            }
        })
    }
}

#[cfg(test)]
mod tests {

    #[tokio::test(flavor = "current_thread")]
    #[ignore]
    async fn test_deepseek() {}
}