final-agent-course / agent.py
jjvelezo's picture
Update agent.py
54b69c4 verified
raw
history blame
5.03 kB
import os
import requests
import urllib.parse
from bs4 import BeautifulSoup
class DuckDuckGoAgent:
def __init__(self):
print("DuckDuckGoAgent initialized.")
self.headers = {
"User-Agent": "Mozilla/5.0"
}
# Support for multiple model backends
self.supported_models = {
"huggingface": self.call_huggingface_llm,
# You can easily extend this dictionary to support:
# "openai": self.call_openai_model,
# "lite_llm": self.call_litellm_model,
# "custom_server": self.call_custom_model,
}
self.default_model = "huggingface"
self.model_config = {
"huggingface": {
"api_key": os.getenv("HF_API_TOKEN"),
"model_name": "mistralai/Mistral-7B-Instruct-v0.1"
}
}
def __call__(self, question: str) -> str:
"""
Main method to process a question. It first tries DuckDuckGo,
then scraping, and finally uses a language model if needed.
"""
print(f"Agent received question: {question[:50]}...")
answer = self.get_duckduckgo_answer(question)
print(f"Agent returning answer: {answer}")
return answer.strip()
def get_duckduckgo_answer(self, query: str) -> str:
"""
Attempt to get an answer from the DuckDuckGo API.
If no abstract text is found, fall back to scraping.
"""
search_query = urllib.parse.quote(query)
url = f"https://api.duckduckgo.com/?q={search_query}&format=json&no_html=1&skip_disambig=1"
try:
response = requests.get(url, timeout=10)
if response.status_code == 200:
data = response.json()
if 'AbstractText' in data and data['AbstractText']:
return data['AbstractText'][:200]
else:
print("No abstract found, falling back to scraping.")
return self.scrape_duckduckgo(query)
else:
print(f"DuckDuckGo API failed with status: {response.status_code}")
return self.scrape_duckduckgo(query)
except Exception as e:
print(f"Error contacting DuckDuckGo API: {e}")
return self.scrape_duckduckgo(query)
def scrape_duckduckgo(self, query: str) -> str:
"""
Fallback to scraping DuckDuckGo search results if API fails or no abstract found.
"""
print("Using fallback: scraping HTML results.")
try:
response = requests.post(
"https://html.duckduckgo.com/html/",
data={"q": query},
headers=self.headers,
timeout=10
)
soup = BeautifulSoup(response.text, "html.parser")
snippets = soup.select(".result__snippet")
for s in snippets:
text = s.get_text().strip()
if text:
return text[:200]
print("No useful snippets found, falling back to language model.")
return self.call_model_backend(query)
except Exception as e:
print(f"Error scraping DuckDuckGo: {e}")
return self.call_model_backend(query)
def call_model_backend(self, prompt: str) -> str:
"""
Dispatch to the selected LLM backend.
"""
if self.default_model in self.supported_models:
return self.supported_models[self.default_model](prompt)
return "No valid model backend configured."
def call_huggingface_llm(self, prompt: str) -> str:
"""
Call Hugging Face Inference API as fallback LLM.
"""
config = self.model_config.get("huggingface", {})
api_key = config.get("api_key")
model = config.get("model_name")
if not api_key or not model:
return "Error: Hugging Face API Token or model not configured."
url = f"https://api-inference.huggingface.co/models/{model}"
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
}
payload = {
"inputs": prompt,
"parameters": {
"max_new_tokens": 200,
"temperature": 0.7
}
}
try:
response = requests.post(url, headers=headers, json=payload, timeout=30)
response.raise_for_status()
output = response.json()
if isinstance(output, list) and "generated_text" in output[0]:
return output[0]["generated_text"].strip()[:200]
elif isinstance(output, dict) and "error" in output:
return f"HF LLM error: {output['error']}"
else:
return "No response generated from Hugging Face LLM."
except Exception as e:
print(f"Error contacting Hugging Face LLM: {e}")
return "Error contacting Hugging Face model."