Upload 2 files

app.py

@@ -1,9 +1,8 @@
 import streamlit as st
-import chromadb
-from llama_index.core import VectorStoreIndex, Settings, StorageContext
-from llama_index.vector_stores.chroma import ChromaVectorStore
+import os
+from llama_index.core import VectorStoreIndex, Settings, SimpleDirectoryReader, StorageContext
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
-# This is the new, correct import for the Gemini model
+# This is the new, correct import for the Gemini model from the modern library
 from llama_index.llms.google import Gemini 
 
 # --- App Configuration ---
@@ -25,58 +24,83 @@ with st.sidebar:
         "Get your free API key from [Google AI Studio](https://aistudio.google.com/)."
     )
 
-# --- Caching to avoid reloading models on every interaction ---
+# --- Define paths for your raw data and the persistent storage ---
+DATA_DIR = "./data"
+STORAGE_DIR = "./storage"
+
+# --- Caching Functions to improve performance ---
+
 @st.cache_resource(show_spinner="Loading embedding model...")
 def load_embed_model():
-    """Loads and caches the embedding model."""
+    """Loads and caches the embedding model from Hugging Face."""
     return HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
 
-@st.cache_resource(show_spinner="Connecting to knowledge base...")
-def load_vector_store():
-    """Loads and caches the ChromaDB vector store and index."""
-    db = chromadb.PersistentClient(path="./chroma_db")
-    chroma_collection = db.get_or_create_collection("genshin_sandbox_guide")
-    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
-    storage_context = StorageContext.from_defaults(vector_store=vector_store)
-    # This correctly loads an index from an existing vector store.
-    return VectorStoreIndex.from_vector_store(vector_store=vector_store, storage_context=storage_context)
+@st.cache_resource(show_spinner="Loading knowledge base... (This may take a moment on first startup)")
+def load_index():
+    """
+    Loads or builds the vector index. 
+    On the very first run, it will build the index from the documents in DATA_DIR.
+    On subsequent runs, it will load the pre-built index from STORAGE_DIR.
+    """
+    Settings.embed_model = load_embed_model()
+
+    if not os.path.exists(STORAGE_DIR):
+        st.info("First time startup: Building the knowledge base from your documents. Please wait.")
+        # Load all your .txt files from the 'data' directory
+        documents = SimpleDirectoryReader(DATA_DIR).load_data()
+        
+        # Create the index from the documents
+        index = VectorStoreIndex.from_documents(documents)
+        
+        # Save the newly built index to disk for future use
+        index.storage_context.persist(persist_dir=STORAGE_DIR)
+        st.success("Knowledge base built successfully!")
+    else:
+        # Load the existing index from the 'storage' directory
+        storage_context = StorageContext.from_defaults(persist_dir=STORAGE_DIR)
+        index = VectorStoreIndex.from_storage(storage_context)
+        
+    return index
 
-# --- Main App Logic ---
-embed_model = load_embed_model()
-index = load_vector_store()
+# --- Main Application Logic ---
+index = load_index()
 
 if "messages" not in st.session_state:
     st.session_state.messages = [
         {"role": "assistant", "content": "Hi! How can I help you build in the Miliastra Sandbox today?"}
     ]
 
-# Display chat messages
+# Display existing chat messages
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
-# Process user input
+# Handle user input
 if prompt := st.chat_input("Your question"):
+    # Ensure the user has entered their API key
     if not google_api_key:
         st.info("Please add your Google API Key to continue.")
         st.stop()
 
+    # Add user message to chat history
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
         st.markdown(prompt)
 
+    # Generate and display the assistant's response
     with st.chat_message("assistant"):
         message_placeholder = st.empty()
         message_placeholder.markdown("Thinking...")
 
         try:
-            # Use the new Gemini class
+            # Configure the LLM with the API key
             Settings.llm = Gemini(model_name="models/gemini-1.5-pro-latest", api_key=google_api_key)
-            Settings.embed_model = embed_model
             
-            query_engine = index.as_query_engine(similarity_top_k=3)
+            # Create a query engine from the loaded index
+            query_engine = index.as_query_engine(similarity_top_k=3) # Get top 3 relevant text chunks
             response = query_engine.query(prompt)
             
+            # Display the final answer
             message_placeholder.markdown(str(response))
             st.session_state.messages.append({"role": "assistant", "content": str(response)})
 

requirements.txt

@@ -1,6 +1,4 @@
 streamlit
 llama-index
 llama-index-embeddings-huggingface
-llama-index-vector-stores-chroma
-llama-index-integrations-llms-google
-chromadb
+llama-index-integrations-llms-google