add unit test for textchat & voicechat, update api files

Signed-off-by: LetongHan <[email protected]>

Author: letonghan

neural_chat/server/restful/finetune_api.py

@@ -20,6 +20,8 @@
 from neural_chat.cli.log import logger
 from neural_chat.server.restful.request import FinetuneRequest
 from neural_chat.server.restful.response import FinetuneResponse
+from neural_chat.config import NeuralChatConfig
+from neural_chat.chatbot import build_chatbot
 
 
 def check_finetune_params(request: FinetuneRequest) -> Optional[str]:
@@ -40,7 +42,7 @@ def set_chatbot(self, chatbot) -> None:
 
     def get_chatbot(self):
         if self.chatbot is None:
-            raise RuntimeError("Finetunebot instance has not been set.")
+            raise RuntimeError("Chatbot instance has not been set.")
         return self.chatbot
 
     def handle_finetune_request(self, request: FinetuneRequest) -> FinetuneResponse:
@@ -50,6 +52,9 @@ def handle_finetune_request(self, request: FinetuneRequest) -> FinetuneResponse:
 
 
 router = FinetuneAPIRouter()
+config = NeuralChatConfig()
+bot = build_chatbot(config)
+router.set_chatbot(bot)
 
 
 @router.post("/v1/finetune")

neural_chat/server/restful/request.py

@@ -44,10 +44,6 @@ class Text2ImageRequest(RequestBaseModel):
     sd_inference_token: Optional[str] = None
 
 
-class VoiceRequest(RequestBaseModel):
-    voice: str
-
-
 class TextRequest(RequestBaseModel):
     text: str
 

neural_chat/server/restful/response.py

@@ -33,14 +33,6 @@ class ImageResponse(ResponseBaseModel):
     image: bytes
 
 
-class TextResponse(ResponseBaseModel):
-    content: str
-
-
-class VoiceResponse(ResponseBaseModel):
-    voice: str
-
-
 class RetrievalResponse(ResponseBaseModel):
     content: str
 

neural_chat/server/restful/retrieval_api.py

@@ -21,6 +21,8 @@
 from neural_chat.cli.log import logger
 from neural_chat.server.restful.request import RetrievalRequest
 from neural_chat.server.restful.response import RetrievalResponse
+from neural_chat.config import NeuralChatConfig
+from neural_chat.chatbot import build_chatbot
 
 
 def check_retrieval_params(request: RetrievalRequest) -> Optional[str]:
@@ -52,6 +54,9 @@ def handle_retrieval_request(self, request: RetrievalRequest) -> RetrievalRespon
 
 
 router = RetrievalAPIRouter()
+config = NeuralChatConfig()
+bot = build_chatbot(config)
+router.set_chatbot(bot)
 
 
 @router.post("/v1/retrieval")

neural_chat/server/restful/textchat_api.py

@@ -26,6 +26,8 @@
     ChatCompletionRequest, ChatCompletionResponseChoice, ChatCompletionResponse, 
     UsageInfo, ModelCard, ModelList, ModelPermission, ChatMessage
 )
+from neural_chat.config import NeuralChatConfig
+from neural_chat.chatbot import build_chatbot
 
 
 # TODO: process request and return params in Dict
@@ -172,18 +174,21 @@ async def handle_chat_completion_request(self, request: ChatCompletionRequest) -
                 for usage_key, usage_value in task_usage.dict().items():
                     setattr(usage, usage_key, getattr(usage, usage_key) + usage_value)
 
-        return ChatCompletionResponse(model=request.model, choices=choices, usage=usage)
+        return ChatCompletionResponse(model=request.model, choices=choices, usage=usage) 
 
 
 router = TextChatAPIRouter()
+config = NeuralChatConfig()
+bot = build_chatbot(config)
+router.set_chatbot(bot)
 
 
 @router.post("/v1/models")
 async def models_endpoint() -> ModelList:
     return await router.handle_models_request()
 
 
-@router.post("/v1/completion")
+@router.post("/v1/completions")
 async def completion_endpoint(request: CompletionRequest) -> CompletionResponse:
     ret = check_completion_request()
     if ret is not None:

neural_chat/server/restful/voicechat_api.py

@@ -15,11 +15,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Union
+from typing import ByteString
 from fastapi import APIRouter
 from neural_chat.cli.log import logger
-from neural_chat.server.restful.request import VoiceRequest, TextRequest
-from neural_chat.server.restful.response import VoiceResponse, TextResponse
+from neural_chat.config import NeuralChatConfig
+from neural_chat.chatbot import build_chatbot
 
 
 class VoiceChatAPIRouter(APIRouter):
@@ -36,28 +36,33 @@ def get_chatbot(self):
             raise RuntimeError("Chatbot instance has not been set.")
         return self.chatbot
 
-    async def handle_voice2text_request(self, request: VoiceRequest) -> TextResponse:
+    async def handle_voice2text_request(self, request: ByteString) -> str:
         # TODO: implement voice to text
         chatbot = self.get_chatbot()
         # TODO: chatbot.voice2text()
-        return TextResponse(content=None)
+        result = chatbot.predict(request.voice)
+        return result
 
