File size: 7,838 Bytes
ba2b9e4 1dcf789 ba2b9e4 ef8ba77 b24beb5 ef8ba77 ba2b9e4 f734451 5513461 b24beb5 ef8ba77 ba2b9e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 |
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) |