fix(api): handle streaming errors and multiple choices

lightllm/server/api_anthropic.py

@@ -8,6 +8,7 @@
 chat_completions_impl and re-emits it as the Anthropic event sequence
 (message_start, content_block_*, message_delta, message_stop).
 """
+
 from __future__ import annotations
 
 import asyncio
@@ -33,6 +34,8 @@
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.log_utils import init_logger
 
+from .api_errors import is_rate_limit_error
+
 logger = init_logger(__name__)
 
 _cached_adapter: Any = None
@@ -780,6 +783,15 @@ def _sse_event(event_type: str, data_obj: Dict[str, Any]) -> bytes:
     return f"event: {event_type}\ndata: {json.dumps(data_obj)}\n\n".encode("utf-8")
 
 
+def _anthropic_error_type(error: Dict[str, Any]) -> str:
+    error_type = error.get("type")
+    if is_rate_limit_error(error):
+        return "rate_limit_error"
+    if error_type == "invalid_request_error":
+        return error_type
+    return "api_error"
+
+
 async def _openai_sse_to_anthropic_events(
     openai_body_iterator,
     requested_model: str,
@@ -842,16 +854,14 @@ async def _openai_sse_to_anthropic_events(
 
             if "error" in chunk and "choices" not in chunk:
                 error = chunk["error"]
-                error_type = error.get("type")
-                if error.get("code") == 429 or error_type in ("RateLimitError", "rate_limit_error"):
-                    error_type = "rate_limit_error"
-                elif error_type != "invalid_request_error":
-                    error_type = "api_error"
                 yield _sse_event(
                     "error",
                     {
                         "type": "error",
-                        "error": {"type": error_type, "message": error.get("message", "generation failed")},
+                        "error": {
+                            "type": _anthropic_error_type(error),
+                            "message": error.get("message", "generation failed"),
+                        },
                     },
                 )
                 return
@@ -1112,7 +1122,7 @@ def _rewrap_openai_error_as_anthropic(resp: JSONResponse) -> JSONResponse:
 async def anthropic_messages_impl(raw_request: Request) -> Response:
     # Lazy imports to avoid pulling in heavy server deps at module import time.
     from .api_models import ChatCompletionRequest, ChatCompletionResponse
-    from .api_openai import chat_completions_impl, prime_pd_master_streaming_response
+    from .api_openai import chat_completions_impl
 
     try:
         raw_body = await raw_request.json()
@@ -1143,21 +1153,19 @@ async def anthropic_messages_impl(raw_request: Request) -> Response:
     downstream = await chat_completions_impl(chat_request, raw_request)
 
     if is_stream:
-        from fastapi.responses import StreamingResponse
+        from .api_stream_obj import CustomStreamingResponse
 
-        if not isinstance(downstream, StreamingResponse):
+        if not isinstance(downstream, CustomStreamingResponse):
             # chat_completions_impl returned an OpenAI-format error — rewrap it.
             if isinstance(downstream, JSONResponse):
                 return _rewrap_openai_error_as_anthropic(downstream)
             return downstream
 
-        downstream = await prime_pd_master_streaming_response(downstream)
-
         message_id = f"msg_{uuid.uuid4().hex[:24]}"
         anthropic_stream = _openai_sse_to_anthropic_events(
             downstream.body_iterator, requested_model=requested_model, message_id=message_id
         )
-        return StreamingResponse(anthropic_stream, media_type="text/event-stream")
+        return CustomStreamingResponse(anthropic_stream, media_type="text/event-stream")
 
     if not isinstance(downstream, ChatCompletionResponse):
         if isinstance(downstream, JSONResponse):

lightllm/server/api_errors.py

@@ -0,0 +1,40 @@
+"""Shared helpers for OpenAI-compatible API error responses."""
+
+from http import HTTPStatus
+from typing import Any, Mapping
+
+from fastapi.responses import JSONResponse
+
+from lightllm.utils.error_utils import ServerBusyError
+
+_RATE_LIMIT_ERROR_TYPES = {"RateLimitError", "rate_limit_error"}
+
+
+def is_rate_limit_error(error: Mapping[str, Any]) -> bool:
+    return error.get("code") == HTTPStatus.TOO_MANY_REQUESTS or error.get("type") in _RATE_LIMIT_ERROR_TYPES
+
+
+def create_error_response(
+    status_code: HTTPStatus, message: str, err_type: str = None, param: str = None
+) -> JSONResponse:
+    if err_type is None:
+        if status_code.value >= 500:
+            err_type = "InternalServerError"
+        elif status_code == HTTPStatus.NOT_FOUND:
+            err_type = "NotFoundError"
+        elif status_code == HTTPStatus.TOO_MANY_REQUESTS:
+            err_type = "RateLimitError"
+        else:
+            err_type = "BadRequestError"
+
+    from .api_http import g_objs
+
+    g_objs.metric_client.counter_inc("lightllm_request_failure")
+    return JSONResponse(
+        {"error": {"message": message, "type": err_type, "param": param, "code": status_code.value}},
+        status_code=status_code.value,
+    )
+
+
+def create_server_busy_response(exc: ServerBusyError) -> JSONResponse:
+    return create_error_response(HTTPStatus(exc.status_code), str(exc))

lightllm/server/api_http.py

@@ -37,7 +37,7 @@
 from typing import Callable
 from lightllm.server import TokenLoad
 from fastapi import BackgroundTasks, FastAPI, Request
-from fastapi.responses import Response, StreamingResponse, JSONResponse
+from fastapi.responses import Response, JSONResponse
 from lightllm.server.core.objs.sampling_params import SamplingParams
 from lightllm.server.core.objs import StartArgs
 from .multimodal_params import MultimodalParams
@@ -52,7 +52,8 @@
 from lightllm.utils.shm_port_args import get_shm_port_args
 from dataclasses import asdict, dataclass, is_dataclass
 
-from .api_openai import chat_completions_impl, completions_impl, prime_pd_master_streaming_response
+from .api_errors import create_error_response, create_server_busy_response
+from .api_openai import chat_completions_impl, completions_impl
 from .api_models import (
     ChatCompletionRequest,
     ChatCompletionResponse,
@@ -153,34 +154,19 @@ async def send_wrapper(message):
 app.add_middleware(_AccessLogMiddleware)
 
 
-def create_error_response(
-    status_code: HTTPStatus, message: str, err_type: str = None, param: str = None
-) -> JSONResponse:
-    if err_type is None:
-        if status_code.value >= 500:
-            err_type = "InternalServerError"
-        elif status_code == HTTPStatus.NOT_FOUND:
-            err_type = "NotFoundError"
-        elif status_code == HTTPStatus.TOO_MANY_REQUESTS:
-            err_type = "RateLimitError"
-        else:
-            err_type = "BadRequestError"
-
-    g_objs.metric_client.counter_inc("lightllm_request_failure")
-    return JSONResponse(
-        {"error": {"message": message, "type": err_type, "param": param, "code": status_code.value}},
-        status_code=status_code.value,
-    )
-
-
-def create_server_busy_response(exc: ServerBusyError) -> JSONResponse:
-    status = HTTPStatus(exc.status_code)
-    return create_error_response(status, str(exc), err_type="RateLimitError")
-
-
 @app.exception_handler(ServerBusyError)
 async def server_busy_exception_handler(request: Request, exc: ServerBusyError) -> JSONResponse:
     logger.warning(str(exc))
+
+    # Streaming responses can raise during their first body iteration, after
+    # the route handler has already returned. Preserve the Anthropic error
+    # envelope for that deferred failure path as well.
+    if request.url.path == "/v1/messages":
+        from .api_anthropic import _anthropic_error_response
+
+        g_objs.metric_client.counter_inc("lightllm_request_failure")
+        return _anthropic_error_response(HTTPStatus(exc.status_code), str(exc))
+
     return create_server_busy_response(exc)
 
 
@@ -327,8 +313,7 @@ async def generate_stream(request: Request) -> Response:
         )
 
     try:
-        response = await g_objs.g_generate_stream_func(request, g_objs.httpserver_manager)
-        return await prime_pd_master_streaming_response(response)
+        return await g_objs.g_generate_stream_func(request, g_objs.httpserver_manager)
     except ServerBusyError as e:
         logger.warning(str(e))
         return create_server_busy_response(e)
@@ -385,7 +370,6 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request)
 
     try:
         resp = await chat_completions_impl(request, raw_request)
-        resp = await prime_pd_master_streaming_response(resp)
     except ValueError as e:
         return create_error_response(HTTPStatus.BAD_REQUEST, str(e))
     except ServerBusyError as e:
@@ -406,7 +390,6 @@ async def completions(request: CompletionRequest, raw_request: Request) -> Respo
 
     try:
         resp = await completions_impl(request, raw_request)
-        resp = await prime_pd_master_streaming_response(resp)
     except ValueError as e:
         return create_error_response(HTTPStatus.BAD_REQUEST, str(e))
     except ServerBusyError as e:
@@ -427,8 +410,7 @@ async def anthropic_messages(raw_request: Request) -> Response:
     from .api_anthropic import _anthropic_error_response, anthropic_messages_impl
 
     try:
-        response = await anthropic_messages_impl(raw_request)
-        return await prime_pd_master_streaming_response(response)
+        return await anthropic_messages_impl(raw_request)
     except ServerBusyError as e:
         logger.warning(str(e))
         g_objs.metric_client.counter_inc("lightllm_request_failure")
@@ -447,8 +429,7 @@ async def openai_responses(raw_request: Request) -> Response:
     from .api_responses import responses_impl
 
     try:
-        response = await responses_impl(raw_request)
-        return await prime_pd_master_streaming_response(response)
+        return await responses_impl(raw_request)
     except ServerBusyError as e:
         logger.warning(str(e))
         return create_server_busy_response(e)

lightllm/server/api_http_rl.py

@@ -7,7 +7,7 @@
        → (多数) RlOpReq → Router → Model RlBackendOps
 
 路由在模块级 ``router`` 上注册，由 ``api_http`` ``include_router`` 挂载。
-``g_objs`` / ``create_error_response`` 在 handler 内懒导入，避免与 api_http 循环依赖。
+``g_objs`` 在 handler 内懒导入，避免与 api_http 循环依赖。
 """
 
 from http import HTTPStatus
