ralph-coder 0.2.1

/// Native DeepSeek provider with optional extended thinking support.
///
/// Default model: deepseek-v4-flash  Context: 1_000_000 tokens
/// When ThinkingMode::On is set, adds "thinking" block to the API request
/// and parses reasoning_content / reasoning_tokens from the response.
use crate::errors::{RalphError, Result};
use crate::providers::{
    ContentPart, LlmProvider, LlmResponse, Message, MessageContent, Role, StopReason, ToolCall,
    ToolDef,
};
use async_trait::async_trait;
use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::HashMap;

pub const DEFAULT_MODEL: &str = "deepseek-v4-flash";
pub const PRO_MODEL: &str = "deepseek-v4-pro";
const CONTEXT_WINDOW: u64 = 1_000_000;
const BASE_URL: &str = "https://api.deepseek.com/v1/chat/completions";
const MAX_TOKENS: u32 = 8192;
const MAX_RETRIES: u32 = 5;

/// Controls extended thinking (reasoning) mode for DeepSeek.
#[derive(Debug, Clone)]
pub enum ThinkingMode {
    Off,
    On { budget_tokens: u32 },
}

pub struct DeepSeekProvider {
    api_key: String,
    model: String,
    thinking: ThinkingMode,
    client: reqwest::Client,
}

impl DeepSeekProvider {
    pub fn new(api_key: String, model: Option<String>, thinking: ThinkingMode) -> Self {
        Self {
            api_key,
            model: model.unwrap_or_else(|| DEFAULT_MODEL.to_string()),
            thinking,
            client: reqwest::Client::new(),
        }
    }
}

// ── Request / response types ──────────────────────────────────────────────────

#[derive(Serialize)]
struct DeepSeekRequest {
    model: String,
    messages: Vec<OaiMessage>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    tools: Vec<OaiTool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    thinking: Option<ThinkingBlock>,
}

#[derive(Serialize)]
struct ThinkingBlock {
    #[serde(rename = "type")]
    kind: String,
    budget_tokens: u32,
}

#[derive(Serialize)]
struct OaiMessage {
    role: String,
    content: Value,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_calls: Option<Vec<OaiToolCallOut>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_call_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    name: Option<String>,
    /// DeepSeek thinking mode: echo reasoning_content back on subsequent calls.
    #[serde(skip_serializing_if = "Option::is_none")]
    reasoning_content: Option<String>,
}

#[derive(Serialize, Deserialize)]
struct OaiToolCallOut {
    id: String,
    #[serde(rename = "type")]
    kind: String,
    function: OaiFunctionCall,
}

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

#[derive(Serialize)]
struct OaiTool {
    #[serde(rename = "type")]
    kind: String,
    function: OaiToolFunction,
}

#[derive(Serialize)]
struct OaiToolFunction {
    name: String,
    description: String,
    parameters: Value,
}

#[derive(Deserialize)]
struct DeepSeekResponse {
    choices: Vec<DsChoice>,
    usage: DsUsage,
}

#[derive(Deserialize)]
struct DsChoice {
    message: DsRespMessage,
    finish_reason: Option<String>,
}

#[derive(Deserialize)]
struct DsRespMessage {
    content: Option<String>,
    #[serde(default)]
    reasoning_content: Option<String>,
    tool_calls: Option<Vec<OaiToolCallOut>>,
}

#[derive(Deserialize)]
struct DsUsage {
    prompt_tokens: u64,
    completion_tokens: u64,
    #[serde(default)]
    completion_tokens_details: Option<DsCompletionTokensDetails>,
}

#[derive(Deserialize)]
struct DsCompletionTokensDetails {
    #[serde(default)]
    reasoning_tokens: u64,
}

// ── SSE streaming types ───────────────────────────────────────────────────────

#[derive(Deserialize)]
struct DsStreamChunk {
    choices: Vec<DsStreamChoice>,
    #[serde(default)]
    usage: Option<DsUsage>,
}

#[derive(Deserialize)]
struct DsStreamChoice {
    delta: DsStreamDelta,
    finish_reason: Option<String>,
}

#[derive(Deserialize, Default)]
struct DsStreamDelta {
    content: Option<String>,
    #[serde(default)]
    reasoning_content: Option<String>,
    tool_calls: Option<Vec<OaiToolCallChunk>>,
}

#[derive(Deserialize)]
struct OaiToolCallChunk {
    index: usize,
    id: Option<String>,
    function: Option<OaiFunctionChunk>,
}

#[derive(Deserialize)]
struct OaiFunctionChunk {
    name: Option<String>,
    arguments: Option<String>,
}

// ── Message conversion ────────────────────────────────────────────────────────

