Add agent_api_root and search_scope configuration options

packages/datacommons-mcp/.env.sample

@@ -75,6 +75,11 @@ DC_TYPE=base
 # Root URL for a non-prod or local Data Commons API (mixer) instance
 # DC_API_ROOT=http://localhost:8081/v2
 
+# Root URL specifically for Data Commons Agent API endpoints (optional)
+# Takes precedence over DC_API_ROOT when DC_USE_AGENT_API=true
+# DC_AGENT_API_ROOT=https://api.datacommons.org/v2
+
 # Root URL for Data Commons API key validation
 # Configure for non-prod environments
 # DC_API_KEY_VALIDATION_ROOT=https://api.datacommons.org
+

packages/datacommons-mcp/datacommons_mcp/agent_api_client.py

@@ -68,15 +68,22 @@ async def wrapper(
 class AgentAPIClient:
     """Async client for interacting with Data Commons agent endpoints."""
 
-    def __init__(self, api_root: str, api_key: str | None = None) -> None:
+    def __init__(
+        self,
+        api_root: str,
+        api_key: str | None = None,
+        search_scope: str | None = None,
+    ) -> None:
         """Initialize the AgentAPIClient.
 
         Args:
             api_root: The base API root URL (e.g. 'https://api.datacommons.org/v2').
             api_key: Optional API key for authentication.
+            search_scope: Optional search scope / target (e.g. 'custom_only', 'base_and_custom').
         """
         self.api_root = api_root.rstrip("/")
         self.api_key = api_key
+        self.search_scope = search_scope
         self.headers = {
             "Content-Type": "application/json",
             "x-surface": f"mcp-{__version__}",
@@ -114,6 +121,9 @@ async def post(self, endpoint: str, payload: dict) -> dict:
             raise AgentAPIError(
                 error_msg, e.response.status_code, e.response.text
             ) from e
+        except httpx.RequestError as e:
+            error_msg = f"Agent API request to {endpoint} failed: {e}"
+            raise AgentAPIError(error_msg, 503, None) from e
 
     async def close(self) -> None:
         """Close the underlying HTTP client if it was initialized."""

packages/datacommons-mcp/datacommons_mcp/agent_api_service.py

@@ -126,7 +126,9 @@ async def search_indicators(
         "parent_place": parent_place,
         "per_search_limit": per_search_limit,
         "include_topics": include_topics,
+        "target": client.search_scope,
     }
+
     return await client.post("agent/search_indicators", payload)
 
 

packages/datacommons-mcp/datacommons_mcp/app.py

@@ -72,9 +72,21 @@ def __init__(self) -> None:
         # Create agent API client only if enabled
         self.agent_api_client = None
         if self.settings.use_agent_api:
-            api_root = self.settings.api_root or "https://api.datacommons.org/v2"
-            api_key = getattr(self.settings, "api_key", None)
-            self.agent_api_client = AgentAPIClient(api_root=api_root, api_key=api_key)
+            api_root = (
+                self.settings.agent_api_root
+                or self.settings.api_root
+                or "https://api.datacommons.org/v2"
+            )
+            api_key = self.settings.api_key
+            search_scope = self.settings.search_scope
+            search_scope_str = (
+                search_scope.value if hasattr(search_scope, "value") else search_scope
+            )
+            self.agent_api_client = AgentAPIClient(
+                api_root=api_root,
+                api_key=api_key,
+                search_scope=search_scope_str,
+            )
 
         # Load Server Instructions
         server_instructions = self._load_instructions(SERVER_INSTRUCTIONS_FILE)

packages/datacommons-mcp/datacommons_mcp/data_models/settings.py

@@ -65,6 +65,18 @@ class DCSettingsBase(BaseSettings):
         description="API root for Data Commons",
     )
 
+    agent_api_root: str | None = Field(
+        default=None,
+        alias="DC_AGENT_API_ROOT",
+        description="API root specifically for Data Commons Agent API endpoints",
+    )
+
+    search_scope: SearchScope | None = Field(
+        default=None,
+        alias="DC_SEARCH_SCOPE",
+        description="Search scope for queries",
+    )
+
 
 class BaseDCSettings(DCSettingsBase):
     """Settings for base Data Commons instance."""

packages/datacommons-mcp/datacommons_mcp/exceptions.py

@@ -54,3 +54,14 @@ def __init__(self, message: str, status_code: int, body: str | None = None) -> N
         super().__init__(message)
         self.status_code = status_code
         self.body = body
+
+    def __str__(self) -> str:
+        base = super().__str__()
+        if self.body:
+            truncated_body = (
+                self.body
+                if len(self.body) <= 500
+                else f"{self.body[:500]}... [truncated]"
+            )
+            return f"{base} - Body: {truncated_body}"
+        return base

packages/datacommons-mcp/tests/test_agent_api.py

@@ -63,6 +63,18 @@ async def test_agent_api_client_post():
     await client.close()
 
 
+@pytest.mark.asyncio
+async def test_agent_api_client_search_scope():
+    """Verify AgentAPIClient stores search_scope attribute."""
+    client = AgentAPIClient(
+        api_root="https://api.datacommons.org/v2",
+        api_key="test-key",
+        search_scope="custom_only",
+    )
+    assert client.search_scope == "custom_only"
+    await client.close()
+
+
 @pytest.mark.asyncio
 async def test_agent_api_service_get_observations():
     """Verify get_observations builds correct payload and invokes agent_api_client."""
@@ -178,6 +190,7 @@ async def test_agent_api_service_search_indicators():
     from datacommons_mcp.app import app
 
     mock_client = AsyncMock()
+    mock_client.search_scope = "custom_only"
     mock_client.post.return_value = {"variables": []}
 
     with patch.object(app, "agent_api_client", mock_client):
@@ -197,6 +210,7 @@ async def test_agent_api_service_search_indicators():
                 "parent_place": "USA",
                 "per_search_limit": 5,
                 "include_topics": False,
+                "target": "custom_only",
             },
         )
 

packages/datacommons-mcp/tests/test_settings.py

@@ -53,13 +53,31 @@ def test_loads_with_env_var_overrides(self):
                 "/path/to/cache2.json",
             ]
 
+    def test_loads_with_agent_api_root_and_search_scope(self):
+        """Tests loading agent_api_root and search_scope from environment variables in BaseDCSettings."""
+        env_vars = {
+            "DC_API_KEY": "test_key",
+            "DC_TYPE": "base",
+            "DC_AGENT_API_ROOT": "https://custom-agent-api.datacommons.org/v2",
+            "DC_SEARCH_SCOPE": "custom_only",
+        }
+        with patch.dict(os.environ, env_vars):
+            settings = get_dc_settings()
+            assert isinstance(settings, BaseDCSettings)
+            assert (
+                settings.agent_api_root == "https://custom-agent-api.datacommons.org/v2"
+            )
+            assert settings.search_scope == SearchScope.CUSTOM_ONLY
+
     def test_default_dc_type_is_base(self):
         """Tests that DC_TYPE defaults to 'base' when not provided."""
         env_vars = {"DC_API_KEY": "test_key"}
         with patch.dict(os.environ, env_vars):
             settings = get_dc_settings()
             assert isinstance(settings, BaseDCSettings)
             assert settings.dc_type == "base"
+            assert settings.agent_api_root is None
+            assert settings.search_scope is None
 
 
 class TestCustomSettings: