Merge pull request #76 from toggle-corp/feature/chat-session-management

Update chat session and chat message

chat/admin.py

@@ -0,0 +1,16 @@
+from django.contrib import admin
+
+from .models import UserChatMessage, UserChatSession
+
+# Register your models here.
+
+
+@admin.register(UserChatSession)
+class UserChatSessionAdmin(admin.ModelAdmin):
+    list_display = ["user_uuid", "platform"]
+
+
+@admin.register(UserChatMessage)
+class UserChatMessageAdmin(admin.ModelAdmin):
+    list_display = ["session", "type", "query", "response", "status"]
+    list_filter = ["status", "type", "session__platform"]

chat/apps.py

@@ -0,0 +1,6 @@
+from django.apps import AppConfig
+
+
+class ChatConfig(AppConfig):
+    default_auto_field = "django.db.models.BigAutoField"
+    name = "chat"

chat/migrations/0001_initial.py

@@ -0,0 +1,45 @@
+# Generated by Django 5.1.4 on 2025-01-17 08:17
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = []
+
+    operations = [
+        migrations.CreateModel(
+            name="UserChatSession",
+            fields=[
+                ("id", models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")),
+                ("user_uuid", models.UUIDField()),
+                ("created_at", models.DateTimeField(auto_now_add=True)),
+                (
+                    "platform",
+                    models.IntegerField(
+                        choices=[(0, "Whatsapp"), (1, "Facebook"), (2, "Streamlit"), (3, "Website"), (4, "Automation")],
+                        default=2,
+                    ),
+                ),
+            ],
+        ),
+        migrations.CreateModel(
+            name="UserChatMessage",
+            fields=[
+                ("id", models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")),
+                ("type", models.IntegerField(choices=[(0, "Bot"), (1, "User")], default=1)),
+                ("question", models.TextField()),
+                (
+                    "status",
+                    models.IntegerField(choices=[(0, "Pending"), (1, "Started"), (2, "Success"), (3, "Failed")], default=0),
+                ),
+                ("answer", models.TextField(blank=True, null=True)),
+                ("created_at", models.DateTimeField(auto_now_add=True)),
+                ("modified_at", models.DateTimeField(auto_now=True)),
+                ("session", models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to="chat.userchatsession")),
+            ],
+        ),
+    ]

chat/migrations/0002_rename_question_userchatmessage_query_and_more.py

@@ -0,0 +1,23 @@
+# Generated by Django 5.1.4 on 2025-01-17 09:39
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("chat", "0001_initial"),
+    ]
+
+    operations = [
+        migrations.RenameField(
+            model_name="userchatmessage",
+            old_name="question",
+            new_name="query",
+        ),
+        migrations.RenameField(
+            model_name="userchatmessage",
+            old_name="answer",
+            new_name="response",
+        ),
+    ]

chat/models.py

@@ -0,0 +1,41 @@
+from django.db import models
+
+# Create your models here.
+
+
+class UserChatSession(models.Model):
+    class Platform(models.IntegerChoices):
+        WHATSAPP = 0, "Whatsapp"
+        FACEBOOK = 1, "Facebook"
+        STREAMLIT = 2, "Streamlit"
+        WEBSITE = 3, "Website"
+        AUTOMATION = 4, "Automation"  # Used for QA team
+
+    user_uuid = (
+        models.UUIDField()
+    )  # TODO We need to retrieve the user ID from the user model. Currently, we are using a UUID generated by Streamlit.
+    created_at = models.DateTimeField(auto_now_add=True)
+    platform = models.IntegerField(choices=Platform.choices, default=Platform.STREAMLIT)
+
+    def __str__(self):
+        return str(self.user_uuid)
+
+
+class UserChatMessage(models.Model):
+    class MessageType(models.IntegerChoices):
+        BOT = 0, "Bot"
+        USER = 1, "User"
+
+    class Status(models.IntegerChoices):
+        PENDING = 0, "Pending"
+        STARTED = 1, "Started"
+        SUCCESS = 2, "Success"
+        FAILED = 3, "Failed"
+
+    session = models.ForeignKey(UserChatSession, on_delete=models.CASCADE)
+    type = models.IntegerField(choices=MessageType.choices, default=MessageType.USER)
+    query = models.TextField()
+    status = models.IntegerField(choices=Status.choices, default=Status.PENDING)
+    response = models.TextField(null=True, blank=True)
+    created_at = models.DateTimeField(auto_now_add=True)
+    modified_at = models.DateTimeField(auto_now=True)

chat/tests.py

@@ -0,0 +1 @@
+# Create your tests here.

chat/views.py

@@ -0,0 +1,21 @@
+from chat.models import UserChatMessage, UserChatSession
+
+
+class UserSession:
+    def create_chat_session(self, data):
+        obj, _ = UserChatSession.objects.get_or_create(
+            user_uuid=data.get("user_id"), defaults={"platform": data.get("platform")}
+        )
+        return obj
+
+    def create_chat_message(self, data, chat_session):
+        user_chat_message = UserChatMessage.objects.create(
+            session=chat_session, query=data.get("query"), status=UserChatMessage.Status.STARTED
+        )
+        return user_chat_message
+
+    def update_chat_message(self, response, user_message):
+        user_message.response = response
+        user_message.status = UserChatMessage.Status.SUCCESS
+        user_message.save(update_fields=["response", "status"])
+        return user_message

content/serializers.py

@@ -1,6 +1,9 @@
 from rest_framework import serializers
 
+from chat.models import UserChatSession
+
 
 class UserQuerySerializer(serializers.Serializer):
     query = serializers.CharField(required=True, allow_null=False, allow_blank=False)
     user_id = serializers.UUIDField(required=True)
+    platform = serializers.IntegerField(default=UserChatSession.Platform.STREAMLIT)

content/views.py

@@ -5,6 +5,7 @@
 from rest_framework.generics import GenericAPIView
 from rest_framework.response import Response
 
+from chat.views import UserSession
 from chatbotcore.llm import OllamaHandler, OpenAIHandler
 from chatbotcore.utils import LLMType
 from content.serializers import UserQuerySerializer
@@ -17,10 +18,13 @@ class UserQuery(GenericAPIView):
         llm = OpenAIHandler()
     else:
         raise Exception("Wrong LLM Type")
+    Session = UserSession()
 
     def post(self, request, *arg, **kwargs):
         serializer = UserQuerySerializer(data=request.data)
-        if serializer.is_valid():
-            result = asyncio.run(self.llm.execute_chain(request.data["user_id"], request.data["query"]))
-            return Response(result)
-        return Response(serializer.errors, 422)
+        serializer.is_valid(raise_exception=True)
+        user_session = self.Session.create_chat_session(request.data)
+        user_message = self.Session.create_chat_message(request.data, user_session)
+        result = asyncio.run(self.llm.execute_chain(request.data["user_id"], request.data["query"]))
+        self.Session.update_chat_message(result, user_message)
+        return Response(result)

main/settings.py

@@ -151,6 +151,7 @@
     "content",
     "organization",
     "rest_framework",
+    "chat",
 ]
 
 MIDDLEWARE = [