lmstudiofixed/tools/lmstudio/openai_reasoning_proxy.py

#!/usr/bin/env python3
import argparse
import json
import time
import urllib.error
import urllib.request
from typing import Any

from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.responses import JSONResponse
import uvicorn


DEFAULT_UPSTREAM = "http://127.0.0.1:7860"
PASS_THROUGH_FIELDS = {
    "temperature",
    "top_p",
    "presence_penalty",
    "frequency_penalty",
    "stop",
    "seed",
}

app = FastAPI(title="LM Studio OpenAI Reasoning Proxy")
app.state.upstream_base = DEFAULT_UPSTREAM
app.state.request_timeout = 1800


def post_json(url: str, payload: dict[str, Any], timeout: int) -> tuple[int, str]:
    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    request = urllib.request.Request(
        url,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return response.status, response.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read().decode("utf-8", errors="replace")


def get_json(url: str, timeout: int) -> tuple[int, str]:
    request = urllib.request.Request(url, method="GET")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return response.status, response.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read().decode("utf-8", errors="replace")


def parse_json_maybe(raw_text: str) -> dict[str, Any] | None:
    try:
        return json.loads(raw_text)
    except json.JSONDecodeError:
        return None


def normalize_reasoning_value(value: Any) -> str | None:
    if isinstance(value, str):
        lowered = value.strip().lower()
        if lowered in {"on", "off"}:
            return lowered
    if isinstance(value, bool):
        return "on" if value else "off"
    return None


def resolve_reasoning_mode(request_payload: dict[str, Any]) -> str:
    normalized = normalize_reasoning_value(request_payload.get("reasoning"))
    if normalized is not None:
        return normalized

    if "enable_thinking" in request_payload:
        enable_thinking = request_payload.get("enable_thinking")
        if isinstance(enable_thinking, bool):
            return "on" if enable_thinking else "off"

    return "on"


def extract_text_content(content: Any) -> str:
    if isinstance(content, str):
        return content

    if isinstance(content, list):
        text_parts: list[str] = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                text = item.get("text")
                if isinstance(text, str):
                    text_parts.append(text)
        return "\n".join(part for part in text_parts if part)

    return ""


def messages_to_native_input(messages: list[dict[str, Any]]) -> tuple[str, str | None]:
    transcript_parts: list[str] = []
    system_parts: list[str] = []

    for message in messages:
        if not isinstance(message, dict):
            continue

        role = str(message.get("role") or "").strip().lower()
        content = extract_text_content(message.get("content"))
        if not content:
            continue

        if role == "system":
            system_parts.append(content)
            continue

        if role == "user":
            transcript_parts.append(f"User: {content}")
            continue

        if role == "assistant":
            transcript_parts.append(f"Assistant: {content}")
            continue

        transcript_parts.append(f"{role.title() or 'Message'}: {content}")

    transcript = "\n\n".join(transcript_parts)
    system_prompt = "\n\n".join(system_parts) if system_parts else None
    return transcript, system_prompt


def build_native_payload(request_payload: dict[str, Any]) -> dict[str, Any]:
    messages = request_payload.get("messages")
    if not isinstance(messages, list) or not messages:
        raise ValueError("messages must be a non-empty list")

    input_text, system_prompt = messages_to_native_input(messages)
    if not input_text:
        raise ValueError("messages must include at least one non-system text message")

    native_payload: dict[str, Any] = {
        "model": request_payload.get("model"),
        "input": input_text,
        "reasoning": resolve_reasoning_mode(request_payload),
        "store": False,
    }

    if system_prompt:
        native_payload["system_prompt"] = system_prompt

    if "max_tokens" in request_payload:
        native_payload["max_output_tokens"] = request_payload["max_tokens"]

    for field in PASS_THROUGH_FIELDS:
        if field in request_payload:
            native_payload[field] = request_payload[field]

    return native_payload


def _collect_output_text(native_response: dict[str, Any], output_type: str) -> list[str]:
    texts: list[str] = []
    for item in native_response.get("output") or []:
        if item.get("type") == output_type:
            content = item.get("content")
            if isinstance(content, str):
                texts.append(content)
    return texts


def translate_native_response(native_response: dict[str, Any]) -> dict[str, Any]:
    message_parts = _collect_output_text(native_response, "message")
    reasoning_parts = _collect_output_text(native_response, "reasoning")
    stats = native_response.get("stats") or {}

    content = "\n".join(part for part in message_parts if part)
    reasoning_content = "\n".join(part for part in reasoning_parts if part)
    choice_message = {
        "role": "assistant",
        "content": content,
        "tool_calls": [],
    }
    if reasoning_content:
        choice_message["reasoning_content"] = reasoning_content

    return {
        "id": native_response.get("id", f"chatcmpl-proxy-{int(time.time() * 1000)}"),
        "object": "chat.completion",
        "created": int(time.time()),
        "model": native_response.get("model"),
        "choices": [
            {
                "index": 0,
                "message": choice_message,
                "logprobs": None,
                "finish_reason": "stop",
            }
        ],
        "usage": {
            "prompt_tokens": stats.get("input_tokens", 0),
            "completion_tokens": stats.get("total_output_tokens", 0),
            "total_tokens": stats.get("input_tokens", 0) + stats.get("total_output_tokens", 0),
            "completion_tokens_details": {
                "reasoning_tokens": stats.get("reasoning_output_tokens", 0),
            },
        },
        "stats": stats,
        "system_fingerprint": native_response.get("model"),
    }


def json_error(status_code: int, message: str, error_type: str = "invalid_request_error") -> JSONResponse:
    return JSONResponse(
        status_code=status_code,
        content={
            "error": {
                "message": message,
                "type": error_type,
            }
        },
    )


@app.get("/healthz")
async def healthz() -> dict[str, str]:
    return {"status": "ok"}


@app.get("/v1/models")
async def list_models() -> Response:
    status_code, raw = get_json(
        f"{app.state.upstream_base}/v1/models",
        timeout=app.state.request_timeout,
    )
    return Response(content=raw, status_code=status_code, media_type="application/json")


@app.post("/v1/chat/completions")
async def chat_completions(request: Request) -> Response:
    try:
        request_payload = await request.json()
    except Exception as exc:  # pragma: no cover
        raise HTTPException(status_code=400, detail=f"invalid JSON: {exc}") from exc

    if request_payload.get("stream") is True:
        return json_error(400, "stream=true is not supported by this proxy")

    try:
        native_payload = build_native_payload(request_payload)
    except ValueError as exc:
        return json_error(400, str(exc))

    status_code, raw = post_json(
        f"{app.state.upstream_base}/api/v1/chat",
        native_payload,
        timeout=app.state.request_timeout,
    )
    parsed = parse_json_maybe(raw)
    if status_code >= 400:
        if parsed is not None:
            return JSONResponse(status_code=status_code, content=parsed)
        return json_error(status_code, "upstream returned a non-JSON error", "upstream_error")

    if parsed is None:
        return json_error(502, "upstream returned non-JSON content", "bad_gateway")

    translated = translate_native_response(parsed)
    return JSONResponse(status_code=200, content=translated)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8001)
    parser.add_argument("--upstream", default=DEFAULT_UPSTREAM)
    parser.add_argument("--timeout", type=int, default=1800)
    args = parser.parse_args()

    app.state.upstream_base = args.upstream.rstrip("/")
    app.state.request_timeout = args.timeout
    uvicorn.run(app, host=args.host, port=args.port)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())