Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,4 +1,4 @@
-# 🫠 Clean and Final Streamlit RAG App (Hugging Face + Local Ready)
+# 🫠 Clean and Final Streamlit RAG App (Three-Agent Architecture)
 
 # --- Environment Setup (Safe for Hugging Face) ---
 import os
@@ -40,54 +40,72 @@ embeddings_model = OpenAIEmbeddings(openai_api_key=API_KEY)
 st.set_page_config(page_title="RAG File Chat", layout="centered")
 st.title("🧠 Chat with Your Uploaded File")
 
-# --- Session State ---
-if "uploaded_file" not in st.session_state:
-    st.session_state.uploaded_file = None
-if "file_uploaded" not in st.session_state:
-    st.session_state.file_uploaded = False
-if "vectorstore" not in st.session_state:
-    st.session_state.vectorstore = None
-if "agent" not in st.session_state:
-    st.session_state.agent = None
-if "file_type" not in st.session_state:
-    st.session_state.file_type = None
-
-# --- File Parsing Functions ---
-def extract_text_from_file(file_content, file_type):
-    if file_type == "pdf":
-        reader = PyPDF2.PdfReader(io.BytesIO(file_content))
-        return "\n".join([page.extract_text() for page in reader.pages if page.extract_text()])
-    elif file_type == "docx":
-        doc = Document(io.BytesIO(file_content))
-        return "\n".join([p.text for p in doc.paragraphs if p.text.strip()])
-    return ""
-
-def create_agent_and_index(file_content, file_type):
+# --- Agent 1: File Ingestion and Indexing ---
+def agent_file_loader(file_content, file_type):
+    st.info("📂 Agent 1: Loading and indexing your file...")
     if file_type == "csv":
         df = pd.read_csv(io.StringIO(file_content.decode("utf-8")))
         llm = OpenAI(openai_api_key=API_KEY)
-        st.session_state.agent = create_pandas_dataframe_agent(llm, df, verbose=False)
-        st.success("🤖 Agent created for CSV.")
+        return create_pandas_dataframe_agent(llm, df, verbose=False), None
     elif file_type == "xlsx":
         df = pd.read_excel(file_content)
         llm = OpenAI(openai_api_key=API_KEY)
-        st.session_state.agent = create_pandas_dataframe_agent(llm, df, verbose=False)
-        st.success("🤖 Agent created for Excel.")
+        return create_pandas_dataframe_agent(llm, df, verbose=False), None
     elif file_type == "json":
         df = pd.DataFrame(json.loads(file_content.decode("utf-8")))
         llm = OpenAI(openai_api_key=API_KEY)
-        st.session_state.agent = create_pandas_dataframe_agent(llm, df, verbose=False)
-        st.success("🤖 Agent created for JSON.")
+        return create_pandas_dataframe_agent(llm, df, verbose=False), None
     elif file_type in ["pdf", "docx"]:
         text = extract_text_from_file(file_content, file_type)
         chunks = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0).split_text(text)
-        st.session_state.vectorstore = FAISS.from_texts(chunks, embeddings_model)
-        st.success("📊 Text embedded into FAISS vectorstore.")
+        vectorstore = FAISS.from_texts(chunks, embeddings_model)
+        return None, vectorstore
     else:
         st.error("❌ Unsupported file type.")
