Update app.py

app.py

@@ -5,79 +5,60 @@ import requests
 import os
 from typing import List, Tuple, Optional
 
-# Retrieve API key from environment variable
-api_key = os.getenv('HF_TOKEN')
-if not api_key:
-    raise ValueError("API key not found. Please set the HFTOKEN environment variable.")
-
-# Initialize the InferenceClient with the specified model and API key
-client = InferenceClient(
-    model="meta-llama/Meta-Llama-3.1-405B-Instruct",
-    token=api_key
-)
+# --- Configuration ---
+MODEL_NAME = "meta-llama/Meta-Llama-3.1-405B-Instruct"
+API_KEY_ENV_VAR = "HF_TOKEN"
+SEARCH_TRIGGER_WORD = "search"
+DEFAULT_NUM_RESULTS = 3
 
-def scrape_yahoo_search(query: str, num_results: int = 3) -> Tuple[str, str]:
-    """
-    Scrapes Yahoo search results for the given query and returns detailed snippets and URLs for the top results.
-    
-    Args:
-        query (str): The search query.
-        num_results (int): Number of results to retrieve.
-    
-    Returns:
-        Tuple[str, str]: The formatted snippets and URLs of the top search results.
-    """
+# --- Utility Functions ---
+
+def get_api_key() -> str:
+    """Retrieves the API key from an environment variable."""
+    api_key = os.getenv(API_KEY_ENV_VAR)
+    if not api_key:
+        raise ValueError(f"API key not found. Please set the {API_KEY_ENV_VAR} environment variable.")
+    return api_key
+
+def scrape_yahoo_search(query: str, num_results: int = DEFAULT_NUM_RESULTS) -> Tuple[Optional[str], Optional[str]]:
+    """Scrapes Yahoo search results and returns formatted snippets and the search URL."""
     search_url = f"https://search.yahoo.com/search?p={query}"
-    headers = {
-        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
-    }
-    
+    headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36'} 
+
     try:
         response = requests.get(search_url, headers=headers)
-        response.raise_for_status()
+        response.raise_for_status() 
+
         soup = BeautifulSoup(response.content, 'html.parser')
-        
-        results = []
         result_elements = soup.find_all('div', {'class': 'dd algo algo-sr Sr'}, limit=num_results)
-        
+
         if result_elements:
-            for result in result_elements:
-                title = result.find('h3').get_text(strip=True) if result.find('h3') else "No title"
-                snippet = result.find('div', {'class': 'compText aAbs'}).get_text(strip=True) if result.find('div', {'class': 'compText aAbs'}) else "No snippet"
-                url = result.find('a')['href'] if result.find('a') else "No URL"
-                results.append((title, snippet, url))
-            
-            formatted_results = "\n\n".join(
-                f"Title: {title}\nSnippet: {snippet}\nURL: {url}" for title, snippet, url in results
-            )
-            return formatted_results, search_url
+            results = [
+                f"**Title:** {res.find('h3').get_text(strip=True)}\n**Snippet:** {res.find('div', {'class': 'compText aAbs'}).get_text(strip=True)}\n**URL:** {res.find('a')['href']}" 
+                for res in result_elements 
+                if res.find('h3') and res.find('div', {'class': 'compText aAbs'}) and res.find('a')
+            ]
+            return "\n\n".join(results), search_url
         else:
-            return "No results found.", search_url
-    
+            return None, search_url 
+
     except requests.RequestException as e:
-        return f"Request error: {str(e)}", search_url
+        return f"Request error: {str(e)}", None
     except Exception as e:
-        return f"Processing error: {str(e)}", search_url
+        return f"Processing error: {str(e)}", None
 
-def extract_search_query(message: str, trigger_word: str) -> Optional[str]:
-    """
-    Extracts the search query from the message based on the trigger word.
-    
-    Args:
-        message (str): The user's input message.
-        trigger_word (str): The word that activates the search feature.
-    
-    Returns:
-        Optional[str]: The extracted search query if found, otherwise None.
-    """
+def extract_search_query(message: str, trigger_word: str = SEARCH_TRIGGER_WORD) -> Optional[str]:
+    """Extracts the search query from the message if the trigger word is present."""
     lower_message = message.lower()
     if trigger_word in lower_message:
-        parts = lower_message.split(trigger_word, 1)
-        if len(parts) > 1:
-            query = parts[1].strip()
-            return query if query else None
+        query = lower_message.split(trigger_word, 1)[1].strip()
+        return query if query else None
     return None
 
+# --- Initialize Inference Client ---
+client = InferenceClient(model=MODEL_NAME, token=get_api_key())
+
+# --- Chatbot Logic ---
 def respond(
     message: str,
     history: List[Tuple[str, str]],
@@ -86,48 +67,22 @@ def respond(
     temperature: float,
     top_p: float,
 ) -> str:
-    """
-    Generates a response from the AI model based on the user's message, chat history, and optional Yahoo search results.
-    
-    Args:
-        message (str): The user's input message.
-        history (List[Tuple[str, str]]): A list of tuples representing the conversation history (user, assistant).
-        system_message (str): A system-level message guiding the AI's behavior.
-        max_tokens (int): The maximum number of tokens for the output.
-        temperature (float): Sampling temperature for controlling the randomness.
-        top_p (float): Top-p (nucleus sampling) for controlling diversity.
-    
-    Returns:
-        str: The AI's response as it is generated, including the source URL if applicable.
-    """
-    
-    # Check for trigger word and activate search feature if present
-    trigger_word = "search"
-    query = extract_search_query(message, trigger_word)
-    
-    if query:
-        snippet, url = scrape_yahoo_search(query, num_results=3)
-        message += f"\n\nWeb Content:\n{snippet}\nSource: {url}"
-    elif query is None:
-        message = "Please provide a search query after the trigger word."
+    """Generates a response from the AI model, incorporating search results if requested."""
+    query = extract_search_query(message)
 
-    # Prepare the conversation history for the API call
-    messages = [{"role": "system", "content": system_message}]
+    if query:
+        search_results, search_url = scrape_yahoo_search(query)
+        if search_results:
+            message += f"\n\n## Web Search Results:\n{search_results}\n**Source:** {search_url}"
+        else:
+            message += "\n\nI couldn't find any relevant web results for your query."
     
-    for user_input, assistant_response in history:
-        if user_input:
-            messages.append({"role": "user", "content": user_input})
-        if assistant_response:
-            messages.append({"role": "assistant", "content": assistant_response})
+    messages = [{"role": "system", "content": system_message}] + \
+               [{"role": role, "content": content} for role, content in history] + \
+               [{"role": "user", "content": message}]
 
-    # Add the latest user message to the conversation
-    messages.append({"role": "user", "content": message})
-
-    # Initialize an empty response
     response = ""
-
     try:
-        # Generate a response from the model with streaming
         for response_chunk in client.chat_completion(
             messages=messages,
             max_tokens=max_tokens,
@@ -135,32 +90,25 @@ def respond(
             temperature=temperature,
             top_p=top_p,
         ):
-            token = response_chunk.choices[0].delta.content
-            response += token
+            response += response_chunk.choices[0].delta.content
     except Exception as e:
-        return f"AI model error: {str(e)}"
+        return f"AI model error: {str(e)}" 
 
     return response
 
-# Define the ChatInterface with additional input components for user customization
+# --- Gradio Interface ---
 demo = gr.ChatInterface(
     fn=respond,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+        gr.Textbox(value="You are a helpful and informative AI assistant.", label="System Message"),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max New Tokens"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (Nucleus Sampling)"),
     ],
-    title="Chatbot Interface",
-    description="A customizable chatbot interface using Hugging Face's Inference API with Yahoo search scraping capabilities.",
+    title="Chatbot with Search",
+    description="Chat and search the web using the power of Meta-Llama!",
 )
 
-# Launch the Gradio interface
+# --- Launch the App ---
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()