Spaces:
Runtime error
Runtime error
| 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) | |
| out = "" | |
| for chunk in answer: | |
| out += chunk | |
| return 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""" | |
| # 🤖 MCP DB Answer | |
| Your mcp server that allow you to talk to any database | |
| Powered by Ibis it support : PostgreSQL , SQLite , MySQL , MSSQL , ClickHouse , BigQuery and many other | |
| Also support .CSV and .parquet files | |
| """, | |
| 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=None) | |
| if __name__ == "__main__": | |
| demo.launch(mcp_server=True) |