From 6d65ab70305364a33e304c6f7f0fc939ee44cd90 Mon Sep 17 00:00:00 2001
From: Anulax <anulax1225@icloud.com>
Date: Mon, 19 May 2025 08:49:36 +0200
Subject: [PATCH] Add realtime response when prompting and fixed swagger docs

---
 api/templates/drf-yasg/redoc.html |  4 +-
 api/urls.py                       |  2 +
 api/utils.py                      | 22 ++++++++
 api/views.py                      | 87 ++++++++++++++++++++-----------
 4 files changed, 84 insertions(+), 31 deletions(-)
 create mode 100644 api/utils.py

diff --git a/api/templates/drf-yasg/redoc.html b/api/templates/drf-yasg/redoc.html
index f9a520b..129f6ae 100644
--- a/api/templates/drf-yasg/redoc.html
+++ b/api/templates/drf-yasg/redoc.html
@@ -4,8 +4,8 @@
 <link rel="icon" type="image/ico" href="/static/favicon.ico"/>
   <style>
     img[alt="Botzilla logo"] {
-
-      max-height: 80px !important;
+      margin: auto;
+      max-height: 60px !important;
       width: auto;
       padding-top: 3px !important;
       padding-bottom: 3px !important;
diff --git a/api/urls.py b/api/urls.py
index eb2a263..f29aa30 100644
--- a/api/urls.py
+++ b/api/urls.py
@@ -22,6 +22,7 @@ from rest_framework_simplejwt.views import (
 )
 from .views import MasterKeyTokenObtainView
 from .views import ConversationView
+from .views import ConversationActions
 from rest_framework.routers import DefaultRouter
 
 router = DefaultRouter()
@@ -30,5 +31,6 @@ urlpatterns = [
     path('token/', MasterKeyTokenObtainView.as_view(), name='token_obtain_pair'),
     path('token/refresh/', TokenRefreshView.as_view(), name='token_refresh'),
     path('token/verify/', TokenVerifyView.as_view(), name='token_verify'),
+    path('conversations/<int:id>/prompt/', ConversationActions.as_view()),
     path('', include(router.urls)),
 ]
\ No newline at end of file
diff --git a/api/utils.py b/api/utils.py
new file mode 100644
index 0000000..b29e94e
--- /dev/null
+++ b/api/utils.py
@@ -0,0 +1,22 @@
+from drf_yasg.utils import swagger_auto_schema
+
+def add_swagger_summaries(viewset_class):
+    """Add standard swagger summaries to a ModelViewSet class"""
+    model_name = viewset_class.serializer_class.Meta.model.__name__.lower()
+    
+    # Apply decorators to the class methods
+    for action, template in {
+        'list': f"List all {model_name}s",
+        'create': f"Create new {model_name}",
+        'retrieve': f"Get specific {model_name}",
+        'update': f"Update {model_name} completely",
+        'partial_update': f"Update {model_name} partially",
+        'destroy': f"Delete {model_name}",
+    }.items():
+        if hasattr(viewset_class, action):
+            method = getattr(viewset_class, action)
+            if not hasattr(method, '_swagger_auto_schema'):
+                setattr(viewset_class, action, 
+                        swagger_auto_schema(operation_summary=template)(method))
+    
+    return viewset_class
diff --git a/api/views.py b/api/views.py
index fb3e369..c4bcc06 100644
--- a/api/views.py
+++ b/api/views.py
@@ -1,6 +1,8 @@
 from drf_yasg.utils import swagger_auto_schema
 from drf_yasg import openapi
 
+import base64
+import json
 from rest_framework_simplejwt.views import TokenObtainPairView
 from rest_framework_simplejwt.tokens import RefreshToken
 from django.views.decorators.http import require_http_methods
@@ -11,15 +13,24 @@ from django.utils import timezone
 from django.contrib.auth import get_user_model
 from .models import MasterKey
 from .models import Conversation
+from django.http import StreamingHttpResponse
 from .models import Message
 from .serializers import ConversationSerializer, MessageSerializer
 from rest_framework import viewsets, status
 from rest_framework.decorators import action
 from rest_framework.response import Response
+from .utils import add_swagger_summaries
+from django.views import View
 import ollama
+from asgiref.sync import sync_to_async
+from django.utils.decorators import method_decorator
+import asyncio
+
+from django.views.decorators.csrf import csrf_exempt
 
 User = get_user_model()
 
+@add_swagger_summaries
 class ConversationView(viewsets.ModelViewSet):
 
     queryset = Conversation.objects.all()
@@ -36,7 +47,8 @@ class ConversationView(viewsets.ModelViewSet):
 
     @swagger_auto_schema(
         method='get',
-        operation_description="Get messages of the conversation",
+        operation_description="Get all the messages of the conversation",
+        operation_summary="Get all messages",
         responses={
             200: openapi.Response('List of messages', MessageSerializer(many=True)),
             400: 'Bad Request',
@@ -57,9 +69,11 @@ class ConversationView(viewsets.ModelViewSet):
         messages = conversation.messages.all()
         return Response(data=list(messages.values()))
 
+@method_decorator(csrf_exempt, name='dispatch')
+class ConversationActions(View):
     @swagger_auto_schema(
-        method='post',
         operation_description="Discutes with the ai",
+        operation_summary="Make a new prompt",
         request_body=openapi.Schema(
             type=openapi.TYPE_OBJECT,
             properties={
@@ -72,43 +86,58 @@ class ConversationView(viewsets.ModelViewSet):
             400: 'Bad Request',
         }
     )
-    @action(detail=True, methods=['post'])
-    def prompt(self, request, pk=None):
-        conversation = self.get_object()
-        messages = [{
-            "role": "system",
-            "content": """
-                You must strictly refuse to engage with ANY of the following:
-                1. Violence or harm (even fictional or hypothetical scenarios)
-                2. ANY explicit, suggestive, or romantic content
-                3. Controversial political topics
-                4. ANY content that could potentially be misused
-                5. Medical, legal, or financial advice
-                6. Personal information or privacy violations
-                7. Anything that could be remotely offensive to anyone
+    @sync_to_async
+    def post(self, request, *args, **kwargs):
+        conversation = Conversation.objects.get(pk=self.kwargs['id'])
+        data = json.loads(request.body)
+        messages = []#{
+        #     "role": "system",
+        #     "content": """
+        #         You must strictly refuse to engage with ANY of the following categories:
+        #         1. Violence or harm (even fictional or hypothetical scenarios)
+        #         2. ANY explicit, suggestive, or romantic content
+        #         3. Controversial political topics
+        #         4. ANY content that could potentially be illegal
+        #         5. Medical, legal, or financial advice
+        #         6. Personal information or privacy violations
+        #         7. Anything that could be remotely offensive to anyone
                 
-                If you detect such content, immediately respond with:
-                "I cannot assist with that request as it appears to be inappropriate. I'm designed to be helpful, but within strict ethical boundaries. Is there something else I can help you with?"
-            """
-        }]
+        #         If you detect such content, immediately respond with:
+        #         "I cannot assist with that request as it appears to be inappropriate. I'm designed to be helpful, but within strict ethical boundaries. Is there something else I can help you with?"
+        #     """
+        # }]
         for message in conversation.messages.all():
             if message:
                 messages.append({
                     "role": message.role,
                     "content": message.content
-                })
+                }) 
         messages.append({
             "role": "user",
-            "content": request.data.get("content", "")
+            "content": data['content'] or ""
         })
-        res = ollama.chat(model="gemma3", messages=messages)
-        Message(role="user", content=request.data.get("content", ""), conversation=conversation).save()
-        Message(role="assistant", content=res['message']['content'], conversation=conversation).save()
-        return Response(data={
-            "content": res['message']['content']
-        })
-
+        Message(role="user", content=data['content'] or "", conversation=conversation).save()
+        ai_message = Message(role="assistant", content="", conversation=conversation)
 
+        @sync_to_async
+        def save_message():
+            ai_message.save()
+            
+        async def chat_event_stream():
+            message = ""
+            try: 
+                stream = ollama.chat(model="llama2-uncensored", messages=messages, stream=True)
+                for chunk in stream:
+                    message = chunk['message']['content']
+                    #print(message, base64.b64encode(message.encode("utf-8")))
+                    ai_message.content += message
+                    yield f"{base64.b64encode(message.encode("utf-8")).decode("utf-8")}"
+            finally: 
+                await save_message()
+        response = StreamingHttpResponse(chat_event_stream(), content_type='text/event-stream')
+        response['Cache-Control'] = 'no-cache'
+        response['X-Accel-Buffering'] = 'no'
+        return response