gemini.py

import psycopg2
import requests
import json
from decimal import Decimal
from dotenv import load_dotenv
import os

# Load environment variables from .env file
load_dotenv()

# Gemini 1.5 Flash API Configuration
API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent"
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
headers = {"Content-Type": "application/json"}

# Database Configuration
DB_CONFIG = {
    "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"),
}

# Global variables to hold database connection
conn = None
cur = None

def get_connection():
    """Get the database connection and cursor."""
    global conn, cur
    if conn is None or cur is None:
        conn = psycopg2.connect(**DB_CONFIG)
        cur = conn.cursor()
    return conn, cur

def close_connection():
    """Close the database connection and cursor."""
    global conn, cur
    if cur:
        cur.close()
    if conn:
        conn.close()

def get_schema():
    """Fetch database schema"""
    try:
        conn, cur = get_connection()

        # Get all tables
        cur.execute("""
            SELECT table_name 
            FROM information_schema.tables 
            WHERE table_schema = 'public'
        """)
        tables = cur.fetchall()

        schema_text = "CREATE TABLE statements:\n\n"

        for table in tables:
            table_name = table[0]
            # Get columns and their details
            cur.execute(f"""
                SELECT 
                    c.column_name, 
                    c.data_type,
                    c.is_nullable,
                    c.column_default,
                    CASE 
                        WHEN tc.constraint_type = 'PRIMARY KEY' THEN 'PRIMARY KEY'
                        ELSE ''
                    END as key_type
                FROM information_schema.columns c
                LEFT JOIN information_schema.table_constraints tc
                    ON c.table_name = tc.table_name
                    AND tc.constraint_type = 'PRIMARY KEY'
                WHERE c.table_name = '{table_name}'
                ORDER BY c.ordinal_position;
            """)
            columns = cur.fetchall()

            # Create CREATE TABLE statement
            create_table = f"CREATE TABLE {table_name} (\n"
            for col in columns:
                name, data_type, nullable, default, key_type = col
                create_table += f"    {name} {data_type}"
                if key_type:
                    create_table += f" {key_type}"
                if nullable == 'NO':
                    create_table += " NOT NULL"
                if default:
                    create_table += f" DEFAULT {default}"
                create_table += ",\n"
            create_table = create_table.rstrip(",\n") + "\n);\n\n"
            schema_text += create_table

        return schema_text

    except Exception as e:
        print(f"Error fetching schema: {e}")
        return None

def get_sql_query(user_query, schema):
    """Get SQL query from Gemini 1.5 Flash model"""
    try:
        # Prepare the prompt
        prompt = f"""### PostgreSQL database schema:
{schema}

### Instructions:
- Generate a single, correct SQL query that answers the user's question.
- Include appropriate table aliases, JOIN conditions, and WHERE clauses.
- Provide the SQL query only, with no additional explanations or questions.
- If you dont find something analyse and correlate with others in the schema 

### User Query: {user_query}

### SQL Query:"""

        # Prepare the API payload
        payload = {
            "contents": [{
                "parts": [{"text": prompt}]
            }]
        }

        # Send the POST request to Gemini 1.5 Flash API
        response = requests.post(f"{API_URL}?key={GEMINI_API_KEY}", headers=headers, json=payload)

        if response.status_code != 200:
            print(f"Error from API: {response.text}")
            return None

        # Extract SQL query from response
        response_data = response.json()
        if 'candidates' not in response_data or not response_data['candidates']:
            print("No candidates in response")
            return None

        # Extract the SQL query text from the first candidate
        sql_query = response_data['candidates'][0]['content']['parts'][0]['text'].strip()

        # Remove any surrounding code block formatting (if any)
        if sql_query.startswith('```sql') and sql_query.endswith('```'):
            sql_query = sql_query[7:-3].strip()

        if not sql_query.endswith(';'):
            sql_query += ';'

        return sql_query

    except Exception as e:
        print(f"Error getting SQL query: {e}")
        return None

def execute_query(sql_query):
    """Execute SQL query and return results"""
    try:
        conn, cur = get_connection()
        cur.execute(sql_query)

        # If SELECT query, fetch results
        if sql_query.strip().upper().startswith('SELECT'):
            columns = [desc[0] for desc in cur.description]
            results = []
            for row in cur.fetchall():
                # Convert Decimal values to float
                row_dict = {columns[i]: float(row[i]) if isinstance(row[i], Decimal) else row[i] for i in range(len(row))}
                results.append(row_dict)
        else:
            results = {"message": "Query executed successfully"}

        return results

    except Exception as e:
        print(f"Error executing query: {e}")
        return None

def main():
    try:
        # Get database schema
        print("Fetching database schema...")
        schema = get_schema()
        if not schema:
            print("Failed to fetch schema")
            return

        while True:
            # Get user query
            user_query = input("\nEnter your query in English (or 'exit' to quit): ").strip()

            if user_query.lower() == 'exit':
                break

            print("\nGenerating SQL query...")

            # Get SQL query from Gemini 1.5 Flash
            sql_query = get_sql_query(user_query, schema)
            if not sql_query:
                print("Failed to generate SQL query")
                continue

            print(f"\nGenerated SQL Query:")
            print(sql_query)

            # Execute the query
            print("\nExecuting query...")
            results = execute_query(sql_query)

            if results:
                print("\nResults:")
                print(json.dumps(results, indent=2))
            else:
                print("No results returned")

    except Exception as e:
        print(f"An error occurred: {e}")
    finally:
        close_connection()

if __name__ == "__main__":
    main()