
from http import HTTPStatus
from typing import Optional

from vllm.entrypoints.openai.protocol import ErrorResponse
from vllm.entrypoints.openai.serving_engine import OpenAIServing, AnyRequest
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.logger import init_logger

logger = init_logger(__name__)

async def vllm__entrypoints__openai__serving_engine__OpenAIServing___check_model(
    self,
    request: AnyRequest,
) -> Optional[ErrorResponse]:
    if self._is_model_supported(request.model):
        return None
    if request.model in [lora.lora_name for lora in self.lora_requests]:
        return None
    if request.model in [
            prompt_adapter.prompt_adapter_name
            for prompt_adapter in self.prompt_adapter_requests
    ]:
        return None
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: when client send a request with model=init/save_scheduler_view,
            scheduler will dump profile data.
    '''
    if request.model == "init_scheduler_view":
        await self.engine_client.init_scheduler_view()
    if request.model == "save_scheduler_view":
        await self.engine_client.save_scheduler_view()
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    return self.create_error_response(
        message=f"The model `{request.model}` does not exist.",
        err_type="NotFoundError",
        status_code=HTTPStatus.NOT_FOUND)


MluHijackObject.apply_hijack(OpenAIServing,
                             OpenAIServing._check_model,
                             vllm__entrypoints__openai__serving_engine__OpenAIServing___check_model)