Add the corpus as an AI agent tool (MCP)

import requests

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

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}
    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()

    results = []
    for r in resp.json()["results"]:
        results.append(
            {
                "actId": r.get("actId"),
                "citation": r.get("citation"),
                "title": r.get("title"),
                "excerpt": r.get("excerpt"),
                "sourceUrl": (
                    r.get("govInfoHtmlUrl")
                    or r.get("stateHtmlUrl")
                    or r.get("htmlUrl")
                ),
            }
        )
    return 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