ic-rig 0.2.0

//! DeepSeek Chat Completions API provider.
//!
//! DeepSeek's API is OpenAI-compatible, so the wire format is nearly identical.
//! The main differences are the base URL and the `reasoning_content` field that
//! the thinking-mode models add to assistant messages (we surface it as a
//! separate field on [`CompletionResponse`] via the `extra` JSON blob, but the
//! agentic loop treats it as an opaque text response).
//!
//! # Usage
//!
//! ```rust,ignore
//! use ic_rig::providers::deepseek::{Client, DEEPSEEK_V4_FLASH};
//!
//! let client = Client::new(my_http, "sk-...");
//!
//! let agent = ic_rig::Agent::builder(client.model(DEEPSEEK_V4_FLASH))
//!     .preamble("You are a helpful assistant.")
//!     .build();
//!
//! let reply = agent.prompt("Hello!").await?;
//! ```

use crate::{
    completion::{
        CompletionError, CompletionModel, CompletionRequest, CompletionResponse, ModelChoice, Usage,
    },
    http::{HttpClient, HttpRequest},
    message::{AssistantContent, Message, ToolCall, UserContent},
    tool::ToolDefinition,
};
use serde::{Deserialize, Serialize};

// ── Model constants ───────────────────────────────────────────────────────────

/// Non-thinking flagship model (~300B total / ~13B activated params, 1M context).
pub const DEEPSEEK_V4_FLASH: &str = "deepseek-v4-flash";
/// Larger 1.6T-param flagship model, 1M context.
pub const DEEPSEEK_V4_PRO: &str = "deepseek-v4-pro";
/// Experimental multimodal (vision) variant of V4 Flash.
pub const DEEPSEEK_V4_FLASH_VISION_EXP: &str = "deepseek-v4-flash-vision-exp";

/// Retired by DeepSeek on 2026-07-24; kept only so old code still compiles.
/// Use [`DEEPSEEK_V4_FLASH`] instead (this alias used to route to it in
/// non-thinking mode).
#[deprecated(note = "retired by DeepSeek on 2026-07-24; use DEEPSEEK_V4_FLASH instead")]
pub const DEEPSEEK_CHAT: &str = "deepseek-chat";
/// Retired by DeepSeek on 2026-07-24; kept only so old code still compiles.
/// Use [`DEEPSEEK_V4_FLASH`] instead (this alias used to route to it in
/// thinking mode).
#[deprecated(note = "retired by DeepSeek on 2026-07-24; use DEEPSEEK_V4_FLASH instead")]
pub const DEEPSEEK_REASONER: &str = "deepseek-reasoner";

const BASE_URL: &str = "https://api.deepseek.com/v1";

// ── Client ────────────────────────────────────────────────────────────────────

/// DeepSeek API client. Use [`model`](Client::model) to get a [`CompletionModel`].
pub struct Client<H> {
    http: H,
    api_key: String,
    base_url: String,
}

impl<H: HttpClient + Clone> Client<H> {
    pub fn new(http: H, api_key: impl Into<String>) -> Self {
        Self {
            http,
            api_key: api_key.into(),
            base_url: BASE_URL.to_owned(),
        }
    }

    /// Override the base URL (e.g. for a local proxy).
    pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
        self.base_url = url.into();
        self
    }

    /// Produce a [`Model`] for the given model name (use the constants above).
    pub fn model(&self, model: impl Into<String>) -> Model<H> {
        Model {
            http: self.http.clone(),
            api_key: self.api_key.clone(),
            base_url: self.base_url.clone(),
            model: model.into(),
        }
    }
}

// ── Model ─────────────────────────────────────────────────────────────────────

/// A DeepSeek model that implements [`CompletionModel`].
pub struct Model<H> {
    http: H,
    api_key: String,
    base_url: String,
    model: String,
}

impl<H: HttpClient> CompletionModel for Model<H> {
    type Error = CompletionError;

