Jacqkues's picture
Upload 11 files
d25ee4b verified
raw
history blame
5.47 kB
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):
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 connect_to_database(db_url):
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=600,
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")
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
)
if __name__ == "__main__":
demo.launch(mcp_server=True)