Add the corpus as an AI agent tool (MCP)

import requests

BASE_URL = "https://api.vaquill.ai/api/v1/us"
HEADERS = {
    "Authorization": "Bearer vq_key_...",
    "Content-Type": "application/json",
}

# Exactly what the model needs to answer and cite, and nothing else.
# An unknown name here is rejected with 422, so a typo cannot silently
# drop a field you were relying on.
TOOL_FIELDS = [
    "title",
    "excerpt",
    "govInfoHtmlUrl",
    "stateHtmlUrl",
    "htmlUrl",
]

def search_statutes(
    query: str,
    corpus_type: str | None = None,
    state: str | None = None,
    limit: int = 5,
) -> list[dict]:
    """Search US statutes and regulations. Returns cited results."""
    payload: dict = {"query": query, "limit": limit, "fields": TOOL_FIELDS}
    if corpus_type:
        payload["corpusType"] = corpus_type
    if state:
        payload["state"] = state

    resp = requests.post(
        f"{BASE_URL}/statutes/search", headers=HEADERS, json=payload
    )
    resp.raise_for_status()

    # best_source_url() is the shared helper from Text Formats and Source
    # URLs, linked below. Copy it once; do not re-derive the field order.
    return [
        {
            "actId": r.get("actId"),
            "citation": r.get("citation"),
            "title": r.get("title"),
            "excerpt": r.get("excerpt"),
            "sourceUrl": best_source_url(r),
        }
        for r in resp.json()["results"]
    ]

SEARCH_STATUTES_TOOL = {
    "name": "search_statutes",
    "description": (
        "Search US statutes and regulations (US Code, CFR, state codes, "
        "constitutions, court rules, and more). Use for any question about "
        "what the law says. Returns citations, excerpts, and source links."
    ),
    "parameters": {
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "Natural language question or keywords.",
            },
            "corpus_type": {
                "type": "string",
                "description": (
                    "Optional corpus filter, e.g. USC, CFR, STATE, "
                    "REGULATION. Omit to search everything."
                ),
            },
            "state": {
                "type": "string",
                "description": "Optional 2-letter lowercase state code, e.g. ca.",
            },
        },
        "required": ["query"],
    },
}

def dispatch_tool_call(name: str, arguments: dict) -> list[dict]:
    if name == "search_statutes":
        return search_statutes(
            query=arguments["query"],
            corpus_type=arguments.get("corpus_type"),
            state=arguments.get("state"),
        )
    raise ValueError(f"Unknown tool: {name}")

# Pseudocode for a tool-calling agent loop:
#
#   response = agent.run(user_message, tools=[SEARCH_STATUTES_TOOL])
#   while response.wants_tool_call:
#       result = dispatch_tool_call(response.tool_name, response.tool_args)
#       response = agent.continue_with_tool_result(result)
#   print(response.text)  # grounded answer with citations