(DRAFT) Text to sql

chat/models.py

@@ -21,3 +21,11 @@ class Message(models.Model):
     deleted = models.BooleanField(default=False)
     tokens = models.IntegerField(default=0)
     created_at = models.DateTimeField(auto_now_add=True)
+
+
+class SQLQuery(models.Model):
+    chat = models.ForeignKey(Chat, on_delete=models.CASCADE, related_name='sql_queries')
+    question = models.TextField()
+    sql_query = models.TextField()
+    results = models.JSONField()
+    created_at = models.DateTimeField(auto_now_add=True)

chat/utils.py

@@ -4,6 +4,45 @@
 import google.generativeai as genai
 
 import markdown
+import psycopg2
+from psycopg2.extras import RealDictCursor
+
+
+def natural_language_to_sql(question, db_schema):
+    GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY', "get_one_from_google")
+    genai.configure(api_key=GOOGLE_API_KEY)
+    model = genai.GenerativeModel('gemini-1.5-pro')
+
+    prompt = f"""
+    Given the following database schema:
+    {db_schema}
+
+    Convert the following natural language question into a SQL query:
+    "{question}"
+
+    Return only the SQL query, without any additional explanation.
+    """
+
+    response = model.generate_content(prompt)
+    return response.text.strip()
+
+
+def execute_sql_query(query):
+    # Use your database connection parameters
+    conn = psycopg2.connect(
+        dbname=os.getenv('DB_NAME'),
+        user=os.getenv('DB_USER'),
+        password=os.getenv('DB_PASSWORD'),
+        host=os.getenv('DB_HOST'),
+        port=os.getenv('DB_PORT')
+    )
+
+    with conn.cursor(cursor_factory=RealDictCursor) as cur:
+        cur.execute(query)
+        results = cur.fetchall()
+
+    conn.close()
+    return results
 
 
 def run_llm(chat_history, context):

chat/views.py

@@ -8,7 +8,7 @@
 from django.db.models import Q, F
 
 
-from .models import Chat, Message
+from .models import Chat, Message, SQLQuery
 from .utils import to_markdown, count_tokens
 from django.conf import settings
 
@@ -68,4 +68,7 @@ def get_context_data(self, **kwargs):
     context["size_count"] = size_count_tokens
     context["size_count_percentage"] = (size_count_tokens/1000000)*100
 
+    sql_queries = SQLQuery.objects.filter(chat=context["chat"]).order_by("-created_at")
+    context["sql_queries"] = sql_queries
+
     return context

chat/viewsets.py

@@ -7,7 +7,7 @@
 
 from chat.models import Chat, Message
 from chat.serializers import ChatSerializer
-from .utils import run_llm, to_markdown
+from .utils import run_llm, to_markdown, natural_language_to_sql, execute_sql_query
 
 
 class ChatViewSet(viewsets.GenericViewSet):
@@ -61,6 +61,27 @@ def star_message(self, request, pk=None, msg_id=None):
         message.save()
         return Response({"message": f"Message starred: {message.starred}"})
 
+    @action(detail=False, methods=['POST'])
+    def sql_query(self, request):
+        question = request.data.get('question')
+        db_schema = """
+            # Insert your database schema here
+            # For example:
+            # Table: sales
+            # Columns: id, date, amount, product_id
+            # Table: products
+            # Columns: id, name, price
+            """
+
+        sql_query = natural_language_to_sql(question, db_schema)
+        results = execute_sql_query(sql_query)
+
+        return Response({
+            'question': question,
+            'sql_query': sql_query,
+            'results': results
+        })
+
 
 def _build_summary(context):
     mes_history = [{"role": "user", "parts": [context]}]