@@ -29,14 +29,14 @@
 )
 from lightllm.utils.log_utils import init_logger
 
+from .api_errors import create_error_response
+
 logger = init_logger(__name__)
 
 router = APIRouter()
 
 
 async def handle_request_common(request_obj, handler):
-    from .api_http import create_error_response
-
     try:
         ret: RlOpRsp = await handler(request_obj)
         if ret.success:
@@ -51,7 +51,7 @@ async def handle_request_common(request_obj, handler):
 @router.post("/abort_request")
 async def abort_request(request: AbortReq, raw_request: Request):
     """Abort a request."""
-    from .api_http import create_error_response, g_objs
+    from .api_http import g_objs
 
     try:
         success, msg = await g_objs.httpserver_manager.abort_request(request)

lightllm/server/api_lightllm.py

@@ -1,10 +1,11 @@
 import collections
 from typing import AsyncGenerator
 from fastapi import BackgroundTasks, Request
-from fastapi.responses import Response, StreamingResponse
+from fastapi.responses import Response
 from lightllm.server.core.objs.sampling_params import SamplingParams
 from .multimodal_params import MultimodalParams
 from .httpserver.manager import HttpServerManager
+from .api_stream_obj import CustomStreamingResponse
 import ujson as json
 
 
@@ -164,6 +165,6 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
     from .api_openai import _safe_stream_wrapper
 
     background_tasks = BackgroundTasks()
-    return StreamingResponse(
+    return CustomStreamingResponse(
         _safe_stream_wrapper(stream_results()), media_type="text/event-stream", background=background_tasks
     )