tools/wiki_tool.py

from smolagents.tools import Tool
import wikipedia

class WikiTool(Tool):
    def __init__(self):
        self.name = "WikiTool"  # Явно задаем имя
        self.description = (
            "Performs Wikipedia lookups. Supports actions: 'summary' and 'is_historical_country'.\n"
            "Usage:\n"
            "- {'action': 'summary', 'topic': 'France'}\n"
            "- {'action': 'is_historical_country', 'country_name': 'Yugoslavia'}"
        )
        self.inputs = {
            "action": {
                "type": "string",
                "description": "The action to perform: 'summary' or 'is_historical_country'.",
                "nullable": True  # Сделать action необязательным
            },
            "topic": {
                "type": "string",
                "description": "The topic to summarize (required for 'summary' action).",
                "nullable": True
            },
            "country_name": {
                "type": "string",
                "description": "The country name to check (required for 'is_historical_country' action).",
                "nullable": True
            }
        }
        self.output_type = "string" 
        super().__init__(
            name=self.name,
            description=self.description,
            inputs=self.inputs,
            output_type=self.output_type
        )

    def fetch_summary(self, topic: str) -> str:
        try:
            summary = wikipedia.summary(topic, sentences=3)
            return summary
        except wikipedia.DisambiguationError as e:
            return f"Disambiguation error: The topic '{topic}' is ambiguous. Suggestions: {e.options[:5]}"
        except wikipedia.PageError:
            return f"Page error: The topic '{topic}' does not exist on Wikipedia."
        except Exception as e:
            return f"An unexpected error occurred: {str(e)}"

    def is_historical_country(self, country_name: str) -> str:
        try:
            summary = wikipedia.summary(country_name, sentences=2).lower()
            keywords = [
                "former country", "no longer exists", "historical country",
                "was a country", "defunct", "dissolved", "existed until",
                "disestablished", "merged into"
            ]
            result = any(keyword in summary for keyword in keywords)
            return "Yes" if result else "No"
        except Exception:
            return "Unknown (error during lookup)"

    def forward(self, action: str = "", topic: str = "", country_name: str = "") -> str:
        if action == "summary":
            return self.fetch_summary(topic)
        elif action == "is_historical_country":
            return self.is_historical_country(country_name)
        else:
            return "Error: Unknown action. Use 'summary' or 'is_historical_country'."