server.py

"""MCP tools for extracting locally stored literature PDFs.

Tools:
  - extract_pdf_local(pdf_path, query)
  - extract_folder_pdfs(folder_path, query)
  - retrieve_pdf_passages(pdf_path, query)
  - extract_perovskite_from_text(text)
  - extract_perovskite_from_pdf(pdf_path)
"""

from typing import Dict, Optional

from fastmcp import FastMCP

from _server_transport import run
from batch import extract_folder_pdfs as _extract_folder_pdfs
from extractor import extract_pdf
from perovskite import extract_perovskite_fields
from retrieval import retrieve_relevant_passages

mcp = FastMCP("extract-mcp")


@mcp.tool
def extract_pdf_local(
    pdf_path: str,
    query: str = "",
    analyze_images: Optional[bool] = None,
    image_query: str = "",
) -> Dict:
    """Extract text, images, and a query-aware snippet from one local PDF."""
    try:
        return extract_pdf(
            pdf_path,
            query=query,
            analyze_images=analyze_images,
            image_query=image_query,
        )
    except Exception as e:
        return {"ok": False, "error": f"extract_failed: {e}", "pdf_path": pdf_path}


@mcp.tool
def extract_folder_pdfs(
    folder_path: str,
    query: str = "",
    recursive: bool = True,
    max_files: Optional[int] = None,
    use_llm_filter: Optional[bool] = None,
    include_full_text: bool = False,
    analyze_images: Optional[bool] = None,
    image_query: str = "",
) -> Dict:
    """Batch extract all PDFs under a local folder."""
    return _extract_folder_pdfs(
        folder_path,
        query=query,
        recursive=recursive,
        max_files=max_files,
        use_llm_filter=use_llm_filter,
        include_full_text=include_full_text,
        analyze_images=analyze_images,
        image_query=image_query,
    )


@mcp.tool
def retrieve_pdf_passages(
    pdf_path: str,
    query: str,
    top_k: int = 5,
    method: Optional[str] = None,
) -> Dict:
    """Retrieve relevant text passages from one local PDF using BM25/embedding."""
    try:
        ext = extract_pdf(pdf_path, query="")
    except Exception as e:
        return {"ok": False, "error": f"extract_failed: {e}", "pdf_path": pdf_path}
    passages = retrieve_relevant_passages(
        ext.get("full_text") or "",
        query,
        top_k=top_k,
        method=method,
    )
    return {
        "ok": True,
        "pdf_path": pdf_path,
        "query": query,
        "method": method,
        "passages": passages,
    }


@mcp.tool
def extract_perovskite_from_text(
    text: str,
    use_llm_filter: Optional[bool] = None,
) -> Dict:
    """Extract perovskite solar-cell fields from arbitrary text."""
    return extract_perovskite_fields(text, use_llm_filter=use_llm_filter)


@mcp.tool
def extract_perovskite_from_pdf(
    pdf_path: str,
    use_llm_filter: Optional[bool] = None,
) -> Dict:
    """Extract perovskite solar-cell fields from one local PDF."""
    try:
        ext = extract_pdf(pdf_path, query="")
    except Exception as e:
        return {"ok": False, "error": f"extract_failed: {e}", "pdf_path": pdf_path}
    if not ext.get("ok"):
        return ext
    fields = extract_perovskite_fields(
        ext.get("full_text") or "",
        use_llm_filter=use_llm_filter,
    )
    return {
        "ok": True,
        "pdf_path": pdf_path,
        "full_text_len": ext.get("full_text_len"),
        "perovskite_fields": fields,
    }


if __name__ == "__main__":
    run(mcp, "extract-mcp")