Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,7 +2,6 @@ import streamlit as st
 import pandas as pd
 import json
 import io
-import os
 
 from langchain.llms import OpenAI
 from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
@@ -21,53 +20,64 @@ _ = load_dotenv(find_dotenv())
 # Get API key from Streamlit secrets
 API_KEY = os.getenv("OPENAI_API_KEY")
 
-# Initialize embedding model and vector store in memory (no disk persistence)
+# Initialize Chroma in-memory
 embeddings_model = OpenAIEmbeddings(openai_api_key=API_KEY)
 vectorstore = Chroma(embedding_function=embeddings_model)
 
-# Session flags
+# Streamlit UI setup
+st.set_page_config(page_title="RAG File Chat", layout="centered")
+st.title("🧠 Chat with Your Data File")
+
+# Session state flag
 if "agent_created" not in st.session_state:
     st.session_state.agent_created = False
 
+# Upload section
+uploaded_file = st.file_uploader("📁 Upload a file", type=["csv", "xlsx", "json", "pdf", "docx"])
+
+def extract_text_from_file(file_content, file_type):
+    """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()])
+    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(file_content, file_type):
-    """Create an agent from file content and index the data."""
+    """Reads, processes, and embeds the file content."""
+    # Load file content into DataFrame or raw text
     if file_type == "csv":
-        df = pd.read_csv(io.StringIO(file_content.decode("utf-8")), header=0)
+        df = pd.read_csv(io.StringIO(file_content.decode("utf-8")))
+        st.success("📄 CSV file loaded into DataFrame.")
     elif file_type == "xlsx":
-        df = pd.read_excel(file_content, header=0)
+        df = pd.read_excel(file_content)
+        st.success("📄 Excel file loaded into DataFrame.")
     elif file_type == "json":
         df = pd.DataFrame(json.loads(file_content.decode("utf-8")))
+        st.success("📄 JSON file loaded into DataFrame.")
     elif file_type in ["pdf", "docx"]:
         text = extract_text_from_file(file_content, file_type)
+        st.success(f"📃 {file_type.upper()} text extracted.")
         text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
         texts = text_splitter.split_text(text)
         df = pd.DataFrame({"text": texts})
-    else:
-        raise ValueError(f"Unsupported file type: {file_type}")
-
-    # Add text chunks to vectorstore
-    if file_type in ["pdf", "docx"]:
+        st.success("✂️ Text split into chunks.")
         vectorstore.add_texts(texts=df['text'].tolist(), metadatas=[{'source': file_type}] * len(df))
-
-    llm = OpenAI(openai_api_key=API_KEY)
-    return create_pandas_dataframe_agent(llm, df, verbose=False)
-
-
-def extract_text_from_file(file_content, file_type):
-    """Extract raw text from supported document formats."""
-    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()])
+        st.success("🧠 Embeddings generated and stored in vector database.")
     else:
-        return ""
+        st.error("❌ Unsupported file type.")
+        return None
 
+    # Create agent
+    llm = OpenAI(openai_api_key=API_KEY)
+    agent = create_pandas_dataframe_agent(llm, df, verbose=False)
+    st.success("🤖 Agent created successfully.")
+    return agent
 
 def query_agent(query):
-    """Query the vectorstore using RAG."""
+    """Query vectorstore via RetrievalQA."""
     qa_chain = RetrievalQA.from_chain_type(
         llm=OpenAI(openai_api_key=API_KEY),
         chain_type="stuff",
@@ -76,40 +86,24 @@ def query_agent(query):
     result = qa_chain({"query": query})
     return result["result"]
 
-
-# --- Streamlit UI ---
-st.set_page_config(page_title="RAG from Upload", layout="centered")
-st.title("🧠 Chat with Your File")
-
-uploaded_file = st.file_uploader("Upload a file", type=["csv", "xlsx", "json", "pdf", "docx"])
-
+# Main Logic
 if uploaded_file is not None:
     st.success(f"✅ File uploaded: `{uploaded_file.name}` ({uploaded_file.size / 1024:.1f} KB)")
-    
     file_content = uploaded_file.read()
     file_type = uploaded_file.name.split(".")[-1]
 
-    st.write("**File type detected:**", file_type.upper())
-
     if not st.session_state.agent_created:
-        with st.spinner("Indexing your file..."):
+        with st.spinner("🔄 Processing and indexing the file..."):
             create_agent(file_content, file_type)
             st.session_state.agent_created = True
-        st.success("📚 File successfully processed and indexed. You can now ask your question below.")
-
 
-    query = st.text_area("Enter your query")
+    query = st.text_area("💬 Ask a question based on the file")
 
-    if st.button("Submit Query", type="primary"):
+    if st.button("Submit Query"):
         if not query.strip():
-            st.warning("Please enter a valid query.")
-            st.stop()
-
-        if not st.session_state.agent_created:
-            create_agent(file_content, file_type)
-            st.session_state.agent_created = True
-            st.success("Data loaded and indexed.")
-
-        response = query_agent(query)
-        st.subheader("📌 Answer")
-        st.write(response)
+            st.warning("⚠️ Please enter a query.")
+        else:
+            with st.spinner("💡 Thinking..."):
+                answer = query_agent(query)
+            st.subheader("📌 Answer")
+            st.write(answer)