Spaces:
Sleeping
Sleeping
| """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") | |
| 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, | |
| ) | |
| 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, | |
| ) | |
| def get_paper_summaries(document_ids: list[str]) -> str: | |
| """Load pre-written summaries for chat context.""" | |
| return get_paper_summaries_tool(document_ids) | |
| 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") | |