| import streamlit as st | |
| import os | |
| 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") | |
| if "qa_chain" not in st.session_state: | |
| st.session_state.qa_chain = None | |
| 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() | |
| vectorstore = create_vectorstore(docs) | |
| st.session_state.qa_chain = build_rag_chain( | |
| vectorstore, | |
| os.getenv("GROQ_API_KEY") | |
| ) | |
| 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("Index a codebase first.") | |
| elif not question.strip(): | |
| st.warning("Enter a question.") | |
| else: | |
| with st.spinner("Thinking..."): | |
| answer = st.session_state.qa_chain.invoke(question) | |
| st.markdown("### Answer") | |
| st.write(answer) | |