app.py

import os
import gradio as gr
import pandas as pd
import duckdb
import openai

# 1) Read your OpenAI key from the Space’s Secrets
openai.api_key = os.getenv("OPENAI_API_KEY")

# 2) Load your synthetic data into DuckDB
df = pd.read_csv('synthetic_profit.csv')
conn = duckdb.connect(':memory:')
conn.register('sap', df)

# 3) Build a one-line schema description for prompts
schema = ", ".join(df.columns)  # e.g. "Region,Product,FiscalYear,FiscalQuarter,Revenue,Profit,ProfitMargin"

# 4) SQL-generation via OpenAI
def generate_sql(question: str) -> str:
    system = (
        f"You are an expert SQL generator for a DuckDB table named `sap` "
        f"with columns: {schema}. "
        "Translate the user’s question into a valid SQL query and return _only_ the SQL."
    )
    messages = [
        {"role": "system", "content": system},
        {"role": "user",   "content": question},
    ]
    resp = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=messages,
        temperature=0.0,
        max_tokens=150,
    )
    sql = resp.choices[0].message.content.strip()
    # strip triple-backticks if present
    if sql.startswith("```") and sql.endswith("```"):
        sql = "\n".join(sql.splitlines()[1:-1])
    return sql

# 5) Core QA function: NL → SQL → execute → format
def answer_profitability(question: str) -> str:
    # a) generate SQL
    sql = generate_sql(question)
    # b) try to run it
    try:
        result_df = conn.execute(sql).df()
    except Exception as e:
        return (
            f"❌ **Error executing SQL**\n\n"
            f"```\n{e}\n```\n\n"
            f"**Generated SQL**\n```sql\n{sql}\n```"
        )
    # c) format the result
    if result_df.empty:
        return f"No rows returned.\n\n**Generated SQL**\n```sql\n{sql}\n```"
    # single-cell → just the value
    if result_df.shape == (1,1):
        return str(result_df.iat[0,0])
    # otherwise, markdown table
    return result_df.to_markdown(index=False)

# 6) Gradio UI
iface = gr.Interface(
    fn=answer_profitability,
    inputs=gr.Textbox(lines=2, placeholder="Ask a question about profitability…"),
    outputs=gr.Markdown(),             # renders errors, code, and tables nicely
    title="SAP Profitability Q&A (OpenAI → SQL → DuckDB)",
    description=(
        "Uses OpenAI’s GPT-3.5-Turbo to translate your question into SQL, "
        "executes it on the `sap` table in DuckDB, and returns the result."
    ),
    allow_flagging="never",
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)