Update app.py

app.py

@@ -1,29 +1,85 @@
+# Combined Imports
 import os
-
 import streamlit as st
 from dotenv import load_dotenv
+from apify_client import ApifyClient
 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
 
+# Environment variables and configuration
 load_dotenv()
-
-website_url = os.environ.get('WEBSITE_URL', 'a website')
-
-st.set_page_config(page_title=f'Chat with {website_url}')
-st.title('Chat with a website')
+WEBSITE_URL = os.environ.get('WEBSITE_URL', 'a website')
+OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
+APIFY_API_TOKEN = os.environ.get('APIFY_API_TOKEN')
+
+# Scraper Functionality
+def scrape_website():
+    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}]}
+    )
+    st.write('Saving data into the vector database. Please wait...')
+    loader = ApifyDatasetLoader(
+        dataset_id=actor_run_info['defaultDatasetId'],
+        dataset_mapping_function=lambda item: Document(
+            page_content=item['text'] or '', metadata={'source': item['url']}
+        ),
+    )
+    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)
+    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
 
 class StreamHandler(BaseCallbackHandler):
@@ -35,27 +91,9 @@ class StreamHandler(BaseCallbackHandler):
         self.text += token
         self.container.markdown(self.text)
 
-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)
+# Main App Flow
+if st.sidebar.button("Scrape a new website"):
+    scrape_website()
 
-    with st.chat_message('assistant'):
-        stream_handler = StreamHandler(st.empty())
-        response = qa_chain.run(user_query, callbacks=[stream_handler])
+if st.sidebar.button("Chat with scraped website"):
+    chat_with_website()