tools/tool_collection_wiki.py

from langchain.tools import StructuredTool

from tools.wiki_search import (
    GetPageSectionContentInput,
    GetPageTitleExcerptSectionsInput,
    SearchWikipediaEn,
    WikipediaOpensearchInput,
)


class ToolsCollection:
    @staticmethod
    def get_tools(needed_tools: list[str]) -> list:
        tools = []

        for nt in needed_tools:
            if nt == "wikipedia_opensearch":
                tools.append(
                    StructuredTool(
                        name="wikipedia_opensearch",
                        func=SearchWikipediaEn.wikipedia_opensearch,
                        description=SearchWikipediaEn.wikipedia_opensearch.__doc__,
                        args_schema=WikipediaOpensearchInput,
                    )
                )
            elif nt == "get_page_title_excerpt_sections":
                tools.append(
                    StructuredTool(
                        name="get_page_title_excerpt_sections",
                        func=SearchWikipediaEn.get_page_title_excerpt_sections,
                        description=SearchWikipediaEn.get_page_title_excerpt_sections.__doc__,
                        args_schema=GetPageTitleExcerptSectionsInput,
                    )
                )
            elif nt == "get_page_section_content":
                tools.append(
                    StructuredTool(
                        name="get_page_section_content",
                        func=SearchWikipediaEn.get_page_section_content,
                        description=SearchWikipediaEn.get_page_section_content.__doc__,
                        args_schema=GetPageSectionContentInput,
                    )
                )

        return tools