Data Commons MCP Server

# Copyright 2025 Google LLC. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import calendar import re from datetime import datetime from enum import Enum from functools import lru_cache from typing import Any from datacommons_client.endpoints.response import ObservationResponse from datacommons_client.models.observation import Observation, OrderedFacet from dateutil.parser import parse from pydantic import BaseModel, Field, field_validator from datacommons_mcp.exceptions import ( InvalidDateFormatError, InvalidDateRangeError, ) # Wrapper to rename datacommons_client ObservationResponse to avoid confusion. ObservationApiResponse = ObservationResponse STANDARDIZED_DATE_FORMAT = "%Y-%m-%d" DEFAULT_DATE = datetime(1970, 1, 1) # UTC start date class ObservationDateType(str, Enum): """Enumeration for special date strings in observation queries.""" ALL = "all" LATEST = "latest" RANGE = "range" # Regex to validate that a string is in YYYY, YYYY-MM, or YYYY-MM-DD format. DATE_FORMAT_REGEX = re.compile(r"^\d{4}(-\d{2})?(-\d{2})?$") class ObservationDate(BaseModel): date: str @field_validator("date") @classmethod def validate_date_format(cls, v: str) -> str: """Validates that the date is a known constant or a valid date format.""" if v.lower() in [member.value for member in ObservationDateType]: return v.lower() if not DATE_FORMAT_REGEX.match(v): raise InvalidDateFormatError( f"Date string '{v}' is not one of the valid constants nor in YYYY, YYYY-MM, or YYYY-MM-DD format." ) try: # After regex validation, parse to catch invalid values like '2023-99'. ObservationDate.parse_date(v) return v except ValueError as e: # This will catch errors from dateutil.parser for invalid dates. raise InvalidDateFormatError( f"Date string '{v}' contains an invalid value" ) from e @staticmethod def parse_date(date_str: str) -> datetime: try: return parse(date_str, default=DEFAULT_DATE) except ValueError as e: raise InvalidDateFormatError(f"for date '{date_str}': {e}") from e class DateRange(BaseModel): "Accepted formats: YYYY or YYYY-MM or YYYY-MM-DD" start_date: datetime | None = None end_date: datetime | None = None def __init__( self, start_date: str | None = None, end_date: str | None = None, **kwargs: dict[str, Any], ) -> None: """Initializes and validates the date range from string inputs.""" super().__init__(**kwargs) range_start = DateRange.get_start_date(start_date) if start_date else None range_end = DateRange.get_end_date(end_date) if end_date else None if range_start and range_end and range_start > range_end: raise InvalidDateRangeError( f"start_date '{start_date}' cannot be after end_date '{end_date}'" ) self.start_date = range_start self.end_date = range_end @property def start_date_str(self) -> str | None: """Returns the start_date as a standardized 'YYYY-MM-DD' string, or None.""" if not self.start_date: return None return self.start_date.strftime(STANDARDIZED_DATE_FORMAT) @property def end_date_str(self) -> str | None: """Returns the end_date as a standardized 'YYYY-MM-DD' string, or None.""" if not self.end_date: return None return self.end_date.strftime(STANDARDIZED_DATE_FORMAT) @staticmethod def get_start_date(date_str: str) -> datetime: """ Converts a partial date string into a full start date. Caches results to avoid re-calculating for the same input string. Examples: >>> DateRange.parse_interval("2022") ('2022-01-01') >>> DateRange.parse_interval("2023-05") ('2023-05-01') >>> DateRange.parse_interval("2024-01-15") ('2024-01-15') Raises: InvalidDateFormatError: If the date string format is invalid. """ return ObservationDate.parse_date(date_str) @staticmethod def get_end_date(date_str: str) -> datetime: """ Converts a partial date string into a full (start, end) date tuple. Caches results to avoid re-calculating for the same input string. Examples: >>> DateRange.parse_interval("2022") ('2022-12-31') >>> DateRange.parse_interval("2023-05") ('2023-05-31') >>> DateRange.parse_interval("2024-01-15") ('2024-01-15') Raises: InvalidDateFormatError: If the date string format is invalid. """ try: parts = date_str.split("-") num_parts = len(parts) if num_parts == 1: year = int(parts[0]) return datetime(year=year, month=12, day=31) if num_parts == 2: year, month = map(int, parts) # This will raise ValueError for an invalid month. datetime(year=year, month=month, day=1) _, last_day = calendar.monthrange(year, month) return datetime(year=year, month=month, day=last_day) if num_parts == 3: year, month, day = map(int, parts) # This will raise ValueError for an invalid year, month, or day. return datetime(year=year, month=month, day=day) # If we reach here, the number of parts is not 1, 2, or 3. raise ValueError( "Date string must be in YYYY, YYYY-MM, or YYYY-MM-DD format." ) except ValueError as e: # Catch multiple potential errors and raise a single, clear custom exception. raise InvalidDateFormatError(f"for date '{date_str}': {e}") from e @staticmethod @lru_cache(maxsize=128) def parse_interval(date_str: str) -> tuple[datetime, datetime]: """ Converts a partial date string into a full (start, end) date tuple. Caches results to avoid re-calculating for the same input string. Examples: >>> DateRange.parse_interval("2022") ('2022-01-01', '2022-12-31') >>> DateRange.parse_interval("2023-05") ('2023-05-01', '2023-05-31') >>> DateRange.parse_interval("2024-01-15") ('2024-01-15', '2024-01-15') Raises: InvalidDateFormatError: If the date string format is invalid. """ return DateRange.get_start_date(date_str), DateRange.get_end_date(date_str) class ObservationRequest(BaseModel): variable_dcid: str place_dcid: str child_place_type_dcid: str | None = None source_ids: list[str] | None = None date_type: ObservationDateType = None date_filter: DateRange | None = None child_place_type: str | None = None class SourceProcessingResult(BaseModel): """Intermediate model for holding the results of source processing.""" class ProcessedPlaceData(BaseModel): """ Holds the filtered observations and facet data for a single place. This is a private inner class as it's only used within SourceProcessingResult. """ facet: OrderedFacet observations: list[Observation] primary_source_id: str | None = Field( default=None, description="The DCID of the selected primary data source (facet).", ) alternative_source_counts: dict[str, int] = Field( default_factory=dict, description="A map of alternative source DCIDs to the number of places they have data for.", ) processed_data_by_place: dict[str, "ProcessedPlaceData"] = Field( default_factory=dict, description="A map where keys are place DCIDs and values contain the facet and filtered observations for that place.", ) @property def has_data(self) -> bool: """Returns True if any data was processed, False otherwise.""" return self.primary_source_id is not None or self.processed_data_by_place TimeSeriesPoint = tuple[str, float] # [date, value] class ToolResponseBaseModel(BaseModel): """A base model to configure all tool responses to exclude None values.""" model_config = {"populate_by_name": True} class Node(ToolResponseBaseModel): """Represents a Data Commons node, with an optional name, type, and dcid.""" dcid: str | None = None name: str | None = None type_of: list[str] | None = Field(default=None, alias="typeOf") class FacetMetadata(BaseModel): """Represents the static metadata for a data source.""" source_id: str import_name: str | None = Field(default=None, alias="importName") measurement_method: str | None = Field(default=None, alias="measurementMethod") observation_period: str | None = Field(default=None, alias="observationPeriod") provenance_url: str | None = Field(default=None, alias="provenanceUrl") unit: str | None = None class AlternativeSource(FacetMetadata): """Represents metadata for an alternative data source.""" places_found_count: int | None = Field( default=None, description=( "The number of places within the current API response for which this" " alternative source has data. Used only in child-place type responses." ), ) class PlaceObservation(ToolResponseBaseModel): """Contains all observation data for a single place.""" place: Node time_series: list[TimeSeriesPoint] = Field( default_factory=list, description="List of observation tuples with the format (date, value)", ) class ObservationToolResponse(ToolResponseBaseModel): """The response from the get_observations tool. It contains observation data organized as a list of places. To save tokens, source information is normalized into a top-level `source_info` dictionary. """ variable: Node resolved_parent_place: Node | None = Field( default=None, description=( "The parent place that was resolved from the request, if a hierarchical" " query was made. This confirms how the tool interpreted the `place_name`." ), ) child_place_type: str | None = Field( default=None, description=( "The common place type for all observations in the response (e.g., 'State', 'County', 'AdministrativeArea2'). " "This is present when a hierarchical query was made." ), ) place_observations: list[PlaceObservation] = Field( default_factory=list, description="A list of observation data, with one entry per place. This list may be empty if no data is found.", ) source_metadata: FacetMetadata = Field( description="Information about the primary data source used for the observations." ) alternative_sources: list[AlternativeSource] = Field( default_factory=list, description="A list of other available data sources for the same variable and places.", )