websearch_agent / app_1.py
palbha's picture
Create app_1.py
32edf4c verified
raw
history blame
1.96 kB
import gradio as gr
from smolagents import CodeAgent, HfApiModel, DuckDuckGoSearchTool, tool
from smolagents.agents import ActionStep
from time import sleep
import helium
from helium import S, Keys
import requests
from bs4 import BeautifulSoup
# Define a function to extract image URLs from DuckDuckGo search
@tool
def search_images_for_presentation(query: str):
"""
Searches for images related to the given query on DuckDuckGo and returns a list of image URLs.
"""
helium.go_to("https://duckduckgo.com/")
search_box = helium.find(S("input[type='text']"))
search_box.write(query + " image")
search_box.press(Keys.ENTER)
# Wait for search results to load
sleep(3)
# Extract image links from the page
soup = BeautifulSoup(helium.get_driver().page_source, "html.parser")
images = [img["src"] for img in soup.find_all("img") if "src" in img.attrs]
return images[:5] # Return the top 5 image URLs
# Initialize agent
def initialize_agent(model):
return CodeAgent(
tools=[DuckDuckGoSearchTool(), search_images_for_presentation],
model=model,
max_steps=20,
verbosity_level=2,
)
# Function to run the agent and return image links
def run_agent(query: str):
model = HfApiModel(model_id="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud/")
agent = initialize_agent(model)
image_urls = agent.run(query)
return image_urls
# Gradio interface function
def gradio_interface(query: str):
image_urls = run_agent(query)
# Format output as clickable links
links = [f'<a href="{url}" target="_blank">{url}</a>' for url in image_urls]
return "\n".join(links)
# Set up Gradio interface
gr.Interface(
fn=gradio_interface,
inputs="text",
outputs="html",
title="Figure, Image & Logo Finder",
description="Enter a query to search for relevant images, logos, or figures for your presentation.",
).launch()