fetch-aggregate-data.tool.ts

import { Tool } from '@modelcontextprotocol/sdk/types.js'
import { z } from 'zod'

import { BaseTool } from './base.tool.js'
import { buildCitation } from '../helpers/citation.js'
import {
  FetchAggregateDataToolSchema,
  TableArgs,
  TableSchema,
} from '../schema/fetch-aggregate-data.schema.js'
import { ToolContent } from '../types/base.types.js'

import {
  datasetValidator,
  validateGeographyArgs,
} from '../schema/validators.js'

export const toolDescription = `
  Fetches statistical data from U.S. Census Bureau datasets including population, demographics, income, housing, employment, and economic indicators. Use this tool when users request Census statistics, demographic breakdowns, or socioeconomic data for specific geographic areas. Requires a dataset identifier, year/vintage, geographic scope (state, county, tract, etc.), and specific variables or table groups. Returns structured data with proper citations for authoritative government statistics.

  Important: ACS 1-year estimates (acs/acs1) are only available for geographic areas with populations of 65,000 or more. For smaller areas, use ACS 5-year estimates (acs/acs5) instead. See: https://www.census.gov/programs-surveys/acs/guidance/estimates.html
`

export class FetchAggregateDataTool extends BaseTool<TableArgs> {
  name = 'fetch-aggregate-data'
  description = toolDescription
  inputSchema: Tool['inputSchema'] = TableSchema as Tool['inputSchema']
  readonly requiresApiKey = true

  get argsSchema() {
    return FetchAggregateDataToolSchema.superRefine((args, ctx) => {
      //Check that the correct tool is used to fetch data
      const identifiedDataset = datasetValidator(args.dataset)

      if (identifiedDataset.tool !== this.name) {
        ctx.addIssue({
          path: ['dataset'],
          code: z.ZodIssueCode.custom,
          message: identifiedDataset.message,
        })
      }

      validateGeographyArgs(args, ctx)
    })
  }

  constructor() {
    super()
    this.handler = this.handler.bind(this)
  }

  validateArgs(input: unknown) {
    return this.argsSchema.safeParse(input)
  }

  async toolHandler(
    args: TableArgs,
    apiKey: string,
  ): Promise<{ content: ToolContent[] }> {
    const baseUrl = `https://api.census.gov/data/${args.year}/${args.dataset}`

    let getParams = ''

    if (args.get.variables || args.get.group) {
      if (args.get.variables) {
        getParams = args.get.variables.join(',')
      }

      if (args.get.group) {
        if (getParams != '') {
          getParams += ','
        }
        getParams += `group(${args.get.group})`
      }
    }

    const query = new URLSearchParams({
      get: getParams,
    })

    if (args.for) {
      query.append('for', args.for)
    }

    if (args.in) {
      query.append('in', args.in)
    }

    if (args.ucgid) {
      query.append('ucgid', args.ucgid)
    }

    if (args.predicates) {
      for (const [key, value] of Object.entries(args.predicates)) {
        query.append(key, value)
      }
    }

    const descriptive = args.descriptive?.toString() ?? 'false'

    query.append('descriptive', descriptive)
    query.append('key', apiKey)

    const url = `${baseUrl}?${query.toString()}`

    try {
      const fetch = (await import('node-fetch')).default
      const res = await fetch(url)

      console.log(`URL Attempted: ${url}`)

      if (!res.ok) {
        if (res.status === 400 && isAcs1YearDataset(args.dataset)) {
          return this.createErrorResponse(
            `Census API error: ${res.status} ${res.statusText}. ` +
            `Note: ACS 1-year estimates (acs/acs1) are only available for geographic areas ` +
            `with populations of 65,000 or more. If the requested area has a smaller population, ` +
            `try using ACS 5-year estimates instead by changing the dataset to "acs/acs5". ` +
            `See: https://www.census.gov/programs-surveys/acs/guidance/estimates.html`,
          )
        }

        return this.createErrorResponse(
          `Census API error: ${res.status} ${res.statusText}`,
        )
      }

      const data = (await res.json()) as string[][]
      const [headers, ...rows] = data

      const output = rows
        .map((row) => headers.map((h, i) => `${h}: ${row[i]}`).join(', '))
        .join('\n')

      const citation = buildCitation(url)

      return this.createSuccessResponse(
        `Response from ${args.dataset}:\n${output}\n${citation}`,
      )
    } catch (err) {
      return this.createErrorResponse(`Fetch failed: ${(err as Error).message}`)
    }
  }
}

function isAcs1YearDataset(dataset: string): boolean {
  const normalized = dataset.toLowerCase().replace(/\s+/g, '')
  return normalized.includes('acs/acs1') || normalized.includes('acs1')
}