app.py

# app.py
import os
import random
import requests
import streamlit as st
from io import BytesIO
from PIL import Image
from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, Tool
from huggingface_hub import login
import warnings

warnings.filterwarnings("ignore")

# Set page configuration
st.set_page_config(
    page_title="SmolaAgents Tools",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialize HF API key from environment variable
@st.cache_resource
def initialize_api():
    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        login(hf_token)
    else:
        st.warning("Warning: HF_API_TOKEN not found in environment variables")
    
    # Initialize the model
    return HfApiModel(model_id="meta-llama/Llama-3.3-70B-Instruct")

model = initialize_api()

# Define all your tools
class RandomComicFetcher(Tool):
    name = "fetch_random_comic_image"
    description = "Fetches the image of a random XKCD comic"
    inputs = {}
    output_type = "string"

    def __init__(self):
        super().__init__()
        self.xkcd_base_url = "https://xkcd.com/"
        self.xkcd_latest_url = f"{self.xkcd_base_url}info.0.json"

    def forward(self):
        try:
            # Fetch the latest comic info to get max comic number
            latest_resp = requests.get(self.xkcd_latest_url)
            latest_resp.raise_for_status()
            latest_data = latest_resp.json()
            max_comic_id = latest_data["num"]

            # Pick a random comic ID between 1 and latest
            selected_id = random.randint(1, max_comic_id)
            comic_info_url = f"{self.xkcd_base_url}{selected_id}/info.0.json"

            # Fetch comic metadata
            comic_resp = requests.get(comic_info_url)
            comic_resp.raise_for_status()
            comic_data = comic_resp.json()
            comic_img_url = comic_data["img"]
            comic_title = comic_data["title"]
            comic_alt = comic_data["alt"]

            # Download and return the comic image
            image_resp = requests.get(comic_img_url)
            image_resp.raise_for_status()
            img = Image.open(BytesIO(image_resp.content))
            
            return img, comic_title, comic_alt, comic_img_url
        except requests.exceptions.RequestException as err:
            st.error(f"Could not fetch comic: {err}")
            return None, "Error fetching comic", "Error details", ""

class PrimeCheckTool(Tool):
    name = "prime_check"
    description = "Checks if a given number is a prime number."
    inputs = {
        "number": {
            "type": "integer",
            "description": "The number to check for primality.",
        }
    }
    output_type = "boolean"

    def forward(self, number: int) -> bool:
        if number < 2:
            return False
        for i in range(2, int(number**0.5) + 1):
            if number % i == 0:
                return False
        return True

# Initialize tools
search_tool = DuckDuckGoSearchTool()
random_comic = RandomComicFetcher()
prime_check_tool = PrimeCheckTool()

# Streamlit App
st.title("SmolaAgents Tools")
st.markdown("Select a tool from the sidebar and interact with AI-powered assistants")

# Sidebar for tool selection
with st.sidebar:
    st.title("Tool Selection")
    tool_choice = st.radio(
        "Choose a tool:",
        ["Search Tool", "XKCD Comic Fetcher", "Prime Number Checker"]
    )
    
    st.divider()
    st.markdown("### About")
    st.markdown("""
    This app demonstrates the capabilities of SmolaAgents, allowing you to:
    - Search the web with an AI assistant
    - Fetch random XKCD comics with AI commentary
    - Check if numbers are prime with creative explanations
    """)

# Main content area
if tool_choice == "Search Tool":
    st.header("AI Web Search")
    st.markdown("Ask any question and the AI will search the web for answers")
    
    query = st.text_input("Enter your search query:", placeholder="What are the latest advancements in renewable energy?")
    
    if st.button("Search"):
        if query:
            with st.spinner("Searching and processing..."):
                agent = CodeAgent(tools=[search_tool], model=model)
                response = agent.run(query)
                st.success("Search complete!")
                st.markdown("### Results")
                st.markdown(response)
        else:
            st.warning("Please enter a search query")

elif tool_choice == "XKCD Comic Fetcher":
    st.header("XKCD Comic Explorer")
    st.markdown("Fetch a random XKCD comic and get AI commentary")
    
    if "comic_fetched" not in st.session_state:
        st.session_state.comic_fetched = False
        st.session_state.comic_img = None
        st.session_state.comic_title = ""
        st.session_state.comic_alt = ""
        st.session_state.comic_url = ""
    
    col1, col2 = st.columns([1, 2])
    
    with col1:
        if st.button("Fetch Random Comic"):
            with st.spinner("Fetching a random comic..."):
                img, title, alt, url = random_comic.forward()
                if img:
                    st.session_state.comic_fetched = True
                    st.session_state.comic_img = img
                    st.session_state.comic_title = title
                    st.session_state.comic_alt = alt
                    st.session_state.comic_url = url
                    st.experimental_rerun()
    
    if st.session_state.comic_fetched:
        with col1:
            st.image(st.session_state.comic_img, caption=st.session_state.comic_title)
            st.caption(f"Alt text: {st.session_state.comic_alt}")
            st.markdown(f"[View on XKCD]({st.session_state.comic_url})")
        
        with col2:
            st.subheader("Ask AI about this comic")
            query = st.text_input("What would you like to know about this comic?", 
                                 placeholder="Explain this comic in a funny way")
            
            if st.button("Ask AI"):
                with st.spinner("Generating response..."):
                    agent = CodeAgent(tools=[random_comic], model=model)
                    response = agent.run(query if query else "Tell me about this XKCD comic.")
                    st.markdown("### AI Response")
                    st.markdown(response)

elif tool_choice == "Prime Number Checker":
    st.header("Prime Number Checker")
    st.markdown("Check if a number is prime and get a creative explanation")
    
    number = st.number_input("Enter a number to check:", min_value=1, step=1, value=23)
    explanation_style = st.selectbox(
        "Choose explanation style:",
        ["Poetic", "Nursery Rhyme", "Scientific", "Humorous", "Historical"]
    )
    
    if st.button("Check Prime"):
        with st.spinner("Checking and generating response..."):
            is_prime = prime_check_tool.forward(int(number))
            
            # Format query based on selected style
            query = f"Check if the number {number} is a prime number and explain it in a {explanation_style.lower()} style."
            
            agent = CodeAgent(tools=[prime_check_tool], model=model)
            response = agent.run(query)
            
            # Display result with some styling
            prime_status = "✅ PRIME" if is_prime else "❌ NOT PRIME"
            st.markdown(f"### Result: {prime_status}")
            
            st.markdown("### AI Explanation")
            st.markdown(response)
            
            # Optional: Show a mathematical representation
            if st.checkbox("Show mathematical details"):
                if is_prime:
                    st.markdown(f"{number} has no factors other than 1 and itself.")
                else:
                    # Find the factors
                    factors = [i for i in range(1, number + 1) if number % i == 0]
                    st.markdown(f"Factors of {number}: {', '.join(map(str, factors))}")

# Add footer
st.sidebar.divider()
st.sidebar.markdown("Made with ❤️ using SmolaAgents")