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..6d316f04 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__}", @@ -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.""" 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..51fb42af 100644 --- a/packages/datacommons-mcp/datacommons_mcp/app.py +++ b/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) diff --git a/packages/datacommons-mcp/datacommons_mcp/data_models/settings.py b/packages/datacommons-mcp/datacommons_mcp/data_models/settings.py index bd4e2d11..988d5986 100644 --- a/packages/datacommons-mcp/datacommons_mcp/data_models/settings.py +++ b/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.""" diff --git a/packages/datacommons-mcp/datacommons_mcp/exceptions.py b/packages/datacommons-mcp/datacommons_mcp/exceptions.py index 039b9223..d40c5b13 100644 --- a/packages/datacommons-mcp/datacommons_mcp/exceptions.py +++ b/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 diff --git a/packages/datacommons-mcp/tests/test_agent_api.py b/packages/datacommons-mcp/tests/test_agent_api.py index 071936f6..1a9e2119 100644 --- a/packages/datacommons-mcp/tests/test_agent_api.py +++ b/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", }, ) diff --git a/packages/datacommons-mcp/tests/test_settings.py b/packages/datacommons-mcp/tests/test_settings.py index 893e450b..8dc3f866 100644 --- a/packages/datacommons-mcp/tests/test_settings.py +++ b/packages/datacommons-mcp/tests/test_settings.py @@ -53,6 +53,22 @@ 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 +76,8 @@ 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: