mistral.py

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

# Load environment variables from .env file
load_dotenv()

# Mistral API Configuration
API_URL = "https://api-inference.huggingface.co/models/mistral-model-name"  # Replace with the specific Mistral model URL
HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACE_API_KEY")
headers = {
    "Authorization": f"Bearer {HUGGINGFACE_API_KEY}",
    "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"),
}

conn = None
cur = None

def get_connection():
    """Establish the database connection."""
    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."""
    global conn, cur
    if cur:
        cur.close()
    if conn:
        conn.close()

def get_schema():
    """Fetch database schema."""
    try:
        conn, cur = get_connection()
        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]
            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_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 Mistral model."""
    try:
        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.

### User Query: {user_query}

### SQL Query:"""

        payload = {"inputs": prompt}

        # Send the POST request to the Hugging Face Mistral API
        response = requests.post(API_URL, headers=headers, json=payload)
        response_data = response.json()

        if "error" in response_data:
            print(f"Error from API: {response_data['error']}")
            return None

        # Extract the SQL query from the response
        sql_query = response_data.get("generated_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 sql_query.strip().upper().startswith('SELECT'):
            columns = [desc[0] for desc in cur.description]
            results = []
            for row in cur.fetchall():
                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