app.py

import gradio as gr
from database import Database
from filesource import FileSource
from agent import run_agent
from services.utils import get_db_scheme_from_uri
source = None


def connect_to_file(file):
    """
    This function is used to connect to a data source file like CSV or parquet

    Parameters:
    -----------
    file : File
        A file-like object (e.g., from a file upload or file picker) that contains structured data.

    Returns:
    --------
    tuple
        A tuple (schema, status) where:
        - schema (str): A string representation of the data schema if successful, or an empty string on failure.
        - status (str): A message indicating the success or failure of the connection.
    
    """
    global source
    try:
        source = FileSource(file.name)
        status = source.connect()
        schema = source._pretify_schema()
        status = "Connection successful!"
    except Exception as e:
        schema = ""
        status = f"Error: {str(e)}"
    return schema, status

def generate_and_run_sql(prompt):
    """
    This function is used to give an answer to a user question related to a database or file source


    Parameters:
    -----------
    prompt : str
        A natural language instruction or question from the user (e.g., "Show me the top 10 customers by revenue").

    Returns:
    --------
    str
        The given response of the user question
    Notes:
    ------
    - Requires a connected `source` (via `connect_to_file` or `connect_to_database`).
    - Uses a language agent to interpret the prompt and generate an appropriate SQL query or explanation.
    """
    if source is None:
        return "Please connect to a database or add a file"
    
    answer = run_agent(source, prompt,False)
    out = ""
    for chunk in answer:
        out += chunk
    return [{"role":"assistant","content":out}]


def connect_to_database(db_url):
    """
    This function is used to connect to a database using the provided database URL and returns its schema.

    Parameters:
    -----------
    db_url : str
        A valid database connection string (e.g., "mysql://user:pass@host:port/dbname").

    Returns:
    --------
    tuple
        A tuple (schema, status) where:
        - schema (str): A string representation of the database schema if the connection is successful, 
                        or an empty string if it fails.
        - status (str): A message indicating the result of the connection attempt.

    Notes:
    ------
    - The dialect is extracted from the URI to determine the database type (e.g., MySQL, PostgreSQL).
    """
    global source
    try:
        dialect = get_db_scheme_from_uri(db_url)
        source = Database(db_url, dialect)
        status = source.connect()
        schema = source._pretify_schema()
        status = "Connection successful!"
    except Exception as e:
        schema = ""
        status = f"Error: {str(e)}"
    return schema, status
# Function to add user message to chat history
def user(user_message, chat_history):
    chat_history.append({"role": "user", "content": user_message})
    return "", chat_history

# Function to generate a bot response
def bot(chat_history):
  
    if source is None:
         chat_history.append({"role":"assistant","content":"please connect to a database  before asking question"})
         yield chat_history
    else:
        answer = run_agent(source,chat_history[-1]['content'])
        chat_history.append({"role":"assistant","content":""})

        for chunk in answer:
            chat_history[-1]['content'] += chunk    
            yield chat_history

# Create the Gradio interface
with gr.Blocks( theme=gr.themes.Default(), css="""
    .gr-button { margin: 5px; border-radius:16px; }
    .gr-textbox, .gr-text-area, .gr-dropdown, .gr-json { border-radius: 8px; }
    .gr-row { gap: 10px; }
    .gr-tab { border-radius: 8px; }
    .status-text { font-size: 0.9em; color: #555; }
    .gr-json { max-height: 300px; overflow-y: auto; } /* Added scrolling for JSON */
""") as demo:
    gr.Markdown(
        f"""
        # ⚡ Ibiza Server  
        Chat with your data
        
        Connect to PostgreSQL, MySQL, SQLite, Snowflake, BigQuery, and more — powered by Ibis, with support for CSV and Parquet too
        
        Powered by **Modal** and Qwen3 Reranking model 🔥
        
        
        (It may take some time to start up the first time if the space hasn't been used before.)
        
       
        
        """,
        elem_classes=["header"]
    )

    with gr.Column(scale=3):
        with gr.Tabs():
            with gr.TabItem("💬 Chat"):
                with gr.Group():
                    main_chat_disp = gr.Chatbot(
                        label=None, height=400,
                        avatar_images=(None, "https://huggingface.co/spaces/Space-Share/bucket/resolve/main/images/pfp.webp"),
                        show_copy_button=True, render_markdown=True, sanitize_html=True, type='messages'
                    )
                    with gr.Row(variant="compact"):
                        user_msg_tb = gr.Textbox(
                            show_label=False, placeholder="Talk with your data...",
                            scale=7, lines=1, max_lines=3
                        )
                        send_btn = gr.Button("Send", variant="primary", scale=1, min_width=100)
            with gr.TabItem("Config"):        
                with gr.Row():
                    # Left column for database configuration.
                    with gr.Column(scale=1):
                        gr.Markdown("## Database Configuration")
                        # Textbox for entering the database URL.
                        db_url_tb = gr.Textbox(
                            show_label=True, label="Database URL", placeholder="Enter the URL to connect to the database..."
                        )
                        # Button to connect to the database.
                        connect_btn = gr.Button("Connect", variant="primary")
                        gr.Markdown("## File Upload")
                        file_uploader = gr.File(
                            label="Upload File", file_types=[".csv", ".parquet", ".xls", ".xlsx"]
                        )
                        # Button to connect to the database.
                        load_btn = gr.Button("Load", variant="primary")

                    # Right column for displaying the database schema and status message.
                    with gr.Column(scale=3):
                        gr.Markdown("## Database Schema")
                        # Textarea to display the database schema.
                        schema_ta = gr.TextArea(
                            show_label=False, placeholder="Database schema will be displayed here...",
                            lines=20, max_lines=50, interactive=False
                        )
                        # Textbox to display the status message.
                        status_tb = gr.Textbox(
                            show_label=False, placeholder="Status message will be displayed here...",
                            lines=1, max_lines=1, interactive=False, elem_classes=["status-text"]
                        )
    connect_btn.click(fn=connect_to_database, inputs=db_url_tb, outputs=[schema_ta, status_tb])
    load_btn.click(fn=connect_to_file, inputs=file_uploader, outputs=[schema_ta, status_tb])
    send_btn.click(fn=user, inputs=[user_msg_tb, main_chat_disp], outputs=[user_msg_tb, main_chat_disp], queue=False).then(
        fn=bot, inputs=main_chat_disp, outputs=main_chat_disp
    )
    invisible_btn = gr.Button(visible=False)
    invisible_btn.click(fn=generate_and_run_sql, inputs=[user_msg_tb], outputs=[main_chat_disp])

if __name__ == "__main__":
    demo.launch(mcp_server=True)