From 4861535919b813ceee8e9f4178cbc9e5f5f7668b Mon Sep 17 00:00:00 2001 From: NOLANeight Date: Thu, 23 Jul 2026 15:24:27 +0800 Subject: [PATCH] feat: add DeepSeek provider with Chat Completions API support Add a new DeepSeek client that uses the standard Chat Completions API (chat.completions.create) via the OpenAICompatibleClient base, matching the pattern used by the existing Doubao provider. Changes: - New deepseek_client.py: DeepSeekClient + DeepSeekProvider (82 lines) - Register DEEPSEEK in LLMProvider enum and LLMClient factory - Add 12 unit tests covering provider config, client init, message handling, tool-calling support, and YAML config integration - Update README with DeepSeek examples and env var docs Closes #406 --- README.md | 7 +- tests/utils/test_deepseek_client.py | 206 ++++++++++++++++++ .../utils/llm_clients/deepseek_client.py | 82 +++++++ trae_agent/utils/llm_clients/llm_client.py | 5 + 4 files changed, 299 insertions(+), 1 deletion(-) create mode 100644 tests/utils/test_deepseek_client.py create mode 100644 trae_agent/utils/llm_clients/deepseek_client.py diff --git a/README.md b/README.md index 0e47f0593..a693e45ca 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,7 @@ For technical details please refer to [our technical report](https://arxiv.org/a ## ✨ Features - 🌊 **Lakeview**: Provides short and concise summarisation for agent steps -- 🤖 **Multi-LLM Support**: Works with OpenAI, Anthropic, Doubao, Azure, OpenRouter, Ollama and Google Gemini APIs +- 🤖 **Multi-LLM Support**: Works with OpenAI, Anthropic, Doubao, DeepSeek, Azure, OpenRouter, Ollama and Google Gemini APIs - 🛠️ **Rich Tool Ecosystem**: File editing, bash execution, sequential thinking, and more - 🎯 **Interactive Mode**: Conversational interface for iterative development - 📊 **Trajectory Recording**: Detailed logging of all agent actions for debugging and analysis @@ -106,6 +106,8 @@ export OPENROUTER_API_KEY="your-openrouter-api-key" export OPENROUTER_BASE_URL="https://openrouter.ai/api/v1" export DOUBAO_API_KEY="your-doubao-api-key" export DOUBAO_BASE_URL="https://ark.cn-beijing.volces.com/api/v3/" +export DEEPSEEK_API_KEY="your-deepseek-api-key" +export DEEPSEEK_BASE_URL="https://api.deepseek.com/v1" ``` ### MCP Services (Optional) @@ -160,6 +162,9 @@ trae-cli run "Refactor the database module" --provider doubao --model doubao-see # Ollama (local models) trae-cli run "Comment this code" --provider ollama --model qwen3 + +# DeepSeek +trae-cli run "Refactor the database module" --provider deepseek --model deepseek-chat ``` ### Advanced Options diff --git a/tests/utils/test_deepseek_client.py b/tests/utils/test_deepseek_client.py new file mode 100644 index 000000000..c861f799a --- /dev/null +++ b/tests/utils/test_deepseek_client.py @@ -0,0 +1,206 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# SPDX-License-Identifier: MIT + +"""Tests for the DeepSeek client. + +Verifies initialization, provider configuration, message handling, +and tool-calling support. Chat integration tests that require a live +API key are skipped by default (set SKIP_DEEPSEEK_TEST=true). +""" + +import os +import unittest +from unittest.mock import patch + +from trae_agent.utils.config import ModelConfig, ModelProvider +from trae_agent.utils.llm_clients.deepseek_client import DeepSeekClient, DeepSeekProvider +from trae_agent.utils.llm_clients.llm_basics import LLMMessage + + +class TestDeepSeekProvider(unittest.TestCase): + """Unit tests for the DeepSeekProvider configuration.""" + + def setUp(self) -> None: + self.provider = DeepSeekProvider() + + def test_service_name(self) -> None: + """Service name should identify DeepSeek for retry logging.""" + self.assertEqual(self.provider.get_service_name(), "DeepSeek") + + def test_provider_name(self) -> None: + """Provider name should match the trajectory-recording identifier.""" + self.assertEqual(self.provider.get_provider_name(), "deepseek") + + def test_extra_headers_is_empty(self) -> None: + """DeepSeek does not require extra HTTP headers.""" + self.assertEqual(self.provider.get_extra_headers(), {}) + + def test_supports_tool_calling(self) -> None: + """All DeepSeek chat models should advertise tool-calling support.""" + self.assertTrue(self.provider.supports_tool_calling("deepseek-chat")) + self.assertTrue(self.provider.supports_tool_calling("deepseek-reasoner")) + self.assertTrue(self.provider.supports_tool_calling("unknown-model")) + + +class TestDeepSeekClientInit(unittest.TestCase): + """Tests for DeepSeekClient instantiation and basic properties.""" + + @staticmethod + def _model_config(**overrides: object) -> ModelConfig: + """Build a minimal ModelConfig for DeepSeek testing.""" + defaults: dict[str, object] = { + "model": "deepseek-chat", + "model_provider": ModelProvider( + provider="deepseek", + api_key="test-api-key", + base_url="https://api.deepseek.com/v1", + api_version=None, + ), + "max_tokens": 4096, + "temperature": 0.5, + "top_p": 1.0, + "top_k": 0, + "parallel_tool_calls": False, + "max_retries": 3, + } + defaults.update(overrides) + return ModelConfig(**defaults) # type: ignore[arg-type] + + def test_init_stores_api_key(self) -> None: + """API key from config should be propagated to the client.""" + client = DeepSeekClient(self._model_config()) + self.assertEqual(client.api_key, "test-api-key") + + def test_init_stores_base_url(self) -> None: + """Base URL from config should be propagated to the client.""" + client = DeepSeekClient(self._model_config()) + self.assertEqual(client.base_url, "https://api.deepseek.com/v1") + + def test_init_without_base_url(self) -> None: + """When no base_url is provided the OpenAI SDK will use its default.""" + client = DeepSeekClient( + self._model_config( + model_provider=ModelProvider( + provider="deepseek", + api_key="test-key", + base_url=None, + api_version=None, + ) + ) + ) + self.assertIsNone(client.base_url) + + def test_set_chat_history(self) -> None: + """set_chat_history should parse and store messages without error.""" + client = DeepSeekClient(self._model_config()) + message = LLMMessage(role="user", content="Hello, DeepSeek!") + client.set_chat_history(messages=[message]) + self.assertTrue(True) # Runnable — no exception raised. + + @patch("trae_agent.utils.llm_clients.openai_compatible_base.openai.OpenAI") + def test_chat_mock(self, mock_openai_class: object) -> None: + """chat() should construct a valid request (mocked API call).""" + mock_client = mock_openai_class.return_value + + # Build a minimal ChatCompletion response. + mock_choice = unittest.mock.MagicMock() + mock_choice.message.content = "Hello from DeepSeek" + mock_choice.message.tool_calls = None + mock_choice.finish_reason = "stop" + mock_choice.index = 0 + + mock_response = unittest.mock.MagicMock() + mock_response.choices = [mock_choice] + mock_response.model = "deepseek-chat" + mock_response.usage.prompt_tokens = 10 + mock_response.usage.completion_tokens = 5 + + mock_client.chat.completions.create.return_value = mock_response + + client = DeepSeekClient(self._model_config(max_retries=1)) + # Override the auto-created OpenAI client with our mock. + client.client = mock_client + + message = LLMMessage(role="user", content="Hi") + result = client.chat( + messages=[message], + model_config=self._model_config(max_retries=1), + tools=None, + reuse_history=False, + ) + + self.assertEqual(result.content, "Hello from DeepSeek") + self.assertIsNone(result.tool_calls) + self.assertEqual(result.model, "deepseek-chat") + + def test_supports_tool_calling(self) -> None: + """The DeepSeekClient should delegate to DeepSeekProvider.""" + client = DeepSeekClient(self._model_config()) + self.assertTrue(client.supports_tool_calling(self._model_config())) + + +class TestDeepSeekConfigIntegration(unittest.TestCase): + """Verify that the DeepSeek provider is registered in LLMClient.""" + + def test_llm_client_creates_deepseek(self) -> None: + """LLMClient should instantiate DeepSeekClient for provider='deepseek'.""" + from trae_agent.utils.llm_clients.llm_client import LLMClient + + model_config = ModelConfig( + model="deepseek-chat", + model_provider=ModelProvider( + provider="deepseek", + api_key="test-key", + base_url="https://api.deepseek.com/v1", + api_version=None, + ), + max_tokens=4096, + temperature=0.5, + top_p=1.0, + top_k=0, + parallel_tool_calls=False, + max_retries=3, + ) + llm_client = LLMClient(model_config) + self.assertIsInstance(llm_client.client, DeepSeekClient) + self.assertEqual(llm_client.provider.value, "deepseek") + + def test_resolve_api_key_from_env(self) -> None: + """Config should fall back to DEEPSEEK_API_KEY when api_key is empty.""" + with unittest.mock.patch.dict( + os.environ, {"DEEPSEEK_API_KEY": "env-key-123"}, clear=False + ): + from trae_agent.utils.config import Config + + yaml_config = """ +agents: + trae_agent: + enable_lakeview: false + model: ds + max_steps: 3 + tools: [bash, task_done] +model_providers: + ds_provider: + api_key: "" + provider: deepseek + base_url: https://api.deepseek.com/v1 +models: + ds: + model_provider: ds_provider + model: deepseek-chat + max_tokens: 4096 + temperature: 0.5 + top_p: 1 + top_k: 0 + max_retries: 3 + parallel_tool_calls: false +""" + config = Config.create(config_string=yaml_config) + config.resolve_config_values() + self.assertEqual( + config.trae_agent.model.model_provider.api_key, "env-key-123" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/trae_agent/utils/llm_clients/deepseek_client.py b/trae_agent/utils/llm_clients/deepseek_client.py new file mode 100644 index 000000000..5cbce23ac --- /dev/null +++ b/trae_agent/utils/llm_clients/deepseek_client.py @@ -0,0 +1,82 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# SPDX-License-Identifier: MIT + +"""DeepSeek client wrapper with tool integrations. + +DeepSeek provides an OpenAI-compatible Chat Completions API at +https://api.deepseek.com/v1. This client uses the standard +chat.completions.create endpoint rather than the newer Responses API. +""" + +import openai + +from trae_agent.utils.config import ModelConfig +from trae_agent.utils.llm_clients.openai_compatible_base import ( + OpenAICompatibleClient, + ProviderConfig, +) + + +class DeepSeekProvider(ProviderConfig): + """DeepSeek provider configuration. + + DeepSeek's API is OpenAI-compatible. The base URL defaults to + https://api.deepseek.com/v1 and can be overridden via config or + the DEEPSEEK_BASE_URL environment variable. + """ + + def create_client( + self, api_key: str, base_url: str | None, api_version: str | None + ) -> openai.OpenAI: + """Create an OpenAI client pointed at DeepSeek's API endpoint.""" + return openai.OpenAI(base_url=base_url, api_key=api_key) + + def get_service_name(self) -> str: + """Get the service name for retry logging.""" + return "DeepSeek" + + def get_provider_name(self) -> str: + """Get the provider name for trajectory recording.""" + return "deepseek" + + def get_extra_headers(self) -> dict[str, str]: + """DeepSeek does not require extra headers.""" + return {} + + def supports_tool_calling(self, model_name: str) -> bool: + """Check if the model supports tool calling. + + All DeepSeek chat models (deepseek-chat, deepseek-reasoner) support + function calling via the OpenAI-compatible API. + """ + return True + + +class DeepSeekClient(OpenAICompatibleClient): + """DeepSeek client wrapper for the OpenAI-compatible Chat Completions API. + + Uses the standard chat.completions.create endpoint, which is compatible + with DeepSeek's API as well as other OpenAI-compatible providers. + + Example YAML configuration:: + + model_providers: + deepseek: + api_key: "" # Reads DEEPSEEK_API_KEY from environment + provider: deepseek + base_url: https://api.deepseek.com/v1 + + models: + my_model: + model_provider: deepseek + model: deepseek-chat + max_tokens: 4096 + temperature: 0.5 + top_p: 1 + top_k: 0 + max_retries: 3 + parallel_tool_calls: false + """ + + def __init__(self, model_config: ModelConfig): + super().__init__(model_config, DeepSeekProvider()) diff --git a/trae_agent/utils/llm_clients/llm_client.py b/trae_agent/utils/llm_clients/llm_client.py index 888266bdc..c394b3c60 100644 --- a/trae_agent/utils/llm_clients/llm_client.py +++ b/trae_agent/utils/llm_clients/llm_client.py @@ -22,6 +22,7 @@ class LLMProvider(Enum): OPENROUTER = "openrouter" DOUBAO = "doubao" GOOGLE = "google" + DEEPSEEK = "deepseek" class LLMClient: @@ -52,6 +53,10 @@ def __init__(self, model_config: ModelConfig): from .doubao_client import DoubaoClient self.client = DoubaoClient(model_config) + case LLMProvider.DEEPSEEK: + from .deepseek_client import DeepSeekClient + + self.client = DeepSeekClient(model_config) case LLMProvider.OLLAMA: from .ollama_client import OllamaClient