From 4c732623aac3007388b87ff76462f6cc4dd9bd28 Mon Sep 17 00:00:00 2001
From: Chris Clark <chris@untrod.com>
Date: Thu, 11 Jan 2024 11:25:25 -0500
Subject: [PATCH] prompt logs

---
 explorer/assistant/models.py          | 21 +++++++++++++++++++
 explorer/assistant/utils.py           |  1 -
 explorer/assistant/views.py           | 27 +++++++++++++++++++++----
 explorer/migrations/0014_promptlog.py | 29 +++++++++++++++++++++++++++
 explorer/models.py                    |  1 +
 5 files changed, 74 insertions(+), 5 deletions(-)
 create mode 100644 explorer/assistant/models.py
 create mode 100644 explorer/migrations/0014_promptlog.py

diff --git a/explorer/assistant/models.py b/explorer/assistant/models.py
new file mode 100644
index 00000000..19b4a565
--- /dev/null
+++ b/explorer/assistant/models.py
@@ -0,0 +1,21 @@
+from django.db import models
+from django.conf import settings
+
+
+class PromptLog(models.Model):
+
+    class Meta:
+        app_label = 'explorer'
+
+    prompt = models.TextField(blank=True)
+    response = models.TextField(blank=True)
+    run_by_user = models.ForeignKey(
+        settings.AUTH_USER_MODEL,
+        null=True,
+        blank=True,
+        on_delete=models.CASCADE
+    )
+    run_at = models.DateTimeField(auto_now_add=True)
+    duration = models.FloatField(blank=True, null=True)  # seconds
+    model = models.CharField(blank=True, max_length=128, default='')
+    error = models.TextField(blank=True, null=True)
diff --git a/explorer/assistant/utils.py b/explorer/assistant/utils.py
index 2ea54637..95f1b451 100644
--- a/explorer/assistant/utils.py
+++ b/explorer/assistant/utils.py
@@ -24,7 +24,6 @@ def do_req(prompt, history=None):
         {"role": "system", "content": prompt['system']},
         {"role": "user", "content": prompt['user']},
     ]
-    print(messages)
     resp = openai_client.chat.completions.create(
       model=OPENAI_MODEL,
       messages=messages
diff --git a/explorer/assistant/views.py b/explorer/assistant/views.py
index 7c9d4137..7543a5dd 100644
--- a/explorer/assistant/views.py
+++ b/explorer/assistant/views.py
@@ -1,7 +1,9 @@
 from django.http import JsonResponse
+from django.utils import timezone
 import json
 
 from sql_metadata import Parser
+from .models import PromptLog
 from explorer.schema import schema_info
 from explorer.assistant.prompts import PROMPT_MAP
 from explorer.assistant.utils import format_prompt, do_req, extract_response
@@ -12,7 +14,7 @@ def tables_from_schema_info(connection, table_names):
     return [table for table in schema if table[0] in table_names]
 
 
-def do_modify_query(request_data):
+def do_modify_query(request_data, user):
     sql = request_data['sql']
     connection = request_data['connection']
     assistant_request = request_data['assistant_request']
@@ -24,8 +26,25 @@ def do_modify_query(request_data):
                                  'description': assistant_request,
                                  'table_structure': table_struct
                             })
-    resp = do_req(prompt)
-    return extract_response(resp)
+    start = timezone.now()
+    pl = PromptLog(
+        prompt=prompt,
+        run_by_user=user,
+        run_at=timezone.now(),
+    )
+    response_text = None
+    try:
+        resp = do_req(prompt)
+        response_text = extract_response(resp)
+        pl.response = response_text
+    except Exception as e:
+        pl.error = str(e)
+    finally:
+        stop = timezone.now()
+        pl.duration = (stop - start).total_seconds()
+        pl.save()
+    return response_text
+
 
 
 def assistant_help(request):
@@ -35,7 +54,7 @@ def assistant_help(request):
             assistant_function = data['assistant_function']
 
             if assistant_function == 'modify_query':
-                resp = do_modify_query(data)
+                resp = do_modify_query(data, request.user)
             else:
                 resp = None
 
diff --git a/explorer/migrations/0014_promptlog.py b/explorer/migrations/0014_promptlog.py
new file mode 100644
index 00000000..38795e34
--- /dev/null
+++ b/explorer/migrations/0014_promptlog.py
@@ -0,0 +1,29 @@
+# Generated by Django 4.2.8 on 2024-01-11 08:22
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('explorer', '0013_querylog_error_querylog_success'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='PromptLog',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('prompt', models.TextField(blank=True)),
+                ('response', models.TextField(blank=True)),
+                ('run_at', models.DateTimeField(auto_now_add=True)),
+                ('duration', models.FloatField(blank=True, null=True)),
+                ('model', models.CharField(blank=True, default='', max_length=128)),
+                ('error', models.TextField(blank=True, null=True)),
+                ('run_by_user', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL)),
+            ],
+        ),
+    ]
diff --git a/explorer/models.py b/explorer/models.py
index 7a93ca8b..8856f466 100644
--- a/explorer/models.py
+++ b/explorer/models.py
@@ -7,6 +7,7 @@
 from django.urls import reverse
 from django.utils.translation import gettext_lazy as _
 
+from assistant.models import PromptLog
 from explorer import app_settings
 from explorer.utils import (
     extract_params, get_params_for_url, get_s3_bucket, get_valid_connection, passes_blacklist, s3_url,