Create app.py

app.py

@@ -0,0 +1,1049 @@
+import asyncio
+import json
+import os
+import time
+from http import HTTPStatus
+from typing import AsyncGenerator, Dict, List, Optional, Tuple, Union
+
+import fastapi
+import uvicorn
+from fastapi import Request, Depends, HTTPException, BackgroundTasks
+from fastapi.exceptions import RequestValidationError
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, StreamingResponse, Response
+from packaging import version
+from pydantic import BaseModel, Field, ValidationError, validator, conint, root_validator
+
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.engine.async_llm_engine import AsyncLLMEngine
+
+from vllm.entrypoints.openai.protocol import (
+    CompletionResponse, CompletionResponseChoice,
+    CompletionResponseStreamChoice, CompletionStreamResponse,
+    ChatCompletionResponse,
+    ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
+    ChatCompletionStreamResponse, ChatMessage, DeltaMessage, ErrorResponse,
+    ModelCard, ModelList, ModelPermission, UsageInfo)
+
+
+from vllm.logger import init_logger
+from vllm.outputs import RequestOutput
+from vllm.sampling_params import SamplingParams
+from vllm.transformers_utils.tokenizer import get_tokenizer
+from vllm.utils import random_uuid
+from vllm import LLM
+
+from huggingface_hub import snapshot_download
+
+
+try:
+    import fastchat
+    from fastchat.conversation import Conversation, SeparatorStyle
+    from fastchat.model.model_adapter import get_conversation_template
+    _fastchat_available = True
+except ImportError:
+    _fastchat_available = False
+
+TIMEOUT_KEEP_ALIVE = 5
+DEFAULT_API_KEY = "your_default_api_key"
+API_KEY = os.environ.get("API_KEY", DEFAULT_API_KEY)
+MODEL_NAME = os.environ.get("SERVED_MODEL", "jnjj/gemma-3-4b-it-qat-int4-quantized-inference-unrestricted-pruned-sf")
+HOST = os.environ.get("HOST", "0.0.0.0")
+PORT = int(os.environ.get("PORT", "7860"))
+MAX_MODEL_LEN_CONFIG = int(os.environ.get("MAX_MODEL_LEN", "8000"))
+GPU_MEMORY_UTILIZATION = float(os.environ.get("GPU_MEMORY_UTILIZATION", "0.0"))
+REQUESTS_PER_MINUTE = int(os.environ.get("REQUESTS_PER_MINUTE", "120"))
+LOG_LEVEL = os.environ.get("LOG_LEVEL", "INFO").upper()
+DOWNLOADED_MODEL_PATH = None
+ENABLE_REQUEST_LOGGING = os.environ.get("ENABLE_REQUEST_LOGGING", "false").lower() == "true"
+MAX_CONCURRENT_DOWNLOADS = int(os.environ.get("MAX_CONCURRENT_DOWNLOADS", "2"))
+QUEUE_SIZE = int(os.environ.get("QUEUE_SIZE", "100"))
+
+logger = init_logger(__name__)
+served_model = MODEL_NAME
+app = fastapi.FastAPI(title="vLLM OpenAI API", description="Concurrent OpenAI Compatible API - vLLM Powered - Advanced, Robust & Optimized", version="1.2.0")
+engine = None
+tokenizer = None
+max_model_len = MAX_MODEL_LEN_CONFIG
+download_semaphore = asyncio.Semaphore(MAX_CONCURRENT_DOWNLOADS)
+request_queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_SIZE)
+
+request_timestamps = []
+
+
+async def rate_limit_dependency(request: Request):
+    current_time = time.monotonic()
+    request_timestamps.append(current_time)
+    request_timestamps[:] = [ts for ts in request_timestamps if current_time - ts <= 60]
+    if len(request_timestamps) > REQUESTS_PER_MINUTE:
+        raise HTTPException(status_code=429, detail="Too Many Requests. Please try again later.")
+    return True
+
+async def queue_dependency():
+    if request_queue.full():
+        raise HTTPException(status_code=429, detail="Queue is full. Please try again later.")
+    await request_queue.put(1)
+    try:
+        yield
+    finally:
+        await request_queue.get(1)
+        request_queue.task_done()
+
+
+class HTTPException(fastapi.HTTPException):
+    pass
+
+
+class ChatCompletionRequest(BaseModel):
+    model: str = Field(default=MODEL_NAME, description="Model name for chat completion")
+    api_key: str = Field(..., description="API Key for authentication")
+    messages: Union[str, List[Dict[str, str]]] = Field(..., description="Conversation messages")
+    temperature: Optional[float] = Field(0.7, description="Sampling temperature")
+    top_p: Optional[float] = Field(1.0, description="Top p sampling parameter")
+    n: Optional[conint(ge=1, le=10)] = Field(1, description="Number of chat completion choices (max 10)")
+    max_tokens: Optional[conint(ge=1, le=max_model_len)] = Field(None, description=f"Max tokens, up to {max_model_len}")
+    stop: Optional[Union[str, List[str]]] = Field(default_factory=list, description="Stop sequences")
+    stream: Optional[bool] = Field(False, description="Enable streaming responses")
+    presence_penalty: Optional[float] = Field(0.0, description="Presence penalty")
+    frequency_penalty: Optional[float] = Field(0.0, description="Frequency penalty")
+    logit_bias: Optional[Dict[str, float]] = Field(None, description="Logit bias map")
+    user: Optional[str] = Field(None, description="User identifier")
+    best_of: Optional[conint(ge=1, le=10)] = Field(None, description="Best of sampling (max 10)")
+    top_k: Optional[conint(ge=-1)] = Field(-1, description="Top k sampling")
+    ignore_eos: Optional[bool] = Field(False, description="Ignore EOS token")
+    use_beam_search: Optional[bool] = Field(False, description="Use beam search (not for chat)")
+    stop_token_ids: Optional[List[int]] = Field(default_factory=list, description="Stop token IDs")
+    skip_special_tokens: Optional[bool] = Field(True, description="Skip special tokens")
+    spaces_between_special_tokens: Optional[bool] = Field(True, description="Spaces between special tokens")
+
+    @validator("messages")
+    def messages_must_be_list_or_str(cls, v):
+        if not isinstance(v, (str, list)):
+            raise ValueError("Messages must be a string or a list of messages")
+        return v
+
+
+class CompletionRequest(BaseModel):
+    model: str = Field(default=MODEL_NAME, description="Model name for text completion")
+    api_key: str = Field(..., description="API Key for authentication")
+    prompt: Union[List[int], List[List[int]], str, List[str]] = Field(..., description="Text prompt for completion")
+    suffix: Optional[str] = Field(None, description="Suffix (not supported)")
+    max_tokens: Optional[conint(ge=1, le=max_model_len)] = Field(16, description=f"Max completion tokens, up to {max_model_len}")
+    temperature: Optional[float] = Field(1.0, description="Sampling temperature")
+    top_p: Optional[float] = Field(1.0, description="Top p sampling")
+    n: Optional[conint(ge=1, le=10)] = Field(1, description="Number of completions (max 10)")
+    stream: Optional[bool] = Field(False, description="Enable streaming responses")
+    echo: Optional[bool] = Field(False, description="Echo prompt (not supported)")
+    stop: Optional[Union[str, List[str]]] = Field(default_factory=list, description="Stop sequences")
+    presence_penalty: Optional[float] = Field(0.0, description="Presence penalty")
+    frequency_penalty: Optional[float] = Field(0.0, description="Frequency penalty")
+    logit_bias: Optional[Dict[str, float]] = Field(None, description="Logit bias map")
+    user: Optional[str] = Field(None, description="User identifier")
+    best_of: Optional[conint(ge=1, le=10)] = Field(None, description="Best of sampling (max 10)")
+    top_k: Optional[conint(ge=-1)] = Field(-1, description="Top k sampling")
+    ignore_eos: Optional[bool] = Field(False, description="Ignore EOS token")
+    use_beam_search: Optional[bool] = Field(False, description="Use beam search (not for completion)")
+    stop_token_ids: Optional[List[int]] = Field(default_factory=list, description="Stop token IDs")
+    skip_special_tokens: Optional[bool] = Field(True, description="Skip special tokens")
+    spaces_between_special_tokens: Optional[bool] = Field(True, description="Spaces between special tokens")
+
+    @validator("prompt")
+    def prompt_must_be_list_or_str(cls, v):
+        if not isinstance(v, (str, list)):
+            raise ValueError("Prompt must be a string or a list of prompts")
+        return v
+
+
+class ModelDownloadResponse(BaseModel):
+    model_name: str = Field(..., description="Name of model downloaded")
+    download_path: Optional[str] = Field(None, description="Local download path")
+    status: str = Field(..., description="Download status")
+    message: Optional[str] = Field(None, description="Download message")
+
+
+def create_error_response(status_code: HTTPStatus, message: str, err_type="invalid_request_error") -> JSONResponse:
+    logger.error(f"Error Response: {status_code.value} - {message} ({err_type})")
+    return JSONResponse(ErrorResponse(message=message, type=err_type).dict(), status_code=status_code.value)
+
+
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request: Request, exc: RequestValidationError):
+    logger.warning(f"Validation Error: {exc}")
+    return create_error_response(HTTPStatus.BAD_REQUEST, str(exc), err_type="validation_error")
+
+
+@app.exception_handler(HTTPException)
+async def http_exception_handler(request: Request, exc: HTTPException):
+    logger.warning(f"HTTP Exception: {exc.detail} Status Code: {exc.status_code}")
+    return create_error_response(exc.status_code, exc.detail, err_type="rate_limit_error" if exc.status_code == 429 else "http_error")
+
+
+async def check_api_key(api_key: str = Depends(lambda request: request.headers.get("Authorization") or request.query_params.get("api_key"))):
+    if api_key is None or api_key.replace("Bearer ", "") != API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid API key.")
+    return True
+
+
+async def check_model(request_model_name: str) -> Optional[JSONResponse]:
+    model_to_check = DOWNLOADED_MODEL_PATH if DOWNLOADED_MODEL_PATH else served_model
+    if request_model_name == model_to_check:
+        return None
+    return create_error_response(
+        HTTPStatus.NOT_FOUND,
+        f"Model '{request_model_name}' not found. Serving: {model_to_check}",
+        err_type="model_not_found"
+    )
+
+
+async def get_gen_prompt(request: ChatCompletionRequest) -> str:
+    if not _fastchat_available:
+        raise ModuleNotFoundError("fastchat not installed. Install to use chat API: pip install fschat")
+    if version.parse(fastchat.__version__) < version.parse("0.2.23"):
+        raise ImportError(f"fastchat version too low: {fastchat.__version__}. Upgrade: pip install -U fschat")
+
+    try:
+        try:
+             conv = get_conversation_template(request.model)
+        except Exception:
+             logger.warning(f"Conversation template for model '{request.model}' not found. Using default template.")
+             if isinstance(request.messages, str):
+                 return request.messages
+             else:
+                 raise ValueError(f"Conversation template for model '{request.model}' not found and messages is not a string.")
+
+        conv_dict = request.dict()
+        conversation_keys = {f.name for f in Conversation.__fields__.values()}
+        filtered_conv_dict = {k: v for k, v in conv_dict.items() if k in conversation_keys}
+        conv = Conversation(**filtered_conv_dict)
+
+
+        if isinstance(request.messages, str):
+            prompt = request.messages
+        else:
+            for message in request.messages:
+                role = message["role"]
+                if role == "system":
+                    conv.system_message = message["content"]
+                elif role == "user":
+                    conv.append_message(conv.roles[0], message["content"])
+                elif role == "assistant":
+                    conv.append_message(conv.roles[1], message["content"])
+                else:
+                    raise ValueError(f"Unknown role: {role}")
+
+            conv.append_message(conv.roles[1], None)
+            prompt = conv.get_prompt()
+        return prompt
+    except ValueError as e:
+        logger.error(f"Prompt generation error: {e}")
+        raise ValueError(f"Failed to generate prompt: {e}")
+    except Exception as e:
+        logger.error(f"An unexpected error occurred during prompt generation: {e}")
+        raise RuntimeError(f"An unexpected error occurred during prompt generation: {e}")
+
+
+async def check_length(request: Union[ChatCompletionRequest, CompletionRequest], prompt: Optional[str] = None, prompt_ids: Optional[List[int]] = None) -> Tuple[List[int], Optional[JSONResponse]]:
+    assert (not (prompt is None and prompt_ids is None) and not (prompt is not None and prompt_ids is not None)), "Provide either prompt or prompt_ids."
+
+    if tokenizer is None:
+         return [], create_error_response(HTTPStatus.INTERNAL_SERVER_ERROR, "Tokenizer not initialized.", err_type="internal_error")
+
+    try:
+        input_ids = prompt_ids if prompt_ids else tokenizer(prompt).input_ids
+    except Exception as e:
+         logger.error(f"Error during tokenization: {e}")
+         return [], create_error_response(HTTPStatus.INTERNAL_SERVER_ERROR, f"Error tokenizing prompt: {e}", err_type="tokenization_error")
+
+    token_num = len(input_ids)
+
+    if request.max_tokens is None:
+        remaining_tokens = max_model_len - token_num
+        if remaining_tokens <= 0:
+             return input_ids, create_error_response(
+                HTTPStatus.BAD_REQUEST,
+                f"Prompt length ({token_num}) exceeds or equals max model length ({max_model_len}). No space for completion.",
+                err_type="context_length_exceeded"
+            )
+        request.max_tokens = remaining_tokens
+
+    if token_num + request.max_tokens > max_model_len:
+        return input_ids, create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            f"Context length exceeded. Max: {max_model_len}, Prompt Tokens: {token_num}, Requested Completion Tokens: {request.max_tokens}, Total: {request.max_tokens + token_num}",
+            err_type="context_length_exceeded"
+        )
+    return input_ids, None
+
+
+@app.get("/health", tags=["System"])
+async def health() -> Response:
+    if engine is None:
+         return Response(status_code=503, content="Engine not initialized")
+    try:
+        await engine.get_model_config()
+        return Response(status_code=200)
+    except Exception as e:
+        logger.error(f"Health check failed: {e}")
+        return Response(status_code=503, content=f"Engine health check failed: {e}")
+
+
+@app.get("/metrics", tags=["System"])
+async def metrics() -> Response:
+    return Response(content="", media_type="text/plain")
+
+
+@app.get("/models", response_model=ModelList, tags=["System"])
+async def show_available_models():
+    model_cards = [
+        ModelCard(id= served_model if DOWNLOADED_MODEL_PATH is None else DOWNLOADED_MODEL_PATH,
+                  root= served_model if DOWNLOADED_MODEL_PATH is None else DOWNLOADED_MODEL_PATH,
+                  permission=[ModelPermission()])
+    ]
+    return ModelList(data=model_cards)
+
+
+@app.get("/model_config", tags=["System"])
+async def get_model_configuration():
+    model_config = {
+        "model_name": served_model if DOWNLOADED_MODEL_PATH is None else DOWNLOADED_MODEL_PATH,
+        "max_model_len_config": MAX_MODEL_LEN_CONFIG,
+        "cpu_only": True,
+        "gpu_memory_utilization": GPU_MEMORY_UTILIZATION,
+    }
+    if engine:
+        try:
+            engine_model_config = await engine.get_model_config()
+            model_config["actual_max_model_len"] = engine_model_config.max_model_len
+            model_config["dtype"] = engine_model_config.dtype
+            model_config["num_layers"] = engine_model_config.num_layers
+            model_config["num_attention_heads"] = engine_model_config.num_attention_heads
+            model_config["hidden_size"] = engine_model_config.hidden_size
+            model_config["vocab_size"] = engine_model_config.vocab_size
+
+        except Exception as e:
+            logger.warning(f"Could not retrieve detailed engine config: {e}")
+            model_config["engine_config_status"] = f"Error retrieving engine config: {e}"
+
+    return model_config
+
+@app.post("/models/download", response_model=ModelDownloadResponse, tags=["Model Management"])
+async def download_model(model_name: str = fastapi.Query(..., description="Model name to download"), background_tasks: BackgroundTasks = BackgroundTasks()):
+    logger.info(f"Download requested for model: {model_name}")
+    if download_semaphore.locked():
+        raise HTTPException(status_code=429, detail="Model download already in progress.")
+
+    global DOWNLOADED_MODEL_PATH
+    previous_downloaded_path = DOWNLOADED_MODEL_PATH
+    DOWNLOADED_MODEL_PATH = None
+
+    background_tasks.add_task(run_model_download, model_name, previous_downloaded_path)
+
+    return ModelDownloadResponse(model_name=model_name, status="pending", message="Model download started. Check logs for progress.")
+
+
+async def run_model_download(model_name: str, previous_downloaded_path: Optional[str]):
+    async with download_semaphore:
+        logger.info(f"Starting background download for model: {model_name}")
+        loop = asyncio.get_running_loop()
+        global DOWNLOADED_MODEL_PATH, engine, tokenizer, max_model_len
+        try:
+            download_path = await loop.run_in_executor(None, snapshot_download, model_name)
+            logger.info(f"Model downloaded to: {download_path}")
+
+            if engine:
+                 logger.info("Shutting down existing engine...")
+                 engine = None
+                 tokenizer = None
+                 max_model_len = MAX_MODEL_LEN_CONFIG
+                 await asyncio.sleep(2)
+                 logger.info("Existing engine dereferenced.")
+
+            await initialize_llm_engine(download_path)
+            DOWNLOADED_MODEL_PATH = download_path
+            logger.info(f"Model '{model_name}' ready from downloaded path: {DOWNLOADED_MODEL_PATH}")
+
+        except Exception as e:
+            logger.error(f"Model download & init error for {model_name}: {e}")
+            DOWNLOADED_MODEL_PATH = previous_downloaded_path
+
+
+async def completion_stream_generator_chat(request: ChatCompletionRequest, result_generator: AsyncGenerator[RequestOutput, None]) -> AsyncGenerator[str, None]:
+    model_name = request.model
+    request_id = f"cmpl-{random_uuid()}"
+    created_time = int(time.time())
+    prompt_token_count = 0
+
+    for i in range(request.n):
+         choice_data = ChatCompletionResponseStreamChoice(index=i, delta=DeltaMessage(role="assistant"))
+         chunk = ChatCompletionStreamResponse(id=request_id, choices=[choice_data], model=model_name, created=created_time)
+         yield f"data: {chunk.json(exclude_unset=True, ensure_ascii=False)}\n\n"
+
+
+    previous_texts = [""] * request.n
+    output_completion_tokens = [0] * request.n
+
+
+    async for res in result_generator:
+        prompt_token_count = len(res.prompt_token_ids)
+        final_res = res
+
+        for output in res.outputs:
+            i = output.index
+            delta_text = output.text[len(previous_texts[i]):]
+            previous_texts[i] = output.text
+
+            output_completion_tokens[i] = len(output.token_ids)
+
+            if delta_text:
+                choice_data = ChatCompletionResponseStreamChoice(index=i, delta=DeltaMessage(content=delta_text))
+                chunk = ChatCompletionStreamResponse(id=request_id, choices=[choice_data], model=model_name, created=created_time)
+                yield f"data: {chunk.json(exclude_unset=True, ensure_ascii=False)}\n\n"
+
+            if output.finish_reason:
+                 choice_data_finish = ChatCompletionResponseStreamChoice(index=i, delta=DeltaMessage(), finish_reason=output.finish_reason)
+                 chunk_finish = ChatCompletionStreamResponse(id=request_id, choices=[choice_data_finish], model=model_name, created=created_time)
+                 yield f"data: {chunk_finish.json(exclude_unset=True, ensure_ascii=False)}\n\n"
+
+    total_completion_tokens = sum(output_completion_tokens)
+
+
+    yield "data: [DONE]\n\n"
+
+
+async def completion_stream_generator_completion(request: CompletionRequest, result_generator: AsyncGenerator[RequestOutput, None]) -> AsyncGenerator[str, None]:
+    model_name = request.model
+    request_id = f"cmpl-{random_uuid()}"
+    created_time = int(time.time())
+    prompt_token_count = 0
+    completion_token_count = 0
+
+    previous_texts = [""] * request.n
+
+    async for res in result_generator:
+        prompt_token_count = len(res.prompt_token_ids)
+
+        for output in res.outputs:
+            i = output.index
+            delta_text = output.text[len(previous_texts[i]):]
+
+            current_output_tokens = len(output.token_ids)
+            tokens_generated_in_chunk = current_output_tokens - len(tokenizer(previous_texts[i]).input_ids)
+
+            logprobs_obj = None
+
+            previous_texts[i] = output.text
+
+            completion_token_count += tokens_generated_in_chunk
+
+            choice_data = CompletionResponseStreamChoice(
+                index=i,
+                text=delta_text,
+                logprobs=None
+            )
+            chunk = CompletionStreamResponse(id=request_id, choices=[choice_data], model=model_name, created=created_time)
+            yield f"data: {chunk.json(exclude_unset=True, ensure_ascii=False)}\n\n"
+
+            if output.finish_reason:
+                 choice_data_finish = CompletionResponseStreamChoice(
+                     index=i,
+                     text="",
+                     logprobs=None,
+                     finish_reason=output.finish_reason
+                 )
+                 chunk_finish = CompletionStreamResponse(id=request_id, choices=[choice_data_finish], model=model_name, created=created_time)
+                 yield f"data: {chunk_finish.json(exclude_unset=True, ensure_ascii=False)}\n\n"
+
+
+    yield "data: [DONE]\n\n"
+
+
+@app.post("/completions", response_model=CompletionResponse, tags=["Completions"], dependencies=[Depends(rate_limit_dependency), Depends(check_api_key), Depends(queue_dependency)])
+async def create_completion(request: CompletionRequest, raw_request: Request):
+    start_time = time.monotonic()
+    if ENABLE_REQUEST_LOGGING:
+        logger.info(f"Completion Request: {request}")
+
+    model_error_check = await check_model(request.model)
+    if model_error_check:
+        if request.stream:
+             error_json_str = json.dumps(json.loads(model_error_check.body))
+             async def error_stream():
+                 yield f"data: {error_json_str}\n\n"
+                 yield "data: [DONE]\n\n"
+             return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=model_error_check.status_code)
+        else:
+            raise HTTPException(status_code=model_error_check.status_code, detail=json.loads(model_error_check.body)['message'])
+
+
+    if request.echo:
+        error_message = "Echo not supported."
+        error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="not_supported")
+        if request.stream:
+            error_json_str = json.dumps(json.loads(error_res.body))
+            async def error_stream():
+                 yield f"data: {error_json_str}\n\n"
+                 yield "data: [DONE]\n\n"
+            return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+            raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+    if request.suffix:
+        error_message = "Suffix not supported."
+        error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="not_supported")
+        if request.stream:
+            error_json_str = json.dumps(json.loads(error_res.body))
+            async def error_stream():
+                 yield f"data: {error_json_str}\n\n"
+                 yield "data: [DONE]\n\n"
+            return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+            raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+    if request.logit_bias and len(request.logit_bias) > 0:
+        error_message = "Logit bias not supported."
+        error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="not_supported")
+        if request.stream:
+            error_json_str = json.dumps(json.loads(error_res.body))
+            async def error_stream():
+                 yield f"data: {error_json_str}\n\n"
+                 yield "data: [DONE]\n\n"
+            return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+            raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+
+    model_name = request.model
+    request_id = f"cmpl-{random_uuid()}"
+
+    use_token_ids = False
+    prompt = request.prompt
+    prompt_processed = None
+    prompt_token_ids_input = None
+
+    if isinstance(prompt, list):
+        if not prompt:
+            error_message = "Provide at least one prompt."
+            error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="invalid_prompt")
+            if request.stream:
+               error_json_str = json.dumps(json.loads(error_res.body))
+               async def error_stream():
+                   yield f"data: {error_json_str}\n\n"
+                   yield "data: [DONE]\n\n"
+               return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+            else:
+               raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+        first_element = prompt[0]
+        if isinstance(first_element, int):
+            use_token_ids = True
+            prompt_token_ids_input = prompt
+        elif isinstance(first_element, list):
+             if len(prompt) > 1:
+                  error_message = "Batch requests are not fully supported for 'prompt' field as List[str] or List[List[int]] > 1."
+                  error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="not_supported")
+                  if request.stream:
+                      error_json_str = json.dumps(json.loads(error_res.body))
+                      async def error_stream():
+                          yield f"data: {error_json_str}\n\n"
+                          yield "data: [DONE]\n\n"
+                      return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+                  else:
+                      raise HTTPException(status_code=error_res.status_code, detail=error_message)
+             if isinstance(first_element, int):
+                 use_token_ids = True
+                 prompt_token_ids_input = prompt
+             elif isinstance(first_element, str):
+                 prompt_processed = prompt[0]
+             elif isinstance(first_element, list) and isinstance(first_element[0], int):
+                 use_token_ids = True
+                 prompt_token_ids_input = first_element
+             else:
+                 error_message = "Invalid format for 'prompt' list."
+                 error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="invalid_prompt")
+                 if request.stream:
+                    error_json_str = json.dumps(json.loads(error_res.body))
+                    async def error_stream():
+                        yield f"data: {error_json_str}\n\n"
+                        yield "data: [DONE]\n\n"
+                    return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+                 else:
+                    raise HTTPException(status_code=error_res.status_code, detail=error_message)
+        elif isinstance(first_element, str):
+             prompt_processed = prompt[0]
+        else:
+            error_message = "Invalid format for 'prompt' list."
+            error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="invalid_prompt")
+            if request.stream:
+               error_json_str = json.dumps(json.loads(error_res.body))
+               async def error_stream():
+                   yield f"data: {error_json_str}\n\n"
+                   yield "data: [DONE]\n\n"
+               return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+            else:
+               raise HTTPException(status_code=error_res.status_code, detail=error_message)
+    elif isinstance(prompt, str):
+         prompt_processed = prompt
+    else:
+         error_message = "Prompt must be a string or a list."
+         error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="invalid_prompt")
+         if request.stream:
+            error_json_str = json.dumps(json.loads(error_res.body))
+            async def error_stream():
+                 yield f"data: {error_json_str}\n\n"
+                 yield "data: [DONE]\n\n"
+            return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+         else:
+            raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+
+    try:
+        if use_token_ids:
+            input_ids, length_error = await check_length(request, prompt_ids=prompt_token_ids_input)
+
+        else:
+            input_ids, length_error = await check_length(request, prompt=prompt_processed)
+
+        if length_error:
+             if request.stream:
+                error_json_str = json.dumps(json.loads(length_error.body))
+                async def error_stream():
+                    yield f"data: {error_json_str}\n\n"
+                    yield "data: [DONE]\n\n"
+                return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=length_error.status_code)
+             else:
+                 raise HTTPException(status_code=length_error.status_code, detail=json.loads(length_error.body)['message'])
+
+    except ValueError as ve:
+        error_message = str(ve)
+        error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="prompt_error")
+        if request.stream:
+           error_json_str = json.dumps(json.loads(error_res.body))
+           async def error_stream():
+               yield f"data: {error_json_str}\n\n"
+               yield "data: [DONE]\n\n"
+           return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+           raise HTTPException(status_code=error_res.status_code, detail=error_message)
+    except Exception as e:
+        error_message = f"Error processing prompt length: {e}"
+        error_res = create_error_response(HTTPStatus.INTERNAL_SERVER_ERROR, error_message, err_type="internal_error")
+        if request.stream:
+           error_json_str = json.dumps(json.loads(error_res.body))
+           async def error_stream():
+               yield f"data: {error_json_str}\n\n"
+               yield "data: [DONE]\n\n"
+           return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+           raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+
+    created_time = int(time.time())
+
+    sampling_params = SamplingParams(**request.dict(
+        exclude={
+            "stream",
+            "api_key",
+            "model",
+            "prompt",
+            "user",
+            "echo",
+            "suffix",
+            "logit_bias",
+        },
+        exclude_none=True
+    ))
+
+    try:
+        if use_token_ids:
+             result_generator = engine.generate(
+                 prompt=None,
+                 sampling_params=sampling_params,
+                 request_id=request_id,
+                 prompt_token_ids=input_ids
+             )
+        else:
+             result_generator = engine.generate(
+                 prompt=prompt_processed,
+                 sampling_params=sampling_params,
+                 request_id=request_id,
+                 prompt_token_ids=input_ids
+             )
+
+    except Exception as e:
+        logger.error(f"Error submitting generation request to engine: {e}")
+        error_message = f"Error submitting generation request: {e}"
+        error_res = create_error_response(HTTPStatus.INTERNAL_SERVER_ERROR, error_message, err_type="engine_error")
+        if request.stream:
+           error_json_str = json.dumps(json.loads(error_res.body))
+           async def error_stream():
+               yield f"data: {error_json_str}\n\n"
+               yield "data: [DONE]\n\n"
+           return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+           raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+
+    try:
+        if request.stream:
+            response = StreamingResponse(completion_stream_generator_completion(request, result_generator), media_type="text/event-stream")
+            return response
+        else:
+            final_res = None
+            async for res in result_generator:
+                final_res = res
+
+            if final_res is None or not final_res.outputs:
+                 error_message = "Engine returned no output."
+                 error_res = create_error_response(HTTPStatus.INTERNAL_SERVER_ERROR, error_message, err_type="engine_output_error")
+                 raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+            choices = [
+                CompletionResponseChoice(
+                    index=output.index,
+                    text=output.text,
+                    logprobs=None,
+                    finish_reason=output.finish_reason
+                ) for output in final_res.outputs
+            ]
+
+            prompt_tokens = len(final_res.prompt_token_ids)
+            completion_tokens = sum(len(output.token_ids) for output in final_res.outputs)
+            total_tokens = prompt_tokens + completion_tokens
+            usage = UsageInfo(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens)
+
+
+            response = CompletionResponse(id=request_id, created=created_time, model=model_name, choices=choices, usage=usage)
+
+            if ENABLE_REQUEST_LOGGING:
+                 logger.info(f"Completion Response (non-stream): {response}")
+
+            return response
+
+    except Exception as e:
+        logger.error(f"Error processing generation result for request {request_id}: {e}")
+        error_message = f"Error processing generation result: {e}"
+        error_res = create_error_response(HTTPStatus.INTERNAL_SERVER_ERROR, error_message, err_type="engine_error")
+        if request.stream:
+           error_json_str = json.dumps(json.loads(error_res.body))
+           async def error_stream():
+               yield f"data: {error_json_str}\n\n"
+               yield "data: [DONE]\n\n"
+           return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+           raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+
+@app.post("/chat/completions", response_model=ChatCompletionResponse, tags=["Chat Completions"], dependencies=[Depends(rate_limit_dependency), Depends(check_api_key), Depends(queue_dependency)])
+async def create_chat_completion(request: ChatCompletionRequest, raw_request: Request):
+    start_time = time.monotonic()
+    if ENABLE_REQUEST_LOGGING:
+        log_request_dict = request.dict()
+        messages = log_request_dict.pop("messages", "N/A")
+        logger.info(f"Chat Completion Request: {log_request_dict}, Messages: {messages}")
+
+    model_error_check = await check_model(request.model)
+    if model_error_check:
+         if request.stream:
+             error_json_str = json.dumps(json.loads(model_error_check.body))
+             async def error_stream():
+                 yield f"data: {error_json_str}\n\n"
+                 yield "data: [DONE]\n\n"
+             return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=model_error_check.status_code)
+         else:
+            raise HTTPException(status_code=model_error_check.status_code, detail=json.loads(model_error_check.body)['message'])
+
+    if request.use_beam_search:
+         error_message = "Beam search not supported for chat completions."
+         error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="not_supported")
+         if request.stream:
+            error_json_str = json.dumps(json.loads(error_res.body))
+            async def error_stream():
+                 yield f"data: {error_json_str}\n\n"
+                 yield "data: [DONE]\n\n"
+            return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+         else:
+            raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+    if request.best_of is not None and request.best_of > 1:
+        error_message = "Best of > 1 not fully supported for chat completions."
+        error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="not_supported")
+        if request.stream:
+            error_json_str = json.dumps(json.loads(error_res.body))
+            async def error_stream():
+                 yield f"data: {error_json_str}\n\n"
+                 yield "data: [DONE]\n\n"
+            return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+            raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+    if request.logit_bias and len(request.logit_bias) > 0:
+        error_message = "Logit bias not supported."
+        error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="not_supported")
+        if request.stream:
+            error_json_str = json.dumps(json.loads(error_res.body))
+            async def error_stream():
+                 yield f"data: {error_json_str}\n\n"
+                 yield "data: [DONE]\n\n"
+            return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+            raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+
+    try:
+        prompt = await get_gen_prompt(request)
+    except ValueError as ve:
+        error_message = str(ve)
+        error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="prompt_generation_error")
+        if request.stream:
+           error_json_str = json.dumps(json.loads(error_res.body))
+           async def error_stream():
+               yield f"data: {error_json_str}\n\n"
+               yield "data: [DONE]\n\n"
+           return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+           raise HTTPException(status_code=error_res.status_code, detail=error_message)
+    except RuntimeError as re:
+        error_message = str(re)
+        error_res = create_error_response(HTTPStatus.INTERNAL_SERVER_ERROR, error_message, err_type="internal_error")
+        if request.stream:
+           error_json_str = json.dumps(json.loads(error_res.body))
+           async def error_stream():
+               yield f"data: {error_json_str}\n\n"
+               yield "data: [DONE]\n\n"
+           return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+           raise HTTPException(status_code=error_res.status_code, detail=error_message)
+    except Exception as e:
+        logger.error(f"An unexpected error occurred during chat prompt generation: {e}")
+        error_message = f"An unexpected error occurred during prompt generation: {e}"
+        error_res = create_error_response(HTTPStatus.INTERNAL_SERVER_ERROR, error_message, err_type="internal_error")
+        if request.stream:
+           error_json_str = json.dumps(json.loads(error_res.body))
+           async def error_stream():
+               yield f"data: {error_json_str}\n\n"
+               yield "data: [DONE]\n\n"
+           return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+           raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+
+    try:
+        input_ids, length_error = await check_length(request, prompt=prompt)
+        if length_error:
+            if request.stream:
+               error_json_str = json.dumps(json.loads(length_error.body))
+               async def error_stream():
+                   yield f"data: {error_json_str}\n\n"
+                   yield "data: [DONE]\n\n"
+               return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=length_error.status_code)
+            else:
+                raise HTTPException(status_code=length_error.status_code, detail=json.loads(length_error.body)['message'])
+    except ValueError as ve:
+        error_message = str(ve)
+        error_res = create_error_response(HTTPStatus.BAD_REQUEST, error_message, err_type="prompt_error")
+        if request.stream:
+           error_json_str = json.dumps(json.loads(error_res.body))
+           async def error_stream():
+               yield f"data: {error_json_str}\n\n"
+               yield "data: [DONE]\n\n"
+           return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+           raise HTTPException(status_code=error_res.status_code, detail=error_message)
+    except Exception as e:
+        error_message = f"Error processing prompt length: {e}"
+        error_res = create_error_response(HTTPStatus.INTERNAL_SERVER_ERROR, error_message, err_type="internal_error")
+        if request.stream:
+           error_json_str = json.dumps(json.loads(error_res.body))
+           async def error_stream():
+               yield f"data: {error_json_str}\n\n"
+               yield "data: [DONE]\n\n"
+           return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+           raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+
+    created_time = int(time.time())
+    request_id = f"chatcmpl-{random_uuid()}"
+
+    sampling_params = SamplingParams(**request.dict(
+        exclude={
+            "stream",
+            "api_key",
+            "model",
+            "messages",
+            "user",
+            "use_beam_search",
+            "logit_bias",
+        },
+         exclude_none=True
+    ))
+
+    try:
+         result_generator = engine.generate(
+             prompt=prompt,
+             sampling_params=sampling_params,
+             request_id=request_id,
+             prompt_token_ids=input_ids
+         )
+    except Exception as e:
+        logger.error(f"Error submitting chat generation request to engine: {e}")
+        error_message = f"Error submitting chat generation request: {e}"
+        error_res = create_error_response(HTTPStatus.INTERNAL_SERVER_ERROR, error_message, err_type="engine_error")
+        if request.stream:
+           error_json_str = json.dumps(json.loads(error_res.body))
+           async def error_stream():
+               yield f"data: {error_json_str}\n\n"
+               yield "data: [DONE]\n\n"
+           return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+           raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+
+    try:
+        if request.stream:
+            response = StreamingResponse(completion_stream_generator_chat(request, result_generator), media_type="text/event-stream")
+            return response
+        else:
+            final_res = None
+            async for res in result_generator:
+                final_res = res
+
+            if final_res is None or not final_res.outputs:
+                 error_message = "Engine returned no output."
+                 error_res = create_error_response(HTTPStatus.INTERNAL_SERVER_ERROR, error_message, err_type="engine_output_error")
+                 raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+            choices = [
+                ChatCompletionResponseChoice(
+                    index=output.index,
+                    message=ChatMessage(role="assistant", content=output.text),
+                    logprobs=None,
+                    finish_reason=output.finish_reason,
+                ) for output in final_res.outputs
+            ]
+
+            prompt_tokens = len(final_res.prompt_token_ids)
+            completion_tokens = sum(len(output.token_ids) for output in final_res.outputs)
+            total_tokens = prompt_tokens + completion_tokens
+            usage = UsageInfo(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens)
+
+
+            response = ChatCompletionResponse(id=request_id, created=created_time, model=model_name, choices=choices, usage=usage)
+
+            if ENABLE_REQUEST_LOGGING:
+                 logger.info(f"Chat Completion Response (non-stream): {response}")
+
+            return response
+
+    except Exception as e:
+        logger.error(f"Error processing chat generation result for request {request_id}: {e}")
+        error_message = f"Error processing generation result: {e}"
+        error_res = create_error_response(HTTPStatus.INTERNAL_SERVER_ERROR, error_message, err_type="engine_error")
+        if request.stream:
+           error_json_str = json.dumps(json.loads(error_res.body))
+           async def error_stream():
+               yield f"data: {error_json_str}\n\n"
+               yield "data: [DONE]\n\n"
+           return StreamingResponse(error_stream(), media_type="text/event-stream", status_code=error_res.status_code)
+        else:
+           raise HTTPException(status_code=error_res.status_code, detail=error_message)
+
+
+async def initialize_llm_engine(model_path_to_load: str):
+    global engine, tokenizer, max_model_len
+    try:
+        logger.info(f"Initializing LLM Engine for CPU with model from: {model_path_to_load}")
+
+        if engine:
+            logger.info("Shutting down existing engine...")
+            engine = None
+            tokenizer = None
+            max_model_len = MAX_MODEL_LEN_CONFIG
+            await asyncio.sleep(2)
+            logger.info("Existing engine dereferenced.")
+
+
+        engine_args = AsyncEngineArgs(
+            model=model_path_to_load,
+            tensor_parallel_size=1, # For CPU
+            dtype="auto", # Let vLLM determine dtype
+            max_model_len=MAX_MODEL_LEN_CONFIG,
+            gpu_memory_utilization=GPU_MEMORY_UTILIZATION, # This might still be used even on CPU for planning
+            swap_space=4, # Swap space in GiB (CPU host memory for KV cache)
+        )
+
+        # Instantiate the AsyncLLMEngine directly
+        # If LLM is preferred, check its init signature for CPU arguments
+        # The error "EngineArgs.__init__() got an unexpected keyword argument 'cpu_only'"
+        # suggests 'cpu_only' should be passed elsewhere or is not a direct EngineArgs param
+        # In recent vLLM, device='cpu' or engine_args.device='cpu' is used.
+        # LLM(cpu_only=True) correctly sets device='cpu' in its underlying EngineArgs.
+        # The error might be from an older vLLM version or a conflict.
+        # Let's try passing device='cpu' to LLM init, which is the modern way.
+
+        llm = LLM(model=model_path_to_load,
+                  device="cpu", # Use device='cpu' instead of cpu_only=True if available
+                  max_model_len=MAX_MODEL_LEN_CONFIG,
+                  enable_chunked_prefill=False,
+                  tensor_parallel_size=1,
+                  swap_space=4
+                  )
+
+
+        engine = llm.llm_engine
+        engine_model_config = await engine.get_model_config()
+        max_model_len = engine_model_config.max_model_len
+
+        tokenizer = get_tokenizer(llm.get_tokenizer_name(),
+                                  tokenizer_mode=llm.get_tokenizer_mode(),
+                                  trust_remote_code=llm.get_tokenizer_trust_remote_code())
+
+        logger.info(f"LLM Engine initialized for CPU with model: {model_path_to_load}. Max model length: {max_model_len}")
+
+    except Exception as e:
+        logger.error(f"LLM Engine initialization failed: {e}", exc_info=True)
+        engine = None
+        tokenizer = None
+        max_model_len = MAX_MODEL_LEN_CONFIG
+        raise RuntimeError(f"LLM Engine initialization failed: {e}") from e
+
+
+@app.on_event("startup")
+async def startup_event():
+    logger.info("Application startup initiated.")
+
+    model_to_load_initially = DOWNLOADED_MODEL_PATH if DOWNLOADED_MODEL_PATH else MODEL_NAME
+    logger.info(f"Initial model to load: {model_to_load_initially}")
+
+    try:
+        await initialize_llm_engine(model_to_load_initially)
+    except RuntimeError as e:
+        logger.error(f"Failed to initialize LLM Engine during startup: {e}")
+
+    logger.info("Application startup complete.")
+
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    logger.info("Application shutdown initiated.")
+
+    global engine, tokenizer
+    if engine:
+        logger.info("Attempting to clean up vLLM engine resources.")
+        engine = None
+        tokenizer = None
+        logger.info("vLLM engine and tokenizer dereferenced.")
+
+    logger.info("Application shutdown complete.")
+
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=False,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+if __name__ == "__main__":
+    uvicorn.run(app, host=HOST, port=PORT, log_level=LOG_LEVEL.lower(), timeout_keep_alive=TIMEOUT_KEEP_ALIVE)