Upload 2 files

app.py

@@ -4,8 +4,23 @@ import os
 import shutil
 import subprocess
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from langchain_community.document_loaders import DirectoryLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter, Language
+from langchain_core.documents import Document
+
+EXTENSION_TO_LANGUAGE = {
+    '.py': Language.PYTHON,
+    '.js': Language.JS,
+    '.ts': Language.JS,
+    '.java': Language.JAVA,
+    '.cpp': Language.CPP,
+    '.c': Language.CPP,
+    '.h': Language.CPP,
+    '.go': Language.GO,
+    '.rs': Language.RUST,
+    '.rb': Language.RUBY,
+    '.html': Language.HTML,
+    '.md': Language.MARKDOWN,
+}
 from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
 from langchain_community.vectorstores import FAISS
 from langchain_core.runnables import RunnablePassthrough
@@ -42,23 +57,65 @@ def setup_vector_db():
     if not os.path.exists('./repo'):
         os.makedirs('./repo')
         
-    loader = DirectoryLoader('./repo', glob="**/*.py", show_progress=True)
-    docs = loader.load()
+    docs_by_language = {}
+    generic_docs = []
+    file_count = 0
     
-    if not docs:
+    for root, _, files in os.walk('./repo'):
+        if '.git' in root:
+            continue
+        for file in files:
+            file_path = os.path.join(root, file)
+            ext = os.path.splitext(file)[1].lower()
+            try:
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+                doc = Document(page_content=content, metadata={"source": file_path})
+                file_count += 1
+                
+                lang = EXTENSION_TO_LANGUAGE.get(ext)
+                if lang:
+                    if lang not in docs_by_language:
+                        docs_by_language[lang] = []
+                    docs_by_language[lang].append(doc)
+                else:
+                    generic_docs.append(doc)
+            except UnicodeDecodeError:
+                pass # Skip binary files
+                
+    if file_count == 0:
         return None, 0
 
-    python_splitter = RecursiveCharacterTextSplitter.from_language(
-        language=Language.PYTHON, 
-        chunk_size=500, 
-        chunk_overlap=50
-    )
-    texts = python_splitter.split_documents(docs)
+    all_splits = []
+    
+    # Split documents by specific language rules
+    for lang, docs in docs_by_language.items():
+        try:
+            splitter = RecursiveCharacterTextSplitter.from_language(
+                language=lang, 
+                chunk_size=500, 
+                chunk_overlap=50
+            )
+            all_splits.extend(splitter.split_documents(docs))
+        except Exception:
+            # Fallback if language is not supported by installed langchain version
+            generic_docs.extend(docs)
+            
+    # Split generic documents
+    if generic_docs:
+        generic_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=500, 
+            chunk_overlap=50
+        )
+        all_splits.extend(generic_splitter.split_documents(generic_docs))
+        
+    if not all_splits:
+        return None, 0
     
     embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-    db = FAISS.from_documents(texts, embeddings)
+    db = FAISS.from_documents(all_splits, embeddings)
     
-    return db, len(docs)
+    return db, file_count
 
 # 3. GLOBAL INITIALIZATION
 print("Initializing models...")
@@ -102,106 +159,130 @@ qa_chain = build_qa_chain(vector_db)
 # 4. INGESTION FUNCTIONS
 def clone_and_index(repo_url):
     global vector_db, file_count, qa_chain
+    if not repo_url or not repo_url.strip():
+        return "⚠️ Please enter a valid GitHub URL."
+
     if os.path.exists('./repo'):
         shutil.rmtree('./repo')
     
     try:
-        subprocess.run(["git", "clone", repo_url, "./repo"], check=True)
+        subprocess.run(["git", "clone", repo_url.strip(), "./repo"], check=True, capture_output=True, text=True)
+    except subprocess.CalledProcessError as e:
+        return f"❌ Failed to clone repo. Error: {e.stderr}"
     except Exception as e:
-        return f"**Repo Status:** Failed to clone repo: {str(e)} ❌"
+        return f"❌ Failed to clone repo: {str(e)}"
         
     vector_db, file_count = setup_vector_db()
     qa_chain = build_qa_chain(vector_db)
     
     if vector_db:
-        return f"**Repo Status:** {file_count} files indexed from `{repo_url}` ✅"
+        return f"✅ Success! {file_count} files indexed from `{repo_url}`"
     else:
-        return f"**Repo Status:** No Python files found in `{repo_url}` ❌"
+        return f"⚠️ Warning: No valid text files found in `{repo_url}`"
 
 def upload_and_index(files):
     global vector_db, file_count, qa_chain
+    if not files:
+        return "⚠️ No files were uploaded."
+
     if os.path.exists('./repo'):
         shutil.rmtree('./repo')
     os.makedirs('./repo', exist_ok=True)
     
-    if not files:
-        return "**Repo Status:** No files uploaded ❌"
-        
     for file in files:
-        dest_path = os.path.join('./repo', os.path.basename(file.name))
-        shutil.copy(file.name, dest_path)
+        # Handle both Gradio 3 (filepath string) and Gradio 4 (File object)
+        file_path = getattr(file, "name", str(file))
+        dest_path = os.path.join('./repo', os.path.basename(file_path))
+        shutil.copy(file_path, dest_path)
         
     vector_db, file_count = setup_vector_db()
     qa_chain = build_qa_chain(vector_db)
     
     if vector_db:
-        return f"**Repo Status:** {file_count} files indexed from local upload ✅"
+        return f"✅ Success! {file_count} files indexed from local upload"
     else:
-        return "**Repo Status:** No Python files found in local upload ❌"
+        return "⚠️ Warning: No valid text files found in the uploaded files"
 
 # 5. CHAT LOGIC
 def respond(message, chat_history):
+    if not message.strip():
+        return "", chat_history
+        
     if not vector_db:
-        bot_message = "👋 Welcome! Please provide a repo link or upload Python files to start chatting."
+        bot_message = "👋 Welcome! Please provide a repo link or upload your code files using the panel on the left to start chatting."
         chat_history.append((message, bot_message))
         return "", chat_history
     
-    # Fetch response from RAG
-    response = qa_chain.invoke(message)
-    answer = response["answer"]
-    sources = response["context"]
-    
-    final_answer = answer
-    
-    if sources:
-        final_answer += "\n\n<details><summary>🔍 View Source Code Referenced</summary>\n\n"
-        for idx, doc in enumerate(sources):
-            source_file = doc.metadata.get("source", "Unknown File")
-            final_answer += f"**Snippet {idx + 1}** from `{source_file}`:\n"
-            final_answer += f"```python\n{doc.page_content}\n```\n\n"
-        final_answer += "</details>"
+    try:
+        # Fetch response from RAG
+        response = qa_chain.invoke(message)
+        answer = response["answer"]
+        sources = response["context"]
+        
+        final_answer = answer
         
-    chat_history.append((message, final_answer))
+        if sources:
+            final_answer += "\n\n<details><summary>🔍 View Source Code Referenced</summary>\n\n"
+            for idx, doc in enumerate(sources):
+                source_file = doc.metadata.get("source", "Unknown File")
+                final_answer += f"**Snippet {idx + 1}** from `{source_file}`:\n"
+                final_answer += f"```python\n{doc.page_content}\n```\n\n"
+            final_answer += "</details>"
+            
+        chat_history.append((message, final_answer))
+    except Exception as e:
+        bot_message = f"❌ An error occurred during processing: {str(e)}"
+        chat_history.append((message, bot_message))
+
     return "", chat_history
 
 # 6. GRADIO UI
 custom_css = """
-.status-box { padding: 10px; border-radius: 8px; background-color: #f0f0f0; margin-bottom: 10px; }
-.dark .status-box { background-color: #1e293b; color: #cbd5e1; }
+.status-box { padding: 15px; border-radius: 8px; background-color: #f0f0f0; margin-bottom: 20px; border-left: 4px solid #007bff;}
+.dark .status-box { background-color: #1e293b; color: #cbd5e1; border-left: 4px solid #3b82f6;}
+.instructions { font-size: 0.95em; color: #555; }
+.dark .instructions { color: #aaa; }
 """
 
 def get_initial_repo_status():
     if vector_db:
-        return f"**Repo Status:** {file_count} files indexed ✅"
-    return "**Repo Status:** Empty ❌\n\nProvide a repo link or upload files to begin analyzing."
+        return f"✅ **Ready!** {file_count} files indexed and loaded."
+    return "❌ **Empty Database.** Provide a codebase below to begin."
 
-with gr.Blocks(title="Codebase Assistant", css=custom_css) as demo:
+with gr.Blocks(title="Codebase Assistant", css=custom_css, theme=gr.themes.Soft()) as demo:
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("# DevAssist AI\nYour personal Qwen-powered codebase expert.")
+            gr.Markdown("# 🦖 RepoRaptor\n**Your personal AI codebase expert.**")
             gr.Markdown("---")
             
             with gr.Column(elem_classes=["status-box"]):
-                gr.Markdown("### System Status")
+                gr.Markdown("### 📊 System Status")
                 gr.Markdown(f"**Hardware:** {device_status}")
                 repo_status = gr.Markdown(get_initial_repo_status())
             
-            gr.Markdown("### Add Codebase")
-            with gr.Tab("GitHub Repo"):
-                repo_url = gr.Textbox(placeholder="https://github.com/user/repo", show_label=False)
-                clone_btn = gr.Button("Clone & Index")
-            with gr.Tab("Local Upload"):
-                local_files = gr.File(file_count="multiple", label="Upload Local Files", file_types=[".py"])
-                upload_btn = gr.Button("Upload & Index")
+            gr.Markdown("### 📂 Ingest Codebase")
+            gr.Markdown("Choose a method to load your codebase into the Vector Database.", elem_classes=["instructions"])
+            
+            with gr.Tabs():
+                with gr.Tab("GitHub Repo"):
+                    gr.Markdown("Clone a public repository directly:")
+                    repo_url = gr.Textbox(placeholder="https://github.com/user/repo", show_label=False)
+                    clone_btn = gr.Button("⬇️ Clone & Index", variant="primary")
+                with gr.Tab("Local Upload"):
+                    gr.Markdown("Upload local codebase files:")
+                    local_files = gr.File(file_count="multiple", label="Upload Files")
+                    upload_btn = gr.Button("📤 Upload & Index", variant="primary")
 
             clone_btn.click(fn=clone_and_index, inputs=[repo_url], outputs=[repo_status])
             upload_btn.click(fn=upload_and_index, inputs=[local_files], outputs=[repo_status])
         
         with gr.Column(scale=3):
-            gr.Markdown("### 💻 Chat with your Codebase\nAsk architecture questions, find bugs, or request code explanations.")
-            chatbot = gr.Chatbot(height=500, show_label=False)
-            msg = gr.Textbox(placeholder="E.g., What does the main function do?", show_label=False)
-            clear = gr.Button("Clear Chat")
+            gr.Markdown("### 💻 Chat Interface\nAsk architecture questions, find bugs, or request code explanations. I will **only** answer questions related to code.")
+            chatbot = gr.Chatbot(height=600, show_label=False, bubble_full_width=False)
+            
+            with gr.Row():
+                msg = gr.Textbox(placeholder="E.g., What does the main function do? (Press Enter to send)", show_label=False, scale=4)
+                clear = gr.Button("🗑️ Clear Chat", scale=1)
 
             msg.submit(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
             clear.click(lambda: None, None, chatbot, queue=False)