Spaces:
Sleeping
Sleeping
File size: 3,102 Bytes
afb4047 |
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 |
import requests
from typing_extensions import Any, Annotated
from langchain_core.tools import tool
from langchain_core.messages import ToolMessage
from langgraph.types import Command
from langgraph.prebuilt import InjectedState
from langchain_core.tools.base import InjectedToolCallId
from src.state import State
@tool
def get_wiki_page_sections(
page_title: str,
tool_call_id: Annotated[str, InjectedToolCallId]
) -> Command:
"""Get sections of a Wikipedia page.
This function retrieves the sections of a Wikipedia page.
It requires the page title as an input parameter.
"""
page_title = page_title.replace(" ", "_")
payload = {
"action": "parse",
"page": page_title,
"prop": "sections",
"format": "json",
}
response = requests.get(
"https://en.wikipedia.org/w/api.php",
params=payload
)
if not response.status_code == 200:
return (f"Error fetching sections for {page_title}: {response.test}")
data = response.json()
sections = data.get("parse", {}).get("sections", [])
sections_map = {}
for section in sections:
section_title = section.get("anchor").lower()
section_number = section.get("index")
if section_title and section_number:
sections_map[section_title] = section_number
sections_text = "The sections of the page are:\n"
for title in sections_map.keys():
sections_text += f"{title}\n"
return Command(
update={
# update the state keys
"wiki_sections": sections_map,
# update the message history
"messages": [
ToolMessage(
sections_text, tool_call_id=tool_call_id
)
],
}
)
@tool
def get_wiki_page_by_section(
page_title: str,
section: str,
state: Annotated[State, InjectedState]
) -> str:
"""Get sections of a Wikipedia page.
This function retrieves the content of a specific section from a Wikipedia page.
It requires the page title and the section name as input parameters.
"""
wiki_sections = state.wiki_sections
if not wiki_sections:
return (f"Error: No sections found for {page_title}. Please run get_page_sections first.")
page_title = page_title.replace(" ", "_")
section = section.replace(" ", "_").lower()
if section not in wiki_sections:
return (f"Error: Section '{section}' not found in {page_title}. Please run get_page_sections first.")
payload = {
"action": "parse",
"page": page_title,
"prop": "wikitext",
"section": wiki_sections[section],
"format": "json",
}
response = requests.get(
"https://en.wikipedia.org/w/api.php",
params=payload
)
if not response.status_code == 200:
return (f"Error fetching sections for {page_title}: {response.test}")
data = response.json()
return data.get("parse", {}).get("wikitext", "No content found.")
|