-        return
-    st.session_state.file_uploaded = True
-    st.session_state.file_type = file_type
+        return None, None
+
+# --- Agent 2: Query Resolution ---
+def agent_query_executor(query, file_type, df_agent=None, vectorstore=None):
+    st.info("🧠 Agent 2: Processing your question...")
+    if file_type in ["pdf", "docx"]:
+        qa_chain = RetrievalQA.from_chain_type(
+            llm=OpenAI(openai_api_key=API_KEY),
+            chain_type="stuff",
+            retriever=vectorstore.as_retriever(search_kwargs={"k": 5}),
+        )
+        result = qa_chain({"query": query})
+        return result["result"]
+    else:
+        return df_agent.run(query)
+
+# --- Agent 3: Response Enhancement ---
+def agent_response_enhancer(response):
+    st.info("🔍 Agent 3: Reviewing and enhancing the response...")
+    enhancement_prompt = f"Improve the clarity and format of the following response:\n{response}"
+    llm = OpenAI(openai_api_key=API_KEY)
+    return llm.invoke(enhancement_prompt)
+
+# --- Helper Function for Text Extraction ---
+def extract_text_from_file(file_content, file_type):
+    if file_type == "pdf":
+        reader = PyPDF2.PdfReader(io.BytesIO(file_content))
+        return "\n".join([page.extract_text() for page in reader.pages if page.extract_text()])
+    elif file_type == "docx":
+        doc = Document(io.BytesIO(file_content))
+        return "\n".join([p.text for p in doc.paragraphs if p.text.strip()])
+    return ""
+
+# --- Session State ---
+if "uploaded_file" not in st.session_state:
+    st.session_state.uploaded_file = None
+if "file_uploaded" not in st.session_state:
+    st.session_state.file_uploaded = False
+if "vectorstore" not in st.session_state:
+    st.session_state.vectorstore = None
+if "agent" not in st.session_state:
+    st.session_state.agent = None
+if "file_type" not in st.session_state:
+    st.session_state.file_type = None
 
 # --- File Upload UI ---
 MAX_SIZE_MB = 50
@@ -106,8 +124,16 @@ if st.session_state.uploaded_file and st.button("📤 Upload File"):
     try:
         content = st.session_state.uploaded_file.read()
         ftype = st.session_state.uploaded_file.name.split(".")[-1].lower()
-        with st.spinner("🔄 Processing file..."):
-            create_agent_and_index(content, ftype)
+        with st.spinner("🔄 Processing file with Agent 1..."):
+            agent, vectorstore = agent_file_loader(content, ftype)
+            if agent or vectorstore:
+                st.session_state.agent = agent
+                st.session_state.vectorstore = vectorstore
+                st.session_state.file_uploaded = True
+                st.session_state.file_type = ftype
+                st.success("✅ File processed successfully.")
+            else:
+                st.error("⚠️ Failed to process file.")
     except Exception as e:
         st.error("❌ Upload failed. Try a smaller file or check connection.")
         st.exception(e)
@@ -121,30 +147,29 @@ if st.session_state.file_uploaded:
         if not query.strip():
             st.warning("⚠️ Please enter a valid question.")
         else:
-            with st.spinner("💡 Thinking..."):
-                if st.session_state.file_type in ["pdf", "docx"]:
-                    qa_chain = RetrievalQA.from_chain_type(
-                        llm=OpenAI(openai_api_key=API_KEY),
-                        chain_type="stuff",
-                        retriever=st.session_state.vectorstore.as_retriever(search_kwargs={"k": 5}),
-                    )
-                    result = qa_chain({"query": query})
-                    response = result["result"]
-                else:
-                    response = st.session_state.agent.run(query)
-
-            st.subheader("📌 Answer")
+            with st.spinner("💡 Sending query to Agent 2..."):
+                raw_response = agent_query_executor(
+                    query,
+                    st.session_state.file_type,
+                    df_agent=st.session_state.agent,
+                    vectorstore=st.session_state.vectorstore,
+                )
+
+            with st.spinner("✨ Enhancing response with Agent 3..."):
+                enhanced_response = agent_response_enhancer(raw_response)
+
+            st.subheader("📌 Final Answer")
             if output_format == "Plain Text":
-                st.text(response)
+                st.text(enhanced_response)
             elif output_format == "Markdown":
-                st.markdown(response)
+                st.markdown(enhanced_response)
             elif output_format == "Tabular View":
-                rows = [line.split("\t") for line in response.split("\n") if "\t" in line]
+                rows = [line.split("\t") for line in enhanced_response.split("\n") if "\t" in line]
                 if not rows or len(rows[0]) == 1:
-                    rows = [line.split(",") for line in response.split("\n") if "," in line]
+                    rows = [line.split(",") for line in enhanced_response.split("\n") if "," in line]
                 try:
                     df = pd.DataFrame(rows[1:], columns=rows[0])
                     st.dataframe(df)
                 except Exception:
                     st.warning("⚠️ Could not render table. Showing raw text.")
-                    st.text(response)
+                    st.text(enhanced_response)