mcp_service.py

import json
from climateqa.engine.chains.retrieve_documents import retrieve_documents
from climateqa.engine.graph_retriever import retrieve_graphs


def make_retrieve_data_mcp(vectorstore, reranker):
    """
    MCP-exposed function: retrieve IPx (IPCC/IPBES/IPOS) documents and figures.
    """
    async def retrieve_data_mcp(
        query: str,
    ) -> str:
        """
        MCP-exposed function: retrieve IPx (IPCC/IPBES/IPOS) documents and figures.

        Args:
            query (str): The user query.

        Returns:
            str: JSON with keys: "documents", "images".
        """

        current_question = {
            "question": query,
            "sources": ["IPCC", "IPBES", "IPOS"],
            "source_type": "IPx",
            "index": 1,
        }

        docs, images = await retrieve_documents(
            current_question=current_question,
            config={"logging": True},
            source_type="IPx",
            vectorstore=vectorstore,
            reranker=reranker,
            search_figures=True,
            rerank_by_question=True,
            k_images_by_question=5,
            k_before_reranking=100,
            k_by_question=5,
            k_summary_by_question=3,
        )

        def _serialize(docs_list):
            return [
                {
                    "content": getattr(d, "page_content", ""),
                    "metadata": getattr(d, "metadata", {}),
                }
                for d in (docs_list or [])
            ]

        payload = {
            "documents": _serialize(docs),
            "images": _serialize(images),
        }

        return json.dumps(payload)

    return retrieve_data_mcp


def make_retrieve_graphs_mcp(vectorstore):
    """
    MCP-exposed function: retrieve graphs (OWID).
    """
    async def retrieve_graphs_mcp(
        query: str,
    ) -> str:
        """
        MCP-exposed function: retrieve graphs (OWID).

        Args:
            query (str): The user query.

        Returns:
            str: JSON with key: "graphs".
        """

        docs = await retrieve_graphs(
            query=query,
            vectorstore=vectorstore,
            # sources=["OWID"],
            threshold=0.2,  # Lowered from 0.5 to get more results
            k_total=10,
        )

        def _serialize(docs_list):
            return [
                {
                    "content": getattr(d, "page_content", ""),
                    "metadata": getattr(d, "metadata", {}),
                }
                for d in (docs_list or [])
            ]

        payload = {
            "graphs": _serialize(docs),
        }

        return json.dumps(payload)

    return retrieve_graphs_mcp