Upload app.py

app.py

@@ -1,7 +1,8 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import pipeline
-import mysql.connector
+import mysql.connector 
+import json
 import os
 from dotenv import load_dotenv
 
@@ -69,20 +70,20 @@ def home():
 @app.post("/query")
 def handle_query(request: QueryRequest):
     try:
+        print("Received query:", request.query)  # Debugging: Print the received query
         text = request.query
         
         # Fetch the database schema
         schema = get_database_schema()
-        schema_str = "\n".join([f"{table}: {', '.join(columns)}" for table, columns in schema.items()])
+        schema_str = json.dumps(schema, indent=4)
+        print("Fetched schema:", schema)  # Debugging: Print the fetched schema
         
-        # Construct the system message
         system_message = f"""
-    You are a helpful, cheerful database assistant. 
-    Use the following dynamically retrieved database schema when creating your answers:
+        You are a helpful, cheerful database assistant. 
+        Use the following dynamically retrieved database schema when creating your answers:
 
-    {schema_str}
-
-    When creating your answers, consider the following:
+        {schema_str}
+         When creating your answers, consider the following:
 
     1. If a query involves a column or value that is not present in the provided database schema, correct it and mention the correction in the summary. If a column or value is missing, provide an explanation of the issue and adjust the query accordingly.
     2. If there is a spelling mistake in the column name or value, attempt to correct it by matching the closest possible column or value from the schema. Mention this correction in the summary to clarify any changes made.
@@ -97,19 +98,18 @@ def handle_query(request: QueryRequest):
     In the preceding JSON response, substitute "your-query" with a MariaDB query to retrieve the requested data.
     In the preceding JSON response, substitute "your-summary" with a summary of the query and any corrections or clarifications made.
     Always include all columns in the table.
-    """
-        
+        """
+
         prompt = f"{system_message}\n\nUser request:\n\n{text}\n\nSQL query:"
         output = pipe(prompt, max_new_tokens=100)
-        
+        print("Generated output:", output)  # Debugging: Print the generated output
+
         generated_text = output[0]['generated_text']
         sql_query = generated_text.split("SQL query:")[-1].strip()
         
-        # Basic validation
         if not sql_query.lower().startswith(('select', 'show', 'describe')):
             raise ValueError("Generated text is not a valid SQL query")
-        
-        # Example: execute the generated SQL query and return the results
+
         conn = get_db_connection()
         cursor = conn.cursor()
         cursor.execute(sql_query)
@@ -120,8 +120,10 @@ def handle_query(request: QueryRequest):
 
         return {"sql": sql_query, "results": results}
     except Exception as e:
+        print("Error occurred:", str(e))  # Debugging: Print the error
         raise HTTPException(status_code=500, detail=str(e))
 
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)