Upload testfile.py

back_end/application/testfile.py

@@ -1,15 +1,7 @@
 import os
 import sqlite3
-import logging
 from groq import Groq
 
-# Configure logging
-logging.basicConfig(
-    filename="healthcare_chatbot.log",
-    level=logging.INFO,
-    format="%(asctime)s - %(levelname)s - %(message)s"
-)
-
 # Load Groq API key from environment variable
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 
@@ -65,11 +57,9 @@ def get_sql_query(natural_language_query):
             model="llama3-8b-8192",
             temperature=0
         )
-        sql_query = completion.choices[0].message.content.strip()
-        logging.info(f"Generated SQL Query: {sql_query}")
-        return sql_query
+        return completion.choices[0].message.content.strip()
+
     except Exception as e:
-        logging.error(f"Error generating SQL query: {str(e)}")
         return f"Error generating SQL query: {str(e)}"
 
 # Function to execute the SQL query
@@ -80,27 +70,29 @@ def execute_sql_query(sql_query):
     database = "veludb.db"  # Ensure this is the correct database
 
     try:
-        logging.info(f"Executing query: {sql_query}")
+        # print("Executing query")
         conn = sqlite3.connect(database)
         cursor = conn.cursor()
         cursor.execute(sql_query)
         rows = cursor.fetchall()
+        # print("Successfully Executing query")
+        # Extract column names
         columns = [desc[0] for desc in cursor.description] if cursor.description else []
+
         conn.close()
 
         if not rows:
-            logging.info("No matching records found.")
             return {"status": "success", "data": "No matching records found."}
 
+        # If the query is a COUNT query, extract the number directly
         if "COUNT" in sql_query.upper():
-            logging.info(f"Count query result: {rows[0][0]}")
             return {"status": "success", "data": f"There are {rows[0][0]} matching records."}
 
+        # Format the data output
         formatted_data = [dict(zip(columns, row)) for row in rows]
-        logging.info(f"Query executed successfully, returned {len(rows)} rows.")
         return {"status": "success", "data": formatted_data}
+
     except Exception as e:
-        logging.error(f"Query execution failed: {str(e)}")
         return {"status": "error", "message": f"Query execution failed: {str(e)}"}
 
 # Function to refine the response using another LLM
@@ -129,23 +121,24 @@ def refine_response(user_query, sql_data):
             model="llama3-8b-8192",
             temperature=0
         )
-        refined_response = completion.choices[0].message.content.strip()
-        logging.info(f"Refined response: {refined_response}")
-        return refined_response
+        return completion.choices[0].message.content.strip()
+
     except Exception as e:
-        logging.error(f"Error refining response: {str(e)}")
         return f"Error refining response: {str(e)}"
 
 # Example Usage
 if __name__ == "__main__":
+    # Example user query
     user_query = "How many male patients are there?"
-    logging.info(f"User Query: {user_query}")
 
+    # Step 1: Convert natural language query to SQL
     sql_query = get_sql_query(user_query)
     print("Generated SQL Query:\n", sql_query)
 
+    # Step 2: Execute SQL query and fetch results
     query_result = execute_sql_query(sql_query)
     print("Query Result:\n", query_result)
 
-    refined_response = refine_response(user_query, query_result["data"])
+    # Step 3: Refine the response using another LLM
+    refined_response = refine_response(user_query, sql_query, query_result["data"])
     print("Refined Response:\n", refined_response)