Update app.py

app.py

@@ -14,6 +14,11 @@ from llama_index.readers.web import FireCrawlWebReader
 from llama_index.core import SummaryIndex
 import streamlit_analytics2 as streamlit_analytics
 import time
+import dotenv
+
+dotenv.load_dotenv()
+# Set page config
+#st.set_page_config(page_title="Talk to Software Documentation", page_icon="📚", layout="wide")
 
 # Initialize session state
 if 'setup_complete' not in st.session_state:
@@ -24,8 +29,15 @@ if 'chat_history' not in st.session_state:
     st.session_state['chat_history'] = []
 if 'index' not in st.session_state:
     st.session_state['index'] = None
+if 'url' not in st.session_state:
+    st.session_state['url'] = ""
+if 'collection_name' not in st.session_state:
+    st.session_state['collection_name'] = ""
+if 'query' not in st.session_state:
+    st.session_state['query'] = ""
 
 os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY")
+os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY")
 
 # Setup functions
 def embed_setup():
@@ -81,58 +93,65 @@ st.title("Talk to Software Documentation")
 
 st.markdown("""
 This tool allows you to chat with software documentation. Here's how to use it:
-1. Enter the URL of the documentation you want to chat about.
-2. Click the "Ingest and Setup" button to crawl the documentation and set up the query engine.
-3. Once setup is complete, enter your query in the text box.
-4. Click "Search" to get a response based on the documentation.
-5. View your chat history in the sidebar.
+1. Enter the URL of the documentation you want to chat about (optional if using an existing collection).
+2. Enter the collection name for the vector store.
+3. Click the "Ingest and Setup" button to crawl the documentation (if URL provided) and set up the query engine.
+4. Once setup is complete, enter your query in the text box.
+5. Click "Search" to get a response based on the documentation.
+6. View your chat history in the sidebar.
 """)
 
 with streamlit_analytics.track():
     # URL input for document ingestion
-    url = st.text_input("Enter URL to crawl and ingest documents:")
+    st.session_state['url'] = st.text_input("Enter URL to crawl and ingest documents (optional):", value=st.session_state['url'])
+    
+    # Collection name input
+    st.session_state['collection_name'] = st.text_input("Enter collection name for vector store:", value=st.session_state['collection_name'])
     
     # Combined Ingest and Setup button
     if st.button("Ingest and Setup"):
-        if url:
-            with st.spinner("Crawling, ingesting documents, and setting up query engine..."):
-                st.session_state['documents'] = ingest_documents(url)
-                embed_setup()
-                client = qdrant_setup()
-                llm = llm_setup()
-                vector_store = QdrantVectorStore(client=client, collection_name="testdoc1")
-                index = VectorStoreIndex.from_documents(st.session_state['documents'], vector_store=vector_store)
-                st.session_state['index'] = index
-                st.session_state['setup_complete'] = True
-            st.success(f"Documents ingested from {url} and query engine setup completed successfully!")
-        else:
-            st.error("Please enter a URL")
+        with st.spinner("Setting up query engine..."):
+            embed_setup()
+            client = qdrant_setup()
+            llm = llm_setup()
+            vector_store = QdrantVectorStore(client=client, collection_name=st.session_state['collection_name'])
+            storage_context = StorageContext.from_defaults(vector_store=vector_store)
+            
+            if st.session_state['url']:
+                st.session_state['documents'] = ingest_documents(st.session_state['url'])
+                st.session_state['index'] = VectorStoreIndex.from_documents(st.session_state['documents'], vector_store=vector_store, storage_context=storage_context)
+                st.success(f"Documents ingested from {st.session_state['url']} and query engine setup completed successfully!")
+            else:
+                st.session_state['index'] = VectorStoreIndex.from_vector_store(vector_store=vector_store, storage_context=storage_context)
+                st.success(f"Query engine setup completed successfully using existing collection: {st.session_state['collection_name']}")
+            
+            st.session_state['setup_complete'] = True
     
     # Query input
-    query = st.text_input("Enter your query:(please click on the search button, do not just press enter)")
+    st.session_state['query'] = st.text_input("Enter your query:", value=st.session_state['query'])
     
     # Search button
     if st.button("Search"):
         if not st.session_state['setup_complete']:
             st.error("Please complete the setup first")
-        elif query:
+        elif st.session_state['query']:
             with st.spinner("Searching..."):
                 try:
                     chat_engine = query_index(st.session_state['index'])
-                    response = chat_engine.chat(query)
+                    response = chat_engine.chat(st.session_state['query'])
                 except Exception as e:
                     st.error(f"An error occurred: {str(e)}")
                     st.info("Retrying in 120 seconds...")
                     time.sleep(120)
                     try:
                         chat_engine = query_index(st.session_state['index'])
-                        response = chat_engine.chat(query)
+                        response = chat_engine.chat(st.session_state['query'])
                     except Exception as e:
                         st.error(f"Retry failed. Error: {str(e)}")
                         st.stop()
 
             # Add the query and response to chat history
-            st.session_state['chat_history'].append(("User", query))
+            st.session_state['chat_history'].append(("User", st.session_state['query']))
             st.session_state['chat_history'].append(("Assistant", str(response.response)))
             
             # Display the most recent response prominently