Skip to content

Commit be8d923

Browse files
keyurvakeyurva
authored
Route MCP tool calls to mixer agent apis (#201)
1 parent 8a6b58a commit be8d923

12 files changed

Lines changed: 780 additions & 20 deletions

File tree

‎packages/datacommons-mcp/.env.sample‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,10 @@ DC_TYPE=base
5757
# - tools/{tool_name}.md
5858
# DC_INSTRUCTIONS_DIR=/path/to/custom/instructions
5959

60+
# Enable using Agent-optimized APIs instead of local processing (optional)
61+
# Default: false
62+
# DC_USE_AGENT_API=false
63+
6064

6165
# =============================================================================
6266
# NON-PROD ROOTS (optional, base DC only)

‎packages/datacommons-mcp/datacommons_mcp/agent_api_client.py‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
# Copyright 2026 Google LLC.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
"""
15+
Client for interacting with the Data Commons Agent API.
16+
"""
17+
18+
import logging
19+
import time
20+
from collections.abc import Callable
21+
from functools import wraps
22+
from typing import Any # noqa: ANN401
23+
24+
import httpx
25+
26+
from datacommons_mcp.exceptions import AgentAPIError
27+
from datacommons_mcp.version import __version__
28+
29+
logger = logging.getLogger(__name__)
30+
31+
32+
def log_api_call(func: Callable[..., Any]) -> Callable[..., Any]: # noqa: ANN401
33+
"""Decorator to log URL, request payload, execution time, and errors for Agent API calls."""
34+
35+
@wraps(func)
36+
async def wrapper(
37+
self: "AgentAPIClient",
38+
endpoint: str,
39+
payload: dict,
40+
*args: object,
41+
**kwargs: object,
42+
) -> Any: # noqa: ANN401
43+
url = f"{self.api_root}/{endpoint}"
44+
logger.info("AgentAPIClient POST request URL: %s, payload: %s", url, payload)
45+
start_time = time.perf_counter()
46+
try:
47+
result = await func(self, endpoint, payload, *args, **kwargs)
48+
elapsed_time = time.perf_counter() - start_time
49+
logger.info(
50+
"AgentAPIClient POST request to %s completed in %.3f seconds",
51+
url,
52+
elapsed_time,
53+
)
54+
return result
55+
except Exception as e:
56+
elapsed_time = time.perf_counter() - start_time
57+
logger.error(
58+
"AgentAPIClient POST request to %s failed after %.3f seconds with error: %s",
59+
url,
60+
elapsed_time,
61+
e,
62+
)
63+
raise
64+
65+
return wrapper
66+
67+
68+
class AgentAPIClient:
69+
"""Async client for interacting with Data Commons agent endpoints."""
70+
71+
def __init__(self, api_root: str, api_key: str | None = None) -> None:
72+
"""Initialize the AgentAPIClient.
73+
74+
Args:
75+
api_root: The base API root URL (e.g. 'https://api.datacommons.org/v2').
76+
api_key: Optional API key for authentication.
77+
"""
78+
self.api_root = api_root.rstrip("/")
79+
self.api_key = api_key
80+
self.headers = {
81+
"Content-Type": "application/json",
82+
"x-surface": f"mcp-{__version__}",
83+
}
84+
if api_key:
85+
self.headers["X-API-Key"] = api_key
86+
self.timeout = 30.0 # 30 seconds default timeout
87+
self._client = None
88+
89+
@property
90+
def client(self) -> httpx.AsyncClient:
91+
"""Lazily initialize the AsyncClient under the active event loop."""
92+
if self._client is None:
93+
self._client = httpx.AsyncClient(headers=self.headers, timeout=self.timeout)
94+
return self._client
95+
96+
@log_api_call
97+
async def post(self, endpoint: str, payload: dict) -> dict:
98+
"""Perform an asynchronous POST request to the specified endpoint.
99+
100+
Args:
101+
endpoint: The API endpoint (e.g. 'agent/get_observations').
102+
payload: The dictionary to send as JSON payload.
103+
104+
Returns:
105+
The parsed JSON response as a dictionary.
106+
"""
107+
url = f"{self.api_root}/{endpoint}"
108+
try:
109+
response = await self.client.post(url, json=payload)
110+
response.raise_for_status()
111+
return response.json()
112+
except httpx.HTTPStatusError as e:
113+
error_msg = f"Agent API call to {endpoint} failed with status {e.response.status_code}"
114+
raise AgentAPIError(
115+
error_msg, e.response.status_code, e.response.text
116+
) from e
117+
118+
async def close(self) -> None:
119+
"""Close the underlying HTTP client if it was initialized."""
120+
if self._client is not None:
121+
await self._client.aclose()
122+
self._client = None

‎packages/datacommons-mcp/datacommons_mcp/agent_api_service.py‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
# Copyright 2026 Google LLC.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
"""
15+
Service layer for calling Agent APIs.
16+
"""
17+
18+
from typing import Any
19+
20+
from datacommons_mcp.agent_api_client import AgentAPIClient
21+
from datacommons_mcp.app import app
22+
23+
24+
def _get_agent_api_client() -> AgentAPIClient:
25+
"""Helper to get the initialized AgentAPIClient, raising RuntimeError if not set."""
26+
if app.agent_api_client is None:
27+
raise RuntimeError(
28+
"Agent API client is not initialized. Ensure DC_USE_AGENT_API is enabled."
29+
)
30+
return app.agent_api_client
31+
32+
33+
async def get_observations(
34+
variable_dcid: str,
35+
place_dcid: str,
36+
child_place_type: str | None = None,
37+
source_override: str | None = None,
38+
date: str | None = None,
39+
date_range_start: str | None = None,
40+
date_range_end: str | None = None,
41+
) -> dict[str, Any]:
42+
"""Fetches observations via the Agent API agent/get_observations endpoint."""
43+
client = _get_agent_api_client()
44+
payload = {
45+
"variable_dcid": variable_dcid,
46+
"place_dcid": place_dcid,
47+
"child_place_type": child_place_type,
48+
"source_override": source_override,
49+
"date": date,
50+
"date_range_start": date_range_start,
51+
"date_range_end": date_range_end,
52+
}
53+
return await client.post("agent/get_observations", payload)
54+
55+
56+
async def search_indicators(
57+
query: str,
58+
places: list[str] | None = None,
59+
parent_place: str | None = None,
60+
per_search_limit: int = 10,
61+
*,
62+
include_topics: bool = True,
63+
) -> dict[str, Any]:
64+
"""Searches for indicators via the Agent API agent/search_indicators endpoint."""
65+
client = _get_agent_api_client()
66+
payload = {
67+
"query": query,
68+
"places": places or [],
69+
"parent_place": parent_place,
70+
"per_search_limit": per_search_limit,
71+
"include_topics": include_topics,
72+
}
73+
return await client.post("agent/search_indicators", payload)

‎packages/datacommons-mcp/datacommons_mcp/agent_api_tools.py‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
# Copyright 2026 Google LLC.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
"""
15+
Tool implementations for the Agent API-based Data Commons MCP server.
16+
"""
17+
18+
from datacommons_mcp.agent_api_service import (
19+
get_observations as agent_api_get_observations,
20+
)
21+
from datacommons_mcp.agent_api_service import (
22+
search_indicators as agent_api_search_indicators,
23+
)
24+
from datacommons_mcp.data_models.observations import ObservationDateType
25+
26+
27+
async def get_observations(
28+
variable_dcid: str,
29+
place_dcid: str,
30+
child_place_type: str | None = None,
31+
source_override: str | None = None,
32+
date: str = ObservationDateType.LATEST.value,
33+
date_range_start: str | None = None,
34+
date_range_end: str | None = None,
35+
) -> dict:
36+
"""Fetches observations for a statistical variable from Data Commons."""
37+
return await agent_api_get_observations(
38+
variable_dcid=variable_dcid,
39+
place_dcid=place_dcid,
40+
child_place_type=child_place_type,
41+
source_override=source_override,
42+
date=date,
43+
date_range_start=date_range_start,
44+
date_range_end=date_range_end,
45+
)
46+
47+
48+
async def search_indicators(
49+
query: str,
50+
places: list[str] | None = None,
51+
parent_place: str | None = None,
52+
per_search_limit: int = 10,
53+
*,
54+
include_topics: bool = True,
55+
) -> dict:
56+
"""Searches for indicators (topics and variables) in Data Commons."""
57+
return await agent_api_search_indicators(
58+
query=query,
59+
places=places,
60+
parent_place=parent_place,
61+
per_search_limit=per_search_limit,
62+
include_topics=include_topics,
63+
)

‎packages/datacommons-mcp/datacommons_mcp/app.py‎

Lines changed: 27 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,15 @@
1717

1818
import json
1919
import logging
20-
from collections.abc import Callable
20+
from collections.abc import AsyncIterator, Callable
2121
from typing import Any
2222

2323
from fastmcp import FastMCP
2424
from fastmcp.tools.tool import Tool
2525
from pydantic import ValidationError
2626

2727
from datacommons_mcp import settings
28+
from datacommons_mcp.agent_api_client import AgentAPIClient
2829
from datacommons_mcp.clients import create_dc_client
2930
from datacommons_mcp.utils import read_external_content, read_package_content
3031
from datacommons_mcp.version import __version__
@@ -54,20 +55,39 @@ def __init__(self) -> None:
5455
logger.error("Settings error: %s", e)
5556
raise
5657

57-
# Create client
58-
try:
59-
self.client = create_dc_client(self.settings)
60-
except Exception as e:
61-
logger.error("Failed to create DC client: %s", e)
62-
raise
58+
# Create client only if agent APIs are NOT enabled (as fallback is not needed)
59+
self.client = None
60+
if not self.settings.use_agent_api:
61+
try:
62+
self.client = create_dc_client(self.settings)
63+
except Exception as e:
64+
logger.error("Failed to create DC client: %s", e)
65+
raise
66+
67+
# Create agent API client only if enabled
68+
self.agent_api_client = None
69+
if self.settings.use_agent_api:
70+
api_root = self.settings.api_root or "https://api.datacommons.org/v2"
71+
api_key = getattr(self.settings, "api_key", None)
72+
self.agent_api_client = AgentAPIClient(api_root=api_root, api_key=api_key)
6373

6474
# Load Server Instructions
6575
server_instructions = self._load_instructions(SERVER_INSTRUCTIONS_FILE)
6676

77+
from contextlib import asynccontextmanager
78+
79+
@asynccontextmanager
80+
async def lifespan(_server: FastMCP) -> AsyncIterator[dict[str, Any]]:
81+
yield {}
82+
if self.agent_api_client:
83+
logger.info("Closing Agent API client...")
84+
await self.agent_api_client.close()
85+
6786
self.mcp = FastMCP(
6887
MCP_SERVER_NAME,
6988
version=__version__,
7089
instructions=server_instructions,
90+
lifespan=lifespan,
7191
)
7292

7393
def _load_instructions(self, filename: str) -> str:

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

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,8 @@ class DCSettingsSelector(BaseSettings):
3535
)
3636

