Update app.py

app.py

@@ -1,34 +1,39 @@
-# Combined Imports
 import os
 import streamlit as st
-from dotenv import load_dotenv
 from apify_client import ApifyClient
+from dotenv import load_dotenv
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.chains import ConversationalRetrievalChain
 from langchain.chat_models import ChatOpenAI
-from langchain.document_loaders import ApifyDatasetLoader
-from langchain.document_loaders.base import Document
 from langchain.embeddings import OpenAIEmbeddings
-from langchain.embeddings.openai import OpenAIEmbeddings
 from langchain.memory import ConversationBufferMemory
 from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
-from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
+from langchain.document_loaders import ApifyDatasetLoader
+from langchain.document_loaders.base import Document
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 
-# Environment variables and configuration
 load_dotenv()
-WEBSITE_URL = os.environ.get('WEBSITE_URL', 'a website')
-OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
+
+st.set_page_config(page_title='Chat with a website')
+
+website_url = st.text_input("Please enter the website URL to scrape:", value="https://www.example.com/")
+st.title(f'Chat with {website_url}')
+
 APIFY_API_TOKEN = os.environ.get('APIFY_API_TOKEN')
 
-# Scraper Functionality
-def scrape_website():
+@st.cache_resource(ttl='1h')
+def get_retriever():
+    embeddings = OpenAIEmbeddings()
+    vectordb = Chroma(persist_directory='db', embedding_function=embeddings)
+    retriever = vectordb.as_retriever(search_type='mmr')
+    return retriever
+
+def scrape_website(website_url: str):
     apify_client = ApifyClient(APIFY_API_TOKEN)
-    st.write(f'Extracting data from "{WEBSITE_URL}". Please wait...')
     actor_run_info = apify_client.actor('apify/website-content-crawler').call(
-        run_input={'startUrls': [{'url': WEBSITE_URL}]}
+        run_input={'startUrls': [{'url': website_url}]}
     )
-    st.write('Saving data into the vector database. Please wait...')
     loader = ApifyDatasetLoader(
         dataset_id=actor_run_info['defaultDatasetId'],
         dataset_mapping_function=lambda item: Document(
@@ -38,49 +43,16 @@ def scrape_website():
     documents = loader.load()
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=100)
     docs = text_splitter.split_documents(documents)
-
-    embedding = OpenAIEmbeddings(api_key=OPENAI_API_KEY)
+    embedding = OpenAIEmbeddings()
     vectordb = Chroma.from_documents(
         documents=docs,
         embedding=embedding,
         persist_directory='db2',
     )
     vectordb.persist()
-    st.write('All done!')
-
-# Chat Functionality
-def chat_with_website():
-    st.set_page_config(page_title=f'Chat with {WEBSITE_URL}')
-    st.title('Chat with a website')
-    retriever = get_retriever()
-    msgs = StreamlitChatMessageHistory()
-    memory = ConversationBufferMemory(memory_key='chat_history', chat_memory=msgs, return_messages=True)
-
-    llm = ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0, streaming=True)
-    qa_chain = ConversationalRetrievalChain.from_llm(
-        llm, retriever=retriever, memory=memory, verbose=False
-    )
 
-    if st.sidebar.button('Clear message history') or len(msgs.messages) == 0:
-        msgs.clear()
-        msgs.add_ai_message(f'Ask me anything about {WEBSITE_URL}!')
-
-    avatars = {'human': 'user', 'ai': 'assistant'}
-    for msg in msgs.messages:
-        st.chat_message(avatars[msg.type]).write(msg.content)
-
-    if user_query := st.chat_input(placeholder='Ask me anything!'):
-        st.chat_message('user').write(user_query)
-        with st.chat_message('assistant'):
-            stream_handler = StreamHandler(st.empty())
-            response = qa_chain.run(user_query, callbacks=[stream_handler])
-
-@st.cache_resource(ttl='1h')
-def get_retriever():
-    embeddings = OpenAIEmbeddings()
-    vectordb = Chroma(persist_directory='db', embedding_function=embeddings)
-    retriever = vectordb.as_retriever(search_type='mmr')
-    return retriever
+if st.button("Start Scraping"):
+    scrape_website(website_url)
 
 class StreamHandler(BaseCallbackHandler):
     def __init__(self, container: st.delta_generator.DeltaGenerator, initial_text: str = ''):
@@ -91,9 +63,25 @@ class StreamHandler(BaseCallbackHandler):
         self.text += token
         self.container.markdown(self.text)
 
-# Main App Flow
-if st.sidebar.button("Scrape a new website"):
-    scrape_website()
+retriever = get_retriever()
+msgs = StreamlitChatMessageHistory()
+memory = ConversationBufferMemory(memory_key='chat_history', chat_memory=msgs, return_messages=True)
+llm = ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0, streaming=True)
+qa_chain = ConversationalRetrievalChain.from_llm(
+    llm, retriever=retriever, memory=memory, verbose=False
+)
+
+if st.sidebar.button('Clear message history') or len(msgs.messages) == 0:
+    msgs.clear()
+    msgs.add_ai_message(f'Ask me anything about {website_url}!')
+
+avatars = {'human': 'user', 'ai': 'assistant'}
+for msg in msgs.messages:
+    st.chat_message(avatars[msg.type]).write(msg.content)
+
+if user_query := st.chat_input(placeholder='Ask me anything!'):
+    st.chat_message('user').write(user_query)
 
-if st.sidebar.button("Chat with scraped website"):
-    chat_with_website()
+    with st.chat_message('assistant'):
+        stream_handler = StreamHandler(st.empty())
+        response = qa_chain.run(user_query, callbacks=[stream_handler])