Update app.py

app.py

@@ -1,7 +1,11 @@
-import streamlit as st
+import gradio as gr
+from selenium import webdriver
+from selenium.webdriver.common.by import By
+from selenium.webdriver.chrome.options import Options
+from selenium.webdriver.chrome.service import Service
+from webdriver_manager.chrome import ChromeDriverManager
 import google.generativeai as genai
-from datetime import datetime
-import requests
+import time
 import os
 from dotenv import load_dotenv
 
@@ -10,172 +14,166 @@ load_dotenv()
 
 # Configure Gemini
 genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
+model = genai.GenerativeModel('gemini-pro')
 
-# Set up the model
-generation_config = {
-    "temperature": 0.9,
-    "top_p": 1,
-    "top_k": 1,
-    "max_output_tokens": 2048,
-}
+# Initialize Selenium WebDriver
+def init_driver():
+    chrome_options = Options()
+    chrome_options.add_argument("--headless")
+    chrome_options.add_argument("--disable-gpu")
+    chrome_options.add_argument("--window-size=1920x1080")
+    chrome_options.add_argument("--no-sandbox")
+    chrome_options.add_argument("--disable-dev-shm-usage")
+    service = Service(ChromeDriverManager().install())
+    driver = webdriver.Chrome(service=service, options=chrome_options)
+    return driver
 
-safety_settings = [
-    {
-        "category": "HARM_CATEGORY_HARASSMENT",
-        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-    },
-    {
-        "category": "HARM_CATEGORY_HATE_SPEECH",
-        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-    },
-    {
-        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-    },
-    {
-        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
-        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-    },
-]
-
-model = genai.GenerativeModel(
-    model_name="gemini-1.5-flash",
-    generation_config=generation_config,
-    safety_settings=safety_settings
-)
-
-# Function to perform web search (using SerpAPI)
-def search_web(query):
+# Enhanced web search with Selenium
+def search_with_selenium(query):
+    driver = init_driver()
     try:
-        api_key = os.getenv("SERPAPI_KEY")
-        if not api_key:
-            return None
-            
-        params = {
-            'q': query,
-            'api_key': api_key,
-            'engine': 'google'
-        }
+        print(f"Searching for: {query}")
+        driver.get(f"https://www.google.com/search?q={query}")
+        time.sleep(2)  # Wait for results to load
         
-        response = requests.get('https://serpapi.com/search', params=params)
-        results = response.json()
+        # Accept cookies if popup appears (for EU users)
+        try:
+            driver.find_element(By.ID, "L2AGLb").click()
+            time.sleep(1)
+        except:
+            pass
         
-        if 'organic_results' in results:
-            return results['organic_results'][:3]  # Return top 3 results
-        return None
-    except Exception as e:
-        st.error(f"Search error: {e}")
-        return None
-
-# Initialize chat history
-if "messages" not in st.session_state:
-    st.session_state.messages = [
-        {
-            "role": "assistant", 
-            "content": "Hello! I'm Gemini AI with web search capabilities. How can I help you today?",
-            "timestamp": datetime.now().strftime("%H:%M")
-        }
-    ]
-
-# App title and sidebar
-st.set_page_config(page_title="Gemini AI Chatbot", page_icon="🤖")
-st.title("💎 Gemini AI Chatbot")
-st.caption("Powered by Gemini 1.5 Flash with web search capabilities")
-
-# Sidebar controls
-with st.sidebar:
-    st.header("Settings")
-    web_search = st.toggle("Enable Web Search", value=True)
-    st.divider()
-    if st.button("New Chat"):
-        st.session_state.messages = [
-            {
-                "role": "assistant", 
-                "content": "Hello! I'm Gemini AI with web search capabilities. How can I help you today?",
-                "timestamp": datetime.now().strftime("%H:%M")
-            }
-        ]
-    st.divider()
-    st.markdown("### About")
-    st.markdown("This chatbot uses Google's Gemini 1.5 Flash model and can perform web searches when enabled.")
-
-# Display chat messages
-for message in st.session_state.messages:
-    avatar = "🤖" if message["role"] == "assistant" else "👤"
-    with st.chat_message(message["role"], avatar=avatar):
-        st.markdown(message["content"])
-        if "search_results" in message:
-            st.divider()
-            st.markdown("**Web Search Results**")
-            for result in message["search_results"]:
-                with st.expander(result.get("title", "No title")):
-                    st.markdown(f"**Snippet:** {result.get('snippet', 'No snippet available')}")
-                    st.markdown(f"**Link:** [{result.get('link', 'No link')}]({result.get('link', '#')})")
-        st.caption(f"{message['timestamp']} • {message['role'].capitalize()}")
-
-# Accept user input
-if prompt := st.chat_input("Ask me anything..."):
-    # Add user message to chat history
-    st.session_state.messages.append({
-        "role": "user", 
-        "content": prompt,
-        "timestamp": datetime.now().strftime("%H:%M")
-    })
-    
-    # Display user message
-    with st.chat_message("user", avatar="👤"):
-        st.markdown(prompt)
-        st.caption(f"{datetime.now().strftime('%H:%M')} • User")
-    
-    # Display assistant response
-    with st.chat_message("assistant", avatar="🤖"):
-        message_placeholder = st.empty()
-        full_response = ""
-        
-        # If web search is enabled, perform search first
-        search_results = None
-        if web_search:
-            with st.spinner("Searching the web..."):
-                search_results = search_web(prompt)
+        # Extract search results
+        results = []
+        search_items = driver.find_elements(By.CSS_SELECTOR, 'div.g')
         
