Spaces:
Sleeping
Sleeping
import streamlit as st | |
from itertools import tee | |
import google.generativeai as genai | |
import sqlite3 | |
def connect_to_database(db_path): | |
# debugging | |
try: | |
conn = sqlite3.connect(db_path) | |
cursor = conn.cursor() | |
# Running a test query | |
cursor.execute("SELECT name FROM sqlite_master WHERE type='table';") | |
tables = cursor.fetchall() | |
if tables: | |
print("Connection successful. Found tables:", tables) | |
else: | |
print("Connection successful but no tables found.") | |
return conn | |
except sqlite3.Error as e: | |
print(f"Error connecting to database: {e}") | |
return None | |
# Connect to the database | |
db_connection = connect_to_database('phones.db') | |
st.set_page_config(layout="wide") | |
MODEL_AVATAR_URL = "./icon.png" | |
GEMINI = st.secrets["GEMINI_API_KEY"] | |
genai.configure(api_key=GEMINI) | |
generation_config = { | |
"temperature": 0.7, | |
"top_p": 0.95, | |
"top_k": 0, | |
"max_output_tokens": 8192, | |
} | |
safety_settings = [ | |
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"}, | |
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"}, | |
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"}, | |
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"}, | |
] | |
model = genai.GenerativeModel(model_name="gemini-1.0-pro", | |
generation_config=generation_config, | |
safety_settings=safety_settings) | |
st.title("Enhancing Customer Experience with Nexus") | |
with open("style.css") as css: | |
st.markdown(f'<style>{css.read()}</style>', unsafe_allow_html=True) | |
if "messages" not in st.session_state: | |
st.session_state["messages"] = [] | |
def clear_chat_history(): | |
st.session_state["messages"] = [] | |
st.button('Clear Chat', on_click=clear_chat_history) | |
# Function to generate SQL query | |
def generate_sql_query(user_input): | |
prompt = f"Generate an SQL query to find information based on the user's question: '{user_input}'. Note: The table name is 'phones' and the columns are 'ProductName', 'Brand', 'Storage', 'Color', 'Price', 'QuantityInStock', and 'Location'. Note: GIve the query without '''sql at the start and end of the query. just give the query text content" | |
response = model.generate_content(prompt) | |
print(response.text) | |
return response.text | |
def execute_sql_query(query): | |
try: | |
cursor = db_connection.cursor() | |
cursor.execute(query) | |
results = cursor.fetchall() | |
cursor.close() | |
print(results) | |
return results | |
except Exception as e: | |
return f"Error executing query: {str(e)}" | |
def format_response(user_input, query_results): | |
if not query_results: | |
return "No data found for your query." | |
result_text = f"Found {len(query_results)} results: " + ', '.join([str(item) for sublist in query_results for item in sublist]) | |
prompt = f"Rephrase this in a more conversational and informative way based on the user's question: '{user_input}'. Here are the details: {result_text}. Answer the user's question in a conversational manner. Note: as its a conversational response, give the response in correct mannser with correct formatting" | |
formatted_response = model.generate_content(prompt) | |
print(formatted_response.text) | |
return formatted_response.text | |
def handle_user_input(user_input): | |
with history: | |
st.session_state["messages"].append({"role": "user", "content": user_input}) | |
with st.chat_message("user"): | |
st.markdown(user_input) | |
# Generate SQL query from user input | |
sql_query = generate_sql_query(user_input) | |
if sql_query: | |
# Execute the generated SQL query | |
query_results = execute_sql_query(sql_query) | |
formatted_answer = format_response(user_input, query_results) | |
with st.chat_message("assistant", avatar=MODEL_AVATAR_URL): | |
st.markdown(formatted_answer) | |
st.session_state["messages"].append({"role": "assistant", "content": formatted_answer}) | |
else: | |
with st.chat_message("assistant", avatar=MODEL_AVATAR_URL): | |
st.markdown("Failed to generate a valid SQL query.") | |
st.session_state["messages"].append({"role": "assistant", "content": "Failed to generate a valid SQL query."}) | |
EXAMPLES = [ | |
"Do you have any iphone?", | |
"How much does the Samsung Galaxy cost?", | |
"Show me all products with more than 64GB of storage", | |
"My budget is around $500, suggest best phones accordingly" | |
] | |
with st.sidebar: | |
with st.container(): | |
st.title("Example Queries") | |
for prompt in EXAMPLES: | |
st.button(prompt, on_click=lambda prompt=prompt: handle_user_input(prompt)) | |
main = st.container() | |
with main: | |
history = st.container(height=400) | |
with history: | |
for message in st.session_state["messages"]: | |
avatar = None | |
if message["role"] == "assistant": | |
avatar = MODEL_AVATAR_URL | |
with st.chat_message(message["role"], avatar=avatar): | |
st.markdown(message["content"]) | |
if prompt := st.chat_input("Type your question:", max_chars=1000): | |
handle_user_input(prompt) |