app.py

import pandas as pd
import gradio as gr
import os
import google.generativeai as genai  # requires GOOGLE_API_KEY set as env var

# 1. Configure Gemini
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
model = genai.GenerativeModel("gemini-2.5-pro")

# 2. Load Excel data
df1 = pd.read_excel(os.path.join("data_source", "OC Onboarding Information.xlsx"))
df2 = pd.read_excel(os.path.join("data_source", "The Alex Ideas Report.xlsx"))

# Build schema info for prompts
def get_schema_info():
    schema1 = f"Report1 columns: {list(df1.columns)}"
    schema2 = f"Report2 columns: {list(df2.columns)}"
    return schema1 + "\n" + schema2

schema_info = get_schema_info()

# 3. Core function
def answer_question(history, message):
    """
    history: chat history (list of [user, assistant] pairs)
    message: latest user message (string)
    """
    # Build prompt for Gemini
    prompt = f"""
You are a data analysis assistant. 
You can ONLY answer questions using the two Excel reports provided. 
Do not hallucinate or use external knowledge.
If the question is irrelevant, respond with:
"I can only answer questions about the provided Excel reports."

The reports have the following schema:
{schema_info}

The user asked:
{message}

Generate Python pandas code that uses df1 and/or df2 to answer.
Return ONLY code, nothing else.
    """

    try:
        # Call Gemini
        response = model.generate_content(prompt)
        code = response.text.strip("```python").strip("```")

        # 4. Execute code safely
        local_vars = {"df1": df1, "df2": df2, "pd": pd}
        try:
            result = eval(code, {"__builtins__": {}}, local_vars)
        except Exception as e:
            exec(code, {"__builtins__": {}}, local_vars)
            result = local_vars.get("result", "No result variable found")

        return str(result)

    except Exception as e:
        return f"Error: {str(e)}"


# 5. Gradio UI
with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox(placeholder="Ask me a question about the reports...")
    clear = gr.ClearButton([msg, chatbot])

    def respond(message, chat_history):
        answer = answer_question(chat_history, message)
        chat_history.append((message, answer))
        return "", chat_history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])

# 6. Run locally (Spaces will call demo.launch() automatically)
if __name__ == "__main__":
    demo.launch()