added fetch references field (#359)

ohcnetwork · Oct 10, 2023 · 7408a64 · 7408a64
1 parent 567993d
commit 7408a64
Show file tree

Hide file tree

Showing 6 changed files with 46 additions and 1 deletion.
diff --git a/ayushma/admin.py b/ayushma/admin.py
@@ -4,6 +4,7 @@
 from simple_history.admin import SimpleHistoryAdmin
 
 from ayushma.models.services import Service, TempToken
+from ayushma.models.testsuite import Feedback, TestQuestion, TestRun, TestSuite
 
 from .models import APIKey, Chat, ChatMessage, Document, Project, User
 
@@ -102,3 +103,26 @@ class TempTokenAdmin(DjangoQLSearchMixin, admin.ModelAdmin):
 @admin.register(Service)
 class ServiceAdmin(DjangoQLSearchMixin, admin.ModelAdmin):
     pass
+
+
+@admin.register(Feedback)
+class FeedbackAdmin(DjangoQLSearchMixin, admin.ModelAdmin):
+    pass
+
+
+@admin.register(TestSuite)
+class TestSuiteAdmin(DjangoQLSearchMixin, admin.ModelAdmin):
+    pass
+
+
+@admin.register(TestQuestion)
+class TestQuestionAdmin(DjangoQLSearchMixin, admin.ModelAdmin):
+    pass
+
+
+@admin.register(TestRun)
+class TestRunAdmin(DjangoQLSearchMixin, admin.ModelAdmin):
+    pass
+
+
+admin.site.site_header = "Ayushma Admin"
diff --git a/ayushma/migrations/0045_testrun_references.py b/ayushma/migrations/0045_testrun_references.py
@@ -0,0 +1,17 @@
+# Generated by Django 4.2.1 on 2023-10-10 16:31
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("ayushma", "0044_alter_project_stt_engine"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="testrun",
+            name="references",
+            field=models.BooleanField(default=True),
+        ),
+    ]
diff --git a/ayushma/models/testsuite.py b/ayushma/models/testsuite.py
@@ -32,6 +32,7 @@ class TestRun(BaseModel):
     test_suite = ForeignKey(TestSuite, on_delete=CASCADE)
     project = ForeignKey(Project, on_delete=CASCADE)
     status = IntegerField(choices=StatusChoices.choices, default=StatusChoices.RUNNING)
+    references = models.BooleanField(default=True)
 
 
 class TestResult(BaseModel):

diff --git a/ayushma/serializers/testsuite.py b/ayushma/serializers/testsuite.py
@@ -90,6 +90,7 @@ class Meta:
             "modified_at",
             "status",
             "test_results",
+            "references",
         )
         read_only_fields = (
             "external_id",

diff --git a/ayushma/tasks/testrun.py b/ayushma/tasks/testrun.py
@@ -71,6 +71,7 @@ def mark_test_run_as_completed(self, test_run_id):
                         user_language=test_question.language + "-IN",
                         stream=False,
                         generate_audio=False,
+                        fetch_references=test_run.references,
                     )
                 )
 

diff --git a/ayushma/utils/openaiapi.py b/ayushma/utils/openaiapi.py
@@ -265,6 +265,7 @@ def converse(
     references=None,
     generate_audio=True,
     noonce=None,
+    fetch_references=True,
 ):
     if not openai_key:
         raise Exception("OpenAI-Key header is required to create a chat or converse")
@@ -288,7 +289,7 @@ def converse(
 
     if references:
         reference = references
-    elif chat.project and chat.project.external_id:
+    elif fetch_references and chat.project and chat.project.external_id:
         try:
             reference = get_reference(
                 english_text, openai_key, str(chat.project.external_id), match_number