chore: Added Tavily search tools.

Author: anirbanbasu

README.md

@@ -42,8 +42,9 @@ There are Dapr related configuration files too.
 
 The following API keys are optional but maybe provided for additional functionality. If not provided, the corresponding functionality will not be available.
 
- - `OLLAMA_API_KEY`: The Ollama API key is required for MCP-based web search, web fetch and cloud hosted models on Ollama. MCP features include 2 tools. Create one from [your Ollama account](https://ollama.com/settings/keys).
+ - `OLLAMA_API_KEY`: The Ollama API key is required for MCP-based web search, web fetch and cloud hosted models on Ollama. MCP features include 2 tools. Create an API key from [your Ollama account](https://ollama.com/settings/keys).
  - `ALPHAVANTAGE_API_KEY`: The API key to MCP-based Alpha Vantage finance related functions. MCP features include 118 tools. Obtain your [free API key from Alpha Vantage](https://www.alphavantage.co/support/#api-key). Note that basic finance tools are available through the [yfinance MCP](https://github.com/narumiruna/yfinance-mcp) even if Alpha Vantage MCP is not loaded.
+ - `TAVILY_API_KEY`: The API key for Tavily search and research services. MCP features include 4 tools. Get your API key [from your Tavily account](https://app.tavily.com/home).
 
 The following environment variables are all optional.
  - `APP_LOG_LEVEL`: The general log level of the DQA app. Defaults to `INFO`.

src/dqa/__init__.py

@@ -51,6 +51,7 @@ class ParsedEnvVars:
     )
 
     API_KEY_ALPHAVANTAGE: str = env.str("ALPHAVANTAGE_API_KEY", default=None)
+    API_KEY_TAVILY: str = env.str("TAVILY_API_KEY", default=None)
     API_KEY_OLLAMA: str = env.str("OLLAMA_API_KEY", default=None)
 
     _instance: ClassVar = None

src/dqa/mcp/primary.py

@@ -48,6 +48,23 @@ def server():
             "No AlphaVantage API key found in environment variable 'API_KEY_ALPHAVANTAGE'. Skipping mounting AlphaVantage remote MCP."
         )
 
+    if ParsedEnvVars().API_KEY_TAVILY:
+        # See: https://github.com/alphavantage/alpha_vantage_mcp
+        tavily_remote_url = (
+            f"https://mcp.tavily.com/mcp/?tavilyApiKey={ParsedEnvVars().API_KEY_TAVILY}"
+        )
+        tavily_remote_proxy = FastMCP.as_proxy(Client(tavily_remote_url))
+        # Prefix is not necessary because Tavily tools are prefixed with "tavily_" already.
+        server.mount(tavily_remote_proxy)
+        logger.info("Mounted Tavily remote MCP.")
+        logger.warning(
+            "Tavily MCP is a rate-limited remote service, expect higher latency."
+        )
+    else:
+        logger.warning(
+            "No AlphaVantage API key found in environment variable 'TAVILY_API_KEY'. Skipping mounting Tavily remote MCP."
+        )
+
     if ParsedEnvVars().API_KEY_OLLAMA:
         server.mount(ollama_mcp(), prefix="ollama")
         logger.info("Mounted Ollama MCP.")

tests/test_primary_mcp.py

@@ -4,6 +4,8 @@
 from dqa.mcp.primary import server as primary_mcp
 from tests.mcp_test_mixin import MCPTestMixin
 
+from itertools import combinations
+
 
 class TestPrimaryMCP(MCPTestMixin):
     # No need to test each MCP exhaustively here, they have their own tests.
@@ -34,17 +36,28 @@ def test_composition(self, mcp_client):
         """
         Test that the primary MCP server has composed the necessary MCPs correctly.
         """
+
+        def possible_tool_counts(base: int, additions: list[int]) -> set[int]:
+            """
+            Given a base count and a list of possible additions, return all possible total counts.
+            """
+
+            counts = {base}
+            for r in range(1, len(additions) + 1):
+                for combo in combinations(additions, r):
+                    counts.add(base + sum(combo))
+            return counts
+
         n_base_tools = 22  # Base tools from local MCPs
         n_ollama_tools = 2
         n_alphavantage_tools = 118
+        n_tavily_tools = 4
         tools = asyncio.run(self.list_tools(mcp_client))
-        assert len(tools) in [
-            # Base
+        assert len(tools) in possible_tool_counts(
             n_base_tools,
-            # With Ollama only
-            n_base_tools + n_ollama_tools,
-            # With AlphaVantage only
-            n_base_tools + n_alphavantage_tools,
-            # With both Ollama and AlphaVantage
-            n_base_tools + n_ollama_tools + n_alphavantage_tools,
-        ]
+            [
+                n_ollama_tools,
+                n_alphavantage_tools,
+                n_tavily_tools,
+            ],
+        )