lmstudiofixed/services/litellm-lmstudio-adapter/litellm_lmstudio_adapter.py

#!/usr/bin/env python3
import argparse
import json
import time
from typing import Any

import httpx
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, Response, StreamingResponse
import uvicorn


DEFAULT_UPSTREAM = "http://127.0.0.1:7860"
PASS_THROUGH_CHAT_FIELDS = {
    "temperature",
    "top_p",
    "presence_penalty",
    "frequency_penalty",
    "stop",
    "seed",
}
STUB_ENDPOINTS = {
    "/v1/audio/speech": "audio.speech",
    "/v1/audio/transcriptions": "audio.transcriptions",
    "/v1/audio/translations": "audio.translations",
    "/v1/images/generations": "images.generations",
    "/v1/images/edits": "images.edits",
    "/v1/images/variations": "images.variations",
    "/v1/moderations": "moderations",
    "/v1/files": "files",
}

app = FastAPI(title="LiteLLM LM Studio Adapter")
app.state.upstream_base = DEFAULT_UPSTREAM
app.state.request_timeout = 1800


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

    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)
        elif role == "user":
            transcript_parts.append(f"User: {content}")
        elif role == "assistant":
            transcript_parts.append(f"Assistant: {content}")
        else:
            transcript_parts.append(f"{role.title() or 'Message'}: {content}")

    return "\n\n".join(transcript_parts), "\n\n".join(system_parts) or None


