Add Wikipedia tool and enable embedding model configuration for tools in `tools.py`

src/gaia_solving_agent/agent.py

@@ -6,8 +6,6 @@ from llama_index.core.prompts import RichPromptTemplate
 from llama_index.llms.nebius import NebiusLLM
 from llama_index.llms.mistralai import MistralAI
 from llama_index.llms.openai import OpenAI
-from llama_index.tools.requests import RequestsToolSpec
-from llama_index.tools.wikipedia import WikipediaToolSpec
 from workflows import Workflow, step, Context
 from workflows.events import StartEvent, Event, StopEvent
 
@@ -15,13 +13,13 @@ from gaia_solving_agent import NEBIUS_API_KEY, MISTRAL_API_KEY, OPENAI_API_KEY
 from gaia_solving_agent.prompts import PLANING_PROMPT, FORMAT_ANSWER
 from gaia_solving_agent.tools import (
     tavily_search_web,
-    load_and_search_tools_from_toolspec,
     simple_web_page_reader_toolspec,
     vllm_ask_image_tool,
     youtube_transcript_reader_toolspec,
     text_content_analysis,
     research_paper_reader_toolspec,
     get_text_representation_of_additional_file,
+    wikipedia_toolspec,
 )
 from gaia_solving_agent.utils import extract_pattern
 
@@ -171,7 +169,7 @@ gaia_solving_agent = FunctionAgent(
         get_text_representation_of_additional_file,
         vllm_ask_image_tool,
         tavily_search_web,
-        *load_and_search_tools_from_toolspec(WikipediaToolSpec()),
+        *wikipedia_toolspec,
         *simple_web_page_reader_toolspec.to_tool_list(),
         *youtube_transcript_reader_toolspec.to_tool_list(),
         *research_paper_reader_toolspec.to_tool_list(),

src/gaia_solving_agent/tools.py

@@ -9,16 +9,33 @@ from llama_index.llms.openai import OpenAI
 from llama_index.readers.web import SimpleWebPageReader
 from llama_index.readers.youtube_transcript import YoutubeTranscriptReader
 from llama_index.readers.papers import ArxivReader
+from llama_index.tools.wikipedia import WikipediaToolSpec
 from tavily import AsyncTavilyClient
 from workflows import Context
 
 from gaia_solving_agent import TAVILY_API_KEY, NEBIUS_API_KEY, MISTRAL_API_KEY, OPENAI_API_KEY
 
+embedding_model = "text-embedding-3-small"
 
-def load_and_search_tools_from_toolspec(tool_spec: BaseToolSpec) -> list[FunctionTool]:
+
+def load_and_search_tools_from_toolspec(tool_spec: BaseToolSpec, **kwargs) -> list[FunctionTool]:
+    """
+    Loads and searches tools based on the provided tool specification, allowing
+    enhanced flexibility and customization through additional parameters.
+
+    Args:
+        tool_spec (BaseToolSpec): The tool specification derived from BaseToolSpec
+            which defines the source tools to be processed.
+        **kwargs: Arbitrary keyword arguments for customizing the tool processing
+            and search behavior.
+
+    Returns:
+        list[FunctionTool]: A list of processed and searched tools derived from
+        the input tool_spec and based on the provided custom parameters.
+    """
     tools_list = []
     for tool in tool_spec.to_tool_list():
-        tools_list.extend(LoadAndSearchToolSpec.from_defaults(tool).to_tool_list())
+        tools_list.extend(LoadAndSearchToolSpec.from_defaults(tool, **kwargs).to_tool_list())
     return tools_list
 
 
@@ -144,6 +161,7 @@ async def vllm_ask_image(query: str, images: ImageDocument | list[ImageDocument]
 
 simple_web_page_reader_tool = OnDemandLoaderTool.from_defaults(
     SimpleWebPageReader(html_to_text=True),
+    index_kwargs={"embed_model": embedding_model},
     name="simple_web_page_reader_tool",
     description="""
 Tool for loading content from a web page and return it as text.
@@ -159,10 +177,14 @@ Do not use this tool for:
 => You have specialized tools for those needs.
 """,
 )
-simple_web_page_reader_toolspec = LoadAndSearchToolSpec.from_defaults(simple_web_page_reader_tool)
+simple_web_page_reader_toolspec = LoadAndSearchToolSpec.from_defaults(
+    simple_web_page_reader_tool,
+    index_kwargs={"embed_model": embedding_model},
+)
 
 youtube_transcript_reader_tool = OnDemandLoaderTool.from_defaults(
     YoutubeTranscriptReader(),
+    index_kwargs={"embed_model": embedding_model},
     name="youtube_transcript_reader_tool",
     description=r"""
 Tool for loading the audio transcript from a youtube video and return it as text.
@@ -176,14 +198,26 @@ Supported formats include:
 If you are provided with a youtube link in the wrong format, make it fit one the supported format.
 """,
 )
-youtube_transcript_reader_toolspec = LoadAndSearchToolSpec.from_defaults(youtube_transcript_reader_tool)
+youtube_transcript_reader_toolspec = LoadAndSearchToolSpec.from_defaults(
+    youtube_transcript_reader_tool,
+    index_kwargs={"embed_model": embedding_model},
+)
 
 
 research_paper_reader_tool = OnDemandLoaderTool.from_defaults(
     ArxivReader(),
+    index_kwargs={"embed_model": embedding_model},
     name="research_paper_reader_tool",
     description=r"""
 Gets a search query, return a list of Documents of the top corresponding scientific papers on Arxiv.
 """,
 )
-research_paper_reader_toolspec = LoadAndSearchToolSpec.from_defaults(research_paper_reader_tool)
+research_paper_reader_toolspec = LoadAndSearchToolSpec.from_defaults(
+    research_paper_reader_tool,
+    index_kwargs={"embed_model": embedding_model},
+)
+
+wikipedia_toolspec = load_and_search_tools_from_toolspec(
+    WikipediaToolSpec(),
+    index_kwargs={"embed_model": embedding_model},
+)