Upload app.py

app.py

@@ -14,7 +14,7 @@ app = FastAPI()
 pipe = pipeline("text-generation", model="defog/llama-3-sqlcoder-8b", pad_token_id=2)
 
 class QueryRequest(BaseModel):
-    text: str
+    query: str
 
 def get_db_connection():
     """Create a new database connection."""
@@ -66,13 +66,10 @@ def get_database_schema():
 def home():
     return {"message": "SQL Generation Server is running"}
 
-@app.post("/generate")
-def generate(request: QueryRequest):
+@app.post("/query")
+def handle_query(request: QueryRequest):
     try:
-        # Log the incoming request text for debugging
-        print(f"Received request with text: {request.text}")
-        
-        text = request.text
+        text = request.query
         
         # Fetch the database schema
         schema = get_database_schema()
@@ -80,13 +77,27 @@ def generate(request: QueryRequest):
         
         # 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}
+    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.
+    3. Ensure that the correct columns and values are used based on the schema provided. Verify the query against the schema to confirm accuracy.
+    4. Include column name headers in the query results for clarity.
 
-        [Additional instructions as in your original code]
-        """
+    Always provide your answer in the JSON format below:
+
+    {{ "summary": "your-summary", "query":  "your-query" }}
+    
+    Output ONLY JSON.
+    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)
@@ -98,11 +109,19 @@ def generate(request: QueryRequest):
         if not sql_query.lower().startswith(('select', 'show', 'describe')):
             raise ValueError("Generated text is not a valid SQL query")
         
-        return {"output": sql_query}
+        # Example: execute the generated SQL query and return the results
+        conn = get_db_connection()
+        cursor = conn.cursor()
+        cursor.execute(sql_query)
+        results = cursor.fetchall()
+
+        cursor.close()
+        conn.close()
+
+        return {"sql": sql_query, "results": results}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)