fix protocol to make sure function call input get accepted correctly.

rubra-ai · Jun 12, 2024 · 4fa44c3 · 4fa44c3
1 parent 49f4197
commit 4fa44c3
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/vllm/entrypoints/openai/protocol.py b/vllm/entrypoints/openai/protocol.py
@@ -41,6 +41,10 @@ class CustomChatCompletionMessageParam(TypedDict, total=False):
     same role.
     """
 
+    tool_calls : Optional[list]
+
+    tool_call_id : str
+
 
 ChatCompletionMessageParam = Union[
     openai.types.chat.ChatCompletionMessageParam,
@@ -129,8 +133,7 @@ class ChatCompletionNamedToolChoiceParam(OpenAIBaseModel):
 class ChatCompletionRequest(OpenAIBaseModel):
     # Ordered by official OpenAI API documentation
     # https://platform.openai.com/docs/api-reference/chat/create
-    # messages: List[ChatCompletionMessageParam] # TODO: figure out why
-    messages: List[dict]
+    messages: List[ChatCompletionMessageParam]
     model: str
     frequency_penalty: Optional[float] = 0.0
     logit_bias: Optional[Dict[str, float]] = None