def build_chat_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_CHAT_FIELDS:
        if field in request_payload:
            native_payload[field] = request_payload[field]
    if request_payload.get("stream") is True:
        native_payload["stream"] = True
    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_chat_response(native_response: dict[str, Any]) -> dict[str, Any]:
    model = native_response.get("model") or native_response.get("model_instance_id")
    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)
    message = {"role": "assistant", "content": content, "tool_calls": []}
    if reasoning_content:
        message["reasoning_content"] = reasoning_content
    return {
        "id": native_response.get("id", f"chatcmpl-{int(time.time() * 1000)}"),
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model,
        "choices": [
            {
                "index": 0,
                "message": 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": model,
    }


def build_responses_response(native_response: dict[str, Any]) -> dict[str, Any]:
    chat_response = translate_chat_response(native_response)
    message = chat_response["choices"][0]["message"]
    output = []
    if message.get("reasoning_content"):
        output.append(
            {
                "id": "rs_reasoning_0",
                "type": "reasoning",
                "summary": [],
                "content": [{"type": "output_text", "text": message["reasoning_content"]}],
            }
        )
    output.append(
        {
            "id": "msg_0",
            "type": "message",
            "role": "assistant",
            "content": [{"type": "output_text", "text": message.get("content", "")}],
        }
    )
    return {
        "id": f"resp_{int(time.time() * 1000)}",
        "object": "response",
        "created_at": int(time.time()),
        "model": chat_response["model"],
        "output": output,
        "usage": chat_response["usage"],
        "status": "completed",
    }


def sse_frame(data: dict[str, Any]) -> str:
    return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"


def parse_sse_event_blocks(raw_text: str):
    for block in raw_text.split("\n\n"):
        block = block.strip()
        if not block:
            continue
        event_name = None
        data_lines: list[str] = []
        for line in block.splitlines():
            if line.startswith("event:"):
                event_name = line[6:].strip()
            elif line.startswith("data:"):
                data_lines.append(line[5:].strip())
        data_raw = "\n".join(data_lines)
        parsed_data = None
        if data_raw and data_raw != "[DONE]":
            try:
                parsed_data = json.loads(data_raw)
            except json.JSONDecodeError:
                parsed_data = None
        yield {
            "event": event_name,
            "data": parsed_data,
            "data_raw": data_raw,
        }


def translate_chat_stream_event(event: dict[str, Any], model: str, chunk_id: str) -> str | None:
    event_type = event.get("type")
    content = event.get("content")
    if not isinstance(content, str):
        return None

    delta: dict[str, Any] = {}
    if event_type in {"reasoning", "reasoning.delta"}:
        delta["reasoning_content"] = content
    elif event_type in {"message", "message.delta"}:
        delta["content"] = content
    else:
        return None

    payload = {
        "id": chunk_id,
        "object": "chat.completion.chunk",
        "created": int(time.time()),
        "model": model,
        "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
    }
    return sse_frame(payload)


def translate_responses_stream_event(event: dict[str, Any], model: str, response_id: str) -> list[str]:
    event_type = event.get("type")
    content = event.get("content")
    if not isinstance(content, str):
        return []

    if event_type in {"reasoning", "reasoning.delta"}:
        return [
            sse_frame(
                {
                    "type": "response.reasoning.delta",
                    "response_id": response_id,
                    "model": model,
                    "delta": content,
                }
            )
        ]

    if event_type in {"message", "message.delta"}:
        return [
            sse_frame(
                {
                    "type": "response.output_text.delta",
                    "response_id": response_id,
                    "model": model,
                    "delta": content,
                }
            )
        ]

    return []


def build_stub_error(feature_name: str) -> JSONResponse:
    return JSONResponse(
        status_code=501,
        content={
            "error": {
                "message": f"{feature_name} is not implemented by this adapter yet",
                "type": "not_implemented",
                "param": None,
                "code": "not_implemented",
            }
        },
    )


async def get_async_client() -> httpx.AsyncClient:
    return httpx.AsyncClient(timeout=app.state.request_timeout)


async def proxy_request(method: str, path: str, body: bytes | None = None, headers: dict[str, str] | None = None) -> Response:
    async with await get_async_client() as client:
        response = await client.request(
            method,
            f"{app.state.upstream_base}{path}",
            content=body,
            headers=headers,
        )
    response_headers = {
        key: value
        for key, value in response.headers.items()
        if key.lower() not in {"content-length", "transfer-encoding", "connection", "content-encoding"}
    }
    return Response(
        content=response.content,
        status_code=response.status_code,
        headers=response_headers,
        media_type=response.headers.get("content-type"),
    )


async def stream_lmstudio_events(native_payload: dict[str, Any], translator, model: str, final_frame: str):
    async with httpx.AsyncClient(timeout=app.state.request_timeout) as client:
        async with client.stream(
            "POST",
            f"{app.state.upstream_base}/api/v1/chat",
            json=native_payload,
            headers={"Accept": "text/event-stream"},
        ) as response:
            if response.status_code >= 400:
                raw = await response.aread()
                payload = {
                    "error": {
                        "message": raw.decode("utf-8", errors="replace"),
                        "type": "upstream_error",
                    }
                }
                yield sse_frame(payload)
                yield "data: [DONE]\n\n"
                return

            buffer = ""
            async for text in response.aiter_text():
                buffer += text
                while "\n\n" in buffer:
                    block, buffer = buffer.split("\n\n", 1)
                    for parsed_block in parse_sse_event_blocks(block + "\n\n"):
                        if parsed_block["data_raw"] == "[DONE]":
                            continue
                        event = parsed_block["data"]
                        if not isinstance(event, dict):
                            continue
                        translated = translator(event, model)
                        if translated is None:
                            continue
                        if isinstance(translated, str):
                            if translated:
                                yield translated
                        else:
                            for frame in translated:
                                yield frame
    if final_frame:
        yield final_frame
    yield "data: [DONE]\n\n"


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


@app.get("/v1/models")
async def list_models() -> Response:
    return await proxy_request("GET", "/v1/models")


@app.post("/v1/chat/completions")
async def chat_completions(request: Request) -> Response:
    payload = await request.json()
    native_payload = build_chat_native_payload(payload)
    model = str(payload.get("model") or "")

    if payload.get("stream") is True:
        chunk_id = f"chatcmpl-{int(time.time() * 1000)}"

        def translator(event: dict[str, Any], event_model: str):
            return translate_chat_stream_event(event, event_model, chunk_id)

        final_payload = {
            "id": chunk_id,
            "object": "chat.completion.chunk",
            "created": int(time.time()),
            "model": model,
            "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
        }
        return StreamingResponse(
            stream_lmstudio_events(native_payload, translator, model, sse_frame(final_payload)),
            media_type="text/event-stream",
        )

    async with await get_async_client() as client:
        response = await client.post(f"{app.state.upstream_base}/api/v1/chat", json=native_payload)
    return JSONResponse(status_code=response.status_code, content=translate_chat_response(response.json()))


@app.post("/v1/responses")
async def responses_api(request: Request) -> Response:
    payload = await request.json()
    messages = payload.get("input")
    if isinstance(messages, str):
        payload = {
            "model": payload.get("model"),
            "messages": [{"role": "user", "content": messages}],
            "reasoning": payload.get("reasoning"),
            "enable_thinking": payload.get("enable_thinking"),
            "stream": payload.get("stream"),
            "max_tokens": payload.get("max_output_tokens") or payload.get("max_tokens"),
            "temperature": payload.get("temperature"),
        }
    elif isinstance(messages, list):
        payload = {
            "model": payload.get("model"),
            "messages": messages,
            "reasoning": payload.get("reasoning"),
            "enable_thinking": payload.get("enable_thinking"),
            "stream": payload.get("stream"),
            "max_tokens": payload.get("max_output_tokens") or payload.get("max_tokens"),
            "temperature": payload.get("temperature"),
        }
    else:
        raise ValueError("responses input must be a string or a message list")

    native_payload = build_chat_native_payload(payload)
    model = str(payload.get("model") or "")
    response_id = f"resp_{int(time.time() * 1000)}"

    if payload.get("stream") is True:
        def translator(event: dict[str, Any], event_model: str):
            return translate_responses_stream_event(event, event_model, response_id)

        final_frame = sse_frame(
            {
                "type": "response.completed",
                "response": {
                    "id": response_id,
                    "model": model,
                    "status": "completed",
                },
            }
        )
        initial_frame = sse_frame(
            {
                "type": "response.created",
                "response": {
                    "id": response_id,
                    "model": model,
                    "status": "in_progress",
                },
            }
        )

        async def generator():
            yield initial_frame
            async for frame in stream_lmstudio_events(native_payload, translator, model, final_frame):
                yield frame

        return StreamingResponse(generator(), media_type="text/event-stream")

    async with await get_async_client() as client:
        response = await client.post(f"{app.state.upstream_base}/api/v1/chat", json=native_payload)
    return JSONResponse(status_code=response.status_code, content=build_responses_response(response.json()))


@app.post("/v1/embeddings")
async def embeddings(request: Request) -> Response:
    body = await request.body()
    return await proxy_request("POST", "/v1/embeddings", body=body, headers={"Content-Type": "application/json"})


@app.post("/v1/completions")
async def completions(request: Request) -> Response:
    body = await request.body()
    return await proxy_request("POST", "/v1/completions", body=body, headers={"Content-Type": "application/json"})


@app.api_route("/api/v1/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE"])
async def native_v1_passthrough(path: str, request: Request) -> Response:
    body = await request.body()
    headers = {"Content-Type": request.headers.get("content-type", "application/json")}
    query = f"?{request.url.query}" if request.url.query else ""
    return await proxy_request(request.method, f"/api/v1/{path}{query}", body=body, headers=headers)


@app.api_route("/api/v0/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE"])
async def native_v0_passthrough(path: str, request: Request) -> Response:
    body = await request.body()
    headers = {"Content-Type": request.headers.get("content-type", "application/json")}
    query = f"?{request.url.query}" if request.url.query else ""
    return await proxy_request(request.method, f"/api/v0/{path}{query}", body=body, headers=headers)


for route_path, feature_name in STUB_ENDPOINTS.items():
    async def _stub(feature=feature_name):
        return build_stub_error(feature)

    app.add_api_route(route_path, _stub, methods=["POST", "GET", "DELETE"])


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8010)
    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())