From 79a2d55807c57d57e8460e931d51aafc3c39e0cd Mon Sep 17 00:00:00 2001
From: root <root@idc708073.jf.intel.com>
Date: Thu, 19 Sep 2024 00:23:09 +0000
Subject: [PATCH] add gateway to GenAIExamples.

---
 ChatQnA/Dockerfile |  1 +
 ChatQnA/chatqna.py |  3 +-
 ChatQnA/gateway.py | 74 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 1 deletion(-)
 create mode 100644 ChatQnA/gateway.py

diff --git a/ChatQnA/Dockerfile b/ChatQnA/Dockerfile
index 4ece0783a..1a8729a6e 100644
--- a/ChatQnA/Dockerfile
+++ b/ChatQnA/Dockerfile
@@ -22,6 +22,7 @@ RUN pip install --no-cache-dir --upgrade pip && \
     pip install --no-cache-dir -r /home/user/GenAIComps/requirements.txt
 
 COPY ./chatqna.py /home/user/chatqna.py
+COPY ./gateway.py /home/user/gateway.py
 
 ENV PYTHONPATH=$PYTHONPATH:/home/user/GenAIComps
 
diff --git a/ChatQnA/chatqna.py b/ChatQnA/chatqna.py
index 09062b5d2..0989c523f 100644
--- a/ChatQnA/chatqna.py
+++ b/ChatQnA/chatqna.py
@@ -3,7 +3,8 @@
 
 import os
 
-from comps import ChatQnAGateway, MicroService, ServiceOrchestrator, ServiceType
+from comps import MicroService, ServiceOrchestrator, ServiceType
+from gateway import ChatQnAGateway
 
 MEGA_SERVICE_HOST_IP = os.getenv("MEGA_SERVICE_HOST_IP", "0.0.0.0")
 MEGA_SERVICE_PORT = int(os.getenv("MEGA_SERVICE_PORT", 8888))
diff --git a/ChatQnA/gateway.py b/ChatQnA/gateway.py
new file mode 100644
index 000000000..611ae548d
--- /dev/null
+++ b/ChatQnA/gateway.py
@@ -0,0 +1,74 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from fastapi import Request
+from fastapi.responses import StreamingResponse
+
+from comps.cores.mega.gateway import Gateway
+from comps.cores.mega.constants import MegaServiceEndpoint
+from comps.cores.proto.api_protocol import (
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionResponseChoice,
+    ChatMessage,
+    UsageInfo,
+)
+from comps.cores.proto.docarray import (
+    LLMParams,
+    RerankerParms,
+    RetrieverParms
+)
+
+
+class ChatQnAGateway(Gateway):
+    def __init__(self, megaservice, host="0.0.0.0", port=8888):
+        super().__init__(
+            megaservice, host, port, str(MegaServiceEndpoint.CHAT_QNA), ChatCompletionRequest, ChatCompletionResponse
+        )
+
+    async def handle_request(self, request: Request):
+        data = await request.json()
+        stream_opt = data.get("stream", True)
+        chat_request = ChatCompletionRequest.parse_obj(data)
+        prompt = self._handle_message(chat_request.messages)
+        parameters = LLMParams(
+            max_new_tokens=chat_request.max_tokens if chat_request.max_tokens else 1024,
+            top_k=chat_request.top_k if chat_request.top_k else 10,
+            top_p=chat_request.top_p if chat_request.top_p else 0.95,
+            temperature=chat_request.temperature if chat_request.temperature else 0.01,
+            repetition_penalty=chat_request.presence_penalty if chat_request.presence_penalty else 1.03,
+            streaming=stream_opt,
+            chat_template=chat_request.chat_template if chat_request.chat_template else None,
+        )
+        retriever_parameters = RetrieverParms(
+            search_type=chat_request.search_type if chat_request.search_type else "similarity",
+            k=chat_request.k if chat_request.k else 4,
+            distance_threshold=chat_request.distance_threshold if chat_request.distance_threshold else None,
+            fetch_k=chat_request.fetch_k if chat_request.fetch_k else 20,
+            lambda_mult=chat_request.lambda_mult if chat_request.lambda_mult else 0.5,
+            score_threshold=chat_request.score_threshold if chat_request.score_threshold else 0.2,
+        )
+        reranker_parameters = RerankerParms(
+            top_n=chat_request.top_n if chat_request.top_n else 1,
+        )
+        result_dict, runtime_graph = await self.megaservice.schedule(
+            initial_inputs={"text": prompt},
+            llm_parameters=parameters,
+            retriever_parameters=retriever_parameters,
+            reranker_parameters=reranker_parameters,
+        )
+        for node, response in result_dict.items():
+            if isinstance(response, StreamingResponse):
+                return response
+        last_node = runtime_graph.all_leaves()[-1]
+        response = result_dict[last_node]["text"]
+        choices = []
+        usage = UsageInfo()
+        choices.append(
+            ChatCompletionResponseChoice(
+                index=0,
+                message=ChatMessage(role="assistant", content=response),
+                finish_reason="stop",
+            )
+        )
+        return ChatCompletionResponse(model="chatqna", choices=choices, usage=usage)