File size: 11,260 Bytes
285af8d e8445f4 285af8d 9319fc8 285af8d 58232ea 285af8d 558b71f 285af8d 558b71f 285af8d 558b71f 285af8d 558b71f 285af8d bf883dc 285af8d bf883dc 285af8d 31a92c4 285af8d 31a92c4 285af8d c471e86 285af8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 | import requests
from bs4 import BeautifulSoup
import re
import wikipedia
import os
from typing import Dict, List, Any, Optional
from duckduckgo_search import DDGS
from base_agent import BaseAgent
class ResearchAgent(BaseAgent):
"""
Agent for handling factual questions by performing web searches, with enhanced structured data scraping.
"""
def __init__(self, model: str):
super().__init__(model=model)
def __call__(self, question: str, file_name: str = None, **kwargs) -> str:
print(f"Research agent processing question: {question[:100]}...")
# First extract the search query using LLM
search_query = self.extract_search_query(question)
# print(f"Extracted search query: {search_query}")
# Then try to search Wikipedia
answer = self.try_structured_wikipedia_scrape(search_query)
# print(f"Structured answer: {answer}")
if not answer:
# If structured answer is not found, perform a web search
search_results = self.search_web(search_query)
answer = self.format_search_results(search_results)
# Now over to LLM again to interpret the result and give final answer
system_prompt = """
You are a research expert tasked with answering factual questions based on search results.
Analyze the search results carefully and extract the information needed to answer the question accurately.
If the search results don't contain enough information, indicate what's missing.
Provide a concise, accurate answer based only on the provided information.
You must provide only the exact answer requested by the user, in the correct type and minimal format, based purely on the provided information.
YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
If you cannot find the answer based on the provided information, reply only with "Unknown".
"""
user_prompt = f"""
Question: {question}
Search results:
{answer}
Based on these search results, provide a concise and accurate answer to the question.
"""
final_answer = self.llm_call(system_prompt, user_prompt)
print(f"Final answer: {final_answer}")
if final_answer.lower() == "unknown":
final_answer = self.ask_james(question)
return final_answer
def try_structured_wikipedia_scrape(self, query: str) -> Optional[str]:
try:
# topic = self.extract_search_query(query) # <-- Use your LLM-based extractor here
# if not topic:
# print("No topic extracted.")
# return None
print(f"Structured scrape attempt for topic: {query}")
page_title = wikipedia.search(query, results=1)
if not page_title:
print("No Wikipedia page found.")
return None
page = wikipedia.page(page_title[0])
page_url = page.url
html = requests.get(page_url).text
soup = BeautifulSoup(html, "html.parser")
target_table = self.find_relevant_table(soup, query)
if not target_table:
print("No relevant table found.")
return None
extracted_items = self.parse_table_to_json(target_table)
if not extracted_items:
return "No relevant structured data found matching the question."
answer = "Here is the structured data extracted from Wikipedia:\n\n"
for item in extracted_items:
answer += f"- {item}\n"
return answer
except Exception as e:
print(f"Structured scrape failed: {e}")
return None
def find_relevant_table(self, soup: BeautifulSoup, question: str) -> Optional[BeautifulSoup]:
"""
Find the most relevant table based on the question type.
"""
print("Finding relevant table in " + question)
target_keywords = []
if "studio albums" in question.lower():
target_keywords = ["Studio albums"]
elif "album" in question.lower():
target_keywords = ["Discography"]
elif "book" in question.lower():
target_keywords = ["Bibliography", "Books"]
elif "battle" in question.lower():
target_keywords = ["Battles", "Military history"]
else:
# fallback to first wikitable
tables = soup.find_all("table", {"class": "wikitable"})
return tables[0] if tables else None
print(f"Target keywords for table search: {target_keywords}")
# headings = soup.find_all(['h2', 'h3'])
headings = soup.find_all(lambda tag: tag.name in ["h2", "h3"] and any(keyword.lower() in tag.get_text(strip=True).lower() for keyword in target_keywords))
# print(f"Found {len(headings)} headings in the page.")
for heading in headings:
heading_text = heading.get_text()
# print(f"Checking heading: {heading_text}")
for keyword in target_keywords:
if keyword.lower() in heading_text.lower():
print(f"Found relevant heading: {heading_text}")
next_sibling = heading.find_next_sibling()
while next_sibling:
# Move to the parent div
heading_div = heading.find_parent("div")
if heading_div:
next_sibling = heading_div.find_next_sibling()
while next_sibling:
if next_sibling.name == "table" and "wikitable" in next_sibling.get("class", []):
print(f"Found relevant table under heading: {heading_text}")
return next_sibling
next_sibling = next_sibling.find_next_sibling()
return None
def parse_table_to_json(self, table: BeautifulSoup) -> List[Dict[str, Any]]:
"""
Generic parser for any Wikipedia 'wikitable' to JSON.
Keys are extracted from <th> headers.
Values from <td> cells.
"""
result = []
# Extract headers
headers = []
header_row = table.find("tr")
if not header_row:
return result
for th in header_row.find_all("th"):
header_text = th.get_text(strip=True)
headers.append(header_text)
# Parse data rows
for row in table.find_all("tr")[1:]: # skip header
cols = row.find_all("td")
if not cols or len(cols) != len(headers):
continue
row_data = {}
for header, col in zip(headers, cols):
value = col.get_text(separator=" ", strip=True)
if header.lower() == "album details":
# Try to get album title
album_link = col.find("a")
if album_link:
album_title = album_link.get_text(strip=True)
else:
# Get plain text up to first <ul> if exists
ul = col.find("ul")
if ul:
# Extract text before the <ul>
raw_text = "".join(
[str(item) for item in col.contents if getattr(item, 'name', None) != "ul"]
).strip()
if '<' in raw_text and '>' in raw_text:
album_title = BeautifulSoup(raw_text, "html.parser").get_text(strip=True)
else:
album_title = raw_text.strip()
else:
album_title = col.get_text(strip=True)
# Extract extras from <ul>
extras = []
ul = col.find("ul")
if ul:
extras = [li.get_text(strip=True) for li in ul.find_all("li")]
# Build structured value
if album_title and extras:
value = {
"title": album_title,
"extra": extras
}
elif album_title:
value = {
"title": album_title
}
elif extras:
value = extras
row_data[header] = value
result.append(row_data)
return result
def search_web(self, query: str, max_results: int = 5) -> List[Dict[str, str]]:
"""
Perform a web search using DuckDuckGo.
Args:
query: The search query
max_results: Maximum number of results to return
Returns:
List of search results with title, content, and url
"""
results = []
DUCKDUCKGO_ENABLED = os.getenv("DUCKDUCKGO_ENABLED")
try:
print("Using DuckDuckGo to conduct a web search")
# Then use DuckDuckGo if enabled and we need more results
if DUCKDUCKGO_ENABLED and len(results) < max_results:
with DDGS() as ddgs:
ddgs_results = list(ddgs.text(query, max_results=max_results))
for r in ddgs_results:
results.append({
"title": r.get("title", ""),
"content": r.get("body", ""),
"url": r.get("href", "")
})
except Exception as e:
print(f"Web search error: {str(e)}")
return results[:max_results]
def format_search_results(self, results: List[Dict[str, str]]) -> str:
"""
Format search results for the LLM.
Args:
results: List of search results
Returns:
Formatted search results as a string
"""
if not results:
return "No search results found."
formatted = ""
for i, result in enumerate(results, 1):
formatted += f"Result {i}:\n"
formatted += f"Title: {result.get('title', 'No title')}\n"
formatted += f"URL: {result.get('url', 'No URL')}\n"
formatted += f"Content: {result.get('content', 'No content')}\n\n"
return formatted |