    async fn complete(
        &self,
        request: CompletionRequest,
    ) -> Result<CompletionResponse, CompletionError> {
        let body = build_request(&self.model, request)?;
        let bytes = serde_json::to_vec(&body)?;

        let http_req = HttpRequest::new(format!("{}/chat/completions", self.base_url))
            .header("Authorization", format!("Bearer {}", self.api_key))
            .json_body(bytes);

        let resp = self
            .http
            .post(http_req)
            .await
            .map_err(|e| CompletionError::Http(e.to_string()))?;

        if !resp.is_success() {
            let message = String::from_utf8_lossy(&resp.body).into_owned();
            return Err(CompletionError::Provider {
                status: resp.status,
                message,
            });
        }

        let api_resp: ApiResponse = resp.json()?;
        parse_response(api_resp)
    }
}

// ── Wire types (DeepSeek JSON format) ─────────────────────────────────────────

#[derive(Serialize)]
struct ApiRequest {
    model: String,
    messages: Vec<serde_json::Value>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    tools: Vec<ApiTool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    thinking: Option<ApiThinking>,
}

#[derive(Serialize)]
struct ApiThinking {
    #[serde(rename = "type")]
    kind: &'static str,
}

#[derive(Serialize)]
struct ApiTool {
    #[serde(rename = "type")]
    kind: &'static str,
    function: ApiFunction,
}

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

#[derive(Deserialize)]
struct ApiResponse {
    choices: Vec<ApiChoice>,
    usage: Option<ApiUsage>,
}

#[derive(Deserialize)]
struct ApiChoice {
    message: ApiMessage,
}

#[derive(Deserialize)]
struct ApiMessage {
    content: Option<String>,
    /// Only present when the model is run in thinking mode.
    #[serde(default)]
    reasoning_content: Option<String>,
    #[serde(default)]
    tool_calls: Vec<ApiToolCall>,
}

#[derive(Deserialize)]
struct ApiToolCall {
    id: String,
    function: ApiToolCallFunction,
}

#[derive(Deserialize)]
struct ApiToolCallFunction {
    name: String,
    /// DeepSeek (like OpenAI) encodes arguments as a JSON **string**.
    arguments: String,
}

#[derive(Deserialize)]
struct ApiUsage {
    prompt_tokens: u32,
    completion_tokens: u32,
}

// ── Conversion helpers ────────────────────────────────────────────────────────

fn build_request(model: &str, req: CompletionRequest) -> Result<ApiRequest, CompletionError> {
    let messages = convert_messages(req.messages)?;
    let tools = req.tools.into_iter().map(convert_tool).collect();

    Ok(ApiRequest {
        model: model.to_owned(),
        messages,
        tools,
        temperature: req.temperature,
        max_tokens: req.max_tokens,
        thinking: req.thinking.map(|enabled| ApiThinking {
            kind: if enabled { "enabled" } else { "disabled" },
        }),
    })
}

