Final_Assignment_Template / agent_tools.py
Facelook's picture
Disabled some prints.
9fde8f2
import requests
import json
from bs4 import BeautifulSoup
def duckduckgo_search(query: str, count: int = 3) -> list:
"""
Perform a search using DuckDuckGo and return the results.
Args:
query: The search query string
count: Maximum number of results to return (default: 3)
Returns:
List of search results
"""
print(f"Performing DuckDuckGo search for: {query}")
try:
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"
}
# Format the query for the URL
formatted_query = query.replace(' ', '+')
# Format the URL with query and parameter to increase snippet size
url = f"https://html.duckduckgo.com/html/?q={formatted_query}&kl=wt-wt"
# Send the request
response = requests.get(url, headers=headers, timeout=10)
response.raise_for_status()
# Parse the HTML response
soup = BeautifulSoup(response.text, 'html.parser')
# Extract search results
results = []
for result in soup.select('.result'):
title_elem = result.select_one('.result__title')
link_elem = result.select_one('.result__url')
snippet_elem = result.select_one('.result__snippet')
if title_elem and link_elem:
title = title_elem.get_text(strip=True)
url = link_elem.get('href') if link_elem.get('href') else link_elem.get_text(strip=True)
snippet = snippet_elem.get_text(strip=True) if snippet_elem else ""
# results.append({
# "title": title,
# "url": url,
# "snippet": snippet
# })
results.append(snippet)
if len(results) >= count:
break
print(f"DuckDuckGo results: {results}")
return results
except Exception as e:
print(f"Error during DuckDuckGo search: {e}")
return []
def langsearch_search(query: str, count: int = 5) -> list:
"""
Perform a search using LangSearch API and return the results.
Args:
query: The search query string
count: Maximum number of results to return (default: 5)
api_key: LangSearch API key (default: None, will look for env variable)
Returns:
List of search results
"""
print(f"Performing LangSearch search for: {query}")
try:
import os
# Use API key from parameters or environment variable
api_key = os.environ.get("LS_TOKEN")
if not api_key:
print("Warning: No LangSearch API key provided. Set LS_TOKEN environment variable.")
return []
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}"
}
payload = json.dumps({
"query": query,
"freshness": "noLimit",
"summary": True,
"count": count
})
url = "https://api.langsearch.com/v1/web-search"
response = requests.post(url, headers=headers, data=payload, timeout=30)
response.raise_for_status()
if response.status_code != 200:
print(f"LangSearch API error: {response.text}")
return []
response = response.json()
results = []
for result in response["data"]["webPages"]["value"]:
results.append(result["summary"])
#print(f"LangSearch results: {results}")
return results
except Exception as e:
print(f"Error during LangSearch search: {e}")
return []
# Dictionary mapping tool names to their functions
TOOLS_MAPPING = {
"duckduckgo_search": duckduckgo_search,
"langsearch_search": langsearch_search
}
# Tool definitions for LLM API
TOOLS_DEFINITION = [
{
"type": "function",
"function": {
"name": "duckduckgo_search",
"description": "Search the web using DuckDuckGo and return relevant results",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query string"
},
"num_results": {
"type": "integer",
"description": "Maximum number of results to return",
"default": 5
}
},
"required": ["query"]
}
}
},
{
"type": "function",
"function": {
"name": "langsearch_search",
"description": "Search the web using LangSearch API for more relevant results with deeper context",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query string"
},
"top_k": {
"type": "integer",
"description": "Maximum number of results to return",
"default": 5
},
"api_key": {
"type": "string",
"description": "LangSearch API key (optional, will use LANGSEARCH_API_KEY env var if not provided)"
}
},
"required": ["query"]
}
}
}
]