canlex/server.py

#!/usr/bin/env python3
"""MCP server exposing CanLex Canadian legislation retrieval as tools.

Runs over stdio. Launched by an MCP client (Claude Code / Claude Desktop);
it requires no API keys -- retrieval is fully local.
"""
import os
import sys
from pathlib import Path
from typing import Optional

# Allow this file to be launched directly by an MCP client, cwd-independent.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from pydantic import BaseModel, ConfigDict, Field
from mcp.server.fastmcp import FastMCP

from canlex.index import LegislationIndex

mcp = FastMCP("canlex_mcp", host="0.0.0.0",
              port=int(os.environ.get("PORT", "8000")))

_READONLY = {
    "readOnlyHint": True,
    "destructiveHint": False,
    "idempotentHint": True,
    "openWorldHint": False,
}

# Prepended to every result so the model grounds, cites, and hedges correctly.
GROUNDING_NOTE = (
    "ANSWERING INSTRUCTIONS: Base the answer only on the material below. Cite "
    "specific provisions and quote key operative words (e.g. 'IRPA s. 34(1)(c)'). "
    "When a result lists related provisions, regulations or D-memoranda, fetch "
    "any that bear on the question -- the definitions section, an exception, a "
    "cross-referenced rule, the regulation that adds detail -- with "
    "canlex_get_section or canlex_search_legislation before answering. "
    "Distinguish the kinds of source: enacted law (Acts and regulations) is binding; "
    "CBSA D-Memoranda are administrative guidance -- persuasive only, not binding, "
    "and a court may disagree with them; collective agreements and the National "
    "Joint Council directives they incorporate are binding employment-terms "
    "instruments for a bargaining unit; court decisions interpret and apply the "
    "law and are binding precedent depending on the court and jurisdiction -- "
    "name the deciding court and the date, and do not assume a decision is still "
    "good law if it may have been overtaken (the canlex_case tool checks a "
    "decision's later treatment on CanLII -- give it the neutral citation). "
    "Always state the date the source is current to, and that the answer "
    "reflects the law only as of that date -- for a time-sensitive matter, tell "
    "the reader to verify no amendment has come into force since. If the material "
    "below does not fully resolve the question -- including where it turns on case "
    "law or facts not present here -- say so explicitly. This is legal information, "
    "not legal advice."
)

HEDGE_THRESHOLD = 0.72   # max semantic similarity below which results are weak

WEAK_MATCH_NOTE = (
    "RETRIEVAL CAUTION: the material below is only a weak match for this query "
    "— CanLex may not contain a provision or decision that directly answers it. "
    "Read it critically; if it does not actually address the question, say so "
    "plainly rather than stretching it to fit, and consider canlex_list_acts to "
    "check what the corpus covers."
)

_INDEX: Optional[LegislationIndex] = None


def _index() -> LegislationIndex:
    """Load and cache the legislation index on first use."""
    global _INDEX
    if _INDEX is None:
        _INDEX = LegislationIndex()
    return _INDEX


