Update app.py

app.py

@@ -8,13 +8,19 @@ import chromadb
 from sentence_transformers import SentenceTransformer
 import google.generativeai as genai
 
+# Initialize Gemini API
 genai.configure(api_key="AIzaSyAxUd2tS-qj9C7frYuHRsv92tziXHgIvLo")
 
+# Initialize ChromaDB
 CHROMA_PATH = "chroma_db"
 chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
 collection = chroma_client.get_or_create_collection(name="formula_1")
 embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
 
+# Initialize session state to track if scraping is complete
+if 'scraped' not in st.session_state:
+    st.session_state.scraped = False
+
 def clean_text(text):
     text = re.sub(r'http\S+', '', text)
     text = re.sub(r'\s+', ' ', text).strip()
@@ -39,6 +45,8 @@ def scrape_text(url):
         text = clean_text(soup.get_text())
         chunks = split_content_into_chunks(text)
         add_chunks_to_db(chunks)
+        # Set scraped state to True when complete
+        st.session_state.scraped = True
         return "Scraping and processing complete. You can now ask questions!"
     except requests.exceptions.RequestException as e:
         return f"Error scraping {url}: {e}"
@@ -47,29 +55,73 @@ def ask_question(query):
     query_embedding = embedding_model.encode(query, convert_to_list=True)
     results = collection.query(query_embeddings=[query_embedding], n_results=2)
     top_chunks = results.get("documents", [[]])[0]
+    
     system_prompt = """
     You are a Formula 1 expert. You answer questions about Formula 1.
     But you only answer based on knowledge I'm providing you. You don't use your internal
     knowledge and you don't make things up.
     If you don't know the answer, just say: I don't know.
     """ + str(top_chunks)
+    
     full_prompt = system_prompt + "\nUser Query: " + query
     model = genai.GenerativeModel('gemini-2.0-flash')
     response = model.generate_content(full_prompt)
     return response.text
 
-st.title("Web Scraping & Chatbot")
+# Main UI
+st.title("Formula 1 Web Scraper & Chatbot")
+
+# Scraping section
+with st.container():
+    st.subheader("Step 1: Scrape a Formula 1 Website")
+    url = st.text_input("Enter a Formula 1 related URL:")
+    
+    if url:
+        if st.button("Scrape & Process"):
+            with st.spinner("Scraping and processing content..."):
+                result = scrape_text(url)
+                st.success(result)
+
+# Q&A section - only appears after scraping is complete
+if st.session_state.scraped:
+    with st.container():
+        st.subheader("Step 2: Ask Questions About Formula 1")
+        st.write("The database contains information scraped from the website. Ask a question about Formula 1:")
+        
+        # Chat history
+        if 'chat_history' not in st.session_state:
+            st.session_state.chat_history = []
+            
+        # Display chat history
+        for message in st.session_state.chat_history:
+            with st.chat_message(message["role"]):
+                st.write(message["content"])
+        
+        # Input for new question
+        user_query = st.chat_input("Ask your Formula 1 question here")
+        
+        if user_query:
+            # Add user question to chat history
+            st.session_state.chat_history.append({"role": "user", "content": user_query})
+            
+            # Display user question
+            with st.chat_message("user"):
+                st.write(user_query)
+            
+            # Get and display answer
+            with st.chat_message("assistant"):
+                with st.spinner("Searching Formula 1 database..."):
+                    answer = ask_question(user_query)
+                    st.write(answer)
+                    
+            # Add answer to chat history
+            st.session_state.chat_history.append({"role": "assistant", "content": answer})
 
-url = st.text_input("Enter a URL:")
-if url:
-    if st.button("Scrape & Process"):
-        result = scrape_text(url)
-        st.success(result)
+else:
+    st.info("Please scrape a Formula 1 website first to populate the database, then you can ask questions!")
 
-if 'scraped' in st.session_state and st.session_state.scraped:
-    st.subheader("Ask a Question")
-    query = st.text_input("Enter your question:")
-    if query:
-        if st.button("Get Answer"):
-            answer = ask_question(query)
-            st.write(answer)
+# Add a button to clear the session and start over
+if st.button("Clear Chat History and Data"):
+    st.session_state.chat_history = []
+    st.session_state.scraped = False
+    st.experimental_rerun()