Update app.py

app.py

@@ -1,43 +1,65 @@
 import streamlit as st
 import os
 
-from ingest import load_repo, ingest_repo
+from ingest import load_repo, ingest_repo, extract_zip
 from vectorstore import create_vectorstore
 from rag_chain import build_rag_chain
 
 st.set_page_config(page_title="Codebase RAG Assistant")
-
 st.title("🧠 Codebase RAG Assistant")
 
-# Session state
 if "qa_chain" not in st.session_state:
     st.session_state.qa_chain = None
 
-# Input repo URL
-repo_url = st.text_input("GitHub Repository URL")
+st.subheader("Input Source")
+
+repo_url = st.text_input("GitHub Repository URL (optional)")
+uploaded_zip = st.file_uploader("Or upload ZIP file", type=["zip"])
+
+if st.button("Index Codebase"):
+    try:
+        with st.spinner("Processing codebase..."):
+
+            if uploaded_zip:
+                repo_path = extract_zip(uploaded_zip)
+                docs = ingest_repo(repo_path)
+
+            elif repo_url:
+                if "/blob/" in repo_url:
+                    st.error("Please provide a repository URL, not a file URL.")
+                    st.stop()
+                repo_path = load_repo(repo_url)
+                docs = ingest_repo(repo_path)
+
+            else:
+                st.warning("Provide GitHub repo URL or upload ZIP.")
+                st.stop()
+
+            if not docs:
+                st.error(
+                    "No supported source files found.\n\n"
+                    "Supported: .py, .js, .java, .cpp, .txt, .ipynb"
+                )
+                st.stop()
 
-if st.button("Index Repository"):
-    if not repo_url:
-        st.warning("Please enter a GitHub repository URL")
-    else:
-        with st.spinner("Cloning and indexing repository..."):
-            repo_path = load_repo(repo_url)
-            docs = ingest_repo(repo_path)
             vectorstore = create_vectorstore(docs)
             st.session_state.qa_chain = build_rag_chain(
                 vectorstore,
                 os.getenv("GROQ_API_KEY")
             )
-        st.success("Repository indexed successfully!")
 
-# Question input
+        st.success("Indexing completed!")
+
+    except Exception as e:
+        st.error(str(e))
+
 question = st.text_input("Ask a question about the codebase")
 
 if st.button("Ask"):
     if st.session_state.qa_chain is None:
-        st.warning("Please index a repository first")
+        st.warning("Index a codebase first.")
     elif not question.strip():
-        st.warning("Please enter a question")
+        st.warning("Enter a question.")
     else:
         with st.spinner("Thinking..."):
             answer = st.session_state.qa_chain.invoke(question)