Spaces:
Running
Running
| import streamlit as st | |
| import os | |
| import shutil | |
| import time | |
| from code_chatbot.ingestion.universal_ingestor import process_source | |
| from code_chatbot.ingestion.indexer import Indexer | |
| from code_chatbot.retrieval.rag import ChatEngine | |
| from code_chatbot.analysis.ast_analysis import ASTGraphBuilder | |
| from code_chatbot.retrieval.graph_rag import GraphEnhancedRetriever | |
| import logging | |
| from dotenv import load_dotenv | |
| # Load Env | |
| load_dotenv() | |
| # Basic Setup | |
| st.set_page_config(page_title="Code Chatbot", page_icon="💻", layout="wide", initial_sidebar_state="collapsed") | |
| logging.basicConfig(level=logging.INFO) | |
| # --- Custom CSS for Premium Slate UI --- | |
| from components import style | |
| style.apply_custom_css() | |
| # Session State | |
| if "messages" not in st.session_state: | |
| st.session_state.messages = [] | |
| if "chat_engine" not in st.session_state: | |
| st.session_state.chat_engine = None | |
| if "processed_files" not in st.session_state: | |
| st.session_state.processed_files = False | |
| # --- Sidebar Configuration --- | |
| from components import sidebar | |
| config = sidebar.render_sidebar() | |
| # Extract config values for easy access in main app | |
| provider = config["provider"] | |
| api_key = config["api_key"] | |
| gemini_model = config["gemini_model"] | |
| use_agent = config["use_agent"] | |
| vector_db_type = config["vector_db_type"] | |
| embedding_provider = config["embedding_provider"] | |
| embedding_api_key = config["embedding_api_key"] | |
| # ============================================================================ | |
| # MAIN 3-PANEL LAYOUT | |
| # ============================================================================ | |
| st.title("🕷️ Code Crawler") | |
| if not st.session_state.processed_files: | |
| # Show onboarding message when no files are processed | |
| # --- Main Ingestion Section --- | |
| st.header("🚀 Import Codebase") | |
| st.caption("Upload your project to get started. Configure advanced settings in the sidebar (open with >).") | |
| if not api_key: | |
| st.warning(f"⚠️ {provider.capitalize()} API Key is missing. Open the sidebar (top-left) to configure it.") | |
| source_type = st.radio("Source Type", ["ZIP File", "GitHub Repository", "Web Documentation"], horizontal=True) | |
| source_input = None | |
| if source_type == "ZIP File": | |
| uploaded_file = st.file_uploader("Upload .zip file", type="zip") | |
| if uploaded_file: | |
| import tempfile | |
| upload_dir = tempfile.gettempdir() | |
| source_input = os.path.join(upload_dir, "uploaded.zip") | |
| with open(source_input, "wb") as f: | |
| f.write(uploaded_file.getbuffer()) | |
| elif source_type == "GitHub Repository": | |
| source_input = st.text_input("GitHub URL", placeholder="https://github.com/owner/repo") | |
| elif source_type == "Web Documentation": | |
| source_input = st.text_input("Web URL", placeholder="https://docs.python.org/3/") | |
| if source_input and not st.session_state.processed_files: | |
| if st.button("🚀 Process & Index", type="primary"): | |
| if not api_key: | |
| st.error(f"Please configure {provider} API Key in the sidebar.") | |
| elif provider == "groq" and not embedding_api_key: | |
| st.error(f"Please configure {embedding_provider} API Key for embeddings in the sidebar.") | |
| else: | |
| # Use the new progress-tracked indexer | |
| from code_chatbot.ingestion.indexing_progress import index_with_progress | |
| chat_engine, success, repo_files, workspace_root = index_with_progress( | |
| source_input=source_input, | |
| source_type=source_type, | |
| provider=provider, | |
| embedding_provider=embedding_provider, | |
| embedding_api_key=embedding_api_key, | |
| vector_db_type=vector_db_type, | |
| use_agent=use_agent, | |
| api_key=api_key, | |
| gemini_model=gemini_model # Pass selected model | |
| ) | |
| if success: | |
| st.session_state.chat_engine = chat_engine | |
| st.session_state.processed_files = True | |
| st.session_state.indexed_files = repo_files | |
| st.session_state.workspace_root = workspace_root | |
| time.sleep(0.5) | |
| st.switch_page("pages/1_⚡_Code_Studio.py") | |
| else: | |
| # Codebase Ready! Redirect to Code Studio | |
| st.switch_page("pages/1_⚡_Code_Studio.py") | |