Update app.py

app.py

@@ -1,11 +1,10 @@
 import os
 import time
 import streamlit as st
-from smolagents import load_tool, CodeAgent, HfApiModel, DuckDuckGoSearchTool
+from smolagents import CodeAgent, HfApiModel, tool
 from huggingface_hub import InferenceClient
-
-# Instantiate the DuckDuckGo search tool from SmolAgents
-search_tool = DuckDuckGoSearchTool()
+import requests
+from bs4 import BeautifulSoup
 
 # Retrieve Hugging Face token
 hf_token = os.getenv("HF_TOKEN")
@@ -15,11 +14,65 @@ if not hf_token:
 # Initialize the Hugging Face Inference client
 client = InferenceClient(token=hf_token)
 
+# Custom tools for SmolAgents
+@tool
+def search_harry_potter_lore(query: str) -> str:
+    """Search for Harry Potter-related lore or facts across the entire Harry Potter Fandom site.
+    
+    Args:
+        query: A specific question or topic about Harry Potter lore.
+    
+    Returns:
+        A concise and informative response based on the query.
+    """
+    headers = {"User-Agent": "Mozilla/5.0"}
+    # Construct the search URL for the Harry Potter Fandom site.
+    search_url = f"https://harrypotter.fandom.com/wiki/Special:Search?query={query}"
+    
+    try:
+        # Fetch the search results page.
+        search_response = requests.get(search_url, headers=headers)
+        if search_response.status_code != 200:
+            return f"Error: Received status code {search_response.status_code} from search."
+        
+        search_soup = BeautifulSoup(search_response.text, 'html.parser')
+        
+        # Look for the first link that appears to be an article.
+        article_url = None
+        for a in search_soup.find_all("a", href=True):
+            href = a["href"]
+            # We want links that start with /wiki/ but skip those that contain "Special:"
+            if href.startswith("/wiki/") and "Special:" not in href:
+                article_url = "https://harrypotter.fandom.com" + href
+                break
+        
+        if not article_url:
+            return "No results found for your query."
+        
+        # Fetch the article page.
+        article_response = requests.get(article_url, headers=headers)
+        if article_response.status_code != 200:
+            return f"Error: Received status code {article_response.status_code} from the article page."
+        
+        article_soup = BeautifulSoup(article_response.text, 'html.parser')
+        
+        # Extract the first meaningful paragraph.
+        paragraphs = article_soup.find_all("p")
+        for p in paragraphs:
+            text = p.get_text().strip()
+            if len(text) > 50:  # A simple threshold to ensure the paragraph is informative.
+                return text
+        
+        return "Couldn't extract detailed lore from the article."
+    
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
+
 # Initialize the SmolAgent model
 model = HfApiModel(model_id="meta-llama/Llama-3.2-3B-Instruct", token=hf_token)
 
-# Create the agent with only the search tool
-agent = CodeAgent(tools=[search_tool], model=model)
+# Create the agent
+agent = CodeAgent(tools=[search_harry_potter_lore], model=model)
 
 def apply_custom_styles():
     st.markdown(
@@ -86,17 +139,12 @@ def main():
         st.session_state["state_reset"] = False
 
     def process_input():
-        """Generate a response as Tom Riddle using efficient web search only if needed."""
+        """Fetch Tom Riddle's response based on user input."""
         user_input = st.session_state["user_input"]
         if user_input.strip():
             try:
-                # Instruct Tom Riddle to use the search tool only when uncertain, efficiently.
                 response = agent.run(
-                    f"You are Tom Riddle, a cunning and enigmatic character from Harry Potter. "
-                    f"Answer the user's query clearly and concisely in your distinct persona. "
-                    f"If you are not certain of your answer, only then use your available DuckDuckGo search tool to quickly gather the necessary insights—doing so efficiently without causing delays. "
-                    f"Finally, produce your answer in your characteristic voice without revealing your internal process. "
-                    f"User Query: {user_input}"
+                    f"You are Tom Riddle, a cunning and enigmatic character from Harry Potter. Respond concisely and pragmatically and remain true to your role in the series: {user_input}"
                 )
                 st.session_state["response"] = response
                 st.session_state["waiting_for_input"] = False
@@ -124,4 +172,4 @@ def main():
         st.rerun()
 
 if __name__ == "__main__":
-    main()
+    main()