()
| 36 | |
| 37 | |
| 38 | def main(): |
| 39 | st.title("Text-to-SQL") |
| 40 | |
| 41 | natural_language_query = st.text_input(label="Ask anything...", label_visibility="hidden", placeholder="Ask anything...") |
| 42 | |
| 43 | if natural_language_query: |
| 44 | with st.spinner(text="Generating SQL..."): |
| 45 | start_time = time.time() |
| 46 | response = requests.post(f"{API_BASE}/text_to_sql", json={"natural_language_query": natural_language_query}) |
| 47 | end_time = time.time() |
| 48 | time_taken = end_time - start_time |
| 49 | if response.status_code == 200: |
| 50 | st.info(f"SQL generated in {time_taken:.2f} seconds") |
| 51 | SQL = f"""```sql |
| 52 | {response.json().get("sql_query")} |
| 53 | """ |
| 54 | st.markdown(SQL) |
| 55 | |
| 56 | RESULT = response.json().get("result", {}) |
| 57 | st.table(RESULT.get("results", [])) |
| 58 | |
| 59 | with st.spinner(text="Generating visualization..."): |
| 60 | start_time = time.time() |
| 61 | response = requests.post(f"{API_BASE}/viz", |
| 62 | json={ |
| 63 | "data": create_viz_data_dict( |
| 64 | RESULT.get("column_names", []), |
| 65 | RESULT.get("column_types", []), |
| 66 | RESULT.get("results", []), |
| 67 | ) |
| 68 | } |
| 69 | ) |
| 70 | end_time = time.time() |
| 71 | time_taken = end_time - start_time |
| 72 | if response.status_code == 200: |
| 73 | st.info(f"Visualization generated in {time_taken:.2f} seconds") |
| 74 | VEGA_LITE_SPEC = response.json().get("vega_lite_spec") |
| 75 | st.vega_lite_chart(data=RESULT.get("results", []), spec=VEGA_LITE_SPEC) |
| 76 | else: |
| 77 | st.error(f"{response.status_code}: {response.reason}") |
| 78 | st.info("Sorry, I couldn't generate a visualization. Please try again.") |
| 79 | else: |
| 80 | st.error(f"{response.status_code}: {response.reason}") |
| 81 | st.info("Sorry, I couldn't answer your question/command. Please try again.") |
| 82 | |
| 83 | |
| 84 | if __name__ == "__main__": |
no test coverage detected