app.py

import streamlit as st
import google.generativeai as genai
from datetime import datetime
import requests
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Configure Gemini
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

# Set up the model
generation_config = {
    "temperature": 0.9,
    "top_p": 1,
    "top_k": 1,
    "max_output_tokens": 2048,
}

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):
    try:
        api_key = os.getenv("SERPAPI_KEY")
        if not api_key:
            return None
            
        params = {
            'q': query,
            'api_key': api_key,
            'engine': 'google'
        }
        
        response = requests.get('https://serpapi.com/search', params=params)
        results = response.json()
        
        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)
        
        # Generate response from Gemini
        with st.spinner("Thinking..."):
            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"
                
                # 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)
                
            except Exception as e:
                full_response = f"Sorry, I encountered an error: {str(e)}"
                message_placeholder.markdown(full_response)
        
        # Display search results if available
        if search_results:
            st.divider()
            st.markdown("**Web Search Results**")
            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', '#')})")
        
        st.caption(f"{datetime.now().strftime('%H:%M')} • Assistant")
    
    # 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
    })