proxy_server.py

import os
import re
import json
import logging
from collections import deque
from pathlib import Path

import yaml
import requests
from flask import Flask, request, Response, stream_with_context

BASE_DIR = Path(__file__).parent
CONFIG_PATH = BASE_DIR / "config.yaml"

app = Flask(__name__)


def load_config():
    with open(CONFIG_PATH, "r") as f:
        raw = f.read()

    def replace_env(match):
        var_name = match.group(1)
        return os.environ.get(var_name, "")

    raw = re.sub(r"\$\{(\w+)\}", replace_env, raw)
    return yaml.safe_load(raw)


config = load_config()

TARGET_BASE_URL = config["target"]["base_url"].rstrip("/")

THINKING_ENABLED = config["thinking"].get("enabled", True)
THINKING_FORMAT = config["thinking"].get("format", "openai")
THINKING_EFFORT = config["thinking"].get("effort", "high")

MODEL_MAPPING = config.get("model_mapping", {})

LOG_LEVEL = config.get("log", {}).get("level", "INFO")
LOG_INPUT = config.get("log", {}).get("log_input", False)
LOG_OUTPUT = config.get("log", {}).get("log_output", False)

REASONING_CACHE = deque(maxlen=200)


def _parse_optional_int(value):
    if value is None:
        return None
    if isinstance(value, bool):
        return None
    if isinstance(value, int):
        return value
    if isinstance(value, str):
        s = value.strip()
        if not s:
            return None
        try:
            return int(s)
        except ValueError:
            return None
    return None


OUTPUT_MAX_TOKENS = _parse_optional_int(config.get("output_length", {}).get("max_tokens"))