/// Convert messages to OAI format.
///
/// DeepSeek enforces consistency across the entire message history: if ANY assistant
/// message carries `reasoning_content`, ALL assistant messages must carry it (using an
/// empty string for turns that didn't produce thinking). Mixing `Some` and `None` in
/// the same request triggers HTTP 400, regardless of whether thinking is currently
/// enabled. We therefore detect the "needs normalization" condition once and apply it
/// uniformly across all messages.
fn messages_to_oai(messages: &[Message]) -> Vec<OaiMessage> {
    // Normalize if any assistant message already has reasoning_content set.
    let normalize = messages
        .iter()
        .any(|m| matches!(m.role, Role::Assistant) && m.reasoning_content.is_some());

    messages
        .iter()
        .map(|m| {
            let role = match m.role {
                Role::System => "system",
                Role::User => "user",
                Role::Assistant => "assistant",
                Role::Tool => "tool",
            };

            if matches!(m.role, Role::Assistant) {
                let reasoning_content = if normalize {
                    // Ensure every assistant message has the field so DeepSeek sees a
                    // consistent history. Turns that produced no thinking get "".
                    Some(m.reasoning_content.clone().unwrap_or_default())
                } else {
                    None
                };

                if let MessageContent::Parts(parts) = &m.content {
                    let text: String = parts
                        .iter()
                        .filter_map(|p| match p {
                            ContentPart::Text { text } => Some(text.as_str()),
                            _ => None,
                        })
                        .collect::<Vec<_>>()
                        .join("\n");

                    let tool_calls: Vec<OaiToolCallOut> = parts
                        .iter()
                        .filter_map(|p| match p {
                            ContentPart::ToolUse { id, name, input } => Some(OaiToolCallOut {
                                id: id.clone(),
                                kind: "function".to_string(),
                                function: OaiFunctionCall {
                                    name: name.clone(),
                                    arguments: serde_json::to_string(input)
                                        .unwrap_or_else(|_| "{}".to_string()),
                                },
                            }),
                            _ => None,
                        })
                        .collect();

                    return OaiMessage {
                        role: "assistant".to_string(),
                        content: if text.is_empty() {
                            json!(null)
                        } else {
                            json!(text)
                        },
                        tool_calls: if tool_calls.is_empty() {
                            None
                        } else {
                            Some(tool_calls)
                        },
                        tool_call_id: None,
                        name: None,
                        reasoning_content,
                    };
                }
                // Text-only assistant message (no tool calls).
                return OaiMessage {
                    role: "assistant".to_string(),
                    content: json!(m.content.as_text()),
                    tool_calls: None,
                    tool_call_id: None,
                    name: None,
                    reasoning_content,
                };
            }

            OaiMessage {
                role: role.to_string(),
                content: json!(m.content.as_text()),
                tool_calls: None,
                tool_call_id: m.tool_call_id.clone(),
                name: m.name.clone(),
                reasoning_content: None,
            }
        })
        .collect()
}

fn build_tool_list(tools: &[ToolDef]) -> Vec<OaiTool> {
    tools
        .iter()
        .map(|t| OaiTool {
            kind: "function".to_string(),
            function: OaiToolFunction {
                name: t.name.clone(),
                description: t.description.clone(),
                parameters: t.parameters.clone(),
            },
        })
        .collect()
}

