Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,14 +1,13 @@
-import streamlit as st
-import pandas as pd
-import json
-import io
+# Fix permission issue on Hugging Face Spaces
 import os
-
 os.environ["STREAMLIT_HOME"] = "/tmp"
 os.environ["XDG_CONFIG_HOME"] = "/tmp"
 os.environ["XDG_DATA_HOME"] = "/tmp"
 
 import streamlit as st
+import pandas as pd
+import json
+import io
 
 from langchain.llms import OpenAI
 from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
@@ -27,22 +26,25 @@ _ = load_dotenv(find_dotenv())
 # Get API key from Streamlit secrets
 API_KEY = os.getenv("OPENAI_API_KEY")
 
-# Initialize embedding model
 embeddings_model = OpenAIEmbeddings(openai_api_key=API_KEY)
 
-# Set up Streamlit app
+# Streamlit app setup
 st.set_page_config(page_title="RAG File Chat", layout="centered")
 st.title("🧠 Chat with Your Uploaded File")
 
-# Session state
+# Initialize session state
 if "vectorstore" not in st.session_state:
     st.session_state.vectorstore = None
 if "agent_created" not in st.session_state:
     st.session_state.agent_created = False
+if "file_type" not in st.session_state:
+    st.session_state.file_type = None
+if "agent" not in st.session_state:
+    st.session_state.agent = None
 
 
 def extract_text_from_file(file_content, file_type):
-    """Extract raw text from supported document formats."""
+    """Extract text from PDF or DOCX."""
     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()])
@@ -53,46 +55,41 @@ def extract_text_from_file(file_content, file_type):
 
 
 def create_agent_and_index(file_content, file_type):
-    """Create a LangChain agent and index file content with FAISS."""
+    """Process and embed file content."""
     if file_type == "csv":
         df = pd.read_csv(io.StringIO(file_content.decode("utf-8")))
         st.success("📄 CSV file loaded into DataFrame.")
+        llm = OpenAI(openai_api_key=API_KEY)
+        st.session_state.agent = create_pandas_dataframe_agent(llm, df, verbose=False)
+        st.success("🤖 Agent created for tabular data.")
     elif file_type == "xlsx":
         df = pd.read_excel(file_content)
         st.success("📄 Excel file loaded into DataFrame.")
+        llm = OpenAI(openai_api_key=API_KEY)
+        st.session_state.agent = create_pandas_dataframe_agent(llm, df, verbose=False)
+        st.success("🤖 Agent created for tabular data.")
     elif file_type == "json":
         df = pd.DataFrame(json.loads(file_content.decode("utf-8")))
         st.success("📄 JSON file loaded into DataFrame.")
+        llm = OpenAI(openai_api_key=API_KEY)
+        st.session_state.agent = create_pandas_dataframe_agent(llm, df, verbose=False)
+        st.success("🤖 Agent created for tabular data.")
     elif file_type in ["pdf", "docx"]:
         text = extract_text_from_file(file_content, file_type)
         st.success(f"📃 Extracted text from {file_type.upper()}.")
         text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
         texts = text_splitter.split_text(text)
-        df = pd.DataFrame({"text": texts})
-        st.success("✂️ Split text into chunks.")
-        st.session_state.vectorstore = FAISS.from_texts(
-            texts=df['text'].tolist(),
-            embedding=embeddings_model
-        )
-        st.success("🧠 Text embedded and stored in FAISS (in-memory).")
+        st.session_state.vectorstore = FAISS.from_texts(texts, embeddings_model)
+        st.success("🧠 Text embedded and stored in FAISS.")
     else:
         st.error("❌ Unsupported file type.")
-        return None
-
-    # Create agent for tabular data
-    if file_type in ["csv", "xlsx", "json"]:
-        llm = OpenAI(openai_api_key=API_KEY)
-        agent = create_pandas_dataframe_agent(llm, df, verbose=False)
-        st.session_state.agent_created = True
-        st.success("🤖 Tabular data agent created.")
-        return agent
-
+        return
     st.session_state.agent_created = True
-    return None
+    st.session_state.file_type = file_type
 
 
 def query_vectorstore(query):
-    """Run RetrievalQA on FAISS vectorstore."""
+    """Query FAISS vectorstore using RAG."""
     qa_chain = RetrievalQA.from_chain_type(
         llm=OpenAI(openai_api_key=API_KEY),
         chain_type="stuff",
@@ -102,18 +99,21 @@ def query_vectorstore(query):
     return result["result"]
 
 
-# UI Logic
-uploaded_file = st.file_uploader("📁 Upload a file", type=["csv", "xlsx", "json", "pdf", "docx"])
+# --- UI Section: File Upload ---
+uploaded_file = st.file_uploader("📁 Browse and select a file", type=["csv", "xlsx", "json", "pdf", "docx"])
 
-if uploaded_file is not None:
-    st.success(f"✅ Uploaded: `{uploaded_file.name}` ({uploaded_file.size / 1024:.1f} KB)")
-    file_content = uploaded_file.read()
-    file_type = uploaded_file.name.split(".")[-1]
+if uploaded_file:
+    st.info(f"✅ File selected: `{uploaded_file.name}` ({uploaded_file.size / 1024:.1f} KB)")
+    if st.button("📤 Upload File"):
+        file_content = uploaded_file.read()
+        file_type = uploaded_file.name.split(".")[-1]
+        with st.spinner("🔄 Uploading and processing..."):
+            create_agent_and_index(file_content, file_type)
 
-    with st.spinner("🔄 Processing file and creating embeddings..."):
-        create_agent_and_index(file_content, file_type)
-        st.session_state.file_type = file_type
+# --- Output Format ---
+output_format = st.selectbox("📋 Select Output Format", ["Plain Text", "Markdown", "Tabular View"])
 
+# --- UI Section: Query ---
 query = st.text_area("💬 Ask a question about your uploaded file")
 
 if st.button("Submit Query"):
@@ -121,12 +121,27 @@ if st.button("Submit Query"):
         st.warning("⚠️ Please enter a valid question.")
     elif not st.session_state.agent_created:
         st.warning("📁 Please upload and process a file first.")
-    elif st.session_state.file_type in ["pdf", "docx"]:
-        with st.spinner("💡 Thinking..."):
-            response = query_vectorstore(query)
-        st.subheader("📌 Answer")
-        st.write(response)
     else:
-        st.warning("⚠️ Vector search is only available for PDF and DOCX files in this demo.")
+        with st.spinner("💡 Thinking..."):
+            if st.session_state.file_type in ["pdf", "docx"]:
+                response = query_vectorstore(query)
+            else:
+                response = st.session_state.agent.run(query)
 
+        st.subheader("📌 Answer")
 
+        if output_format == "Plain Text":
+            st.text(response)
+        elif output_format == "Markdown":
+            st.markdown(response)
+        elif output_format == "Tabular View":
+            # Try to parse tabular response (tab or comma-separated)
+            rows = [line.split("\t") for line in 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]
+            try:
+                df = pd.DataFrame(rows[1:], columns=rows[0])
+                st.dataframe(df)
+            except Exception:
+                st.warning("⚠️ Could not render a table. Showing raw output instead.")
+                st.text(response)