Function lightllm_generate_stream

lightllm/server/api_lightllm.py:109–152

(request: Request, httpserver_manager: HttpServerManager)

async def lightllm_generate_stream(request: Request, httpserver_manager: HttpServerManager) -> Response:

    request_dict = await request.json()
    prompt = request_dict.pop("inputs")
    sample_params_dict = request_dict["parameters"]
    _ = sample_params_dict.pop("return_details", False)
    sampling_params = SamplingParams()
    sampling_params.init(tokenizer=httpserver_manager.tokenizer, **sample_params_dict)
    sampling_params.verify()
    if sampling_params.best_of != 1:
        raise Exception("stream api only support best_of == 1")

    multimodal_params_dict = request_dict.get("multimodal_params", {})
    multimodal_params = MultimodalParams(**multimodal_params_dict)
    results_generator = httpserver_manager.generate(prompt, sampling_params, multimodal_params, request=request)

    # Streaming case
    async def stream_results() -> AsyncGenerator[bytes, None]:
        # input_usage 只会在第一个metadata中出现，所以需要保存下来
        input_usage = None
        async for _, request_output, metadata, finish_status in results_generator:
            if input_usage is None:
                input_usage = metadata.get("input_usage", None)

            ret = {
                "token": {
                    "id": metadata.get("id", None),
                    "text": request_output,
                    "logprob": metadata.get("logprob", None),
                    "special": metadata.get("special", False),
                    "count_output_tokens": metadata.get("count_output_tokens", 0),
                    "prompt_tokens": metadata.get("prompt_tokens", 0),
                },
                "generated_text": None,
                "finished": finish_status.is_finished(),
                "finish_reason": finish_status.get_finish_reason(),
                "details": None,
                "input_usage": input_usage,
            }

            yield ("data:" + json.dumps(ret, ensure_ascii=False) + "\n\n").encode("utf-8")

    background_tasks = BackgroundTasks()
    return StreamingResponse(stream_results(), media_type="text/event-stream", background=background_tasks)