Fix List types in KServe openai classes

logicalclocks · Dec 10, 2024 · 5591d35 · 5591d35
1 parent 3b8c621
commit 5591d35
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 6 deletions.
diff --git a/python/huggingfaceserver/huggingfaceserver/generative_model.py b/python/huggingfaceserver/huggingfaceserver/generative_model.py
@@ -26,6 +26,7 @@
     TypedDict,
     Union,
     cast,
+    List,
 )
 
 import torch
@@ -388,7 +389,7 @@ def apply_chat_template(
         self,
         messages: Iterable[ChatCompletionRequestMessage],
         chat_template: Optional[str] = None,
-        tools: Optional[list[ChatCompletionTool]] = None,
+        tools: Optional[List[ChatCompletionTool]] = None,
     ) -> ChatPrompt:
         """
         Given a list of chat completion messages, convert them to a prompt.

diff --git a/python/huggingfaceserver/huggingfaceserver/vllm/vllm_completions.py b/python/huggingfaceserver/huggingfaceserver/vllm/vllm_completions.py
@@ -360,7 +360,7 @@ def apply_chat_template(
         self,
         messages: Iterable[ChatCompletionRequestMessage],
         chat_template: Optional[str] = None,
-        tools: Optional[list[ChatCompletionTool]] = None,
+        tools: Optional[List[ChatCompletionTool]] = None,
     ):
         return self.tokenizer.apply_chat_template(
             conversation=messages,

diff --git a/python/huggingfaceserver/huggingfaceserver/vllm/vllm_model.py b/python/huggingfaceserver/huggingfaceserver/vllm/vllm_model.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import AsyncIterator, Iterable, Optional, Union
+from typing import AsyncIterator, Iterable, Optional, Union, List
 
 import torch
 from vllm.entrypoints.logger import RequestLogger
@@ -71,7 +71,7 @@ def apply_chat_template(
         self,
         messages: Iterable[ChatCompletionRequestMessage],
         chat_template: Optional[str] = None,
-        tools: Optional[list[ChatCompletionTool]] = None,
+        tools: Optional[List[ChatCompletionTool]] = None,
     ) -> ChatPrompt:
         """
         Given a list of chat completion messages, convert them to a prompt.

diff --git a/python/kserve/kserve/protocol/rest/openai/openai_chat_adapter_model.py b/python/kserve/kserve/protocol/rest/openai/openai_chat_adapter_model.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from abc import abstractmethod
-from typing import AsyncIterator, Iterable, Union, cast, Optional
+from typing import AsyncIterator, Iterable, Union, cast, Optional, List
 
 from kserve.protocol.rest.openai.types import (
     ChatCompletion,
@@ -33,6 +33,8 @@
     TopLogprob,
 )
 
+from kserve.protocol.rest.openai.types.openapi import ChatCompletionTool
+
 from ....errors import InvalidInput
 from .openai_model import (
     OpenAIModel,
@@ -56,6 +58,7 @@ def apply_chat_template(
         self,
         messages: Iterable[ChatCompletionRequestMessage],
         chat_template: Optional[str] = None,
+        tools: Optional[List[ChatCompletionTool]] = None,
     ) -> ChatPrompt:
         """
         Given a list of chat completion messages, convert them to a prompt.

diff --git a/python/kserve/test/test_openai.py b/python/kserve/test/test_openai.py
@@ -88,7 +88,7 @@ def apply_chat_template(
         self,
         messages: Iterable[ChatCompletionRequestMessage],
         chat_template: Optional[str] = None,
-        tools: Optional[list[ChatCompletionTool]] = None,
+        tools: Optional[List[ChatCompletionTool]] = None,
     ) -> ChatPrompt:
         return ChatPrompt(prompt="hello")