Stability Enhancements and Bug Fixes (#510)

* Stability Enhancements and Bug Fixes

* Fix ai name removal

ayushma/migrations/0052_document_failed.py

@@ -0,0 +1,17 @@
+# Generated by Django 4.2.6 on 2024-04-17 13:56
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("ayushma", "0051_project_tts_engine"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="document",
+            name="failed",
+            field=models.BooleanField(default=False),
+        ),
+    ]

ayushma/models/document.py

@@ -19,6 +19,7 @@ class Document(BaseModel):
         "TestQuestion", on_delete=models.PROTECT, null=True, blank=True
     )
     uploading = models.BooleanField(default=True)
+    failed = models.BooleanField(default=False)
 
     def __str__(self) -> str:
         return f"{self.title} in {self.project.title}"

ayushma/serializers/users.py

@@ -16,6 +16,7 @@ class Meta:
             "allow_key",
             "is_staff",
             "is_reviewer",
+            "date_joined",
         )
 
 
@@ -56,6 +57,7 @@ class Meta:
             "is_staff",
             "is_reviewer",
             "password",
+            "date_joined",
         )
         read_only_fields = (
             "external_id",
@@ -64,6 +66,7 @@ class Meta:
             "allow_key",
             "is_staff",
             "is_reviewer",
+            "date_joined",
         )
 
     def update(self, instance, validated_data):

ayushma/tasks/__init__.py

@@ -0,0 +1,18 @@
+from celery import current_app
+from celery.schedules import crontab
+
+from ayushma.tasks.stale_cleanup import clean_stale_test_runs, clean_stale_upsert_doc
+
+
+@current_app.on_after_finalize.connect
+def setup_periodic_tasks(sender, **kwargs):
+    sender.add_periodic_task(
+        crontab(minute="*/30"),  # Every 30 minutes
+        clean_stale_test_runs.s(),
+        name="clean_stale_test_runs",
+    )
+    sender.add_periodic_task(
+        crontab(minute="*/30"),  # Every 30 minutes
+        clean_stale_upsert_doc.s(),
+        name="clean_stale_upsert_doc",
+    )

ayushma/tasks/stale_cleanup.py

@@ -0,0 +1,52 @@
+from datetime import timedelta
+
+from celery import shared_task
+from django.utils.timezone import now
+
+from ayushma.models.document import Document
+from ayushma.models.enums import StatusChoices
+from ayushma.models.testsuite import TestRun
+
+
+@shared_task(bind=True)
+def clean_stale_test_runs(self):
+    try:
+        # Get testRuns that are created over 6 hours ago and are still in RUNNING state
+        test_runs = TestRun.objects.filter(
+            created_at__lt=now() - timedelta(hours=6),
+            status=StatusChoices.RUNNING,
+        )
+
+        # Cancel the testRuns
+        for test_run in test_runs:
+            print(
+                f"Cleaning stale test run: {test_run.id}; Created at: {test_run.created_at}"
+            )
+            test_run.status = StatusChoices.FAILED
+            test_run.save()
+    except Exception as e:
+        print(f"Error occurred while cleaning stale test runs: {e}")
+        raise e
+
+
+@shared_task(bind=True)
+def clean_stale_upsert_doc(self):
+    try:
+        # Get stale Document objects that are still in UPLOADING state after 6 hours
+        documents = Document.objects.filter(
+            created_at__lt=now() - timedelta(hours=6),
+            uploading=True,
+        )
+
+        # Set the documents to failed state
+        for document in documents:
+            print(
+                f"Cleaning stale document: {document.id}; Created at: {document.created_at}"
+            )
+            document.failed = True
+            document.uploading = False
+            document.save()
+
+    except Exception as e:
+        print(f"Error occurred while cleaning stale test runs: {e}")
+        raise e

ayushma/utils/converse.py