def _format_section(c: dict, related=None) -> str:
    """Render one chunk (legislation, D-Memo, or agreement) as cited Markdown."""
    doc_type = c.get("doc_type", "legislation")
    header = f"### {c['citation']} — {c['marginal_note']}".rstrip(" —")
    location = " > ".join(p for p in (c["part"], c["division"]) if p)
    lines = [header]
    if location:
        lines.append(f"*{location}*")
    if doc_type == "memorandum":
        lines.append("_CBSA administrative guidance — persuasive, not binding law._")
        lines.append(f"(modified {c['current_to'] or 'n/a'})")
    elif doc_type == "agreement":
        lines.append("_Treasury Board collective agreement — a binding contract for "
                     "this bargaining unit._")
        lines.append(f"(in force to {c['current_to'] or 'n/a'})")
    elif doc_type == "directive":
        lines.append("_National Joint Council directive — forms part of collective "
                     "agreements; binding for the matters it covers._")
        lines.append(f"(effective {c['current_to'] or 'n/a'})")
    elif doc_type == "caselaw":
        if "Immigration and Refugee Board" in c["part"]:
            lines.append("_Immigration and Refugee Board jurisprudential guide "
                         "— IRB members apply its reasoning to similar cases or "
                         "explain why not; persuasive, and subject to revocation "
                         "or to review by the Federal Court._")
        elif "Board" in c["part"]:
            lines.append("_Labour-board decision — a federal administrative "
                         "tribunal's ruling; persuasive within the board's own "
                         "jurisprudence, and subject to judicial review by the "
                         "Federal Court of Appeal._")
        else:
            lines.append("_Court decision — binding precedent depending on the "
                         "court and jurisdiction; confirm it has not been "
                         "overturned on appeal or overtaken by later authority._")
        lines.append(f"(decided {c['current_to'] or 'n/a'})")
        if c["heading"]:
            lines.append(f"Subject: {c['heading']}")
    elif doc_type == "delegation":
        lines.append("_Instrument of delegation and designation — it records "
                     "which officials the Minister has delegated powers to, or "
                     "designated for functions, under IRPA and the IRPR. "
                     "Administrative; confirm it is still the current version._")
        lines.append(f"(dated {c['current_to'] or 'n/a'})")
    else:
        meta = [f"in force; text current to {c['current_to'] or 'n/a'}"]
        if c["last_amended"]:
            meta.append(f"last amended {c['last_amended']}")
        lines.append(f"**Currency:** {'; '.join(meta)}. Does not reflect any "
                     f"amendment that came into force after the 'current to' date.")
    hl = c.get("highlight")
    if hl:
        label, snippet = hl
        lines.append(f"**Most on point for this query:** "
                     f"{c['citation']}{label} — {snippet}")
    lines.append("")
    lines.append(c["text"])
    lines.append("")
    if related:
        provisions = related.get("provisions")
        if provisions:
            refs = "; ".join(f"s. {s} ({n})" if n else f"s. {s}"
                             for s, n in provisions)
            lines.append(f"Related provisions in this Act: {refs}")
        regs = related.get("regulations")
        if regs:
            lines.append("Regulations made under this Act: "
                         + "; ".join(f"{n} ({s})" for s, n in regs))
        enabling = related.get("enabling_act")
        if enabling:
            lines.append(f"Made under: {enabling[1]} ({enabling[0]})")
        memos = related.get("memoranda")
        if memos:
            lines.append("CBSA D-memoranda citing this section (guidance, not "
                         "binding): " + ", ".join(memos))
    if c["history"]:
        if doc_type == "caselaw":
            lines.append(f"Also reported: {c['history']}")
        elif doc_type == "legislation":
            lines.append(f"Amendment history: {c['history']}")
        else:
            lines.append(f"History: {c['history']}")
    lines.append(f"Source: {c['source_url']}")
    return "\n".join(lines)


class SearchInput(BaseModel):
    """Input for canlex_search_legislation."""
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")

    query: str = Field(
        ...,
        description="Natural-language legal question or keywords, e.g. "
        "'detention review timelines' or 'inadmissibility for serious criminality'.",
        min_length=2, max_length=500,
    )
    top_k: int = Field(
        default=6,
        description="Number of sections to return (1-20). Use more for broad questions.",
        ge=1, le=20,
    )
    act: Optional[str] = Field(
        default=None,
        description="Optional filter to a single Act, by short name or code "
        "(e.g. 'IRPA' or 'I-2.5'). Omit to search every loaded Act.",
    )
    doc_type: Optional[str] = Field(
        default=None,
        description="Optional filter by source type: 'legislation' (Acts and "
        "regulations), 'memorandum' (CBSA D-Memoranda), 'agreement' (collective "
        "agreements), 'directive' (NJC directives), 'caselaw' (court and "
        "tribunal decisions), or 'delegation' (IRPA/IRPR delegation and "
        "designation instruments). Omit to search all.",
    )


class GetSectionInput(BaseModel):
    """Input for canlex_get_section."""
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")

    act: str = Field(..., description="Act short name or code, e.g. 'IRPA' or 'I-2.5'.",
                     min_length=1, max_length=60)
    section: str = Field(..., description="Section number exactly as cited, e.g. '34', '20.1'.",
                         min_length=1, max_length=20)


@mcp.tool(name="canlex_search_legislation",
          annotations={"title": "Search Canadian Legislation", **_READONLY})
def canlex_search_legislation(params: SearchInput) -> str:
    """Search Canadian federal law, CBSA D-Memoranda, agreements, NJC directives,
    and leading court decisions.

    The CanLex corpus has six kinds of source: 31 federal Acts and regulations
    (immigration, customs, criminal, drugs, food/health, labour, privacy and more);
    CBSA D-Memoranda (the Canada Border Services Agency's administrative guidance on
    how it applies customs and border law); Treasury Board collective agreements
    (currently the FB / Border Services group); National Joint Council directives
    (travel, relocation, isolated posts and more); leading decisions of the
    courts and federal tribunals: the Supreme Court, Federal Court of Appeal and
    Federal Court, the Immigration and Refugee Board, and the FPSLREB and CIRB
    labour boards; and instruments of delegation and designation under IRPA and
    the IRPR (which officials the Minister has authorized to exercise which powers). Use this for ANY question about that material. It ranks results by relevance and returns
    their full text so the answer can cite the actual wording; an explicit section
    reference (e.g. "section 34") is always surfaced. Each result is marked with its
    source type.

    Args:
        params (SearchInput): Validated input containing:
            - query (str): Legal question or keywords to search for.
            - top_k (int): How many sections to return, 1-20 (default 6).
            - act (Optional[str]): Restrict to one Act by short name/code, or omit for all.
            - doc_type (Optional[str]): 'legislation', 'memorandum', 'agreement',
              'directive', 'caselaw', or 'delegation' to restrict to one source
              type; omit for all.

    Returns:
        str: Markdown with answering instructions followed by the matching sections.
        Each section block contains its citation, Part/Division, the 'current to' and
        'last amended' dates, the full provision text, amendment history, and a source
        URL. Returns a "No ... matched" message if nothing is found.

    Examples:
        - "What are the security grounds for inadmissibility?" -> query about s. 34.
        - "detention review 48 hours" -> finds the detention-review provisions.
        - Don't use to look up a known section number verbatim -- use canlex_get_section.
    """
    index = _index()
    results = index.search(params.query, top_k=params.top_k, act=params.act,
                           doc_type=params.doc_type)
    if not results:
        scope = f" in '{params.act}'" if params.act else ""
        return (f"No results matched '{params.query}'{scope}. "
                f"Try broader or different keywords, or call canlex_list_acts to see "
                f"what is currently loaded.")
    blocks = []
    weak = results[0].get("confidence")
    if weak is not None and weak < HEDGE_THRESHOLD:
        blocks += [WEAK_MATCH_NOTE, ""]
    blocks += [GROUNDING_NOTE, "",
               f'{len(results)} relevant section(s) for: "{params.query}"']
    for c in results:
        blocks.append("")
        blocks.append("---")
        blocks.append("")
        blocks.append(_format_section(c, index.related(c)))
    return "\n".join(blocks)


@mcp.tool(name="canlex_get_section",
          annotations={"title": "Get a Legislation Section", **_READONLY})
def canlex_get_section(params: GetSectionInput) -> str:
    """Retrieve one specific section of legislation verbatim, by Act and section number.

    Use this when the exact citation is already known (e.g. the user asks about
    "IRPA section 34" or a provision cross-references "subsection 25.1(1)").

    Args:
        params (GetSectionInput): Validated input containing:
            - act (str): Act short name or code, e.g. 'IRPA' or 'I-2.5'.
            - section (str): Section number exactly as cited, e.g. '34' or '20.1'.

    Returns:
        str: Markdown with answering instructions followed by the section's citation,
        Part/Division, 'current to' and 'last amended' dates, full provision text,
        amendment history, and source URL. Returns an actionable error message (listing
        the loaded Acts) if the Act or section is not found.

    Examples:
        - "Show me IRPA s. 34" -> act='IRPA', section='34'.
        - "What does subsection 20.1 say" -> act='IRPA', section='20.1'.
    """
    index = _index()
    section = index.get_section(params.act, params.section)
    if section is None:
        acts = sorted({c["act_short"] for c in index.chunks})
        return (f"Error: no section '{params.section}' found in '{params.act}'. "
                f"Loaded Acts: {', '.join(acts) or 'none'}. Check the section number, "
                f"or use canlex_search_legislation to locate the provision by topic.")
    return GROUNDING_NOTE + "\n\n" + _format_section(section, index.related(section))


@mcp.tool(name="canlex_list_acts",
          annotations={"title": "List Loaded Legislation", **_READONLY})
