"""The ``aeat_corpus_search`` grounding tool: hybrid search over the legal corpus.
ADR R3's grounding surface reaches the model as a read-only console tool that
searches the bundled BOE/AEAT corpus and returns grounded hits — each carrying
its ``corpus_ref``, title, a verbatim snippet, the fused relevance score, and
an ``aeat://corpus/{ref}`` URI a resources-capable client can read to pull the
full verbatim text. An exact citation id short-circuits straight to the
resolved authoritative text.
Like ``_harness_tools`` / ``_resources``, this module is SDK-independent pure
functions over typed models: :func:`corpus_search_payload_from_response` and
:func:`render_corpus_search_text` carry no protocol detail and are unit-tested
directly, while :func:`build_corpus_search_tool` lazily adapts onto the MCP
SDK's ``Tool`` type so the module still imports (and the server refuses
gracefully) when the ``aeat-cli[agent]`` extra is absent. The retrieval itself is
owned by the application service (:func:`~application.corpus_search.search_corpus`),
consumed through the package facade per ``service-imports-via-top-level-reexports``.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from pydantic import BaseModel, ConfigDict, Field
from ...application.corpus_search import RetrievalMode, RetrievalResponse, search_corpus
if TYPE_CHECKING:
# Typing-only: the MCP SDK is an optional runtime dependency (``aeat-cli[agent]``);
# the real import stays deferred to inside the function body below.
from mcp.types import Tool
#: The grounding tool's MCP name (the ``corpus.search`` verb, per the
#: ``aeat_<key>`` convention).
CORPUS_SEARCH_TOOL = "aeat_corpus_search"
_CORPUS_URI_PREFIX = "aeat://corpus/"
_SNIPPET_MAX = 280
_DEFAULT_LIMIT = 8
_MAX_LIMIT = 50
_STRICT_FROZEN = ConfigDict(frozen=True, strict=True, validate_assignment=True, extra="forbid")
[docs]
def corpus_uri(corpus_ref: str) -> str:
"""Render the ``aeat://corpus/<ref>`` URI for a corpus reference."""
return f"{_CORPUS_URI_PREFIX}{corpus_ref}"
[docs]
class CorpusSearchResultRow(BaseModel):
"""One grounded corpus hit surfaced to the model."""
model_config = _STRICT_FROZEN
corpus_ref: str = Field(min_length=1)
title: str = Field(min_length=1)
snippet: str = Field(min_length=1)
score: float = Field(ge=0.0)
uri: str = Field(min_length=1)
[docs]
class CorpusCitationResult(BaseModel):
"""The short-circuit result when the query is an exact citation id."""
model_config = _STRICT_FROZEN
citation_id: str = Field(min_length=1)
document_id: str = Field(min_length=1)
permalink: str = Field(min_length=1)
uri: str = Field(min_length=1)
snippet: str = Field(min_length=1)
[docs]
class CorpusSearchPayload(BaseModel):
"""The grounding tool's structured result.
``mode`` records how the response was produced (``citation`` / ``hybrid`` /
``lexical_only``). For a citation query ``citation`` is populated and
``results`` is empty; otherwise ``results`` carries the ranked hits.
"""
model_config = _STRICT_FROZEN
query: str = Field(min_length=1)
mode: RetrievalMode
results: tuple[CorpusSearchResultRow, ...] = ()
citation: CorpusCitationResult | None = None
def _snippet(text: str) -> str:
collapsed = " ".join(text.split())
if len(collapsed) <= _SNIPPET_MAX:
return collapsed
return collapsed[: _SNIPPET_MAX - 1].rstrip() + "…"
[docs]
def corpus_search_payload_from_response(response: RetrievalResponse) -> CorpusSearchPayload:
"""Map a :class:`RetrievalResponse` to the tool's typed payload.
Returns:
A :class:`CorpusSearchPayload`.
"""
if response.mode is RetrievalMode.CITATION and response.citation is not None:
citation = response.citation
return CorpusSearchPayload(
query=response.query,
mode=response.mode,
citation=CorpusCitationResult(
citation_id=citation.citation_id,
document_id=citation.document_id,
permalink=citation.permalink,
uri=corpus_uri(citation.citation_id),
snippet=_snippet(citation.verbatim_text),
),
)
rows = tuple(
CorpusSearchResultRow(
corpus_ref=hit.corpus_ref,
title=hit.doc_title,
snippet=_snippet(hit.text),
score=hit.score,
uri=corpus_uri(hit.corpus_ref),
)
for hit in response.hits
)
return CorpusSearchPayload(query=response.query, mode=response.mode, results=rows)
[docs]
def build_corpus_search_payload(query: str, *, limit: int = _DEFAULT_LIMIT) -> CorpusSearchPayload:
"""Run grounding retrieval for ``query`` and return the tool payload.
Returns:
A :class:`CorpusSearchPayload`.
"""
response = search_corpus(query, limit=limit)
return corpus_search_payload_from_response(response)
[docs]
def render_corpus_search_text(payload: CorpusSearchPayload) -> str:
"""Render the payload as markdown for the tool's text content."""
if payload.citation is not None:
citation = payload.citation
return "\n".join(
[
f"# corpus citation: {citation.citation_id}",
f"- document: {citation.document_id}",
f"- permalink: {citation.permalink}",
f"- resource: {citation.uri}",
"",
citation.snippet,
]
)
if not payload.results:
return f"No corpus results for '{payload.query}'."
lines = [f"# corpus results for '{payload.query}' ({payload.mode.value})", ""]
for index, row in enumerate(payload.results, start=1):
lines += [
f"{index}. {row.title} (score {row.score:.4f})",
f" {row.uri}",
f" {row.snippet}",
]
return "\n".join(lines)
__all__ = [
"CORPUS_SEARCH_TOOL",
"CorpusCitationResult",
"CorpusSearchPayload",
"CorpusSearchResultRow",
"build_corpus_search_payload",
"build_corpus_search_tool",
"corpus_search_payload_from_response",
"corpus_uri",
"render_corpus_search_text",
]