src/custom_tools/wikipedia.py

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.")