#[async_trait]
impl LlmProvider for DeepSeekProvider {
    async fn chat(&self, messages: &[Message], tools: &[ToolDef]) -> Result<LlmResponse> {
        let oai_tools = build_tool_list(tools);

        let thinking_block = match &self.thinking {
            ThinkingMode::Off => None,
            ThinkingMode::On { budget_tokens } => Some(ThinkingBlock {
                kind: "enabled".to_string(),
                budget_tokens: *budget_tokens,
            }),
        };

        let body = DeepSeekRequest {
            model: self.model.clone(),
            messages: messages_to_oai(messages),
            tools: oai_tools,
            max_tokens: Some(MAX_TOKENS),
            thinking: thinking_block,
        };

        let mut last_decode_err = String::new();
        for attempt in 0..MAX_RETRIES {
            if attempt > 0 {
                // Exponential backoff: 2s, 4s, 8s, 16s
                let secs = (2_u64).pow(attempt).min(30);
                tokio::time::sleep(tokio::time::Duration::from_secs(secs)).await;
            }

            let resp = self
                .client
                .post(BASE_URL)
                .bearer_auth(&self.api_key)
                .json(&body)
                .send()
                .await?;

            let status = resp.status();
            if status == 401 {
                return Err(RalphError::LlmAuth {
                    provider: "deepseek".to_string(),
                });
            }
            if status == 429 {
                return Err(RalphError::LlmRateLimit {
                    provider: "deepseek".to_string(),
                    attempts: 1,
                });
            }
            if !status.is_success() {
                let err_body = resp.text().await.unwrap_or_default();
                // Retry on transient server errors (5xx); fail fast on client errors
                if status.as_u16() >= 500 {
                    last_decode_err = format!("HTTP {}: {}", status, err_body);
                    continue;
                }
                return Err(RalphError::LlmApi {
                    provider: "deepseek".to_string(),
                    message: format!("HTTP {}: {}", status, err_body),
                });
            }

            let parsed: DeepSeekResponse = match resp.json().await {
                Ok(p) => p,
                Err(e) => {
                    last_decode_err = e.to_string();
                    continue;
                }
            };

            let choice = parsed.choices.into_iter().next().ok_or_else(|| {
                RalphError::LlmResponseParse("No choices in response".to_string())
            })?;

            let stop_reason = match choice.finish_reason.as_deref() {
                Some("tool_calls") => StopReason::ToolUse,
                Some("stop") => StopReason::Stop,
                Some("length") => StopReason::MaxTokens,
                _ => StopReason::EndTurn,
            };

            // For thinking-mode responses, always store Some(...) — even empty — so that
            // subsequent requests can include it consistently without mixing Some/None.
            let reasoning_content = match &self.thinking {
                ThinkingMode::On { .. } => {
                    Some(choice.message.reasoning_content.clone().unwrap_or_default())
                }
                ThinkingMode::Off => choice.message.reasoning_content.clone(),
            };

            let tool_calls = choice
                .message
                .tool_calls
                .unwrap_or_default()
                .into_iter()
                .map(|tc| {
                    let args: Value = serde_json::from_str(&tc.function.arguments)
                        .unwrap_or(Value::Object(Default::default()));
                    ToolCall {
                        id: tc.id,
                        name: tc.function.name,
                        arguments: args,
                    }
                })
                .collect();

            let input_tokens = parsed.usage.prompt_tokens;
            let reasoning_tokens = parsed
                .usage
                .completion_tokens_details
                .as_ref()
                .map(|d| d.reasoning_tokens)
                .unwrap_or(0);
            let output_tokens = parsed.usage.completion_tokens;

            return Ok(LlmResponse {
                text: choice.message.content,
                tool_calls,
                input_tokens,
                output_tokens,
                reasoning_tokens,
                reasoning_content,
                tokens_used: input_tokens + output_tokens + reasoning_tokens,
                stop_reason,
            });
        } // end retry loop
        Err(RalphError::LlmResponseParse(last_decode_err))
    }

