From 91eafcbcd3b41a2b48b9467eb169f2c39ab22a8c Mon Sep 17 00:00:00 2001
From: Oren Zhang <orenzhang@ovinc.cn>
Date: Thu, 3 Oct 2024 22:02:53 +0800
Subject: [PATCH] feat: json mode

---
 apps/chat/consumers_async.py | 16 ++++++++++++++++
 apps/chat/views.py           | 18 ++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/apps/chat/consumers_async.py b/apps/chat/consumers_async.py
index 7e2443c..8976f74 100644
--- a/apps/chat/consumers_async.py
+++ b/apps/chat/consumers_async.py
@@ -145,3 +145,19 @@ def get_model_client(self, model: AIModel) -> Type[BaseClient]:
                 return MidjourneyClient
             case _:
                 raise UnexpectedProvider()
+
+
+class JSONModeConsumer(AsyncConsumer):
+    """
+    JSON Mode Consumer
+    """
+
+    def __init__(self, key: str):
+        super().__init__("", key)
+        self.message = ""
+
+    async def send(self, text_data: str):
+        self.message += json.loads(text_data).get("data", "")
+
+    async def close(self):
+        return
diff --git a/apps/chat/views.py b/apps/chat/views.py
index 88d5fee..4d091fa 100644
--- a/apps/chat/views.py
+++ b/apps/chat/views.py
@@ -15,6 +15,7 @@
 from rest_framework.response import Response
 
 from apps.chat.constants import MESSAGE_CACHE_KEY
+from apps.chat.consumers_async import JSONModeConsumer
 from apps.chat.models import AIModel, ChatLog, SystemPreset
 from apps.chat.permissions import AIModelPermission
 from apps.chat.serializers import (
@@ -92,6 +93,23 @@ def load_model_map(self) -> dict:
         models = AIModel.objects.all()
         return {model.model: model.name for model in models}
 
+    @action(methods=["POST"], detail=False, permission_classes=[AIModelPermission])
+    async def json(self, request, *args, **kwargs):
+        """
+        JSON Mode
+        """
+
+        # pre check
+        pre_response = await self.pre_check(request, *args, **kwargs)
+        data = pre_response.data
+
+        # chat
+        consumer = JSONModeConsumer(data["key"])
+        await consumer.chat()
+
+        # response
+        return Response(data={"data": consumer.message})
+
 
 # pylint: disable=R0901
 class AIModelViewSet(ListMixin, MainViewSet):