/// Convert ic-rig messages into the DeepSeek flat-list format.
///
/// Key wire-format rules:
/// - Tool results become individual `role: "tool"` messages, one per result.
/// - Assistant tool calls live in `tool_calls`; arguments are a JSON **string**.
/// - Multiple user text parts are merged into a single content string.
fn convert_messages(messages: Vec<Message>) -> Result<Vec<serde_json::Value>, CompletionError> {
    let mut out = Vec::new();

    for msg in messages {
        match msg {
            Message::System { content } => {
                out.push(serde_json::json!({ "role": "system", "content": content }));
            }

            Message::User { content } => {
                let mut text_parts: Vec<String> = Vec::new();

                for part in content {
                    match part {
                        UserContent::Text(t) => {
                            text_parts.push(t.text);
                        }
                        UserContent::ToolResult(r) => {
                            // Flush accumulated text before emitting a tool result.
                            if !text_parts.is_empty() {
                                let merged = text_parts.join("\n");
                                text_parts.clear();
                                out.push(
                                    serde_json::json!({ "role": "user", "content": merged }),
                                );
                            }
                            out.push(serde_json::json!({
                                "role": "tool",
                                "tool_call_id": r.call_id,
                                "content": r.content,
                            }));
                        }
                    }
                }

                if !text_parts.is_empty() {
                    let merged = text_parts.join("\n");
                    out.push(serde_json::json!({ "role": "user", "content": merged }));
                }
            }

            Message::Assistant { content } => {
                let mut text: Option<String> = None;
                let mut tool_calls: Vec<serde_json::Value> = Vec::new();

                for part in content {
                    match part {
                        AssistantContent::Text(t) => {
                            text = Some(t.text);
                        }
                        AssistantContent::ToolCall(c) => {
                            // Arguments must be serialised to a JSON string.
                            let arguments = serde_json::to_string(&c.arguments)?;
                            tool_calls.push(serde_json::json!({
                                "id": c.id,
                                "type": "function",
                                "function": { "name": c.name, "arguments": arguments },
                            }));
                        }
                    }
                }

                let mut msg = serde_json::json!({ "role": "assistant", "content": text });
                if !tool_calls.is_empty() {
                    msg["tool_calls"] = serde_json::json!(tool_calls);
                }
                out.push(msg);
            }
        }
    }

    Ok(out)
}

fn convert_tool(def: ToolDefinition) -> ApiTool {
    ApiTool {
        kind: "function",
        function: ApiFunction {
            name: def.name,
            description: def.description,
            parameters: def.parameters,
        },
    }
}

fn parse_response(resp: ApiResponse) -> Result<CompletionResponse, CompletionError> {
    let choice = resp
        .choices
        .into_iter()
        .next()
        .ok_or_else(|| CompletionError::Response("no choices in response".into()))?;

    // Captured up front (cloned) so it's available via `reasoning` even when
    // the branch below also consumes it as a last-resort answer.
    let reasoning = choice.message.reasoning_content.clone();

    let model_choice = if !choice.message.tool_calls.is_empty() {
        let calls = choice
            .message
            .tool_calls
            .into_iter()
            .map(|c| {
                let arguments: serde_json::Value = serde_json::from_str(&c.function.arguments)?;
                Ok(ToolCall {
                    id: c.id,
                    name: c.function.name,
                    arguments,
                })
            })
            .collect::<Result<Vec<_>, serde_json::Error>>()?;
        ModelChoice::ToolCall(calls)
    } else {
        // In thinking mode, prefer the final `content` over `reasoning_content`.
        // The reasoning trace is surfaced here only when there is no regular content.
        let text = choice
            .message
            .content
            .filter(|s| !s.is_empty())
            .or(choice.message.reasoning_content)
            .ok_or_else(|| {
                CompletionError::Response("no content and no tool_calls".into())
            })?;
        ModelChoice::Message(text)
    };

    let usage = resp.usage.map(|u| Usage {
        prompt_tokens: u.prompt_tokens,
        completion_tokens: u.completion_tokens,
    });

    Ok(CompletionResponse {
        choice: model_choice,
        reasoning,
        usage,
    })
}

