"""MCP retrieval tools used by the agent loop (summaries).""" from __future__ import annotations import os from dotenv import load_dotenv load_dotenv(override=True) from mcp.server.fastmcp import FastMCP from mcp_tools.local_rag_tools import build_paper_context_tool, get_paper_summaries_tool from mcp_tools.metadata_tools import search_metadata_tool, select_papers_from_search_tool mcp = FastMCP("waterscope-retrieval") @mcp.tool() def search_metadata( query: str, top_k: int = 10, year_from: int | None = None, year_to: int | None = None, ) -> str: """Search title, abstract, and keywords with optional publication-year bounds.""" return search_metadata_tool( query, top_k=top_k, year_from=year_from, year_to=year_to, ) @mcp.tool() def select_papers_from_search( query: str, ranks: list[int], top_k: int = 10, year_from: int | None = None, year_to: int | None = None, ) -> str: """Return document_ids for 1-based search result ranks.""" return select_papers_from_search_tool( query, ranks, top_k=top_k, year_from=year_from, year_to=year_to, ) @mcp.tool() def get_paper_summaries(document_ids: list[str]) -> str: """Load pre-written summaries for chat context.""" return get_paper_summaries_tool(document_ids) @mcp.tool() def build_paper_context(document_ids: list[str], question: str = "") -> str: """Assemble truncated multi-paper context for WaterScope /agent_run.""" return build_paper_context_tool(document_ids, question=question) if __name__ == "__main__": transport = os.getenv("MCP_TRANSPORT", "stdio") if transport == "streamable-http": mcp.run( transport="streamable-http", host=os.getenv("MCP_HOST", "127.0.0.1"), port=int(os.getenv("MCP_PORT", "8001")), ) else: mcp.run(transport="stdio")