Check server input

lmdeploy/model.py

@@ -3,7 +3,7 @@
 import json
 import uuid
 from abc import abstractmethod
-from typing import List, Literal, Optional
+from typing import List, Literal, Optional, Union
 
 from mmengine import Registry
 
@@ -18,6 +18,13 @@ def random_uuid() -> str:
     return str(uuid.uuid4().hex)
 
 
+def get_text(content: Union[str, List[dict]]):
+    """extract text from content."""
+    if isinstance(content, str):
+        return content
+    return content[0]['text']
+
+
 @dataclasses.dataclass
 class ChatTemplateConfig:
     """Parameters for chat template.
@@ -219,7 +226,7 @@ def messages2prompt(self, messages, sequence_start=True, **kwargs):
                 ret += f'{self.system}{self.meta_instruction}{self.eosys}'
         for message in messages:
             role = message['role']
-            content = message['content']
+            content = get_text(message['content'])
             ret += f'{box_map[role]}{content}{eox_map[role]}'
         if len(messages) and messages[-1]['role'] == 'assistant':
             return ret[:-len(eox_map['assistant'])]  # prefix of response
@@ -509,7 +516,7 @@ def messages2prompt(self,
             messages.insert(insert_index, tools_prompt)
         for message in messages:
             role = message['role']
-            content = message['content']
+            content = get_text(message['content'])
             if role == 'assistant' and message.get('tool_calls',
                                                    None) is not None:
                 for tool_call in message['tool_calls']:
@@ -861,7 +868,7 @@ def messages2prompt(self,
                     ret += f'{self.system}{self.knowledge}{self.tools}{tool_prompt}{self.eotools}{self.meta_instruction}{self.eosys}'
         for message in messages:
             role = message['role']
-            content = message['content']
+            content = get_text(message['content'])
             if role == 'assistant' and ('<|python_tag|>' in content
                                         or '</function>' in content):
                 ret += f'{box_map[role]}{content}<|eom_id|>'
@@ -1037,7 +1044,7 @@ def messages2prompt(self, messages, sequence_start=True, **kwargs):
         count = 0
         for message in messages:
             role = message['role']
-            content = message['content']
+            content = get_text(message['content'])
             if role == 'user':
                 count += 1
                 ret += f'[Round {count}]\n\n'

lmdeploy/serve/openai/api_server.py

@@ -12,22 +12,23 @@
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, Response, StreamingResponse
 from fastapi.security.http import HTTPAuthorizationCredentials, HTTPBearer
+from pydantic import ValidationError
 
 from lmdeploy.archs import get_task
 from lmdeploy.messages import (GenerationConfig, LogitsProcessor,
                                PytorchEngineConfig, TurbomindEngineConfig)
 from lmdeploy.model import ChatTemplateConfig
 from lmdeploy.serve.async_engine import AsyncEngine
 from lmdeploy.serve.openai.protocol import (  # noqa: E501
-    ChatCompletionRequest, ChatCompletionResponse,
-    ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
-    ChatCompletionStreamResponse, ChatCompletionTokenLogprob, ChatMessage,
-    ChoiceLogprobs, CompletionRequest, CompletionResponse,
-    CompletionResponseChoice, CompletionResponseStreamChoice,
-    CompletionStreamResponse, DeltaMessage, EmbeddingsRequest, EncodeRequest,
-    EncodeResponse, ErrorResponse, FunctionResponse, GenerateRequest,
-    GenerateResponse, LogProbs, ModelCard, ModelList, ModelPermission,
-    ToolCall, TopLogprob, UsageInfo)
+    ChatCompletionRequest, ChatCompletionRequestMessage,
+    ChatCompletionResponse, ChatCompletionResponseChoice,
+    ChatCompletionResponseStreamChoice, ChatCompletionStreamResponse,
+    ChatCompletionTokenLogprob, ChatMessage, ChoiceLogprobs, CompletionRequest,
+    CompletionResponse, CompletionResponseChoice,
+    CompletionResponseStreamChoice, CompletionStreamResponse, DeltaMessage,
+    EmbeddingsRequest, EncodeRequest, EncodeResponse, ErrorResponse,
+    FunctionResponse, GenerateRequest, GenerateResponse, LogProbs, ModelCard,
+    ModelList, ModelPermission, ToolCall, TopLogprob, UsageInfo)
 from lmdeploy.tokenizer import DetokenizeState, Tokenizer
 from lmdeploy.utils import get_logger
 
@@ -334,6 +335,13 @@ async def chat_completions_v1(request: ChatCompletionRequest,
     error_check_ret = await check_request(request)
     if error_check_ret is not None:
         return error_check_ret
+    try:
+        _ = ChatCompletionRequestMessage(messages=request.messages)
+    except ValidationError as error:
+        return create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            f'Invalid input, input={request.messages}, error={error}')
+
     if VariableInterface.async_engine.id2step.get(str(request.session_id),
                                                   0) != 0:
         return create_error_response(

lmdeploy/serve/openai/protocol.py

@@ -114,7 +114,7 @@ class ChatCompletionRequest(BaseModel):
     temperature: Optional[float] = 0.7
     top_p: Optional[float] = 1.0
     tools: Optional[List[Tool]] = Field(default=None, examples=[None])
-    tool_choice: Union[ToolChoice, Literal['auto', 'required','none']] = Field(default='auto', examples=['none'])  # noqa
+    tool_choice: Union[ToolChoice, Literal['auto', 'required', 'none']] = Field(default='auto', examples=['none'])  # noqa
     logprobs: Optional[bool] = False
     top_logprobs: Optional[int] = None
     n: Optional[int] = 1
@@ -139,6 +139,45 @@ class ChatCompletionRequest(BaseModel):
     seed: Optional[int] = None
 
 
+class ContentPartTextParam(BaseModel):
+    type: Literal['text']
+    text: str
+
+
+class ImageURL(BaseModel):
+    url: str
+
+
+class ContentPartImageParam(BaseModel):
+    type: Literal['image_url']
+    image_url: ImageURL
+
+
+ContentPartParam = Union[ContentPartTextParam, ContentPartImageParam]
+
+
+class AssistantMessage(BaseModel):
+    role: Literal['assistant']
+    content: Union[str, List[ContentPartTextParam]]
+
+
+class SystemMessage(BaseModel):
+    role: Literal['system']
+    content: Union[str, List[ContentPartTextParam]]
+
+
+class UserMessage(BaseModel):
+    role: Literal['user']
+    content: Union[str, List[ContentPartParam]]
+
+
+ChatCompletionMessage = Union[SystemMessage, UserMessage, AssistantMessage]
+
+
+class ChatCompletionRequestMessage(BaseModel):
+    messages: Union[str, List[ChatCompletionMessage]]
+
+
 class FunctionResponse(BaseModel):
     """Function response."""
     name: str
@@ -240,7 +279,6 @@ class CompletionRequest(BaseModel):
     stream_options: Optional[StreamOptions] = Field(default=None,
                                                     examples=[None])
     top_p: Optional[float] = 1.0
-    logprobs: Optional[int] = None
     echo: Optional[bool] = False
     presence_penalty: Optional[float] = 0.0
     frequency_penalty: Optional[float] = 0.0

tests/test_lmdeploy/test_server_input.py

@@ -0,0 +1,37 @@
+import pytest
+from pydantic import ValidationError
+
+from lmdeploy.serve.openai.protocol import ChatCompletionRequestMessage
+
+
+def test_single_str_input():
+    messages = 'hello'
+    _ = ChatCompletionRequestMessage(messages=messages)
+
+
+@pytest.mark.parametrize('role', ['system', 'user', 'assistant'])
+def test_list_str_input(role):
+    content = 'hello'
+    messages = [dict(role=role, content=content)]
+    _ = ChatCompletionRequestMessage(messages=messages)
+
+
+@pytest.mark.parametrize('role', ['system', 'user', 'assistant'])
+def test_list_content_input(role):
+    content = [dict(type='text', text='hello')]
+    messages = [dict(role=role, content=content)]
+    _ = ChatCompletionRequestMessage(messages=messages)
+
+
+def test_user_image_input():
+    content = [dict(type='image_url', image_url=dict(url='xxx'))]
+    messages = [dict(role='user', content=content)]
+    _ = ChatCompletionRequestMessage(messages=messages)
+
+
+@pytest.mark.parametrize('role', ['system', 'assistant'])
+def test_system_assistant_image_input(role):
+    content = [dict(type='image_url', image_url=dict(url='xxx'))]
+    messages = [dict(role=role, content=content)]
+    with pytest.raises(ValidationError):
+        _ = ChatCompletionRequestMessage(messages=messages)