// ── Tests ─────────────────────────────────────────────────────────────────────

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

    #[test]
    fn thinking_toggle_serializes_expected_shape() {
        let mut on = CompletionRequest::new(vec![Message::user("hi")]);
        on.thinking = Some(true);
        let json = serde_json::to_value(build_request(DEEPSEEK_V4_FLASH, on).unwrap()).unwrap();
        assert_eq!(json["thinking"]["type"], "enabled");

        let mut off = CompletionRequest::new(vec![Message::user("hi")]);
        off.thinking = Some(false);
        let json = serde_json::to_value(build_request(DEEPSEEK_V4_FLASH, off).unwrap()).unwrap();
        assert_eq!(json["thinking"]["type"], "disabled");

        let unset = CompletionRequest::new(vec![Message::user("hi")]);
        let json = serde_json::to_value(build_request(DEEPSEEK_V4_FLASH, unset).unwrap()).unwrap();
        assert!(json.get("thinking").is_none());
    }

    fn make_response(json: &str) -> ApiResponse {
        serde_json::from_str(json).expect("test fixture must deserialise")
    }

    #[test]
    fn parse_simple_text_response() {
        let resp = make_response(r#"{
            "choices": [{
                "message": { "role": "assistant", "content": "Hello, world!" },
                "finish_reason": "stop",
                "index": 0
            }],
            "usage": { "prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15 }
        }"#);

        let result = parse_response(resp).unwrap();
        assert!(matches!(result.choice, ModelChoice::Message(ref s) if s == "Hello, world!"));
        let usage = result.usage.unwrap();
        assert_eq!(usage.prompt_tokens, 10);
        assert_eq!(usage.completion_tokens, 5);
    }

    #[test]
    fn parse_tool_call_response() {
        let resp = make_response(r#"{
            "choices": [{
                "message": {
                    "role": "assistant",
                    "content": "",
                    "tool_calls": [{
                        "id": "call_abc123",
                        "type": "function",
                        "index": 0,
                        "function": { "name": "add", "arguments": "{\"x\":2,\"y\":5}" }
                    }]
                },
                "finish_reason": "tool_calls",
                "index": 0
            }],
            "usage": { "prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30 }
        }"#);

        let result = parse_response(resp).unwrap();
        match result.choice {
            ModelChoice::ToolCall(calls) => {
                assert_eq!(calls.len(), 1);
                assert_eq!(calls[0].id, "call_abc123");
                assert_eq!(calls[0].name, "add");
                assert_eq!(calls[0].arguments["x"], 2);
                assert_eq!(calls[0].arguments["y"], 5);
            }
            other => panic!("expected ToolCall, got {other:?}"),
        }
    }

    #[test]
    fn parse_reasoner_falls_back_to_reasoning_content() {
        // deepseek-reasoner may return empty `content` with only `reasoning_content`.
        let resp = make_response(r#"{
            "choices": [{
                "message": {
                    "role": "assistant",
                    "content": "",
                    "reasoning_content": "I think therefore I am."
                },
                "finish_reason": "stop",
                "index": 0
            }],
            "usage": { "prompt_tokens": 5, "completion_tokens": 8, "total_tokens": 13 }
        }"#);

        let result = parse_response(resp).unwrap();
        assert!(
            matches!(result.choice, ModelChoice::Message(ref s) if s == "I think therefore I am.")
        );
    }

    #[test]
    fn parse_no_choices_returns_error() {
        let resp = make_response(r#"{ "choices": [], "usage": null }"#);
        assert!(parse_response(resp).is_err());
    }

    #[test]
    fn convert_messages_merges_user_text() {
        use crate::message::{Text, UserContent};

        let messages = vec![Message::User {
            content: vec![
                UserContent::Text(Text { text: "first".into() }),
                UserContent::Text(Text { text: "second".into() }),
            ],
        }];

        let out = convert_messages(messages).unwrap();
        assert_eq!(out.len(), 1);
        assert_eq!(out[0]["content"], "first\nsecond");
    }

    #[test]
    fn convert_messages_splits_tool_results() {
        use crate::message::{Text, ToolResult, UserContent};

        let messages = vec![Message::User {
            content: vec![
                UserContent::ToolResult(ToolResult {
                    call_id: "call_1".into(),
                    name: "my_tool".into(),
                    content: "result".into(),
                }),
                UserContent::Text(Text { text: "follow up".into() }),
            ],
        }];

        let out = convert_messages(messages).unwrap();
        assert_eq!(out.len(), 2);
        assert_eq!(out[0]["role"], "tool");
        assert_eq!(out[0]["tool_call_id"], "call_1");
        assert_eq!(out[1]["role"], "user");
        assert_eq!(out[1]["content"], "follow up");
    }
}