Revert "feat: add the endpoint compatible with OpenAI client protocols"

This reverts commit bcea8fe.

Author: baixiac

app/api/routers/generative.py

@@ -1,24 +1,16 @@
-import json
 import logging
-import time
-import uuid
 import app.api.globals as cms_globals
 
-from typing import Union, Iterable, AsyncGenerator
 from typing_extensions import Annotated
 from fastapi import APIRouter, Depends, Request, Body, Query
-from fastapi.encoders import jsonable_encoder
-from fastapi.responses import PlainTextResponse, StreamingResponse, JSONResponse
-from starlette.status import HTTP_200_OK, HTTP_400_BAD_REQUEST
-from app.domain import Tags, OpenAIChatRequest, OpenAIChatResponse, PromptMessage, PromptRole
+from fastapi.responses import PlainTextResponse, StreamingResponse
+from app.domain import Tags
 from app.model_services.base import AbstractModelService
-from app.utils import get_settings, get_prompt_from_messages
+from app.utils import get_settings
 from app.api.utils import get_rate_limiter
-from app.api.dependencies import validate_tracking_id
 
 PATH_GENERATE = "/generate"
 PATH_GENERATE_ASYNC = "/stream/generate"
-PATH_OPENAI_COMPLETIONS = "/v1/chat/completions"
 
 router = APIRouter()
 config = get_settings()
@@ -39,8 +31,6 @@ def generate_text(
     request: Request,
     prompt: Annotated[str, Body(description="The prompt to be sent to the model", media_type="text/plain")],
     max_tokens: Annotated[int, Query(description="The maximum number of tokens to generate", gt=0)] = 512,
-    temperature: Annotated[float, Query(description="The temperature of the generated text", gt=0.0, lt=1.0)] = 0.7,
-    tracking_id: Union[str, None] = Depends(validate_tracking_id),
     model_service: AbstractModelService = Depends(cms_globals.model_service_dep)
 ) -> PlainTextResponse:
     """
@@ -50,27 +40,13 @@ def generate_text(
         request (Request): The request object.
         prompt (str): The prompt to be sent to the model.
         max_tokens (int): The maximum number of tokens to generate.
-        temperature (float): The temperature of the generated text.
-        tracking_id (Union[str, None]): An optional tracking ID of the requested task.
         model_service (AbstractModelService): The model service dependency.
 
     Returns:
         PlainTextResponse: A response containing the generated text.
     """
 