-    async def handle_text2voice_request(self, request: TextRequest) -> VoiceResponse:
+    async def handle_text2voice_request(self, text: str) -> ByteString:
         # TODO: implement text to voice
         chatbot = self.get_chatbot()
         # TODO: chatbot.text2voice()
-        return VoiceResponse(content=None)
+        result = chatbot.predict(text)
+        return result
 
 
 router = VoiceChatAPIRouter()
+config = NeuralChatConfig()
+bot = build_chatbot(config)
+router.set_chatbot(bot)
 
 # voice to text
 @router.post("/v1/voice/asr")
-async def voice2text(requst: VoiceRequest) -> TextResponse:
-    return await router.handle_voice2text_request(requst)
+async def voice2text(request: ByteString) -> str:
+    return await router.handle_voice2text_request(request)
 
 
 # text to voice
 @router.post("/v1/voice/tts")
-async def voice2text(requst: TextRequest) -> VoiceResponse:
+async def voice2text(requst: str) -> ByteString:
     return await router.handle_text2voice_request(requst)

neural_chat/tests/restful/config.py

@@ -0,0 +1,22 @@
+import os
+
+# Get the host and port from the environment variables
+host = os.environ.get('MY_HOST')
+port = os.environ.get('MY_PORT')
+
+# Check if the environment variables are set and not empty
+if host and port:
+    # Combine the host and port to form the full URL
+    HOST = f"http://{host}:{port}"
+    API_COMPLETION = '/v1/completions'
+    API_CHAT_COMPLETION = '/v1/chat/completions'
+    API_ASR = '/v1/voice/asr'
+    API_TTS = '/v1/voice/tts'
+
+    print("HOST URL:", HOST)
+    print("Completions Endpoint:", API_COMPLETION)
+    print("Chat completions Endpoint:", API_CHAT_COMPLETION)
+    print("Voice ASR Endpoint:", API_ASR)
+    print("Voice TTS Endpoint:", API_TTS)
+else:
+    raise("Please set the environment variables MY_HOST and MY_PORT.")

neural_chat/tests/restful/unit_test/test_textchat_api

@@ -0,0 +1,42 @@
+#!/usr/bin/env python
+import requests
+import unittest
+from neural_chat.tests.restful.config import HOST, API_COMPLETION, API_CHAT_COMPLETION
+from neural_chat.server.restful.openai_protocol import CompletionRequest, ChatCompletionRequest
+from neural_chat.cli.log import logger
+
+
+class UnitTest(unittest.TestCase):
+
+    def __init__(self, *args):
+        super(UnitTest, self).__init__(*args)
+        self.host = HOST
+
+    def test_completions(self):
+        logger.info(f'Testing POST request: {self.host+API_COMPLETION}')
+        request = CompletionRequest(
+            model="mpt-7b-chat",
+            prompt="This is a test."
+        )
+        response = requests.post(self.host+API_COMPLETION, data=request)
+        logger.info('Response status code: {}'.format(response.status_code))
+        logger.info('Response text: {}'.format(response.choices.text))
+        self.assertEqual(response.status_code, 200, msg="Abnormal response status code.")
+
+    def test_chat_completions(self):
+        logger.info(f'Testing POST request: {self.host+API_CHAT_COMPLETION}')
+        request = ChatCompletionRequest(
+            model="mpt-7b-chat",
+            messages=[
+                {"role": "system","content": "You are a helpful assistant."},
+                {"role": "user","content": "Hello!"}
+            ]
+        )
+        response = requests.post(self.host+API_CHAT_COMPLETION, data=request)
+        logger.info('Response status code: {}'.format(response.status_code))
+        logger.info('Response text: {}'.format(response.choices.message.content))
+        self.assertEqual(response.status_code, 200, msg="Abnormal response status code.")
+
+
+if __name__ == "__main__":
+    unittest.main()

neural_chat/tests/restful/unit_test/test_voicechat_api

@@ -0,0 +1,33 @@
+#!/usr/bin/env python
+import requests
+import unittest
+from neural_chat.tests.restful.config import HOST, API_ASR, API_TTS
+from neural_chat.cli.log import logger
+from neural_chat.pipeline.plugins.audio.asr import AudioSpeechRecognition
+
+
+class UnitTest(unittest.TestCase):
+
+    def __init__(self, *args):
+        super(UnitTest, self).__init__(*args)
+        self.host = HOST
+
+    def test_asr(self):
+        logger.info(f'Testing POST request: {self.host+API_ASR}')
+        request = ""
+        response = requests.post(self.host+API_ASR, json=request)
+        logger.info('Response status code: {}'.format(response.status_code))
+        logger.info('Response text: {}'.format(response.text))
+        self.assertEqual(response.status_code, 200, msg="Abnormal response status code.")
+
+    def test_tts(self):
+        logger.info(f'Testing POST request: {self.host+API_TTS}')
+        request = "Hello, nice to meet you!"
+        response = requests.post(self.host+API_TTS, data=request)
+        logger.info('Response status code: {}'.format(response.status_code))
+        logger.info('Response voice: {}'.format(response.voice))
+        self.assertEqual(response.status_code, 200, msg="Abnormal response status code.")
+
+
+if __name__ == "__main__":
+    unittest.main()