logging.basicConfig(
    level=getattr(logging, LOG_LEVEL.upper(), logging.INFO),
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger("proxy")


def _pretty_log_input(body: dict):
    logger.info("=" * 60)
    logger.info("[INPUT] Request Body:")
    logger.info(json.dumps(body, ensure_ascii=False, indent=2))
    logger.info("=" * 60)


def _pretty_log_output(body: dict):
    logger.info("=" * 60)
    logger.info("[OUTPUT] Response Body:")
    logger.info(json.dumps(body, ensure_ascii=False, indent=2))
    logger.info("=" * 60)


def inject_thinking_params(body: dict) -> dict:
    if not THINKING_ENABLED:
        return body

    if THINKING_FORMAT == "openai":
        body["thinking"] = {"type": "enabled"}
        if "reasoning_effort" not in body:
            body["reasoning_effort"] = THINKING_EFFORT

    elif THINKING_FORMAT == "anthropic":
        body["thinking"] = {"type": "enabled"}
        if "output_config" not in body:
            body["output_config"] = {"effort": THINKING_EFFORT}

    return body


def expand_model_name(body: dict) -> dict:
    model = body.get("model", "")
    if model in MODEL_MAPPING:
        expanded = MODEL_MAPPING[model]
        body["model"] = expanded
        logger.info(f"Model alias: {model} -> {expanded}")
    return body


def apply_output_length_override(body: dict) -> dict:
    if OUTPUT_MAX_TOKENS is None:
        return body

    keys = ("max_tokens", "max_completion_tokens", "max_output_tokens")
    overridden = []
    for k in keys:
        if k in body:
            body[k] = OUTPUT_MAX_TOKENS
            overridden.append(k)

    if not overridden:
        body["max_tokens"] = OUTPUT_MAX_TOKENS
        overridden.append("max_tokens")

    logger.info(f"Output length override: {OUTPUT_MAX_TOKENS} ({', '.join(overridden)})")
    return body


def restore_reasoning_content(body: dict) -> dict:
    messages = body.get("messages", [])
    matched_count = 0
    fifo_count = 0
    placeholder_count = 0

    for msg in messages:
        if msg.get("role") != "assistant":
            continue
        rc = msg.get("reasoning_content")
        if rc is not None and rc != "":
            continue

        cached = None
        content = msg.get("content")
        if isinstance(content, str) and content:
            for item in REASONING_CACHE:
                if item.get("content") == content:
                    cached = item
                    break

        if cached is not None:
            REASONING_CACHE.remove(cached)
            msg["reasoning_content"] = cached.get("reasoning") or " "
            matched_count += 1
        elif REASONING_CACHE:
            item = REASONING_CACHE.popleft()
            msg["reasoning_content"] = item.get("reasoning") or " "
            fifo_count += 1
        else:
            msg["reasoning_content"] = " "
            placeholder_count += 1

    if matched_count or fifo_count or placeholder_count:
        parts = []
        if matched_count:
            parts.append(f"{matched_count} matched")
        if fifo_count:
            parts.append(f"{fifo_count} FIFO")
        if placeholder_count:
            parts.append(f"{placeholder_count} placeholder")
        logger.info(f"Restored reasoning_content: {', '.join(parts)}")

    return body


def cache_reasoning_content(reasoning: str, content: str | None = None):
    REASONING_CACHE.append({"content": content or "", "reasoning": reasoning})
    logger.info(f"Cached reasoning_content ({len(reasoning)} chars) [{len(REASONING_CACHE)} in queue]")


def normalize_path(path: str) -> str:
    chat_completions = "/v1/chat/completions"

    full_url_match = re.match(r"https?://[^/]+(/.*)?", path)
    if full_url_match:
        extracted = full_url_match.group(1) or "/"
        logger.info(f"Path stripped domain: {path} -> {extracted}")
        path = extracted

    normalized = path.strip("/")
    if normalized in ("v1", ""):
        logger.info(f"Path rewritten: {path} -> {chat_completions}")
        return chat_completions
    if not normalized.startswith("/"):
        normalized = "/" + normalized
    return normalized


def build_upstream_url(path: str) -> str:
    return f"{TARGET_BASE_URL}{path}"


def build_headers():
    headers = {}
    for key, value in request.headers:
        if key.lower() == "host":
            continue
        headers[key] = value
    headers.pop("Content-Length", None)
    headers.pop("Transfer-Encoding", None)
    headers.pop("Accept-Encoding", None)
    return headers


@app.route("/", defaults={"path": ""}, methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"])
@app.route("/<path:path>", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"])
def proxy(path: str):
    path = normalize_path(path)
    upstream_url = build_upstream_url(path)

    headers = build_headers()
    body = None
    is_stream = False

    if request.method in ("POST", "PUT", "PATCH"):
        raw_body = request.get_data(as_text=True)
        if raw_body:
            try:
                body = json.loads(raw_body)
            except json.JSONDecodeError:
                body = raw_body

        if isinstance(body, dict):
            if LOG_INPUT:
                _pretty_log_input(body)
            is_stream = body.get("stream", False)
            body = restore_reasoning_content(body)
            body = expand_model_name(body)
            body = apply_output_length_override(body)
            body = inject_thinking_params(body)

    query_params = request.args.to_dict()
    logger.info(f"[{request.method}] {upstream_url} | stream={is_stream}")

    if is_stream:
        return _handle_stream(upstream_url, headers, body, query_params)
    else:
        return _handle_normal(upstream_url, headers, body, query_params)


def _handle_normal(upstream_url, headers, body, query_params):
    try:
        if isinstance(body, dict):
            resp = requests.request(
                method=request.method,
                url=upstream_url,
                headers=headers,
                json=body,
                params=query_params,
                timeout=600,
            )
        else:
            resp = requests.request(
                method=request.method,
                url=upstream_url,
                headers=headers,
                data=body,
                params=query_params,
                timeout=600,
            )

        excluded = {"content-encoding", "content-length", "transfer-encoding", "connection"}
        resp_headers = [
            (k, v) for k, v in resp.headers.items() if k.lower() not in excluded
        ]

        resp_data = resp.content
        try:
            data = json.loads(resp_data)
            msg = data["choices"][0]["message"]
            if "reasoning_content" in msg:
                cache_reasoning_content(msg["reasoning_content"], msg.get("content"))
            if LOG_OUTPUT:
                _pretty_log_output(data)
        except Exception:
            if LOG_OUTPUT:
                logger.info(f"[OUTPUT] Non-JSON response ({len(resp_data)} bytes)")

        return Response(
            resp_data,
            status=resp.status_code,
            headers=resp_headers,
            content_type=resp.headers.get("Content-Type", "application/json"),
        )
    except requests.exceptions.RequestException as e:
        logger.error(f"Upstream error: {e}")
        error_body = ""
        if hasattr(e, "response") and e.response is not None:
            error_body = e.response.text[:2000]
            logger.error(f"Upstream response body: {error_body}")
        return Response(
            json.dumps({"error": str(e)}),
            status=502,
            content_type="application/json",
        )


def _attempt_cache_from_chunk(chunk: bytes, acc_content: list, acc_reasoning: list):
    try:
        text = chunk.decode(errors="ignore")
        for line in text.split("\n"):
            line = line.strip()
            if not line.startswith("data:"):
                continue
            data_str = line[5:].strip()
            if data_str == "[DONE]":
                continue
            try:
                data = json.loads(data_str)
            except json.JSONDecodeError:
                continue
            delta = data.get("choices", [{}])[0].get("delta", {})
            if "content" in delta and delta["content"] is not None:
                acc_content.append(delta["content"])
            if "reasoning_content" in delta and delta["reasoning_content"] is not None:
                acc_reasoning.append(delta["reasoning_content"])
    except Exception:
        pass


def _handle_stream(upstream_url, headers, body, query_params):
    try:
        resp = requests.request(
            method=request.method,
            url=upstream_url,
            headers=headers,
            json=body if isinstance(body, dict) else None,
            data=body if not isinstance(body, dict) else None,
            params=query_params,
            stream=True,
            timeout=600,
        )
        resp.raise_for_status()
    except requests.exceptions.RequestException as e:
        logger.error(f"Stream upstream connection error: {e}")
        error_body = ""
        if hasattr(e, "response") and e.response is not None:
            error_body = e.response.text[:2000]
            logger.error(f"Upstream response body: {error_body}")
        return Response(
            json.dumps({"error": str(e)}),
            status=502,
            content_type="application/json",
        )

    acc_content = []
    acc_reasoning = []

    def generate():
        try:
            for chunk in resp.iter_content(chunk_size=None):
                if chunk:
                    _attempt_cache_from_chunk(chunk, acc_content, acc_reasoning)
                    yield chunk
        except requests.exceptions.RequestException as e:
            logger.error(f"Stream read error: {e}")
            yield f'data: {{"error": "{str(e)}"}}\\n\\n'.encode()
        finally:
            full_content = "".join(acc_content).strip()
            full_reasoning = "".join(acc_reasoning).strip()
            if full_reasoning:
                cache_reasoning_content(full_reasoning, full_content)
            if LOG_OUTPUT:
                if full_reasoning:
                    logger.info("-" * 40)
                    logger.info(f"[STREAM OUTPUT] reasoning_content ({len(full_reasoning)} chars):")
                    logger.info(full_reasoning[:2000])
                if full_content:
                    logger.info("-" * 40)
                    logger.info(f"[STREAM OUTPUT] content ({len(full_content)} chars):")
                    logger.info(full_content[:2000])
                logger.info("=" * 60)

    excluded = {"content-encoding", "content-length", "transfer-encoding", "connection"}
    resp_headers = [
        (k, v) for k, v in resp.headers.items() if k.lower() not in excluded
    ]

    return Response(
        stream_with_context(generate()),
        status=resp.status_code,
        headers=resp_headers,
    )


def main():
    host = config["server"]["host"]
    port = config["server"]["port"]

    logger.info("=" * 50)
    logger.info(f"DeepSeek API Proxy starting on {host}:{port}")
    logger.info(f"Upstream: {TARGET_BASE_URL}")
    logger.info(f"Thinking mode: {'ON' if THINKING_ENABLED else 'OFF'}")
    if THINKING_ENABLED:
        logger.info(f"  Format: {THINKING_FORMAT}")
        logger.info(f"  Effort: {THINKING_EFFORT}")
    logger.info(f"Log input: {'ON' if LOG_INPUT else 'OFF'}")
    logger.info(f"Log output: {'ON' if LOG_OUTPUT else 'OFF'}")
    logger.info("=" * 50)

    app.run(host=host, port=port, debug=False)


if __name__ == "__main__":
    main()