def canlex_list_acts() -> str:
    """List what the CanLex corpus contains -- Acts and regulations, CBSA
    D-Memoranda, collective agreements, NJC directives, leading cases, and
    delegation instruments.

    Use this to learn the scope and currency of the corpus before searching, or to
    report it to the user.

    Returns:
        str: Markdown grouped by source type.
    """
    index = _index()
    acts: dict[str, dict] = {}
    agreements: dict[str, dict] = {}
    directives: dict[str, dict] = {}
    cases: dict[str, dict] = {}
    delegations: dict[str, dict] = {}
    memo_numbers: set[str] = set()
    memo_chunks = 0
    memo_date = ""
    for c in index.chunks:
        doc_type = c.get("doc_type", "legislation")
        if doc_type == "memorandum":
            memo_numbers.add(c["section"])
            memo_chunks += 1
            memo_date = max(memo_date, c["current_to"] or "")
        elif doc_type == "agreement":
            entry = agreements.setdefault(c["act_code"], {
                "short": c["act_short"], "name": c["act_name"],
                "current_to": c["current_to"], "count": 0,
            })
            entry["count"] += 1
        elif doc_type == "directive":
            entry = directives.setdefault(c["act_code"], {
                "short": c["act_short"], "current_to": c["current_to"], "count": 0,
            })
            entry["count"] += 1
        elif doc_type == "caselaw":
            entry = cases.setdefault(c["act_code"], {
                "name": c["act_name"], "decided": c["current_to"], "count": 0,
            })
            entry["count"] += 1
        elif doc_type == "delegation":
            entry = delegations.setdefault(c["act_code"], {
                "short": c["act_short"], "name": c["act_name"],
                "current_to": c["current_to"], "count": 0,
            })
            entry["count"] += 1
        else:
            entry = acts.setdefault(c["act_code"], {
                "short": c["act_short"], "name": c["act_name"],
                "code": c["act_code"], "current_to": c["current_to"], "count": 0,
            })
            entry["count"] += 1
    lines = ["# CanLex corpus", "", "## Enacted law"]
    for a in sorted(acts.values(), key=lambda x: x["short"]):
        lines.append(f"- **{a['short']}** — {a['name']} ({a['code']}): "
                     f"{a['count']} sections, current to {a['current_to'] or 'n/a'}")
    if memo_numbers:
        lines += ["", "## CBSA guidance",
                  f"- **D-Memoranda** — {len(memo_numbers)} memoranda "
                  f"({memo_chunks} sections), modified up to {memo_date or 'n/a'}. "
                  f"Administrative guidance on customs and border law; "
                  f"persuasive, not binding."]
    if agreements:
        lines += ["", "## Collective agreements"]
        for a in sorted(agreements.values(), key=lambda x: x["short"]):
            lines.append(f"- **{a['short']}** — {a['name']}: {a['count']} articles, "
                         f"in force to {a['current_to'] or 'n/a'}")
    if directives:
        lines += ["", "## NJC directives"]
        for a in sorted(directives.values(), key=lambda x: x["short"]):
            lines.append(f"- **{a['short']}**: {a['count']} sections, "
                         f"effective {a['current_to'] or 'n/a'}")
    if cases:
        lines += ["", "## Case law"]
        for cite, a in sorted(cases.items(), key=lambda kv: kv[1]["decided"]):
            lines.append(f"- **{a['name']}**, {cite}: {a['count']} excerpts, "
                         f"decided {a['decided'] or 'n/a'}")
    if delegations:
        lines += ["", "## Delegation instruments"]
        for a in sorted(delegations.values(), key=lambda x: x["short"]):
            lines.append(f"- **{a['short']}** — {a['name']}: {a['count']} items, "
                         f"dated {a['current_to'] or 'n/a'}")
    lines += ["", "Search with canlex_search_legislation; filter by doc_type "
              "(legislation / memorandum / agreement / directive / caselaw / "
              "delegation). Fetch a known provision with canlex_get_section, or "
              "a case's citations with canlex_case."]
    return "\n".join(lines)


_CITATOR = None


def _citator():
    """Load and cache the CanLII citator on first use."""
    global _CITATOR
    if _CITATOR is None:
        from canlex.citator import Citator
        _CITATOR = Citator()
    return _CITATOR


class CaseInput(BaseModel):
    """Input for canlex_case."""
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")

    case_url: str = Field(
        ...,
        description="A Canadian case, given either as a full canlii.org URL or "
        "-- for a Supreme Court, Federal Court of Appeal or Federal Court "
        "decision -- its neutral citation (e.g. '2019 SCC 65' or '2016 FCA 93'). "
        "For other courts, supply the canlii.org URL; find it by web search if "
        "you only have the case name.",
        min_length=8, max_length=400,
    )


