From 8115fe680238658b141955d8411523fe749fe99e Mon Sep 17 00:00:00 2001 From: Keyur Shah Date: Fri, 24 Jul 2026 00:19:04 -0700 Subject: [PATCH 1/3] Add agent_api_root and search_scope configuration options --- packages/datacommons-mcp/.env.sample | 5 +++++ .../datacommons_mcp/agent_api_client.py | 14 ++++++++++++- .../datacommons_mcp/agent_api_service.py | 2 ++ .../datacommons-mcp/datacommons_mcp/app.py | 19 ++++++++++++++++-- .../datacommons_mcp/data_models/settings.py | 13 ++++++++++++ .../datacommons_mcp/exceptions.py | 7 +++++++ .../datacommons-mcp/tests/test_agent_api.py | 16 +++++++++++++++ .../datacommons-mcp/tests/test_settings.py | 20 +++++++++++++++++++ 8 files changed, 93 insertions(+), 3 deletions(-) diff --git a/packages/datacommons-mcp/.env.sample b/packages/datacommons-mcp/.env.sample index 039df52a..02d5de17 100644 --- a/packages/datacommons-mcp/.env.sample +++ b/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 + diff --git a/packages/datacommons-mcp/datacommons_mcp/agent_api_client.py b/packages/datacommons-mcp/datacommons_mcp/agent_api_client.py index 70d6ae06..728eb7f7 100644 --- a/packages/datacommons-mcp/datacommons_mcp/agent_api_client.py +++ b/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__}", @@ -86,6 +93,7 @@ def __init__(self, api_root: str, api_key: str | None = None) -> None: self.timeout = 30.0 # 30 seconds default timeout self._client = None + @property def client(self) -> httpx.AsyncClient: """Lazily initialize the AsyncClient under the active event loop.""" @@ -114,6 +122,10 @@ 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.""" diff --git a/packages/datacommons-mcp/datacommons_mcp/agent_api_service.py b/packages/datacommons-mcp/datacommons_mcp/agent_api_service.py index 2098ec2a..64728701 100644 --- a/packages/datacommons-mcp/datacommons_mcp/agent_api_service.py +++ b/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) diff --git a/packages/datacommons-mcp/datacommons_mcp/app.py b/packages/datacommons-mcp/datacommons_mcp/app.py index fdea9a79..6a626a41 100644 --- a/packages/datacommons-mcp/datacommons_mcp/app.py +++ b/packages/datacommons-mcp/datacommons_mcp/app.py @@ -72,9 +72,24 @@ 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_root = ( + self.settings.agent_api_root + or 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) + search_scope = getattr(self.settings, "search_scope", None) + 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) diff --git a/packages/datacommons-mcp/datacommons_mcp/data_models/settings.py b/packages/datacommons-mcp/datacommons_mcp/data_models/settings.py index bd4e2d11..35f6a604 100644 --- a/packages/datacommons-mcp/datacommons_mcp/data_models/settings.py +++ b/packages/datacommons-mcp/datacommons_mcp/data_models/settings.py @@ -65,6 +65,19 @@ 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.""" diff --git a/packages/datacommons-mcp/datacommons_mcp/exceptions.py b/packages/datacommons-mcp/datacommons_mcp/exceptions.py index 039b9223..eaf229d7 100644 --- a/packages/datacommons-mcp/datacommons_mcp/exceptions.py +++ b/packages/datacommons-mcp/datacommons_mcp/exceptions.py @@ -54,3 +54,10 @@ 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: + return f"{base} - Body: {self.body}" + return base + diff --git a/packages/datacommons-mcp/tests/test_agent_api.py b/packages/datacommons-mcp/tests/test_agent_api.py index 071936f6..fa45e00e 100644 --- a/packages/datacommons-mcp/tests/test_agent_api.py +++ b/packages/datacommons-mcp/tests/test_agent_api.py @@ -63,6 +63,19 @@ 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 +191,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,10 +211,12 @@ async def test_agent_api_service_search_indicators(): "parent_place": "USA", "per_search_limit": 5, "include_topics": False, + "target": "custom_only", }, ) + @pytest.mark.asyncio async def test_agent_api_tools_execution(): """Verify agent_api_tools functions delegate to agent_api_service.""" diff --git a/packages/datacommons-mcp/tests/test_settings.py b/packages/datacommons-mcp/tests/test_settings.py index 893e450b..6f660eee 100644 --- a/packages/datacommons-mcp/tests/test_settings.py +++ b/packages/datacommons-mcp/tests/test_settings.py @@ -53,6 +53,23 @@ 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"} @@ -60,6 +77,9 @@ def test_default_dc_type_is_base(self): 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: From 4c95c4f8b7bf3d9bc4a49d7f9e34eb9af48050e2 Mon Sep 17 00:00:00 2001 From: Keyur Shah Date: Fri, 24 Jul 2026 00:21:30 -0700 Subject: [PATCH 2/3] lint --- packages/datacommons-mcp/datacommons_mcp/agent_api_client.py | 2 -- packages/datacommons-mcp/datacommons_mcp/app.py | 5 +---- .../datacommons-mcp/datacommons_mcp/data_models/settings.py | 1 - packages/datacommons-mcp/datacommons_mcp/exceptions.py | 1 - packages/datacommons-mcp/tests/test_agent_api.py | 2 -- packages/datacommons-mcp/tests/test_settings.py | 4 +--- 6 files changed, 2 insertions(+), 13 deletions(-) diff --git a/packages/datacommons-mcp/datacommons_mcp/agent_api_client.py b/packages/datacommons-mcp/datacommons_mcp/agent_api_client.py index 728eb7f7..6d316f04 100644 --- a/packages/datacommons-mcp/datacommons_mcp/agent_api_client.py +++ b/packages/datacommons-mcp/datacommons_mcp/agent_api_client.py @@ -93,7 +93,6 @@ def __init__( self.timeout = 30.0 # 30 seconds default timeout self._client = None - @property def client(self) -> httpx.AsyncClient: """Lazily initialize the AsyncClient under the active event loop.""" @@ -126,7 +125,6 @@ async def post(self, endpoint: str, payload: dict) -> dict: 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.""" if self._client is not None: diff --git a/packages/datacommons-mcp/datacommons_mcp/app.py b/packages/datacommons-mcp/datacommons_mcp/app.py index 6a626a41..e9b46451 100644 --- a/packages/datacommons-mcp/datacommons_mcp/app.py +++ b/packages/datacommons-mcp/datacommons_mcp/app.py @@ -80,9 +80,7 @@ def __init__(self) -> None: api_key = getattr(self.settings, "api_key", None) search_scope = getattr(self.settings, "search_scope", None) search_scope_str = ( - search_scope.value - if hasattr(search_scope, "value") - else search_scope + search_scope.value if hasattr(search_scope, "value") else search_scope ) self.agent_api_client = AgentAPIClient( api_root=api_root, @@ -90,7 +88,6 @@ def __init__(self) -> None: search_scope=search_scope_str, ) - # Load Server Instructions server_instructions = self._load_instructions(SERVER_INSTRUCTIONS_FILE) diff --git a/packages/datacommons-mcp/datacommons_mcp/data_models/settings.py b/packages/datacommons-mcp/datacommons_mcp/data_models/settings.py index 35f6a604..988d5986 100644 --- a/packages/datacommons-mcp/datacommons_mcp/data_models/settings.py +++ b/packages/datacommons-mcp/datacommons_mcp/data_models/settings.py @@ -78,7 +78,6 @@ class DCSettingsBase(BaseSettings): ) - class BaseDCSettings(DCSettingsBase): """Settings for base Data Commons instance.""" diff --git a/packages/datacommons-mcp/datacommons_mcp/exceptions.py b/packages/datacommons-mcp/datacommons_mcp/exceptions.py index eaf229d7..d7374690 100644 --- a/packages/datacommons-mcp/datacommons_mcp/exceptions.py +++ b/packages/datacommons-mcp/datacommons_mcp/exceptions.py @@ -60,4 +60,3 @@ def __str__(self) -> str: if self.body: return f"{base} - Body: {self.body}" return base - diff --git a/packages/datacommons-mcp/tests/test_agent_api.py b/packages/datacommons-mcp/tests/test_agent_api.py index fa45e00e..1a9e2119 100644 --- a/packages/datacommons-mcp/tests/test_agent_api.py +++ b/packages/datacommons-mcp/tests/test_agent_api.py @@ -75,7 +75,6 @@ async def test_agent_api_client_search_scope(): 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.""" @@ -216,7 +215,6 @@ async def test_agent_api_service_search_indicators(): ) - @pytest.mark.asyncio async def test_agent_api_tools_execution(): """Verify agent_api_tools functions delegate to agent_api_service.""" diff --git a/packages/datacommons-mcp/tests/test_settings.py b/packages/datacommons-mcp/tests/test_settings.py index 6f660eee..8dc3f866 100644 --- a/packages/datacommons-mcp/tests/test_settings.py +++ b/packages/datacommons-mcp/tests/test_settings.py @@ -65,8 +65,7 @@ def test_loads_with_agent_api_root_and_search_scope(self): settings = get_dc_settings() assert isinstance(settings, BaseDCSettings) assert ( - settings.agent_api_root - == "https://custom-agent-api.datacommons.org/v2" + settings.agent_api_root == "https://custom-agent-api.datacommons.org/v2" ) assert settings.search_scope == SearchScope.CUSTOM_ONLY @@ -81,7 +80,6 @@ def test_default_dc_type_is_base(self): assert settings.search_scope is None - class TestCustomSettings: """Test suite for loading CustomDCSettings.""" From 6734fd64bad779c6f21f8b020d7d6b389c935f4e Mon Sep 17 00:00:00 2001 From: Keyur Shah Date: Fri, 24 Jul 2026 00:28:58 -0700 Subject: [PATCH 3/3] Address gemini review --- packages/datacommons-mcp/datacommons_mcp/app.py | 4 ++-- packages/datacommons-mcp/datacommons_mcp/exceptions.py | 7 ++++++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/packages/datacommons-mcp/datacommons_mcp/app.py b/packages/datacommons-mcp/datacommons_mcp/app.py index e9b46451..51fb42af 100644 --- a/packages/datacommons-mcp/datacommons_mcp/app.py +++ b/packages/datacommons-mcp/datacommons_mcp/app.py @@ -77,8 +77,8 @@ def __init__(self) -> None: or self.settings.api_root or "https://api.datacommons.org/v2" ) - api_key = getattr(self.settings, "api_key", None) - search_scope = getattr(self.settings, "search_scope", None) + 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 ) diff --git a/packages/datacommons-mcp/datacommons_mcp/exceptions.py b/packages/datacommons-mcp/datacommons_mcp/exceptions.py index d7374690..d40c5b13 100644 --- a/packages/datacommons-mcp/datacommons_mcp/exceptions.py +++ b/packages/datacommons-mcp/datacommons_mcp/exceptions.py @@ -58,5 +58,10 @@ def __init__(self, message: str, status_code: int, body: str | None = None) -> N def __str__(self) -> str: base = super().__str__() if self.body: - return f"{base} - Body: {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