vtcode-llm 0.136.1

LLM provider abstraction, client implementations, and streaming for VT Code
use serde_json::{Map, Value};
use vtcode_config::constants::{env_vars, models, urls};
use vtcode_config::core::PromptCachingConfig;

use super::extract_reasoning_trace;
use super::openai_compat::{OpenAiCompatCore, OpenAiCompatSpec, SystemPromptPlacement, impl_openai_compat_provider};
use crate::provider::{LLMError, LLMRequest};

pub struct DeepSeekSpec;

fn deepseek_reasoning(message: &Value, choice: &Value) -> Option<String> {
    message
        .get("reasoning_content")
        .and_then(extract_reasoning_trace)
        .or_else(|| message.get("reasoning").and_then(extract_reasoning_trace))
        .or_else(|| choice.get("reasoning_content").and_then(extract_reasoning_trace))
}

impl OpenAiCompatSpec for DeepSeekSpec {
    const NAME: &'static str = "DeepSeek";
    const KEY: &'static str = "deepseek";
    const API_KEY_ENV: &'static str = "DEEPSEEK_API_KEY";
    const DEFAULT_MODEL: &'static str = models::deepseek::DEFAULT_MODEL;
    const DEFAULT_BASE_URL: &'static str = urls::DEEPSEEK_API_BASE;
    const BASE_URL_ENV: Option<&'static str> = Some(env_vars::DEEPSEEK_BASE_URL);
    const LISTED_MODELS: &'static [&'static str] = models::deepseek::SUPPORTED_MODELS;
    const VALIDATION_ALLOWLIST: Option<&'static [&'static str]> = Some(models::deepseek::SUPPORTED_MODELS);

    const SYSTEM_PROMPT: SystemPromptPlacement = SystemPromptPlacement::TopLevelField;
    const STREAM_OPTIONS_INCLUDE_USAGE: bool = true;
    const INCLUDE_USER_ID: bool = true;
    const RESPONSE_REASONING_EXTRACTOR: Option<super::openai_compat::ReasoningExtractor> = Some(deepseek_reasoning);

    fn prompt_cache_enabled(prompt_cache: Option<&PromptCachingConfig>) -> bool {
        prompt_cache.is_some_and(|cfg| {
            let settings = &cfg.providers.deepseek;
            cfg.enabled && settings.enabled && settings.surface_metrics
        })
    }

    fn response_cache_metrics(core: &OpenAiCompatCore<Self>) -> bool {
        core.prompt_cache_enabled
    }

    fn stream_cache_metrics(_core: &OpenAiCompatCore<Self>) -> bool {
        true
    }

    fn insert_reasoning(
        _core: &OpenAiCompatCore<Self>,
        request: &LLMRequest,
        payload: &mut Map<String, Value>,
    ) -> Result<(), LLMError> {
        if let Some(effort) = request.reasoning_effort {
            if effort == vtcode_config::types::ReasoningEffortLevel::None {
                payload.insert("thinking".to_owned(), serde_json::json!({"type": "disabled"}));
            } else {
                use crate::rig_adapter::RigProviderCapabilities;
                use vtcode_config::models::Provider;
                if let Some(params) =
                    RigProviderCapabilities::new(Provider::DeepSeek, &request.model).reasoning_parameters(effort)
                    && let Some(obj) = params.as_object()
                {
                    for (k, v) in obj {
                        payload.insert(k.clone(), v.clone());
                    }
                }
            }
        }
        Ok(())
    }
}

impl_openai_compat_provider!(DeepSeekProvider, DeepSeekSpec, {
    fn supports_reasoning(&self, model: &str) -> bool {
        let requested = if model.trim().is_empty() {
            &self.core.model
        } else {
            model
        };

        // Codex-inspired robustness: Setting model_supports_reasoning to false
        // does NOT disable it for known reasoning models.
        self.core
            .model_behavior
            .as_ref()
            .and_then(|b| b.model_supports_reasoning)
            .unwrap_or(false)
            || requested == models::deepseek::DEEPSEEK_V4_PRO
    }

    fn supports_reasoning_effort(&self, _model: &str) -> bool {
        // Same robustness logic for reasoning effort
        self.core
            .model_behavior
            .as_ref()
            .and_then(|b| b.model_supports_reasoning_effort)
            .unwrap_or(false)
    }

    async fn get_balance(&self) -> Result<Option<vtcode_commons::llm::BalanceInfo>, LLMError> {
        // Strip /v1 suffix to get the root API URL for the balance endpoint.
        let base = self.core.base_url.trim_end_matches('/');
        let root = base.strip_suffix("/v1").unwrap_or(base);
        let url = format!("{root}/user/balance");

        let response = self
            .core
            .http_client
            .get(&url)
            .bearer_auth(&self.core.api_key)
            .send()
            .await
            .map_err(|e| LLMError::Network {
                message: crate::error_display::format_llm_error(
                    <DeepSeekSpec as OpenAiCompatSpec>::NAME,
                    &format!("balance request failed: {e}"),
                ),
                metadata: None,
            })?;

        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            return Err(LLMError::Provider {
                message: crate::error_display::format_llm_error(
                    <DeepSeekSpec as OpenAiCompatSpec>::NAME,
                    &format!("balance API returned {status}: {body}"),
                ),
                metadata: None,
            });
        }

        let balance_resp: vtcode_commons::llm::DeepSeekBalanceResponse =
            response.json().await.map_err(|e| LLMError::Provider {
                message: crate::error_display::format_llm_error(
                    <DeepSeekSpec as OpenAiCompatSpec>::NAME,
                    &format!("failed to parse balance response: {e}"),
                ),
                metadata: None,
            })?;

        Ok(Some(balance_resp.into()))
    }
});

#[cfg(test)]
mod tests {
    use super::DeepSeekProvider;
    use crate::provider::{LLMRequest, Message, ToolChoice};
    use std::sync::Arc;
    use vtcode_config::constants::models;
    use vtcode_config::types::ReasoningEffortLevel;

    fn base_request() -> LLMRequest {
        LLMRequest {
            messages: vec![Message::user("hello".to_string())].into(),
            system_prompt: Some(Arc::new("system guidance".to_string())),
            model: models::deepseek::DEFAULT_MODEL.to_string(),
            max_tokens: Some(512),
            temperature: Some(0.5),
            top_p: Some(0.25),
            stream: true,
            tool_choice: Some(ToolChoice::Auto),
            metadata: Some(serde_json::json!({"user_id": "user-42"})),
            ..Default::default()
        }
    }

    #[test]
    fn golden_payload_basic_shape() {
        let provider = DeepSeekProvider::new("test-key".to_string());
        let payload = provider.core.convert_request(&base_request()).unwrap();

        assert_eq!(payload["model"], models::deepseek::DEFAULT_MODEL);
        let messages = payload["messages"].as_array().unwrap();
        assert_eq!(messages.len(), 1);
        assert_eq!(messages[0]["role"], "user");
        // DeepSeek sends the system prompt as a top-level field, not a message.
        assert_eq!(payload["system"], "system guidance");
        assert_eq!(payload["max_tokens"], 512);
        assert_eq!(payload["temperature"], 0.5);
        assert_eq!(payload["top_p"], 0.25);
        assert_eq!(payload["stream"], true);
        assert_eq!(payload["stream_options"]["include_usage"], true);
        assert_eq!(payload["tool_choice"], "auto");
        assert_eq!(payload["user_id"], "user-42");
        assert!(payload.get("thinking").is_none());
    }

    #[test]
    fn golden_payload_thinking_disabled_and_sampling_suppression() {
        let provider = DeepSeekProvider::new("test-key".to_string());

        let mut request = base_request();
        request.reasoning_effort = Some(ReasoningEffortLevel::None);
        let payload = provider.core.convert_request(&request).unwrap();
        assert_eq!(payload["thinking"]["type"], "disabled");
        assert_eq!(payload["temperature"], 0.5);
        assert_eq!(payload["top_p"], 0.25);

        let mut request = base_request();
        request.reasoning_effort = Some(ReasoningEffortLevel::High);
        let payload = provider.core.convert_request(&request).unwrap();
        assert_eq!(payload["thinking"]["type"], "enabled");
        assert_eq!(payload["reasoning_effort"], "high");
        assert!(payload.get("temperature").is_none());
        assert!(payload.get("top_p").is_none());

        let mut request = base_request();
        request.reasoning_effort = Some(ReasoningEffortLevel::Max);
        let payload = provider.core.convert_request(&request).unwrap();
        assert_eq!(payload["reasoning_effort"], "max");
    }
}