-        # Generate response from Gemini
-        with st.spinner("Thinking..."):
+        for item in search_items[:5]:  # Get top 5 results
             try:
-                # Prepare the prompt
-                chat_prompt = prompt
-                if search_results:
-                    chat_prompt += "\n\nHere are some web search results to help with your response:\n"
-                    for i, result in enumerate(search_results, 1):
-                        chat_prompt += f"\n{i}. {result.get('title', 'No title')}\n{result.get('snippet', 'No snippet')}\n"
+                title = item.find_element(By.CSS_SELECTOR, 'h3').text
+                link = item.find_element(By.CSS_SELECTOR, 'a').get_attribute('href')
                 
-                # Get response from Gemini
-                response = model.generate_content(chat_prompt)
-                
-                # Stream the response
-                for chunk in response.text.split(" "):
-                    full_response += chunk + " "
-                    message_placeholder.markdown(full_response + "▌")
-                message_placeholder.markdown(full_response)
+                # Try different selectors for snippet
+                snippet = ""
+                try:
+                    snippet = item.find_element(By.CSS_SELECTOR, 'div.IsZvec').text
+                except:
+                    try:
+                        snippet = item.find_element(By.CSS_SELECTOR, '.VwiC3b').text
+                    except:
+                        pass
                 
+                if title and link:
+                    results.append({
+                        'title': title,
+                        'link': link,
+                        'snippet': snippet[:200] + "..." if snippet else "No description available"
+                    })
             except Exception as e:
-                full_response = f"Sorry, I encountered an error: {str(e)}"
-                message_placeholder.markdown(full_response)
+                print(f"Error extracting result: {e}")
+                continue
+                
+        return results
+    except Exception as e:
+        print(f"Search error: {e}")
+        return []
+    finally:
+        driver.quit()
+
+# Chatbot function with enhanced error handling
+def chatbot(message, history, use_web_search):
+    try:
+        search_results = []
+        context = message
+        
+        # Perform web search if enabled
+        if use_web_search:
+            with gr.Blocks(analytics_enabled=False):
+                with gr.Row():
+                    gr.Markdown("🔍 Searching the web...")
+            
+            search_results = search_with_selenium(message)
+            if search_results:
+                context += "\n\nWeb search results:\n"
+                for i, result in enumerate(search_results, 1):
+                    context += f"{i}. {result['title']}\n{result['snippet']}\n\n"
         
-        # Display search results if available
+        # Generate response
+        response = model.generate_content(context)
+        bot_message = response.text
+        
+        # Format response with search results
         if search_results:
-            st.divider()
-            st.markdown("**Web Search Results**")
+            bot_message += "\n\n📚 Sources:\n"
             for result in search_results:
-                with st.expander(result.get("title", "No title")):
-                    st.markdown(f"**Snippet:** {result.get('snippet', 'No snippet available')}")
-                    st.markdown(f"**Link:** [{result.get('link', 'No link')}]({result.get('link', '#')})")
+                bot_message += f"- [{result['title']}]({result['link']})\n"
+        
+        return bot_message
+    except Exception as e:
+        return f"⚠️ Error: {str(e)}"
+
+# Gradio interface with enhanced UI
+with gr.Blocks(
+    title="Gemini AI Chatbot with Web Search",
+    theme=gr.themes.Soft(),
+    css=".gradio-container {max-width: 800px !important}"
+) as demo:
+    
+    gr.Markdown("""
+    # 💎 Gemini AI Chatbot
+    **Powered by Google Gemini Pro + Real-time Web Search**
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot(
+                height=500,
+                bubble_full_width=False,
+                avatar_images=(
+                    "user.png", 
+                    "bot.png"
+                )
+            )
+            msg = gr.Textbox(
+                label="Your Message",
+                placeholder="Ask me anything...",
+                lines=2
+            )
+            
+            with gr.Row():
+                submit_btn = gr.Button("Send", variant="primary")
+                clear_btn = gr.Button("Clear")
+                web_search = gr.Checkbox(
+                    label="Enable Web Search",
+                    value=True,
+                    interactive=True
+                )
         
-        st.caption(f"{datetime.now().strftime('%H:%M')} • Assistant")
+        with gr.Column(scale=1):
+            gr.Markdown("""
+            ### 🌐 Web Search Features
+            - Real-time Google results
+            - Source citations
+            - Toggle on/off
+            
+            ### ⚙️ Controls
+            - Press Enter or click Send
+            - Clear resets conversation
+            """)
+    
+    # Event handlers
+    def respond(message, chat_history, use_search):
+        bot_message = chatbot(message, chat_history, use_search)
+        chat_history.append((message, bot_message))
+        return "", chat_history
     
-    # Add assistant response to chat history
-    st.session_state.messages.append({
-        "role": "assistant", 
-        "content": full_response,
-        "timestamp": datetime.now().strftime("%H:%M"),
-        "search_results": search_results if search_results else None
-    })
+    msg.submit(respond, [msg, chatbot, web_search], [msg, chatbot])
+    submit_btn.click(respond, [msg, chatbot, web_search], [msg, chatbot])
+    clear_btn.click(lambda: None, None, chatbot, queue=False)
+
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        favicon_path="favicon.ico"
+    )