Spaces:
Sleeping
Sleeping
| import google.generativeai as genai | |
| from dotenv import load_dotenv | |
| import streamlit as st | |
| import os | |
| import sqlite3 | |
| import pandas as pd | |
| import io | |
| from langchain.prompts import PromptTemplate | |
| from utils.extarct_db import extract_name_and_colums , read_excel_query , read_sql_query | |
| from typing import List | |
| import numpy as np | |
| import plotly.express as px | |
| import matplotlib.pyplot as plt | |
| load_dotenv() | |
| key = os.getenv("GOOGLE_API_KEY") | |
| genai.configure(api_key=key) | |
| model=genai.GenerativeModel('gemini-pro') | |
| # Function to load gemini model | |
| def get_gemini_response(question: str, table_name: str, column_names: List[str]) -> str: | |
| # Convert column_names list to a comma-separated string | |
| columns_str = ', '.join(column_names) | |
| # Defining the prompt | |
| prompt = PromptTemplate( | |
| input_variables=["table_name", "columns"], | |
| template=""" | |
| You are an expert in converting English questions to SQL query! | |
| The SQL database has the name {table_name} and has the following columns - | |
| {columns}\n\nFor example,\nExample 1 - How many entries of records are present?, | |
| the SQL command will be something like this SELECT COUNT(*) FROM {table_name} ; | |
| \nExample 2 - Tell me all the students studying in Data Science class?, | |
| the SQL command will be something like this SELECT * FROM {table_name} | |
| also the sql code should not have ``` in beginning or end and sql word in output | |
| """ | |
| ) | |
| # Format the prompt with the table_name and columns_str | |
| formatted_prompt = prompt.format(table_name=table_name, columns=columns_str) | |
| # Assuming genai is correctly imported and configured | |
| model = genai.GenerativeModel('gemini-pro') | |
| response = model.generate_content([formatted_prompt, question]) | |
| return response.text | |
| st.header("ConverseDB") | |
| uploaded_file = st.sidebar.file_uploader("Choose a file", type=None) | |
| if uploaded_file is not None: | |
| file_type = uploaded_file.type | |
| allowed_file_extensions = ["application/octet-stream", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"] | |
| if file_type in allowed_file_extensions: | |
| print(uploaded_file.name) | |
| question=st.text_input("Input Prompt: " , key="input") | |
| submit = st.button("Query") | |
| if file_type == "application/octet-stream": | |
| file_path = os.path.join(os.getcwd(), uploaded_file.name) | |
| with open(file_path, "wb") as f: | |
| f.write(uploaded_file.getbuffer()) | |
| db_info = extract_name_and_colums(uploaded_file.name) | |
| table_name = db_info['table_name'][0] | |
| column_names = db_info['colum_names']['STUDENT'] | |
| print(column_names) | |
| if submit: | |
| response = get_gemini_response(question, table_name, column_names) | |
| print(response) | |
| response = read_sql_query(response, file_path) | |
| # print(response) | |
| formatted_response=model.generate_content(f"Format this {response} in the table format") | |
| response_text = formatted_response.candidates[0].content.parts[0].text | |
| # Splitting the response text into lines | |
| lines = response_text.strip().split('\n') | |
| print(lines) | |
| # Extracting column names and data | |
| columns = [col.strip() for col in lines[0].split('|') if col.strip()] | |
| data = [dict(zip(columns, [item.strip() for item in line.split('|') if item.strip()])) for line in lines[2:]] | |
| print(columns,data) | |
| # Creating DataFrame | |
| df = pd.DataFrame(data) | |
| st.subheader("The Response is ") | |
| st.table(df) | |
| try: | |
| # Example chart with Plotly | |
| fig = px.bar(df, x=columns[0], y=columns[1], title="Visualize") | |
| st.plotly_chart(fig) | |
| except: | |
| st.success("No Enough data available to do analysis") | |
| if file_type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": | |
| # converting excel file to db format then deleting excel file | |
| db_name = "data.db" | |
| file_path = os.path.join(os.getcwd(), uploaded_file.name) | |
| with open(file_path, "wb") as f: | |
| f.write(uploaded_file.getbuffer()) | |
| df = pd.read_excel(file_path) | |
| conn = sqlite3.connect(db_name) | |
| df.to_sql("excel_data", conn, index=False, if_exists="replace") | |
| conn.close() | |
| # os.remove(file_path) | |
| # extracting DB Name & colum names | |
| db_info = extract_name_and_colums(db_name) | |
| table_name = db_info['table_name'][0] | |
| column_names = db_info['colum_names'] | |
| print(column_names) | |
| if submit: | |
| response = get_gemini_response(question, table_name, column_names) | |
| print(response) | |
| response = read_sql_query(response, db_name) | |
| # print(response) | |
| formatted_response=model.generate_content(f"Format this {response} in the table format") | |
| response_text = formatted_response.candidates[0].content.parts[0].text | |
| # Splitting the response text into lines | |
| lines = response_text.strip().split('\n') | |
| # Extracting column names and data | |
| columns = [col.strip() for col in lines[0].split('|') if col.strip()] | |
| data = [dict(zip(columns, [item.strip() for item in line.split('|') if item.strip()])) for line in lines[2:]] | |
| # Creating DataFrame | |
| df = pd.DataFrame(data) | |
| st.subheader("The Response is ") | |
| st.table(df) | |
| try: | |
| # Example chart with Plotly | |
| fig = px.bar(df, x=columns[0], y=columns[1], title="Visualize") | |
| st.plotly_chart(fig) | |
| except: | |
| st.success("No Enough data available to do analysis") | |
| else: | |
| st.error("File type is not allowed. Please upload a .db or .xlsx file.") | |