-    tracking_id = tracking_id or str(uuid.uuid4())
-    if prompt:
-        return PlainTextResponse(
-            model_service.generate(prompt, max_tokens=max_tokens, temperature=temperature),
-            headers={"x-cms-tracking-id": tracking_id},
-            status_code=HTTP_200_OK,
-        )
-    else:
-        return PlainTextResponse(
-            _empty_prompt_error(),
-            headers={"x-cms-tracking-id": tracking_id},
-            status_code=HTTP_400_BAD_REQUEST,
-        )
+    return PlainTextResponse(model_service.generate(prompt, max_tokens=max_tokens))
 
 
 @router.post(
@@ -84,8 +60,6 @@ async def generate_text_stream(
     request: Request,
     prompt: Annotated[str, Body(description="The prompt to be sent to the model", media_type="text/plain")],
     max_tokens: Annotated[int, Query(description="The maximum number of tokens to generate", gt=0)] = 512,
-    temperature: Annotated[float, Query(description="The temperature of the generated text", gt=0.0, lt=1.0)] = 0.7,
-    tracking_id: Union[str, None] = Depends(validate_tracking_id),
     model_service: AbstractModelService = Depends(cms_globals.model_service_dep)
 ) -> StreamingResponse:
     """
@@ -95,121 +69,13 @@ async def generate_text_stream(
         request (Request): The request object.
         prompt (str): The prompt to be sent to the model.
         max_tokens (int): The maximum number of tokens to generate.
-        temperature (float): The temperature of the generated text.
-        tracking_id (Union[str, None]): An optional tracking ID of the requested task.
         model_service (AbstractModelService): The model service dependency.
 
     Returns:
         StreamingResponse: A streaming response containing the text generated in near real-time.
     """
 
-    tracking_id = tracking_id or str(uuid.uuid4())
-    if prompt:
-        return StreamingResponse(
-            model_service.generate_async(prompt, max_tokens=max_tokens, temperature=temperature),
-            media_type="text/event-stream",
-            headers={"x-cms-tracking-id": tracking_id},
-            status_code=HTTP_200_OK,
-        )
-    else:
-        return StreamingResponse(
-            _empty_prompt_error(),
-            media_type="text/event-stream",
-            headers={"x-cms-tracking-id": tracking_id},
-            status_code=HTTP_400_BAD_REQUEST,
-        )
-
-
-@router.post(
-    "/v1/chat/completions",
-    tags=[Tags.Generative.name],
-    response_model=None,
-    dependencies=[Depends(cms_globals.props.current_active_user)],
-    description="Generate chat response based on messages, similar to OpenAI's /v1/chat/completions",
-)
-def generate_chat_completions(
-    request: Request,
-    request_data: Annotated[OpenAIChatRequest, Body(
-        description="OpenAI-like completion request", media_type="application/json"
-    )],
-    tracking_id: Union[str, None] = Depends(validate_tracking_id),
-    model_service: AbstractModelService = Depends(cms_globals.model_service_dep)
-) -> Union[StreamingResponse, JSONResponse]:
-    """
-    Generate chat response based on messages, mimicking OpenAI's /v1/chat/completions endpoint.
-
-    Args:
-        request (Request): The request object.
-        request_data (OpenAIChatRequest): The request data containing model, messages, and stream.
-        tracking_id (Union[str, None]): An optional tracking ID of the requested task.
-        model_service (AbstractModelService): The model service dependency.
-
-    Returns:
-        StreamingResponse: A OpenAI-like response containing the text generated in near real-time.
-        JSONResponse: A response containing an error message if the prompt messages are empty.
-    """
-
-    messages = request_data.messages
-    stream = request_data.stream
-    max_tokens = request_data.max_tokens
-    temperature = request_data.temperature
-
-    if not messages:
-        error_response = {
-            "error": {
-                "message": "No prompt messages provided",
-                "type": "invalid_request_error",
-                "param": "messages",
-                "code": "missing_field",
-            }
-        }
-        return JSONResponse(content=error_response, status_code=HTTP_400_BAD_REQUEST)
-
-    async def _stream(p: str, mt: int, t: float) -> AsyncGenerator:
-        data = {
-            "id": tracking_id or str(uuid.uuid4()),
-            "object": "chat.completion.chunk",
-            "choices": [{"delta": {"role": PromptRole.ASSISTANT.value}}],
-        }
-        yield f"data: {json.dumps(data)}\n\n"
-        async for chunk in model_service.generate_async(p, max_tokens=mt, temperature=t):
-            data = {
-                "choices": [
-                    {
-                        "delta": {"content": chunk}
-                    }
-                ],
-                "object": "chat.completion.chunk",
-            }
-            yield f"data: {json.dumps(data)}\n\n"
-        yield "data: [DONE]\n\n"
-
-    prompt = get_prompt_from_messages(model_service.tokenizer, messages)    # type: ignore
-    if stream:
-        return StreamingResponse(
-            _stream(prompt, max_tokens, temperature),
-            media_type="text/event-stream"
-        )
-    else:
-        generated_text = model_service.generate(prompt, max_tokens=max_tokens, temperature=temperature)
-        completion = OpenAIChatResponse(
-            id=str(uuid.uuid4()),
-            object="chat.completion",
-            created=int(time.time()),
-            model=model_service.model_name,
-            choices=[
-                {
-                    "index": 0,
-                    "message": PromptMessage(
-                        role=PromptRole.ASSISTANT,
-                        content=generated_text,
-                    ),
-                    "finish_reason": "stop",
-                }
-            ]
-        )
-        return JSONResponse(content=jsonable_encoder(completion))
-
-
-def _empty_prompt_error() -> Iterable[str]:
-    yield "ERROR: No prompt text provided\n"
+    return StreamingResponse(
+        model_service.generate_async(prompt, max_tokens=max_tokens),
+        media_type="text/event-stream"
+    )

app/api/utils.py

@@ -286,7 +286,6 @@ async def init_vllm_engine(app: FastAPI,
     """
 
     try:
-        # Import necessary vLLM components
         from vllm.utils import FlexibleArgumentParser
         from vllm.engine.arg_utils import AsyncEngineArgs
         from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
@@ -299,19 +298,16 @@ async def init_vllm_engine(app: FastAPI,
         )
         from vllm import SamplingParams, TokensPrompt
     except ImportError:
-        # Raise a custom exception if vLLM is not installed
-        raise ConfigurationException("Cannot import the vLLM engine. Please install it with `pip install vllm`.")
+        logger.error("Cannot import the vLLM engine. Please install it with `pip install cms[vllm]`.")
 
     parser = FlexibleArgumentParser()
     parser = make_arg_parser(parser)
     args = parser.parse_args([])
     validate_parsed_serve_args(args)
-
     args.model = model_dir_path
     args.dtype = "float16"
     args.served_model_name = [model_name]
-    args.max_model_len = 2048 # The default batched length (2048) needs to be higher than max_model_len.
-    # args.tokenizer = model_dir_path # Uncomment if your tokenizer is in a different path or needs explicit setting.
+    # args.tokenizer = model_dir_path
     args.log_level = log_level
 
     exit_stack = contextlib.AsyncExitStack()
@@ -321,44 +317,37 @@ async def init_vllm_engine(app: FastAPI,
             disable_frontend_multiprocessing=True,
         )
     )
-
     tokenizer = await engine.get_tokenizer()
     vllm_config = await engine.get_vllm_config()
     model_config = await engine.get_model_config()
-
     await init_app_state(engine, vllm_config, app.state, args)
 
     async def generate_text(
         request: Request,
         prompt: Annotated[str, Body(description="The prompt to be sent to the model", media_type="text/plain")],
         max_tokens: Annotated[int, Query(description="The maximum number of tokens to generate", gt=0)] = 512
     ) -> StreamingResponse:
-        """
-        Custom endpoint for streaming text generation.
-        This endpoint takes a raw text prompt and streams back the generated text.
-        It applies a chat template to the prompt internally for model compatibility.
-        """
         messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
 
         params = SamplingParams(max_tokens=max_tokens)
-
         conversation, _ = parse_chat_messages(messages, model_config, tokenizer, content_format="string")   # type: ignore
-        prompt_tokens = apply_hf_chat_template(    # type: ignore
-            tokenizer,
-            conversation=conversation,
-            tools=None,
-            add_generation_prompt=True,
-            continue_final_message=False,
-            chat_template="{% for message in messages %}\n{% if message['role'] == 'user' %}\nUser: {{ message['content'] }}\n{% elif message['role'] == 'assistant' %}\nAssistant: {{ message['content'] }}\n{% endif %}\n{% endfor %}\nAssistant:",
-            tokenize=True,
+        prompt = TokensPrompt(
+            prompt_token_ids=apply_hf_chat_template(    # type: ignore
+                tokenizer,
+                conversation=conversation,
+                tools=None,
+                add_generation_prompt=True,
+                continue_final_message=False,
+                chat_template="{% for message in messages %}\n{% if message['role'] == 'user' %}\nUser: {{ message['content'] }}\n{% elif message['role'] == 'assistant' %}\nAssistant: {{ message['content'] }}\n{% endif %}\n{% endfor %}\nAssistant:",
+                tokenize=True,
+            )
         )
-        prompt_obj = TokensPrompt(prompt_token_ids=prompt_tokens)
 
         async def _stream() -> AsyncGenerator[bytes, None]:
             start = 0
-            async for output in engine.generate(request_id=uuid.uuid4().hex, prompt=prompt_obj, sampling_params=params):
+            async for output in engine.generate(request_id=uuid.uuid4().hex, prompt=prompt, sampling_params=params):
                 text = output.outputs[0].text
-                yield text[start:].encode("utf-8")
+                yield text[start:]  # type: ignore
                 start = len(text)
 
         return StreamingResponse(_stream(), media_type="text/event-stream")

app/domain.py

@@ -167,30 +167,3 @@ class Doc(BaseModel):
     text: str = Field(description="The text from which the entities are extracted")
     ents: List[Entity] = Field(description="The list of extracted entities")
     title: Optional[str] = Field(default=None, description="The headline of the text")
-
-
-class PromptRole(Enum):
-    SYSTEM = "system"
-    USER = "user"
-    ASSISTANT = "assistant"
-    TOOL = "tool"
-
-
-class PromptMessage(BaseModel):
-    role: PromptRole = Field(description="The role who generates the message")
-    content: str = Field(description="The actual text of the message")
-
-
-class OpenAIChatRequest(BaseModel):
-    messages: List[PromptMessage] = Field(..., description="A list of messages to be sent to the model")
-    stream: bool = Field(..., description="Whether to stream the response")
-    max_tokens: int = Field(512, description="The maximum number of tokens to generate", gt=0)
-    temperature: float = Field(0.7, description="The temperature of the generated text", ge=0.0, le=1.0)
-
-
-class OpenAIChatResponse(BaseModel):
-    id: str
-    object: str
-    created: int
-    model: str
-    choices: List

app/model_services/huggingface_llm_model.py

@@ -168,8 +168,6 @@ def init_model(self) -> None:
             logger.warning("Model service is already initialised and can be initialised only once")
         else:
             self._model, self._tokenizer = self.load_model(self._model_pack_path)
-            if non_default_device_is_available(get_settings().DEVICE):
-                self._model.to(get_settings().DEVICE)
             if self._enable_trainer:
                 logger.error("Trainers are not yet implemented for HuggingFace Generative models")
 
@@ -193,20 +191,13 @@ def annotate(self, text: str) -> List[Annotation]:
     def batch_annotate(self, texts: List[str]) -> List[List[Annotation]]:
         raise NotImplementedError("Batch annotation is not yet implemented for HuggingFace Generative models")
 
-    def generate(
-        self,
-        prompt: str,
-        max_tokens: int = 512,
-        temperature: float = 0.7,
-        **kwargs: Any
-    ) -> str:
+    def generate(self, prompt: str, max_tokens: int = 512, **kwargs: Any) -> str:
         """
         Generates text based on the prompt.
 
         Args:
             prompt (str): The prompt for the text generation
             max_tokens (int): The maximum number of tokens to generate. Defaults to 512.
-            temperature (float): The temperature for the text generation. Defaults to 0.7.
             **kwargs (Any): Additional keyword arguments to be passed to this method.
 
         Returns:
@@ -223,8 +214,8 @@ def generate(
             inputs=inputs.input_ids,
             attention_mask=inputs.attention_mask,
             max_new_tokens=max_tokens,
-            do_sample=False,
-            temperature=temperature,
+            do_sample=True,
+            temperature=0.7,
             top_p=0.9,
         )
 
@@ -236,20 +227,13 @@ def generate(
 
         return generated_text
 
-    async def generate_async(
-        self,
-        prompt: str,
-        max_tokens: int = 512,
-        temperature: float = 0.7,
-        **kwargs: Any
-    ) -> AsyncIterable:
+    async def generate_async(self, prompt: str, max_tokens: int = 512, **kwargs: Any) -> AsyncIterable:
         """
         Asynchronously generates text stream based on the prompt.
 
         Args:
             prompt (str): The prompt for the text generation.
             max_tokens (int): The maximum number of tokens to generate. Defaults to 512.
-            temperature (float): The temperature for the text generation. Defaults to 0.7.
             **kwargs (Any): Additional keyword arguments to be passed to the model loader.
 
         Returns:
@@ -273,7 +257,7 @@ async def generate_async(
             streamer=streamer,
             max_new_tokens=max_tokens,
             do_sample=True,
-            temperature=temperature,
+            temperature=0.7,
             top_p=0.9,
         )