    async fn chat_streaming(
        &self,
        messages: &[Message],
        tools: &[ToolDef],
        token_tx: &tokio::sync::mpsc::UnboundedSender<String>,
    ) -> Result<LlmResponse> {
        let oai_tools = build_tool_list(tools);

        let thinking_block = match &self.thinking {
            ThinkingMode::Off => None,
            ThinkingMode::On { budget_tokens } => Some(json!({
                "type": "enabled",
                "budget_tokens": budget_tokens
            })),
        };

        let mut body = json!({
            "model": self.model,
            "messages": messages_to_oai(messages),
            "stream": true,
            "stream_options": { "include_usage": true },
            "max_tokens": MAX_TOKENS,
        });
        if !oai_tools.is_empty() {
            body["tools"] = json!(oai_tools);
        }
        if let Some(tb) = thinking_block {
            body["thinking"] = tb;
        }

        let mut last_decode_err = String::new();
        for attempt in 0..MAX_RETRIES {
            if attempt > 0 {
                // Exponential backoff: 2s, 4s, 8s, 16s
                let secs = (2_u64).pow(attempt).min(30);
                tokio::time::sleep(tokio::time::Duration::from_secs(secs)).await;
            }

            let resp = self
                .client
                .post(BASE_URL)
                .bearer_auth(&self.api_key)
                .json(&body)
                .send()
                .await?;

            let status = resp.status();
            if status == 401 {
                return Err(RalphError::LlmAuth {
                    provider: "deepseek".to_string(),
                });
            }
            if status == 429 {
                return Err(RalphError::LlmRateLimit {
                    provider: "deepseek".to_string(),
                    attempts: 1,
                });
            }
            if !status.is_success() {
                let err_body = resp.text().await.unwrap_or_default();
                // Retry on transient server errors (5xx); fail fast on client errors
                if status.as_u16() >= 500 {
                    last_decode_err = format!("HTTP {}: {}", status, err_body);
                    continue;
                }
                return Err(RalphError::LlmApi {
                    provider: "deepseek".to_string(),
                    message: format!("HTTP {}: {}", status, err_body),
                });
            }

            let mut stream = resp.bytes_stream();
            let mut buf = String::new();
            let mut text_parts: Vec<String> = Vec::new();
            let mut reasoning_parts: Vec<String> = Vec::new();
            let mut tool_chunks: HashMap<usize, (String, String, String)> = HashMap::new();
            let mut input_tokens: u64 = 0;
            let mut output_tokens: u64 = 0;
            let mut reasoning_tokens: u64 = 0;
            let mut finish_reason: Option<String> = None;
            let mut decode_failed = false;

            while let Some(chunk) = stream.next().await {
                let bytes = match chunk {
                    Ok(b) => b,
                    Err(e) => {
                        last_decode_err = e.to_string();
                        decode_failed = true;
                        break;
                    }
                };
                buf.push_str(&String::from_utf8_lossy(&bytes));

                loop {
                    let Some(pos) = buf.find('\n') else { break };
                    let line = buf[..pos].trim().to_string();
                    buf.drain(..pos + 1);

                    if !line.starts_with("data: ") {
                        continue;
                    }
                    let data = line[6..].trim();
                    if data == "[DONE]" {
                        break;
                    }

                    let Ok(parsed) = serde_json::from_str::<DsStreamChunk>(data) else {
                        continue;
                    };
                    if let Some(usage) = parsed.usage {
                        input_tokens = usage.prompt_tokens;
                        output_tokens = usage.completion_tokens;
                        reasoning_tokens = usage
                            .completion_tokens_details
                            .as_ref()
                            .map(|d| d.reasoning_tokens)
                            .unwrap_or(0);
                    }
                    for choice in parsed.choices {
                        if let Some(r) = choice.finish_reason {
                            finish_reason = Some(r);
                        }
                        // Accumulate reasoning_content silently — must be echoed back to API.
                        if let Some(rc) = choice.delta.reasoning_content {
                            if !rc.is_empty() {
                                reasoning_parts.push(rc);
                            }
                        }
                        if let Some(content) = choice.delta.content {
                            if !content.is_empty() {
                                let _ = token_tx.send(content.clone());
                                text_parts.push(content);
                            }
                        }
                        if let Some(tcs) = choice.delta.tool_calls {
                            for tc in tcs {
                                let entry = tool_chunks.entry(tc.index).or_insert_with(|| {
                                    (String::new(), String::new(), String::new())
                                });
                                if let Some(id) = tc.id {
                                    entry.0 = id;
                                }
                                if let Some(func) = tc.function {
                                    if let Some(name) = func.name {
                                        entry.1 = name;
                                    }
                                    if let Some(args) = func.arguments {
                                        entry.2.push_str(&args);
                                    }
                                }
                            }
                        }
                    }
                }
            }

            let text = if text_parts.is_empty() {
                None
            } else {
                Some(text_parts.join(""))
            };
            // For thinking-mode responses, always store Some(...) — even empty — so that
            // subsequent requests can include it consistently without mixing Some/None.
            let reasoning_content = match &self.thinking {
                ThinkingMode::On { .. } => Some(reasoning_parts.join("")),
                ThinkingMode::Off => {
                    if reasoning_parts.is_empty() {
                        None
                    } else {
                        Some(reasoning_parts.join(""))
                    }
                }
            };

            let mut sorted: Vec<_> = tool_chunks.into_iter().collect();
            sorted.sort_by_key(|(i, _)| *i);
            let tool_calls = sorted
                .into_iter()
                .map(|(_, (id, name, args))| {
                    let arguments: Value =
                        serde_json::from_str(&args).unwrap_or(Value::Object(Default::default()));
                    ToolCall {
                        id,
                        name,
                        arguments,
                    }
                })
                .collect();

            let stop_reason = match finish_reason.as_deref() {
                Some("tool_calls") => StopReason::ToolUse,
                Some("stop") => StopReason::Stop,
                Some("length") => StopReason::MaxTokens,
                _ => StopReason::EndTurn,
            };

            if decode_failed {
                continue;
            }

            return Ok(LlmResponse {
                text,
                tool_calls,
                input_tokens,
                output_tokens,
                reasoning_tokens,
                reasoning_content,
                tokens_used: input_tokens + output_tokens + reasoning_tokens,
                stop_reason,
            });
        } // end retry loop
        Err(RalphError::LlmApi {
            provider: "deepseek".to_string(),
            message: last_decode_err,
        })
    }

    fn supports_streaming(&self) -> bool {
        true
    }

    fn name(&self) -> &str {
        "deepseek"
    }
    fn context_window(&self) -> u64 {
        CONTEXT_WINDOW
    }
    fn default_model(&self) -> &str {
        &self.model
    }
}