diff --git a/serve/openai_server.py b/serve/openai_server.py index d34acdc..1e41bf2 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -26,9 +26,13 @@ import threading import time import uuid from pathlib import Path +from typing import TYPE_CHECKING, Any # Module level so FastAPI can resolve the string annotation `request: Request`. -from fastapi import Request +from fastapi import FastAPI, Request + +if TYPE_CHECKING: + from deepinfer._C.engine import Sequence REPO_ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(REPO_ROOT / "examples")) @@ -123,7 +127,7 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m def mm_data(): return dict(image_bytes=[], image_sizes=[]) if engine.multi_modal_processor is not None else None - def enqueue(rid, seq, num_cached_tokens): + def enqueue(rid: str, seq: Sequence, num_cached_tokens: int) -> None: # Same as LLMEngine.add_request's _enqueue: the processor must get the sequence's own multimodal dict # (overwrite_prompt wrote the text/image segments into it), plus the cached prefix length. engine.scheduler.add_sequence(seq) @@ -228,7 +232,7 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m class IncrementalDetokenizer: """vLLM-style prefix/read offsets so each delta decodes only a short tail.""" - def __init__(self, tokenizer) -> None: + def __init__(self, tokenizer: Any) -> None: self.tok = tokenizer self.ids: list[int] = [] self.num_tokens = 0 @@ -300,12 +304,12 @@ class Frontend: else: self.loop.call_soon_threadsafe(self._deliver_batch, payload) - def _deliver(self, rid: str, item) -> None: + def _deliver(self, rid: str, item: tuple) -> None: q = self.streams.get(rid) if q is not None: q.put_nowait(item) - def _deliver_batch(self, batch) -> None: + def _deliver_batch(self, batch: list[tuple]) -> None: for rid, token_id, finished, reason, prompt_len in batch: q = self.streams.get(rid) if q is not None: @@ -375,7 +379,7 @@ class Frontend: self.stats["finished"] += 1 -def build_app(fe: Frontend): +def build_app(fe: Frontend) -> FastAPI: from fastapi import FastAPI from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse @@ -469,7 +473,7 @@ def build_app(fe: Frontend): done = False try: yield chunk({"role": "assistant", "content": ""}) - async for text, finish in tokens(): + async for text, finish in tokens(): # noqa: B007 (the last finish reason is reported below) if text or finish: yield chunk({"content": text} if text else {}, finish) done = True @@ -484,7 +488,7 @@ def build_app(fe: Frontend): done = False try: parts, finish = [], None - async for text, finish in tokens(): + async for text, finish in tokens(): # noqa: B007 (the last finish reason is reported below) parts.append(text) done = True finally: