import gradio as gr import uuid import logging from typing import List, Tuple from src.main import RAGSystem import asyncio # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Initialize the RAG system rag = RAGSystem() def create_session() -> str: """Create a new session ID""" return str(uuid.uuid4()) def index_website(url: str, session_id: str) -> Tuple[bool, str]: """Index a website for a given session""" try: result = rag.crawl_and_index(session_id, url) if result["status"] == "success": return True, f"Successfully indexed {len(result.get('urls_processed', []))} pages" return False, result.get("message", "Unknown error during indexing") except Exception as e: logger.error(f"Indexing error: {str(e)}") return False, f"Error during indexing: {str(e)}" def chat_response( session_id: str, message: str, model_choice: str, ollama_url: str, gemini_api_key: str, chat_history: List[dict] ) -> Tuple[List[dict], str]: """Generate a chat response with proper error handling""" if not session_id: chat_history.append({"role": "user", "content": "🧑 " + message}) chat_history.append({"role": "assistant", "content": "🤖 Please index a website first or enter a valid session ID"}) return chat_history, "" try: response = asyncio.run(rag.chat( session_id=session_id, question=message, model=model_choice.lower(), ollama_url=ollama_url if model_choice == "mistral" else None, gemini_api_key=gemini_api_key if model_choice == "gemini" else None )) if response["status"] == "success": answer = response["response"] sources = "\n\nSources:\n" + "\n".join( f"- {src['source_url']}" for src in response.get("sources", []) ) if response.get("sources") else "" full_response = answer + sources else: full_response = f"Error: {response.get('message', 'Unknown error')}" chat_history.append({"role": "user", "content": "🧑 " + message}) chat_history.append({"role": "assistant", "content": "🤖 " + full_response}) return chat_history, "" except Exception as e: logger.error(f"Chat error: {str(e)}") chat_history.append({"role": "user", "content": "🧑 " + message}) chat_history.append({"role": "assistant", "content": f"🤖 System error: {str(e)}"}) return chat_history, "" def toggle_model_inputs(model_choice: str) -> List[gr.update]: """Show/hide model-specific inputs""" if model_choice == "mistral": return [gr.update(visible=True), gr.update(visible=False)] return [gr.update(visible=False), gr.update(visible=True)] def load_session(existing_session_id: str) -> Tuple[str, str]: """Load an existing session""" if existing_session_id: # Here you might want to add validation if the session exists return existing_session_id, f"Loaded existing session: {existing_session_id}" return "", "Please enter a valid session ID" def get_session(self, session_id: str): # If session exists in memory, return it if session_id in self.sessions: return self.sessions[session_id] # If not, check if Qdrant collection exists and has documents collection_name = self.get_collection_name(session_id) try: results = self.qdrant_client.scroll(collection_name=collection_name, limit=1) if results and results[0]: # Rehydrate session in memory self.sessions[session_id] = { "documents": [], # Optionally, you can fetch all docs if needed "history": [] } return self.sessions[session_id] except Exception as e: logger.warning(f"Session {session_id} not found in Qdrant: {e}") # If not found, return None or raise raise ValueError("No documents indexed for this session") # Custom CSS for better styling custom_css = """ .gradio-container { max-width: 1200px !important; margin: 0 auto !important; } .dark .gradio-container { background: #1e1e2e !important; } #chatbot { min-height: 500px; border-radius: 12px !important; } .message.user { border-left: 4px solid #4f46e5 !important; } .message.assistant { border-left: 4px solid #10b981 !important; } .btn-primary { background: linear-gradient(to right, #4f46e5, #7c3aed) !important; border: none !important; } .btn-primary:hover { background: linear-gradient(to right, #4338ca, #6d28d9) !important; } .prose { max-width: 100% !important; } """ with gr.Blocks(title="RAG Chat with Mistral/Gemini", css=custom_css, theme="soft") as demo: # Header section with gr.Row(): gr.Markdown(""" # 🌐 RAG Chat Assistant ### Chat with any website using Mistral or Gemini """) # Session state session_id = gr.State("") with gr.Tabs(): with gr.TabItem("📚 Index Website"): with gr.Row(): with gr.Column(): gr.Markdown("### Step 1: Configure and Index") with gr.Group(): url_input = gr.Textbox( label="Website URL to index", placeholder="https://example.com", interactive=True, lines=1 ) with gr.Row(): model_choice = gr.Radio( choices=["mistral", "gemini"], label="Select Model", value="mistral", interactive=True ) index_btn = gr.Button( "🚀 Index Website", variant="primary", scale=0 ) with gr.Accordion("🔐 Model Settings", open=False): ollama_url = gr.Textbox( label="Ollama URL (required for Mistral)", placeholder="http://localhost:11434", visible=True ) gemini_api_key = gr.Textbox( label="Gemini API Key (required for Gemini)", placeholder="your-api-key-here", visible=False, type="password" ) status_output = gr.Textbox( label="Status", interactive=False, elem_classes="prose" ) gr.Markdown(""" **Instructions:** 1. Enter a website URL 2. Select your preferred model 3. Configure model settings if needed 4. Click 'Index Website' """) with gr.TabItem("💬 Chat"): with gr.Row(): with gr.Column(scale=2): # New session ID input for resuming sessions with gr.Accordion("🔍 Resume Previous Session", open=False): existing_session_input = gr.Textbox( label="Enter existing Session ID", placeholder="Paste your session ID here...", interactive=True ) load_session_btn = gr.Button( "🔁 Load Session", variant="secondary" ) session_status = gr.Textbox( label="Session Status", interactive=False ) chatbot = gr.Chatbot( label="Chat History", height=500, avatar_images=(None, None), show_copy_button=True, type="messages" # Use OpenAI-style messages ) with gr.Row(): message_input = gr.Textbox( label="Type your message", placeholder="Ask about the website content...", interactive=True, container=False, scale=7, autofocus=True ) send_btn = gr.Button( "Send", variant="primary", scale=1, min_width=100 ) # Event handlers model_choice.change( fn=toggle_model_inputs, inputs=model_choice, outputs=[ollama_url, gemini_api_key] ) index_btn.click( fn=create_session, outputs=session_id ).success( fn=index_website, inputs=[url_input, session_id], outputs=[status_output] ) # New handler for loading existing sessions load_session_btn.click( fn=load_session, inputs=[existing_session_input], outputs=[session_id, session_status] ) send_btn.click( fn=chat_response, inputs=[session_id, message_input, model_choice, ollama_url, gemini_api_key, chatbot], outputs=[chatbot, message_input] ) # Allow submitting with Enter key message_input.submit( fn=chat_response, inputs=[session_id, message_input, model_choice, ollama_url, gemini_api_key, chatbot], outputs=[chatbot, message_input] ) if __name__ == "__main__": demo.launch( server_name="0.0.0.0", server_port=7860, #favicon_path="assets/favicon.ico" # Local path, not URL )