3737

38-
class DCSettings(BaseSettings):
39-
"""Settings for base Data Commons instance."""
38+
class DCSettingsBase(BaseSettings):
39+
"""Base settings class containing shared configurations."""
4040

4141
model_config = _MODEL_CONFIG
4242

@@ -53,8 +53,20 @@ class DCSettings(BaseSettings):
5353
description="Directory containing custom instruction files (markdown overrides)",
5454
)
5555

56+
use_agent_api: bool = Field(
57+
default=False,
58+
alias="DC_USE_AGENT_API",
59+
description="Use the new Agent-optimized APIs instead of local processing",
60+
)
61+
62+
api_root: str | None = Field(
63+
default=None,
64+
alias="DC_API_ROOT",
65+
description="API root for Data Commons",
66+
)
67+
5668

57-
class BaseDCSettings(DCSettings):
69+
class BaseDCSettings(DCSettingsBase):
5870
"""Settings for base Data Commons instance."""
5971

6072
def __init__(self, **kwargs: dict[str, Any]) -> None:
@@ -76,19 +88,14 @@ def __init__(self, **kwargs: dict[str, Any]) -> None:
7688
alias="DC_BASE_ROOT_TOPIC_DCIDS",
7789
description="List of root topic DCIDs for base DC",
7890
)
79-
api_root: str | None = Field(
80-
default=None,
81-
alias="DC_API_ROOT",
82-
description="API root for local api instance",
83-
)
8491

8592
@field_validator("topic_cache_paths", "base_root_topic_dcids", mode="before")
8693
@classmethod
8794
def parse_list_like_parameter(cls, v: str) -> list[str] | None:
8895
return _parse_list_like_parameter(v)
8996

9097

91-
class CustomDCSettings(DCSettings):
98+
class CustomDCSettings(DCSettingsBase):
9299
"""Settings for custom Data Commons instance."""
93100

94101
model_config = _MODEL_CONFIG

0 commit comments

Comments
 (0)