/ integrations / deepseek.py
deepseek.py
  1  """DeepSeek provider implementation of the LLMProvider protocol."""
  2  
  3  __all__ = ["DeepSeekClient"]
  4  
  5  import json
  6  import logging
  7  import threading
  8  from collections.abc import Sequence
  9  
 10  import openai
 11  from openai import OpenAI
 12  from pydantic import BaseModel
 13  from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential
 14  
 15  from exceptions import APIError, UsageError
 16  from models.llm import Attachment
 17  
 18  logger = logging.getLogger(__name__)
 19  
 20  _DEEPSEEK_BASE_URL = "https://api.deepseek.com/v1"
 21  
 22  
 23  def _is_transient(exc: BaseException) -> bool:
 24      if isinstance(exc, openai.APIStatusError):
 25          return exc.status_code == 429 or exc.status_code >= 500
 26      return isinstance(exc, (openai.APITimeoutError, openai.APIConnectionError))
 27  
 28  
 29  class DeepSeekClient:
 30      """LLMProvider backed by the OpenAI SDK pointed at the DeepSeek base URL."""
 31  
 32      def __init__(self, api_key: str, model: str) -> None:
 33          self._client = OpenAI(api_key=api_key, base_url=_DEEPSEEK_BASE_URL)
 34          self._model = model
 35          self._tls = threading.local()
 36  
 37      @retry(
 38          retry=retry_if_exception(_is_transient),
 39          wait=wait_exponential(multiplier=1, min=2, max=60),
 40          stop=stop_after_attempt(5),
 41          reraise=True,
 42      )
 43      def complete(
 44          self,
 45          system: str,
 46          user: str,
 47          *,
 48          temperature: float | None = None,
 49          seed: int | None = None,
 50          response_schema: type[BaseModel] | None = None,
 51          attachments: Sequence[Attachment] | None = None,
 52      ) -> str:
 53          """Send a prompt to DeepSeek and return the response.
 54  
 55          Args:
 56              system: System prompt text.
 57              user: User message text.
 58              temperature: Sampling temperature, or None to use the model's default.
 59              seed: Passed to the API when provided.
 60              response_schema: If provided, requests JSON output via json_object mode.
 61              attachments: Not supported — raises ``APIError`` if provided.
 62  
 63          Returns:
 64              Response text or JSON string.
 65  
 66          Raises:
 67              APIError: On non-retriable HTTP errors (4xx).
 68              UsageError: If attachments are provided (not supported by DeepSeek).
 69          """
 70          if attachments is not None:
 71              raise UsageError("DeepSeek file attachment support not yet implemented")
 72  
 73          logger.debug("Calling deepseek model=%s", self._model)
 74  
 75          if response_schema is not None:
 76              schema_str = json.dumps(response_schema.model_json_schema(), indent=2)
 77              system_content = (
 78                  system + f"\n\nRespond with valid JSON conforming to this schema:\n{schema_str}"
 79              )
 80          else:
 81              system_content = system
 82          messages = [
 83              {"role": "system", "content": system_content},
 84              {"role": "user", "content": user},
 85          ]
 86  
 87          try:
 88              kwargs: dict[str, object] = {
 89                  "model": self._model,
 90                  "messages": messages,
 91              }
 92              if response_schema is not None:
 93                  kwargs["response_format"] = {"type": "json_object"}
 94              if temperature is not None:
 95                  kwargs["temperature"] = temperature
 96              if seed is not None:
 97                  kwargs["seed"] = seed
 98              completion = self._client.chat.completions.create(**kwargs)  # type: ignore[call-overload]
 99              result = completion.choices[0].message.content or ""
100          except openai.APIStatusError as exc:
101              if exc.status_code == 429 or exc.status_code >= 500:
102                  raise
103              raise APIError("deepseek", exc.status_code, str(exc)) from exc
104          except (openai.APITimeoutError, openai.APIConnectionError):
105              raise
106  
107          self._tls.input_tokens = (
108              completion.usage.prompt_tokens if completion.usage else 0
109          )
110          self._tls.output_tokens = (
111              completion.usage.completion_tokens if completion.usage else 0
112          )
113  
114          logger.info("Response received (%d chars)", len(result))
115          return result
116  
117      @property
118      def last_usage(self) -> tuple[int, int]:
119          """(input_tokens, output_tokens) from the most recent successful call."""
120          return (
121              getattr(self._tls, "input_tokens", 0),
122              getattr(self._tls, "output_tokens", 0),
123          )
124  
125      def ping(self, temperature: float | None = None) -> None:
126          """Send a minimal request to verify the provider and model are reachable."""
127          self.complete(
128              system="You are a helpful assistant.", user="Say: OK", temperature=temperature
129          )