Update app.py

app.py

@@ -1,65 +1,79 @@
 import streamlit as st
 import os
-from langchain_community.document_loaders import TextLoader  # ✅ FIXED
+
+# HF Spaces secret
+os.environ["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"] = st.secrets["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"]
+
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain_community.document_loaders import TextLoader
 from langchain.text_splitter import CharacterTextSplitter
-from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 from langchain.vectorstores import FAISS
 from langchain.prompts import ChatPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 
-# HF Spaces secrets
-os.environ["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"] = st.secrets["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"]
-
 @st.cache_resource
-def load_retriever():
+def get_chatbot():
+    # Load or create vector store
     if not os.path.exists("faiss_index"):
-        st.info("Creating vector index... (first run)")
-        loader = TextLoader("knowledge.txt")
-        docs = loader.load()
-        splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        chunks = splitter.split_documents(docs)
+        with open("knowledge.txt", "r") as f:
+            content = f.read()
+        docs = [content]  # Simple doc
+        
+        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        splits = text_splitter.create_documents([content])
+        
         embeddings = OpenAIEmbeddings()
-        vectorstore = FAISS.from_documents(chunks, embeddings)
+        vectorstore = FAISS.from_documents(splits, embeddings)
         vectorstore.save_local("faiss_index")
     
+    # Load retriever
     embeddings = OpenAIEmbeddings()
     vectorstore = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
-    return vectorstore.as_retriever(search_kwargs={"k": 3})
-
-# Load components
-retriever = load_retriever()
-llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
-prompt = ChatPromptTemplate.from_template(
-    "Use this context only: {context}\n\nQuestion: {question}\nAnswer:"
-)
-
-def get_response(query):
-    context_docs = retriever.invoke(query)
-    context = "\n".join([doc.page_content for doc in context_docs])
-    chain = (
-        {"context": lambda _: context, "question": lambda _: query} 
-        | prompt 
-        | llm 
-        | StrOutputParser()
+    retriever = vectorstore.as_retriever()
+    
+    # LLM
+    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+    
+    prompt = ChatPromptTemplate.from_template(
+        "Answer using ONLY this context:\n{context}\n\nQuestion: {question}\n\nAnswer:"
     )
-    return chain.invoke(query)
+    
+    def rag_chain(query):
+        context_docs = retriever.invoke(query)
+        context = "\n".join([doc.page_content for doc in context_docs])
+        chain = (
+            {"context": lambda x: context, "question": lambda x: query} 
+            | prompt 
+            | llm 
+            | StrOutputParser()
+        )
+        return chain.invoke(query)
+    
+    return rag_chain
 
 # UI
 st.title("🧠 RAG Chatbot")
+chatbot = get_chatbot()
+
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-for msg in st.session_state.messages:
-    with st.chat_message(msg["role"]):
-        st.markdown(msg["content"])
+# Show chat history
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
 
-if prompt := st.chat_input("Ask about your documents..."):
+# Chat input
+if prompt := st.chat_input("Ask about your knowledge base..."):
+    # Add user message
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
         st.markdown(prompt)
     
+    # Generate response
     with st.chat_message("assistant"):
-        with st.spinner("Thinking..."):
-            response = get_response(prompt)
+        with st.spinner("Searching knowledge base..."):
+            response = chatbot(prompt)
             st.markdown(response)
+    
     st.session_state.messages.append({"role": "assistant", "content": response})