model_connector.py

import pandas as pd

from config import client, OPENAI_MODEL


def generate_sql_from_nl(prompt, schema_info):
    """
    Generate an SQL query from a natural language prompt.

    Args:
        prompt (str): The natural language prompt provided by the user.
        schema_info (dict): A dictionary containing the database schema information.

    Returns:
        str: The generated SQL query.
    """
    schema_description = "The database has the following tables:\n"
    for table, cols in schema_info.items():
        schema_description += f"- {table}({', '.join(cols)})\n"

    system_message = (
        "You are a helpful assistant that converts natural language into SQL queries. "
        "The SQL should be valid for the provided database schema. "
        "Don't include explanation text, just provide a single SQL query."
    )

    user_message = f"{schema_description}\nUser query: {prompt}\nReturn just the SQL without explanation."

    # Call the new chat completion method
    response = client.chat.completions.create(
        model=OPENAI_MODEL,
        messages=[
            {"role": "system", "content": system_message},
            {"role": "user", "content": user_message},
        ],
        temperature=0,
    )

    sql = response.choices[0].message.content.strip()
    sql = sql.replace("LIKE", "ILIKE")
    return sql


def format_response(data, sql):
    """
    Format the response for the user.
    :param data:    The data returned from the SQL query.
    :param sql:     The SQL query.
    :return:    The formatted response.
    """
    if len(data) == 0:
        nl_response = "I found no matching results."
        return f"{nl_response}\n\nSQL Query:\n{sql}"
    elif len(data) == 1:
        row = data[0]
        nl_response = "Here is the single matching result:\n"
        for k, v in row.items():
            nl_response += f"{k}: {v}\n"
        return f"{nl_response}\nSQL Query:\n{sql}"
    else:
        df = pd.DataFrame(data)
        table_view = df.to_markdown(index=False)
        return f"Multiple results found:\n{table_view}\n\nSQL Query:\n{sql}"