Update app.py

app.py

@@ -9,8 +9,7 @@ from langchain_community.document_loaders import PyPDFLoader, TextLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
-from langchain.retrievers import BM25Retriever
-from langchain.retrievers import EnsembleRetriever
+from langchain.retrievers import BM25Retriever, EnsembleRetriever
 from langchain.schema import Document
 from langchain.docstore.document import Document as LangchainDocument
 
@@ -21,13 +20,18 @@ HF_TOKEN = st.secrets["HF_TOKEN"]
 st.set_page_config(page_title="DigiTwin RAG", page_icon="📂", layout="centered")
 st.title("📂 DigiTs the Twin")
 
-# --- Upload Files Sidebar ---
+# --- Sidebar ---
 with st.sidebar:
     st.header("📄 Upload Knowledge Files")
     uploaded_files = st.file_uploader("Upload PDFs or .txt files", accept_multiple_files=True, type=["pdf", "txt"])
     hybrid_toggle = st.checkbox("🔀 Enable Hybrid Search", value=True)
+    clear_chat = st.button("🧹 Clear Chat History")
 
-# --- Model Loading ---
+# --- Session State ---
+if "messages" not in st.session_state or clear_chat:
+    st.session_state.messages = []
+
+# --- Load Model + Tokenizer ---
 @st.cache_resource
 def load_model():
     model_id = "tiiuae/falcon-7b-instruct"
@@ -37,7 +41,7 @@ def load_model():
 
 tokenizer, model = load_model()
 
-# --- Document Processing ---
+# --- Process Documents ---
 def process_documents(files):
     documents = []
     for file in files:
@@ -53,55 +57,62 @@ def chunk_documents(documents):
     splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
     return splitter.split_documents(documents)
 
-# --- Embedding and Retrieval ---
+# --- Build Hybrid Retriever ---
 def build_retrievers(chunks):
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     faiss_vectorstore = FAISS.from_documents(chunks, embeddings)
     faiss_retriever = faiss_vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 5})
-
     bm25_retriever = BM25Retriever.from_documents([LangchainDocument(page_content=d.page_content) for d in chunks])
     bm25_retriever.k = 5
+    hybrid = EnsembleRetriever(retrievers=[faiss_retriever, bm25_retriever], weights=[0.5, 0.5])
+    return faiss_retriever, hybrid
 
-    ensemble = EnsembleRetriever(retrievers=[faiss_retriever, bm25_retriever], weights=[0.5, 0.5])
-    return faiss_retriever, ensemble
-
-# --- Inference ---
-def generate_answer(query, retriever):
-    docs = retriever.get_relevant_documents(query)
-    context = "\n".join([doc.page_content for doc in docs])
-
-    system_prompt = (
-        "You are DigiTwin, an expert advisor in asset integrity, reliability, inspection, and maintenance "
-        "of topside piping, structural, mechanical systems, floating units, pressure vessels (VII), and pressure safety devices (PSD's). "
-        "Use the context below to answer professionally.\n\nContext:\n" + context + "\n\nQuery: " + query + "\nAnswer:"
-    )
-
+# --- Inference with Streaming ---
+def generate_stream_response(system_prompt):
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     inputs = tokenizer(system_prompt, return_tensors="pt").to(model.device)
     generation_kwargs = dict(**inputs, streamer=streamer, max_new_tokens=300)
-
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
-
-    answer = ""
+    partial_output = ""
     for token in streamer:
-        answer += token
-        yield answer
+        partial_output += token
+        yield partial_output
 
-# --- Main App ---
+# --- Main App Logic ---
 if uploaded_files:
     with st.spinner("Processing documents..."):
         docs = process_documents(uploaded_files)
         chunks = chunk_documents(docs)
         faiss_retriever, hybrid_retriever = build_retrievers(chunks)
-        st.success("Documents processed successfully.")
-
-    query = st.text_input("🔍 Ask a question based on the uploaded documents")
-    if query:
-        st.subheader("📤 Answer")
         retriever = hybrid_retriever if hybrid_toggle else faiss_retriever
-        response_placeholder = st.empty()
-        full_response = ""
-        for partial_response in generate_answer(query, retriever):
-            full_response = partial_response
-            response_placeholder.markdown(full_response)
+        st.success("Knowledge base ready. Ask your question below.")
+
+    for msg in st.session_state.messages:
+        with st.chat_message(msg["role"]):
+            st.markdown(msg["content"])
+
+    user_input = st.chat_input("💬 Ask DigiTwin something...")
+    if user_input:
+        st.chat_message("user").markdown(user_input)
+        st.session_state.messages.append({"role": "user", "content": user_input})
+
+        with st.chat_message("assistant"):
+            context_docs = retriever.get_relevant_documents(user_input)
+            context_text = "\n".join([doc.page_content for doc in context_docs])
+
+            system_prompt = (
+                "You are DigiTwin, an expert advisor in asset integrity, reliability, inspection, and maintenance "
+                "of topside piping, structural, mechanical systems, floating units, pressure vessels (VII), and pressure safety devices (PSD's).\n\n"
+                f"Context:\n{context_text}\n\n"
+                f"User: {user_input}\nAssistant:"
+            )
+
+            full_response = ""
+            response_area = st.empty()
+            for partial_output in generate_stream_response(system_prompt):
+                full_response = partial_output
+                response_area.markdown(full_response)
+            st.session_state.messages.append({"role": "assistant", "content": full_response})
+else:
+    st.info("👈 Upload one or more PDFs or .txt files to begin.")