Echo-Tavern/services/deepseek_client.py

from __future__ import annotations

import json
import time
from typing import Any

import requests


class DeepSeekClient:
    def __init__(
        self,
        api_key: str,
        base_url: str = "https://api.deepseek.com",
        model: str = "deepseek-chat",
        timeout: float = 25.0,
        connect_timeout: float | None = None,
        read_timeout: float | None = None,
        retries: int = 2,
        extra_headers: dict[str, Any] | None = None,
        force_json_mode: bool = True,
    ):
        self.api_key = api_key.strip()
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.timeout = float(timeout)
        self.connect_timeout = float(connect_timeout) if connect_timeout is not None else 8.0
        self.read_timeout = float(read_timeout) if read_timeout is not None else float(timeout)
        self.retries = retries
        self.force_json_mode = bool(force_json_mode)
        self.extra_headers = self._normalize_headers(extra_headers)

    def _normalize_headers(self, headers: dict[str, Any] | None) -> dict[str, str]:
        if not isinstance(headers, dict):
            return {}
        normalized: dict[str, str] = {}
        for key, value in headers.items():
            safe_key = str(key or "").strip()
            safe_value = str(value or "").strip()
            if safe_key and safe_value:
                normalized[safe_key] = safe_value
        return normalized

    def _chat_url(self) -> str:
        if self.base_url.endswith("/chat/completions"):
            return self.base_url
        if self.base_url.endswith("/v1"):
            return f"{self.base_url}/chat/completions"
        return f"{self.base_url}/v1/chat/completions"

    def _request_headers(self) -> dict[str, str]:
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        headers.update(self.extra_headers)
        if "openrouter.ai" in self.base_url:
            headers.setdefault("HTTP-Referer", "http://localhost")
            headers.setdefault("X-Title", "Echo Tavern")
        return headers

    def _payload(self, messages: list[dict[str, str]], temperature: float, max_tokens: int, use_json_mode: bool) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "model": self.model,
            "messages": messages,
        }
        payload["temperature"] = temperature
        payload["max_tokens"] = max_tokens
        if use_json_mode:
            payload["response_format"] = {"type": "json_object"}
        return payload

    def _build_payload(
        self,
        messages: list[dict[str, str]],
        temperature: float,
        max_tokens: int,
        *,
        use_json_mode: bool,
        token_field: str,
        include_temperature: bool,
    ) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "model": self.model,
            "messages": messages,
        }
        if include_temperature:
            payload["temperature"] = temperature
        if token_field == "max_tokens":
            payload["max_tokens"] = max_tokens
        elif token_field == "max_completion_tokens":
            payload["max_completion_tokens"] = max_tokens
        if use_json_mode:
            payload["response_format"] = {"type": "json_object"}
        return payload

    def _is_response_format_unsupported(self, status_code: int, raw_text: str) -> bool:
        if status_code not in {400, 404, 415, 422}:
            return False
        lower = str(raw_text or "").lower()
        if "response_format" in lower:
            return True
        if "json_object" in lower and "support" in lower:
            return True
        return "json schema" in lower and "support" in lower

    def _extract_error_message(self, body: Any, raw_text: str) -> str:
        if isinstance(body, dict):
            error = body.get("error")
            if isinstance(error, dict):
                msg = error.get("message")
                if isinstance(msg, str) and msg.strip():
                    return msg.strip()
            if isinstance(error, str) and error.strip():
                return error.strip()
            msg = body.get("message")
            if isinstance(msg, str) and msg.strip():
                return msg.strip()
        text = str(raw_text or "").strip()
        if not text:
            return ""
        if len(text) > 1000:
            return text[:1000]
        return text

    def _detect_token_field_switch(self, status_code: int, error_message: str, current_field: str) -> str | None:
        if status_code not in {400, 404, 415, 422}:
            return None
        lower = str(error_message or "").lower()
        if current_field == "max_tokens":
            if "max_tokens" in lower and any(word in lower for word in ("unsupported", "unknown", "invalid", "not allowed", "unrecognized")):
                return "max_completion_tokens"
        if current_field == "max_completion_tokens":
            if "max_completion_tokens" in lower and any(word in lower for word in ("unsupported", "unknown", "invalid", "not allowed", "unrecognized")):
                return "max_tokens"
        if current_field != "none":
            if "maximum context length" in lower and "max_tokens" in lower:
                return "none"
            if "unexpected parameter" in lower and ("max_tokens" in lower or "max_completion_tokens" in lower):
                return "none"
        return None

    def _is_temperature_unsupported(self, status_code: int, error_message: str) -> bool:
        if status_code not in {400, 404, 415, 422}:
            return False
        lower = str(error_message or "").lower()
        if "temperature" not in lower:
            return False
        return any(word in lower for word in ("unsupported", "unknown", "invalid", "not allowed", "unrecognized"))

    def _content_to_text(self, value: Any) -> str:
        if isinstance(value, str):
            return value.strip()
        if isinstance(value, dict):
            for key in ("text", "content", "message"):
                text = value.get(key)
                if isinstance(text, str) and text.strip():
                    return text.strip()
            return ""
        if isinstance(value, list):
            parts: list[str] = []
            for row in value:
                if isinstance(row, str):
                    text = row.strip()
                    if text:
                        parts.append(text)
                    continue
                if not isinstance(row, dict):
                    continue
                text = row.get("text")
                if not isinstance(text, str):
                    nested = row.get("content")
                    if isinstance(nested, str):
                        text = nested
                if isinstance(text, str):
                    clean = text.strip()
                    if clean:
                        parts.append(clean)
            return "\n".join(parts).strip()
        return ""

    def _extract_content(self, body: dict[str, Any]) -> str:
        choices = body.get("choices", [])
        if isinstance(choices, list) and choices:
            first = choices[0] if isinstance(choices[0], dict) else {}
            message = first.get("message", {})
            content = ""
            if isinstance(message, dict):
                content = self._content_to_text(message.get("content"))
                if not content:
                    tool_calls = message.get("tool_calls", [])
                    if isinstance(tool_calls, list) and tool_calls:
                        call = tool_calls[0] if isinstance(tool_calls[0], dict) else {}
                        func = call.get("function", {}) if isinstance(call, dict) else {}
                        if isinstance(func, dict):
                            content = self._content_to_text(func.get("arguments"))
                if not content:
                    content = self._content_to_text(message.get("reasoning"))
            if not content:
                content = self._content_to_text(first.get("text"))
            if not content and isinstance(first.get("delta"), dict):
                content = self._content_to_text(first.get("delta", {}).get("content"))
            if content:
                return content

        output_text = self._content_to_text(body.get("output_text"))
        if output_text:
            return output_text

        output = body.get("output", [])
        if isinstance(output, list):
            blocks: list[str] = []
            for item in output:
                if not isinstance(item, dict):
                    continue
                content_rows = item.get("content", [])
                if not isinstance(content_rows, list):
                    continue
                for row in content_rows:
                    if isinstance(row, dict):
                        text = self._content_to_text(row.get("text"))
                        if text:
                            blocks.append(text)
            merged = "\n".join(blocks).strip()
            if merged:
                return merged
        return ""

    def chat_json(
        self,
        messages: list[dict[str, str]],
        *,
        temperature: float = 0.6,
        max_tokens: int = 800,
    ) -> dict[str, Any]:
        if not self.api_key:
            return {"ok": False, "error": "missing_api_key", "content": None, "detail": "api_key_empty"}

        url = self._chat_url()
        headers = self._request_headers()

        last_error = "request_failed"
        last_detail = ""
        attempt = 0
        use_json_mode = self.force_json_mode
        token_field = "max_tokens"
        include_temperature = True
        while attempt <= self.retries:
            try:
                payload = self._build_payload(
                    messages,
                    temperature,
                    max_tokens,
                    use_json_mode=use_json_mode,
                    token_field=token_field,
                    include_temperature=include_temperature,
                )
                response = requests.post(
                    url,
                    headers=headers,
                    json=payload,
                    timeout=(self.connect_timeout, self.read_timeout),
                )
                raw_text = response.text or ""
                if response.status_code >= 500:
                    last_error = f"server_error_{response.status_code}"
                    raise requests.HTTPError(last_error)

                if response.status_code >= 400:
                    body_error: Any = None
                    try:
                        body_error = response.json()
                    except (json.JSONDecodeError, ValueError):
                        body_error = None
                    error_message = self._extract_error_message(body_error, raw_text)
                    if use_json_mode and self._is_response_format_unsupported(response.status_code, raw_text):
                        use_json_mode = False
                        continue
                    next_token_field = self._detect_token_field_switch(response.status_code, error_message, token_field)
                    if next_token_field is not None and next_token_field != token_field:
                        token_field = next_token_field
                        continue
                    if include_temperature and self._is_temperature_unsupported(response.status_code, error_message):
                        include_temperature = False
                        continue
                    return {
                        "ok": False,
                        "error": f"http_{response.status_code}",
                        "content": None,
                        "detail": error_message,
                        "raw": body_error if isinstance(body_error, dict) else raw_text,
                    }

                try:
                    body = response.json()
                except json.JSONDecodeError:
                    return {
                        "ok": False,
                        "error": "invalid_json_response",
                        "content": None,
                        "detail": "",
                        "raw": raw_text,
                    }

                content = self._extract_content(body)
                if not content:
                    return {"ok": False, "error": "empty_content", "content": None, "detail": "", "raw": body}

                return {
                    "ok": True,
                    "error": None,
                    "content": content,
                    "detail": "",
                    "raw": body,
                    "usage": body.get("usage"),
                }

            except requests.Timeout as exc:
                last_error = "request_timeout"
                last_detail = str(exc)
            except requests.ConnectionError as exc:
                detail = str(exc)
                lower = detail.lower()
                if "nodename nor servname provided" in lower or "name or service not known" in lower:
                    last_error = "dns_resolution_failed"
                elif "connection refused" in lower:
                    last_error = "connection_refused"
                elif "read timed out" in lower:
                    last_error = "request_timeout"
                else:
                    last_error = "network_connection_error"
                last_detail = detail
            except (requests.HTTPError, ValueError):
                if last_error == "request_failed":
                    last_error = "invalid_response"
                last_detail = ""

            attempt += 1
            if attempt <= self.retries:
                time.sleep(0.4 * attempt)

        return {"ok": False, "error": last_error, "content": None, "detail": last_detail}