Spaces:
Sleeping
Sleeping
| # version 2: added custom prompts. | |
| import os | |
| import json | |
| import sqlite3 | |
| from datetime import datetime | |
| import streamlit as st | |
| from langchain_huggingface import HuggingFaceEmbeddings | |
| from langchain_chroma import Chroma | |
| from langchain_groq import ChatGroq | |
| from langchain.memory import ConversationBufferMemory | |
| from langchain.chains import ConversationalRetrievalChain | |
| from langchain.prompts import PromptTemplate | |
| from vectorize_documents import embeddings # If needed elsewhere | |
| # Load config | |
| working_dir = os.path.dirname(os.path.abspath(__file__)) | |
| config_data = json.load(open(f"{working_dir}/config.json")) | |
| GROQ_API_KEY = config_data["GROQ_API_KEY"] | |
| os.environ["GROQ_API_KEY"] = GROQ_API_KEY | |
| # Set up the database | |
| def setup_db(): | |
| conn = sqlite3.connect("chat_history.db", check_same_thread=False) | |
| cursor = conn.cursor() | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS chat_histories ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| username TEXT, | |
| timestamp TEXT, | |
| day TEXT, | |
| user_message TEXT, | |
| assistant_response TEXT | |
| ) | |
| """) | |
| conn.commit() | |
| return conn | |
| # Set up vectorstore | |
| def setup_vectorstore(): | |
| embeddings = HuggingFaceEmbeddings() | |
| vectorstore = Chroma(persist_directory="Vector_db", embedding_function=embeddings) | |
| return vectorstore | |
| # Custom prompt template | |
| custom_prompt_template = PromptTemplate.from_template(""" | |
| You are a helpful assistant that helps users choose laptops. | |
| 1. Analyze the user's query, take information from vectordb and then give top 3 laptops to user from Relevent information that is context. | |
| 2. Keep suggestions clear and concise with names, specs, and reasons only from relevant information context. | |
| Relevant Information: | |
| {context} | |
| Chat History: | |
| {chat_history} | |
| User Query: | |
| {question} | |
| Assistant Response: | |
| """) | |
| # Set up the chatbot chain with a specific model | |
| def chat_chain(vectorstore, model_name): | |
| llm = ChatGroq(model=model_name, temperature=0.3) | |
| retriever = vectorstore.as_retriever() | |
| memory = ConversationBufferMemory( | |
| llm=llm, | |
| output_key="answer", | |
| memory_key="chat_history", | |
| return_messages=True | |
| ) | |
| chain = ConversationalRetrievalChain.from_llm( | |
| llm=llm, | |
| retriever=retriever, | |
| memory=memory, | |
| combine_docs_chain_kwargs={"prompt": custom_prompt_template}, | |
| return_source_documents=True, | |
| verbose=True | |
| ) | |
| return chain | |
| # Streamlit UI setup | |
| st.set_page_config(page_title="ByteX-Ai", page_icon="🤖AI", layout="centered") | |
| st.title("🤖 ByteX-Ai") | |
| st.subheader("Hey! Get your Laptop!!") | |
| # Initialize DB connection | |
| if "conn" not in st.session_state: | |
| st.session_state.conn = setup_db() | |
| # Prompt user to log in | |
| if "username" not in st.session_state: | |
| username = st.text_input("Enter your name to proceed:") | |
| if username: | |
| with st.spinner("Loading chatbot interface... Please wait."): | |
| st.session_state.username = username | |
| st.session_state.chat_history = [] | |
| st.session_state.vectorstore = setup_vectorstore() | |
| st.success(f"Welcome, {username}! Now select a model to start chatting.") | |
| else: | |
| username = st.session_state.username | |
| # Model selection options | |
| model_options = [ | |
| "gemma2-9b-it", | |
| "llama-3.1-8b-instant", | |
| "llama3-70b-8192", | |
| "llama3-8b-8192" | |
| ] | |
| selected_model = st.selectbox("Choose a model:", model_options) | |
| # Ensure vectorstore exists | |
| if "vectorstore" not in st.session_state: | |
| st.session_state.vectorstore = setup_vectorstore() | |
| # Set or update the selected model | |
| if "selected_model" not in st.session_state: | |
| st.session_state.selected_model = selected_model | |
| # Reset conversational_chain if model changes or not yet initialized | |
| if ("conversational_chain" not in st.session_state) or (st.session_state.selected_model != selected_model): | |
| st.session_state.selected_model = selected_model | |
| st.session_state.conversational_chain = chat_chain(st.session_state.vectorstore, selected_model) | |
| st.session_state.chat_history = [] | |
| # Reset chat manually | |
| if st.button("🔄 Reset Chat"): | |
| st.session_state.chat_history = [] | |
| st.session_state.conversational_chain = chat_chain(st.session_state.vectorstore, st.session_state.selected_model) | |
| st.success("Chat reset!") | |
| # Show chat UI | |
| if "username" in st.session_state: | |
| st.subheader(f"Hello {username}, start your query below!") | |
| if st.session_state.chat_history: | |
| for message in st.session_state.chat_history: | |
| if message['role'] == 'user': | |
| with st.chat_message("user"): | |
| st.markdown(message["content"]) | |
| elif message['role'] == 'assistant': | |
| with st.chat_message("assistant"): | |
| st.markdown(message["content"]) | |
| user_input = st.chat_input("Ask AI....") | |
| if user_input: | |
| with st.spinner("Processing your query... Please wait."): | |
| st.session_state.chat_history.append({"role": "user", "content": user_input}) | |
| with st.chat_message("user"): | |
| st.markdown(user_input) | |
| with st.chat_message("assistant"): | |
| response = st.session_state.conversational_chain({"question": user_input}) | |
| assistant_response = response["answer"] | |
| st.markdown(assistant_response) | |
| st.session_state.chat_history.append({"role": "assistant", "content": assistant_response}) | |
| # Version 1: working properly but there is no prompt refinement. | |
| # import os | |
| # import json | |
| # import sqlite3 | |
| # from datetime import datetime | |
| # import streamlit as st | |
| # from langchain_huggingface import HuggingFaceEmbeddings | |
| # from langchain_chroma import Chroma | |
| # from langchain_groq import ChatGroq | |
| # from langchain.memory import ConversationBufferMemory | |
| # from langchain.chains import ConversationalRetrievalChain | |
| # from vectorize_documents import embeddings # If needed elsewhere | |
| # # Load config | |
| # working_dir = os.path.dirname(os.path.abspath(__file__)) | |
| # config_data = json.load(open(f"{working_dir}/config.json")) | |
| # GROQ_API_KEY = config_data["GROQ_API_KEY"] | |
| # os.environ["GROQ_API_KEY"] = GROQ_API_KEY | |
| # # Set up the database | |
| # def setup_db(): | |
| # conn = sqlite3.connect("chat_history.db", check_same_thread=False) | |
| # cursor = conn.cursor() | |
| # cursor.execute(""" | |
| # CREATE TABLE IF NOT EXISTS chat_histories ( | |
| # id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| # username TEXT, | |
| # timestamp TEXT, | |
| # day TEXT, | |
| # user_message TEXT, | |
| # assistant_response TEXT | |
| # ) | |
| # """) | |
| # conn.commit() | |
| # return conn | |
| # # Set up vectorstore | |
| # def setup_vectorstore(): | |
| # embeddings = HuggingFaceEmbeddings() | |
| # vectorstore = Chroma(persist_directory="Vector_db", embedding_function=embeddings) | |
| # return vectorstore | |
| # # Set up the chatbot chain with a specific model | |
| # def chat_chain(vectorstore, model_name): | |
| # llm = ChatGroq(model=model_name, temperature=0) | |
| # retriever = vectorstore.as_retriever() | |
| # memory = ConversationBufferMemory( | |
| # llm=llm, | |
| # output_key="answer", | |
| # memory_key="chat_history", | |
| # return_messages=True | |
| # ) | |
| # chain = ConversationalRetrievalChain.from_llm( | |
| # llm=llm, | |
| # retriever=retriever, | |
| # chain_type="stuff", | |
| # memory=memory, | |
| # verbose=True, | |
| # return_source_documents=True | |
| # ) | |
| # return chain | |
| # # Streamlit UI setup | |
| # st.set_page_config(page_title="ByteX-Ai", page_icon="🤖AI", layout="centered") | |
| # st.title("🤖 ByteX-Ai") | |
| # st.subheader("Hey! Get your Laptop!!") | |
| # # Initialize DB connection | |
| # if "conn" not in st.session_state: | |
| # st.session_state.conn = setup_db() | |
| # # Prompt user to log in | |
| # if "username" not in st.session_state: | |
| # username = st.text_input("Enter your name to proceed:") | |
| # if username: | |
| # with st.spinner("Loading chatbot interface... Please wait."): | |
| # st.session_state.username = username | |
| # st.session_state.chat_history = [] | |
| # st.session_state.vectorstore = setup_vectorstore() | |
| # st.success(f"Welcome, {username}! Now select a model to start chatting.") | |
| # else: | |
| # username = st.session_state.username | |
| # # Model selection options | |
| # model_options = [ | |
| # "gemma2-9b-it", | |
| # "llama-3.1-8b-instant", | |
| # "llama3-70b-8192", | |
| # "llama3-8b-8192" | |
| # ] | |
| # # Model dropdown | |
| # selected_model = st.selectbox("Choose a model:", model_options) | |
| # # Ensure vectorstore exists | |
| # if "vectorstore" not in st.session_state: | |
| # st.session_state.vectorstore = setup_vectorstore() | |
| # # Set or update the selected model | |
| # if "selected_model" not in st.session_state: | |
| # st.session_state.selected_model = selected_model | |
| # # Reset conversational_chain if model changes or not yet initialized | |
| # if ("conversational_chain" not in st.session_state) or (st.session_state.selected_model != selected_model): | |
| # st.session_state.selected_model = selected_model | |
| # st.session_state.conversational_chain = chat_chain(st.session_state.vectorstore, selected_model) | |
| # st.session_state.chat_history = [] | |
| # # Reset chat manually | |
| # if st.button("🔄 Reset Chat"): | |
| # st.session_state.chat_history = [] | |
| # st.session_state.conversational_chain = chat_chain(st.session_state.vectorstore, st.session_state.selected_model) | |
| # st.success("Chat reset!") | |
| # # Show chat UI | |
| # if "username" in st.session_state: | |
| # st.subheader(f"Hello {username}, start your query below!") | |
| # if st.session_state.chat_history: | |
| # for message in st.session_state.chat_history: | |
| # if message['role'] == 'user': | |
| # with st.chat_message("user"): | |
| # st.markdown(message["content"]) | |
| # elif message['role'] == 'assistant': | |
| # with st.chat_message("assistant"): | |
| # st.markdown(message["content"]) | |
| # user_input = st.chat_input("Ask AI....") | |
| # if user_input: | |
| # with st.spinner("Processing your query... Please wait."): | |
| # st.session_state.chat_history.append({"role": "user", "content": user_input}) | |
| # with st.chat_message("user"): | |
| # st.markdown(user_input) | |
| # with st.chat_message("assistant"): | |
| # response = st.session_state.conversational_chain({"question": user_input}) | |
| # assistant_response = response["answer"] | |
| # st.markdown(assistant_response) | |
| # st.session_state.chat_history.append({"role": "assistant", "content": assistant_response}) | |