Address CodeRabbit feedback for Wikipedia tool

Signed-off-by: tamohannes <hovhannes.tamoyan@gmail.com>

Author: tamohannes

docs/agentic_inference/tool_calling.md

@@ -372,4 +372,3 @@ For vLLM, you may need to specify tool calling arguments:
 - [`nemo_skills.mcp.servers.python_tool.PythonTool`](https://github.com/NVIDIA-NeMo/Skills/tree/main/nemo_skills/mcp/servers/python_tool.py) - Python code execution
 - [`nemo_skills.mcp.servers.exa_tool.ExaTool`](https://github.com/NVIDIA-NeMo/Skills/tree/main/nemo_skills/mcp/servers/exa_tool.py) - Web search via Exa API
 - [`nemo_skills.mcp.servers.wikipedia_tool.WikipediaSearchTool`](https://github.com/NVIDIA-NeMo/Skills/tree/main/nemo_skills/mcp/servers/wikipedia_tool.py) - Direct Wikipedia article search and retrieval (no API key required)
-- [`nemo_skills.mcp.servers.wikipedia_tool.WikipediaSearchTool`](https://github.com/NVIDIA-NeMo/Skills/tree/main/nemo_skills/mcp/servers/wikipedia_tool.py) - Wikipedia article search and retrieval (no API key required)

nemo_skills/mcp/servers/wikipedia_tool.py

@@ -72,6 +72,7 @@
     "WIKIPEDIA_RATE_LIMIT_LOCK",
     os.path.join(tempfile.gettempdir(), "nemo_skills_wikipedia_api_rate_limit.lock"),
 )
+EXPECTED_HTTP_ERRORS = (httpx.HTTPError, RuntimeError, json.JSONDecodeError)
 
 _cache: dict[str, str] = {}
 _request_lock = asyncio.Lock()
@@ -270,7 +271,7 @@ async def wikipedia_search(
     async with httpx.AsyncClient(headers=headers) as client:
         try:
             data = await _http_get_json(client, ACTION_BASE, params=params)
-        except Exception as e:
+        except EXPECTED_HTTP_ERRORS as e:
             return f"Wikipedia search failed: {e}"
 
     hits = (data.get("query") or {}).get("search") or []
@@ -328,7 +329,7 @@ async def wikipedia_page(
     async with httpx.AsyncClient(headers=headers) as client:
         try:
             data = await _http_get_json(client, ACTION_BASE, params=params)
-        except Exception as e:
+        except EXPECTED_HTTP_ERRORS as e:
             return f"Wikipedia lookup failed: {e}"
 
     pages = (data.get("query") or {}).get("pages") or []
@@ -383,7 +384,7 @@ async def wikipedia_summary(
     async with httpx.AsyncClient(headers=headers) as client:
         try:
             data = await _http_get_json(client, ACTION_BASE, params=params)
-        except Exception as e:
+        except EXPECTED_HTTP_ERRORS as e:
             return f"Wikipedia summary failed: {e}"
     pages = (data.get("query") or {}).get("pages") or []
     if not pages or pages[0].get("missing"):
@@ -427,7 +428,7 @@ async def wikipedia_sections(
                     "formatversion": "2",
                 },
             )
-        except Exception as e:
+        except EXPECTED_HTTP_ERRORS as e:
             return f"Wikipedia section listing failed: {e}"
     sections = (data.get("parse") or {}).get("sections") or []
     if not sections:
@@ -457,7 +458,7 @@ async def wikipedia_query_summary(
         return cached
     try:
         rendered_title, url, extract = await _page_extract(t)
-    except Exception as e:
+    except EXPECTED_HTTP_ERRORS as e:
         return f"Wikipedia query summary failed: {e}"
     if rendered_title is None:
         return extract
@@ -498,7 +499,7 @@ async def wikipedia_key_facts(
         return cached
     try:
         rendered_title, url, extract = await _page_extract(t)
-    except Exception as e:
+    except EXPECTED_HTTP_ERRORS as e:
         return f"Wikipedia key facts failed: {e}"
     if rendered_title is None:
         return extract
@@ -554,7 +555,7 @@ async def wikipedia_section(
                     "formatversion": "2",
                 },
             )
-        except Exception as e:
+        except EXPECTED_HTTP_ERRORS as e:
             return f"Wikipedia section listing failed: {e}"
         sections = (sec_data.get("parse") or {}).get("sections") or []
         if not sections:
@@ -585,7 +586,7 @@ async def wikipedia_section(
                     "formatversion": "2",
                 },
             )
-        except Exception as e:
+        except EXPECTED_HTTP_ERRORS as e:
             return f"Wikipedia section fetch failed: {e}"
 
         # Prefer the rendered HTML stripped (cleaner than raw wikitext).
@@ -619,7 +620,7 @@ async def _suggest_titles(query: str, n: int = 5) -> list[str]:
                 return []
             data = r.json()
             return data[1] if isinstance(data, list) and len(data) > 1 else []
-    except Exception:
+    except EXPECTED_HTTP_ERRORS:
         return []
 
 
@@ -641,8 +642,24 @@ def default_config(self) -> dict[str, Any]:
         return dict(self._config)
 
     def configure(self, overrides: dict[str, Any] | None = None, context: dict[str, Any] | None = None) -> None:
-        if overrides:
-            self._config.update(overrides)
+        if not overrides:
+            return
+
+        unknown = set(overrides) - set(self._config)
+        if unknown:
+            raise ValueError(f"Unknown WikipediaSearchTool override(s): {sorted(unknown)}")
+
+        if "num_results" in overrides:
+            num_results = int(overrides["num_results"])
+            if num_results < 1 or num_results > 5:
+                raise ValueError("num_results must be between 1 and 5.")
+            self._config["num_results"] = num_results
+        if "max_chars" in overrides:
+            self._config["max_chars"] = max(200, min(int(overrides["max_chars"]), 1500))
+        if "max_sections" in overrides:
+            self._config["max_sections"] = max(1, min(int(overrides["max_sections"]), 50))
+        if "count" in overrides:
+            self._config["count"] = max(1, min(int(overrides["count"]), MAX_FACTS))
 
     async def list_tools(self) -> list[dict[str, Any]]:
         return [
@@ -696,11 +713,14 @@ async def list_tools(self) -> list[dict[str, Any]]:
             },
             {
                 "name": "wikipedia-query-summary",
-                "description": "Search for a query, then summarize the top result.",
+                "description": "Return a compact snippet around a query inside a Wikipedia article.",
                 "input_schema": {
                     "type": "object",
-                    "properties": {"query": {"type": "string", "description": "Search query."}},
-                    "required": ["query"],
+                    "properties": {
+                        "title": {"type": "string", "description": "Article title."},
+                        "query": {"type": "string", "description": "Term or phrase to locate within the article."},
+                    },
+                    "required": ["title", "query"],
                 },
             },
             {
@@ -717,24 +737,22 @@ async def list_tools(self) -> list[dict[str, Any]]:
     async def execute(self, tool_name: str, arguments: dict[str, Any], extra_args: dict[str, Any] | None = None):
         arguments = dict(arguments or {})
         if tool_name == "wikipedia-search":
-            arguments.setdefault("num_results", self._config.get("num_results", 3))
+            arguments.setdefault("num_results", self._config["num_results"])
             return await wikipedia_search(**arguments)
         if tool_name == "wikipedia-page":
-            arguments.setdefault("max_chars", self._config.get("max_chars", EXTRACT_LIMIT))
             return await wikipedia_page(**arguments)
         if tool_name == "wikipedia-summary":
-            arguments.setdefault("max_chars", self._config.get("max_chars", 900))
+            arguments.setdefault("max_chars", self._config["max_chars"])
             return await wikipedia_summary(**arguments)
         if tool_name == "wikipedia-sections":
-            arguments.setdefault("max_sections", self._config.get("max_sections", 40))
+            arguments.setdefault("max_sections", self._config["max_sections"])
             return await wikipedia_sections(**arguments)
         if tool_name == "wikipedia-section":
-            arguments.setdefault("max_chars", self._config.get("max_chars", EXTRACT_LIMIT))
             return await wikipedia_section(**arguments)
         if tool_name == "wikipedia-query-summary":
-            arguments.setdefault("max_chars", self._config.get("max_chars", 900))
+            arguments.setdefault("max_chars", self._config["max_chars"])
             return await wikipedia_query_summary(**arguments)
         if tool_name == "wikipedia-key-facts":
-            arguments.setdefault("count", self._config.get("count", 5))
+            arguments.setdefault("count", self._config["count"])
             return await wikipedia_key_facts(**arguments)
         return f"Error: unknown tool '{tool_name}'"

tests/test_mcp_clients.py

@@ -1043,3 +1043,36 @@ async def test_wikipedia_direct_list_tools(self):
         search_tool = next(t for t in tools if t["name"] == "wikipedia-search")
         assert "query" in search_tool["input_schema"]["properties"]
         assert "num_results" not in search_tool["input_schema"]["properties"]
+
+        query_summary_tool = next(t for t in tools if t["name"] == "wikipedia-query-summary")
+        assert {"title", "query"} <= set(query_summary_tool["input_schema"]["properties"])
+        assert set(query_summary_tool["input_schema"]["required"]) == {"title", "query"}
+
+    @pytest.mark.asyncio
+    async def test_wikipedia_execute_dispatch_contracts(self, monkeypatch):
+        from nemo_skills.mcp.servers import wikipedia_tool
+        from nemo_skills.mcp.servers.wikipedia_tool import WikipediaSearchTool
+
+        async def fake_page(title):
+            return f"page:{title}"
+
+        async def fake_section(title, section):
+            return f"section:{title}:{section}"
+
+        async def fake_query_summary(title, query, max_chars=700):
+            return f"query-summary:{title}:{query}:{max_chars}"
+
+        monkeypatch.setattr(wikipedia_tool, "wikipedia_page", fake_page)
+        monkeypatch.setattr(wikipedia_tool, "wikipedia_section", fake_section)
+        monkeypatch.setattr(wikipedia_tool, "wikipedia_query_summary", fake_query_summary)
+
+        tool = WikipediaSearchTool()
+        assert await tool.execute("wikipedia-page", {"title": "Hydrogen"}) == "page:Hydrogen"
+        assert (
+            await tool.execute("wikipedia-section", {"title": "Hydrogen", "section": "Isotopes"})
+            == "section:Hydrogen:Isotopes"
+        )
+        assert (
+            await tool.execute("wikipedia-query-summary", {"title": "Hydrogen", "query": "isotope"})
+            == "query-summary:Hydrogen:isotope:2500"
+        )