add stream_options for openai api (#1046)

Author: shihaobai

lightllm/server/api_models.py

@@ -48,6 +48,10 @@ class ToolChoice(BaseModel):
     type: Literal["function"] = Field(default="function", examples=["function"])
 
 
+class StreamOptions(BaseModel):
+    include_usage: Optional[bool] = False
+
+
 class CompletionRequest(BaseModel):
     model: str
     # prompt: string or tokens
@@ -58,6 +62,7 @@ class CompletionRequest(BaseModel):
     top_p: Optional[float] = 1.0
     n: Optional[int] = 1
     stream: Optional[bool] = False
+    stream_options: Optional[StreamOptions] = None
     logprobs: Optional[int] = None
     echo: Optional[bool] = False
     stop: Optional[Union[str, List[str]]] = None
@@ -82,6 +87,7 @@ class ChatCompletionRequest(BaseModel):
     top_p: Optional[float] = 1.0
     n: Optional[int] = 1
     stream: Optional[bool] = False
+    stream_options: Optional[StreamOptions] = None
     stop: Optional[Union[str, List[str]]] = None
     max_tokens: Optional[int] = 16
     presence_penalty: Optional[float] = 0.0
@@ -170,6 +176,7 @@ class ChatCompletionStreamResponse(BaseModel):
     created: int = Field(default_factory=lambda: int(time.time()))
     model: str
     choices: List[ChatCompletionStreamResponseChoice]
+    usage: Optional[UsageInfo] = None
 
     @field_validator("id", mode="before")
     def ensure_id_is_str(cls, v):
@@ -216,6 +223,7 @@ class CompletionStreamResponse(BaseModel):
     created: int = Field(default_factory=lambda: int(time.time()))
     model: str
     choices: List[CompletionStreamChoice]
+    usage: Optional[UsageInfo] = None
 
     @field_validator("id", mode="before")
     def ensure_id_is_str(cls, v):

lightllm/server/api_openai.py

@@ -222,7 +222,11 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
         finish_reason = None
         from .req_id_generator import convert_sub_id_to_group_id
 
+        prompt_tokens = 0
+        completion_tokens = 0
         async for sub_req_id, request_output, metadata, finish_status in results_generator:
+            prompt_tokens = metadata["prompt_tokens"]
+            completion_tokens += 1
             if request.tool_choice != "none" and request.tools:
                 delta = request_output
                 group_request_id = convert_sub_id_to_group_id(sub_req_id)
@@ -309,6 +313,22 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
                     choices=[stream_choice],
                 )
                 yield ("data: " + json.dumps(stream_resp.dict(), ensure_ascii=False) + "\n\n").encode("utf-8")
+                # Additional usage chunk
+
+        if request.stream_options and request.stream_options.include_usage:
+            usage = UsageInfo(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                total_tokens=prompt_tokens + completion_tokens,
+            )
+            usage_chunk = ChatCompletionStreamResponse(
+                id=group_request_id,
+                created=created_time,
+                choices=[],  # Empty choices array as per OpenAI spec
+                model=request.model,
+                usage=usage,
+            )
+            yield f"data: {usage_chunk.model_dump_json()}\n\n"
 
     background_tasks = BackgroundTasks()
     return StreamingResponse(stream_results(), media_type="text/event-stream", background=background_tasks)
@@ -453,9 +473,13 @@ async def _handle_streaming_completion(
     async def stream_results() -> AsyncGenerator[bytes, None]:
         from .req_id_generator import convert_sub_id_to_group_id
 
+        prompt_tokens = 0
+        completion_tokens = 0
+
         async for sub_req_id, request_output, metadata, finish_status in results_generator:
             group_request_id = convert_sub_id_to_group_id(sub_req_id)
-
+            prompt_tokens = metadata["prompt_tokens"]
+            completion_tokens += 1
             current_finish_reason = None
             if finish_status.is_finished():
                 current_finish_reason = finish_status.get_finish_reason()
@@ -483,6 +507,21 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
 
         yield "data: [DONE]\n\n".encode("utf-8")
 
+        if request.stream_options and request.stream_options.include_usage:
+            usage = UsageInfo(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                total_tokens=prompt_tokens + completion_tokens,
+            )
+            usage_chunk = CompletionStreamResponse(
+                id=group_request_id,
+                created=created_time,
+                choices=[],  # Empty choices array as per OpenAI spec
+                model=request.model,
+                usage=usage,
+            )
+            yield f"data: {usage_chunk.model_dump_json()}\n\n"
+
     background_tasks = BackgroundTasks()
     return StreamingResponse(stream_results(), media_type="text/event-stream", background=background_tasks)
 

lightllm/server/core/objs/out_token_circlequeue.py

@@ -2,7 +2,7 @@
 import ctypes
 from typing import Tuple
 
-LIGHTLLM_TOKEN_MAX_BYTES = int(os.getenv("LIGHTLLM_TOKEN_MAX_BYTES", 696))
+LIGHTLLM_TOKEN_MAX_BYTES = int(os.getenv("LIGHTLLM_TOKEN_MAX_BYTES", 1280))
 LIGHTLLM_OUT_TOKEN_QUEUE_SIZE = int(os.getenv("LIGHTLLM_OUT_TOKEN_QUEUE_SIZE", 8))