@@ -47,7 +47,7 @@ def converse_api(
         if not open_ai_key:
             open_ai_key = (
                 request.headers.get("OpenAI-Key")
-                or (chat.project and chat.project.openai_key)
+                or (chat.project and chat.project.open_ai_key)
                 or (user.allow_key and settings.OPENAI_API_KEY)
             )
         noonce = request.data.get("noonce")

ayushma/utils/openaiapi.py

@@ -345,7 +345,7 @@ def converse(
         response = lang_chain_helper.get_response(
             english_text, reference, chat_history, documents
         )
-        chat_response = response.replace("Ayushma: ", "")
+        chat_response = response.replace(f"{AI_NAME}:", "").lstrip()
         stats["response_end_time"] = time.time()
         translated_chat_response, url, chat_message = handle_post_response(
             chat_response,
@@ -391,8 +391,6 @@ def converse(
                     documents,
                 )
                 chat_response = ""
-                skip_token = len(f"{AI_NAME}: ")
-
                 while True:
                     if token_queue.empty():
                         continue
@@ -427,10 +425,9 @@ def converse(
                             ayushma_voice=url,
                         )
                         break
-                    if skip_token > 0:
-                        skip_token -= 1
-                        continue
+
                     chat_response += next_token[0]
+                    chat_response = chat_response.replace(f"{AI_NAME}:", "").lstrip()
                     yield create_json_response(
                         local_translated_text,
                         chat.external_id,

ayushma/utils/speech_to_text.py

@@ -79,7 +79,7 @@ def recognize(self, audio):
                 return ""
             response = response.json()
             return response["data"]["transcription"].strip()
-        except Exception as e:
+        except Exception:
             raise ValueError(
                 "[Speech to Text] Failed to recognize speech with Self Hosted engine"
             )
@@ -99,16 +99,17 @@ def speech_to_text(engine_id, audio, language_code):
     engine_class = engines.get(engine_name)
 
     if not engine_class:
-        raise ValueError(f"[Speech to Text] Engine with ID {engine_id} not found")
+        print(f"Invalid Speech to Text engine: {engine_name}")
+        raise ValueError("The selected Speech to Text engine is not valid")
 
     try:
         engine = engine_class(api_key, language_code)
         recognized_text = engine.recognize(audio)
         if not recognized_text:
-            raise ValueError("[Speech to Text] No text recognized")
+            raise ValueError("No text recognized in the audio")
         return recognized_text
     except Exception as e:
         print(f"Failed to transcribe speech with {engine_name} engine: {e}")
         raise ValueError(
-            f"[Speech to Text] Failed to transcribe speech with {engine_name} engine"
+            f"[Speech to Text] Failed to transcribe speech with {engine_name} engine: {e}"
         )

ayushma/views/chat.py

@@ -16,6 +16,8 @@
 from rest_framework.response import Response
 
 from ayushma.models import Chat, ChatFeedback, Project
+from ayushma.models.chat import ChatMessage
+from ayushma.models.enums import ChatMessageType
 from ayushma.permissions import IsTempTokenOrAuthenticated
 from ayushma.serializers import (
     ChatDetailSerializer,
@@ -131,11 +133,25 @@ def speech_to_text(self, *args, **kwarg):
             stats["transcript_end_time"] = time.time()
             translated_text = transcript
         except Exception as e:
-            print(f"Failed to transcribe speech with {stt_engine} engine: {e}")
+            print(f"Failed to transcribe speech with {stt_engine} engine:\n{e}")
+
+            error_msg = (
+                f"[Transcribing] Something went wrong in getting transcription.\n{e}"
+            )
+            chat = Chat.objects.get(external_id=kwarg["external_id"])
+            chat.title = "Transcription Error"
+            chat.save()
+            ChatMessage.objects.create(
+                message=error_msg,
+                original_message=error_msg,
+                chat=chat,
+                messageType=ChatMessageType.SYSTEM,
+                language=language,
+                meta={},
+            )
+
             return Response(
-                {
-                    "error": "[Transcribing] Something went wrong in getting transcription, please try again later"
-                },
+                {"error": error_msg},
                 status=status.HTTP_500_INTERNAL_SERVER_ERROR,
             )
 

ayushma/views/project.py

@@ -26,8 +26,12 @@ class ProjectViewSet(
     CreateModelMixin,
     DestroyModelMixin,
 ):
-    queryset = Project.objects.all()
-    filter_backends = (filters.SearchFilter, DjangoFilterBackend)
+    queryset = Project.objects.all().order_by("-is_default")
+    filter_backends = (
+        filters.SearchFilter,
+        DjangoFilterBackend,
+        filters.OrderingFilter,
+    )
     search_fields = ("title",)
     filterset_fields = ("archived",)
     serializer_class = ProjectSerializer
@@ -56,7 +60,7 @@ def get_queryset(self):
         if self.action == "list":
             if not self.request.user.is_staff:
                 queryset = self.queryset.filter(is_default=True)
-        return queryset
+        return queryset.order_by("-is_default")
 
     def perform_create(self, serializer):
         serializer.save(creator=self.request.user)

ayushma/views/users.py

@@ -1,8 +1,7 @@
 from django_filters.rest_framework import DjangoFilterBackend
 from drf_spectacular.utils import extend_schema
-from rest_framework import permissions
+from rest_framework import filters, permissions
 from rest_framework.decorators import action
-from rest_framework.filters import SearchFilter
 
 from ayushma.models import User
 from ayushma.serializers.users import (
@@ -15,7 +14,12 @@
 
 
 class UserViewSet(FullBaseModelViewSet):
-    queryset = User.objects.all()
+    queryset = User.objects.all().order_by("-date_joined")
+    filter_backends = (
+        filters.SearchFilter,
+        DjangoFilterBackend,
+        filters.OrderingFilter,
+    )
     serializer_class = UserDetailSerializer
     permission_classes = (permissions.IsAdminUser,)
     serializer_action_classes = {
@@ -29,7 +33,6 @@ class UserViewSet(FullBaseModelViewSet):
         "partial_update_me": (permissions.IsAuthenticated(),),
     }
     lookup_field = "username"
-    filter_backends = [SearchFilter, DjangoFilterBackend]
     search_fields = ["full_name"]
     filterset_fields = ["is_staff", "is_reviewer", "allow_key"]