Wikipedia redirect
This commit is contained in:
+28
-14
@@ -76,10 +76,11 @@ async def _search(title: str, limit: int = 5, fallback: bool = False) -> str:
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
async def _fetch_content(title: str, section_index: int) -> Optional[str]:
|
||||
async def _fetch_content(title: str, section_index: int) -> tuple[Optional[str], Optional[str]]:
|
||||
"""
|
||||
Retrieves raw Wikitext for a specific section.
|
||||
section_index=0 returns the lead section.
|
||||
Returns a tuple of (content, final_title).
|
||||
"""
|
||||
params = {
|
||||
"action": "query",
|
||||
@@ -94,24 +95,26 @@ async def _fetch_content(title: str, section_index: int) -> Optional[str]:
|
||||
|
||||
pages = data.get("query", {}).get("pages", {})
|
||||
if not pages:
|
||||
return None
|
||||
return None, None
|
||||
|
||||
page_id = next(iter(pages))
|
||||
page = pages[page_id]
|
||||
final_title = page.get("title")
|
||||
|
||||
if "missing" in page:
|
||||
return None
|
||||
return None, None
|
||||
|
||||
revisions = page.get("revisions", [])
|
||||
if not revisions:
|
||||
return None
|
||||
return None, final_title
|
||||
|
||||
content = revisions[0].get("*")
|
||||
return clean_wikitext(content) if content else None
|
||||
return (clean_wikitext(content) if content else None), final_title
|
||||
|
||||
async def _fetch_toc(title: str) -> Optional[str]:
|
||||
async def _fetch_toc(title: str) -> tuple[Optional[str], Optional[str]]:
|
||||
"""
|
||||
Retrieves the Table of Contents data and formats it hierarchically.
|
||||
Returns a tuple of (toc, final_title).
|
||||
"""
|
||||
params = {
|
||||
"action": "parse",
|
||||
@@ -124,14 +127,15 @@ async def _fetch_toc(title: str) -> Optional[str]:
|
||||
|
||||
parse_data = data.get("parse")
|
||||
if not parse_data:
|
||||
return None
|
||||
return None, None
|
||||
|
||||
final_title = parse_data.get("title")
|
||||
toc_data = parse_data.get("tocdata", {})
|
||||
sections = toc_data.get("sections", [])
|
||||
if not sections:
|
||||
return "No table of contents found for this page."
|
||||
return "No table of contents found for this page.", final_title
|
||||
|
||||
lines = [f"Table of Contents for \"{parse_data.get('title', title)}\":"]
|
||||
lines = [f"Table of Contents for \"{final_title}\":"]
|
||||
for s in sections:
|
||||
level = s.get("tocLevel", 1)
|
||||
index = s.get("index")
|
||||
@@ -139,7 +143,7 @@ async def _fetch_toc(title: str) -> Optional[str]:
|
||||
indent = " " * (level - 1)
|
||||
lines.append(f"{indent}[{index}] {line}")
|
||||
|
||||
return "\n".join(lines)
|
||||
return "\n".join(lines), final_title
|
||||
|
||||
async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
@@ -154,18 +158,24 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
section_index = args.get("section_index")
|
||||
search_limit = args.get("search_limit", 5)
|
||||
|
||||
final_title = None
|
||||
|
||||
if mode == "search":
|
||||
result = await _search(title, search_limit, fallback=False)
|
||||
elif mode == "summary":
|
||||
# Lead section + ToC is the default 'summary' to guide the AI's next steps
|
||||
content = await _fetch_content(title, 0)
|
||||
content, title_from_content = await _fetch_content(title, 0)
|
||||
final_title = title_from_content
|
||||
if content is None:
|
||||
result = await _search(title, search_limit, fallback=True)
|
||||
else:
|
||||
toc = await _fetch_toc(title)
|
||||
toc, title_from_toc = await _fetch_toc(title)
|
||||
if title_from_toc:
|
||||
final_title = title_from_toc
|
||||
result = f"{content}\n\n---\n\n{toc}"
|
||||
elif mode == "toc":
|
||||
toc = await _fetch_toc(title)
|
||||
toc, title_from_toc = await _fetch_toc(title)
|
||||
final_title = title_from_toc
|
||||
if toc is None:
|
||||
result = await _search(title, search_limit, fallback=True)
|
||||
else:
|
||||
@@ -173,7 +183,8 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
elif mode == "section":
|
||||
if section_index is None:
|
||||
raise ToolError("Missing required parameter 'section_index' for mode='section'")
|
||||
content = await _fetch_content(title, int(section_index))
|
||||
content, title_from_content = await _fetch_content(title, int(section_index))
|
||||
final_title = title_from_content
|
||||
if content is None:
|
||||
result = await _search(title, search_limit, fallback=True)
|
||||
else:
|
||||
@@ -181,4 +192,7 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
else:
|
||||
raise ToolError(f"Invalid mode '{mode}'. Supported modes: summary, toc, section, search")
|
||||
|
||||
if final_title and final_title != title:
|
||||
result = f"Redirected to \"{final_title}\"\n\n{result}"
|
||||
|
||||
return [{"type": "text", "text": result}]
|
||||
|
||||
Reference in New Issue
Block a user