Add Gateway for Translation (opea-project#169)

* add translation gateway Signed-off-by: zehao-intel <[email protected]> * fix import Signed-off-by: zehao-intel <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: zehao-intel <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
pallavijaini0525 · Jun 14, 2024 · fecbe4c · fecbe4c
1 parent 959b53d
commit fecbe4c
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 1 deletion.
diff --git a/comps/__init__.py b/comps/__init__.py
@@ -25,7 +25,14 @@
 from comps.cores.mega.orchestrator import ServiceOrchestrator
 from comps.cores.mega.orchestrator_with_yaml import ServiceOrchestratorWithYaml
 from comps.cores.mega.micro_service import MicroService, register_microservice, opea_microservices
-from comps.cores.mega.gateway import Gateway, ChatQnAGateway, CodeGenGateway, CodeTransGateway, DocSumGateway
+from comps.cores.mega.gateway import (
+    Gateway,
+    ChatQnAGateway,
+    CodeGenGateway,
+    CodeTransGateway,
+    DocSumGateway,
+    TranslationGateway,
+)
 
 # Telemetry
 from comps.cores.telemetry.opea_telemetry import opea_telemetry

diff --git a/comps/cores/mega/gateway.py b/comps/cores/mega/gateway.py
@@ -230,6 +230,51 @@ async def handle_request(self, request: Request):
         return ChatCompletionResponse(model="codetrans", choices=choices, usage=usage)
 
 
+class TranslationGateway(Gateway):
+    def __init__(self, megaservice, host="0.0.0.0", port=8888):
+        super().__init__(
+            megaservice, host, port, str(MegaServiceEndpoint.TRANSLATION), ChatCompletionRequest, ChatCompletionResponse
+        )
+
+    async def handle_request(self, request: Request):
+        data = await request.json()
+        language_from = data["language_from"]
+        language_to = data["language_to"]
+        source_language = data["source_language"]
+        prompt_template = """
+            Translate this from {language_from} to {language_to}:
+
+            {language_from}:
+            {source_language}
+
+            {language_to}:
+        """
+        prompt = prompt_template.format(
+            language_from=language_from, language_to=language_to, source_language=source_language
+        )
+        result_dict = await self.megaservice.schedule(initial_inputs={"query": prompt})
+        for node, response in result_dict.items():
+            # Here it suppose the last microservice in the megaservice is LLM.
+            if (
+                isinstance(response, StreamingResponse)
+                and node == list(self.megaservice.services.keys())[-1]
+                and self.megaservice.services[node].service_type == ServiceType.LLM
+            ):
+                return response
+        last_node = self.megaservice.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="translation", choices=choices, usage=usage)
+
+
 class DocSumGateway(Gateway):
     def __init__(self, megaservice, host="0.0.0.0", port=8888):
         super().__init__(