def _format_case(report: dict) -> str:
    """Render a CanLII citator report as Markdown."""
    meta = report["meta"]
    lines = [f"## {meta.get('title', '(untitled case)')}",
             meta.get("citation", "").strip()]
    facts = []
    if meta.get("decisionDate"):
        facts.append(f"decided {meta['decisionDate']}")
    if meta.get("docketNumber"):
        facts.append(f"docket {meta['docketNumber']}")
    if facts:
        lines.append("(" + "; ".join(facts) + ")")
    if meta.get("keywords"):
        lines.append(f"**Keywords:** {meta['keywords']}")
    if meta.get("topics"):
        lines.append(f"**Topics:** {meta['topics']}")
    if meta.get("url"):
        lines.append(f"CanLII: {meta['url']}")
    lines.append("")
    lines.append("_CanLII citator data (live). This is case metadata and the "
                 "citation graph only -- not the judgment text; follow the CanLII "
                 "link to read the decision. A case is binding precedent depending "
                 "on the court and jurisdiction._")

    def case_list(label, block):
        rows = [f"\n### {label}: {block['total']}"]
        for item in block["items"]:
            rows.append(f"- {item.get('title', '')} — {item.get('citation', '')}")
        extra = block["total"] - len(block["items"])
        if extra > 0:
            rows.append(f"  ...and {extra} more.")
        return rows

    lines += case_list("Cited by (later cases citing this one)", report["citingCases"])
    lines += case_list("Cites (authorities this case relies on)", report["citedCases"])

    legis = report["citedLegislations"]
    lines.append(f"\n### Legislation cited: {legis['total']}")
    for item in legis["items"]:
        kind = item.get("type", "")
        lines.append(f"- {item.get('title', '')} — {item.get('citation', '')}"
                     + (f" [{kind}]" if kind else ""))
    extra = legis["total"] - len(legis["items"])
    if extra > 0:
        lines.append(f"  ...and {extra} more.")
    return "\n".join(lines)


@mcp.tool(name="canlex_case",
          annotations={"title": "CanLII Case Citator", "readOnlyHint": True,
                       "destructiveHint": False, "idempotentHint": True,
                       "openWorldHint": True})
def canlex_case(params: CaseInput) -> str:
    """Look up a Canadian case on CanLII and return its citation graph.

    Returns the case's metadata plus its citator: the cases it cites, the cases
    that cite it (its treatment and how leading it is), and the legislation it
    cites -- live from the CanLII API. Use it to gauge whether a decision is
    still good law -- how heavily and how recently it has been cited.

    Supply either a canlii.org URL or, for a Supreme Court / Federal Court of
    Appeal / Federal Court decision, its neutral citation (e.g. '2019 SCC 65') --
    the citation a canlex_search_legislation result already shows. This returns
    metadata and the citation graph only, NOT the judgment text -- follow the
    CanLII link for that. A call takes ~15-20 seconds (the API is rate-limited).

    Args:
        params (CaseInput): contains case_url -- a canlii.org URL or a neutral
            citation.

    Returns:
        str: Markdown -- the case's title, neutral citation, date, docket and
        topics; how many cases cite it (with examples); how many it cites; and the
        legislation it cites. Returns an error message if the URL is not a
        recognized canlii.org case URL, or if the CanLII API is unavailable.
    """
    try:
        citator = _citator()
    except Exception as exc:
        return (f"The case citator is unavailable: {exc} "
                f"It needs a CanLII API key in canlii_key.txt.")
    try:
        report = citator.case_report(params.case_url)
    except Exception as exc:
        return f"CanLII lookup failed: {type(exc).__name__}: {exc}"
    if "error" in report:
        return f"Error: {report['error']}"
    return _format_case(report)


if __name__ == "__main__":
    try:
        idx = _index()
        print(f"CanLex MCP: {len(idx.chunks)} sections loaded.", file=sys.stderr)
    except Exception as exc:  # surfaced in the MCP client's server logs
        print(f"CanLex MCP: failed to load legislation index: {exc}\n"
              f"Run 'py -m canlex.ingest' first.", file=sys.stderr)
        sys.exit(1)
    # CANLEX_HTTP switches on the remote (web) transport; default is local stdio.
    if os.environ.get("CANLEX_HTTP"):
        print("CanLex MCP: serving over streamable-HTTP.", file=sys.stderr)
        mcp.run(transport="streamable-http")
    else:
        mcp.run()