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)