Biomni/mcp_generated/mcp_bbmap/app/bbmap_server.py

import subprocess
import tempfile
from pathlib import Path
from typing import Optional, List

# This is a placeholder for the MCP decorator.
# In a real MCP environment, this would be provided by the MCP framework.
def tool(func):
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper

mcp = type("mcp", (), {"tool": tool})


@mcp.tool
def bbmap(
    in_file: Path,
    ref_file: Path,
    out_file: Path,
    in2_file: Optional[Path] = None,
    out_unmapped: Optional[Path] = None,
    java_memory: Optional[str] = None,
    threads: Optional[int] = None,
    overwrite: bool = False,
    interleaved: str = "auto",
    reads: int = -1,
    samplerate: float = 1.0,
    bamscript: Optional[Path] = None,
    scafstats: Optional[Path] = None,
    covstats: Optional[Path] = None,
    fast: bool = False,
    slow: bool = False,
    vslow: bool = False,
    max_indel: int = 16000,
    min_identity: float = 0.76,
    min_hits: int = 1,
    local_alignment: bool = False,
    perfect_mode: bool = False,
    ambiguous_reads_mode: str = "best",
    sam_version: str = "1.4",
    secondary_alignments: bool = True,
    max_sites: int = 5,
    pairlen: int = 32000,
    qtrim: str = "f",
    untrim: bool = False,
    trimq: int = 6,
    min_avg_quality: int = 0,
    nodisk: bool = False,
) -> dict:
    """
    Maps sequencing reads to a reference genome using BBMap.

    BBMap is a fast and accurate splice-aware read mapper. This tool wraps the
    bbmap.sh script, providing a comprehensive set of parameters for fine-tuning
    the mapping process.

    Args:
        in_file: Path to the primary input file (fasta, fastq, sam; compressed or uncompressed).
        ref_file: Path to the reference genome file (fasta).
        out_file: Path for the output SAM file.
        in2_file: Path to the second input file for paired-end reads.
        out_unmapped: Path to write unmapped reads to.
        java_memory: Java heap size (e.g., '8g', '1000m'). Passed as -Xmx to the JVM.
        threads: Number of threads to use. If not set, BBMap will auto-detect.
        overwrite: If True, allows overwriting of existing output files.
        interleaved: Set to 't' or 'f' if reads are interleaved; 'auto' detects.
        reads: If positive, quit after processing this many reads.
        samplerate: Subsample reads to this fraction (0.0 to 1.0).
        bamscript: Write a shell script to this path for SAM to BAM conversion.
        scafstats: Write statistics on coverage of each reference scaffold.
        covstats: Write coverage statistics to a file.
        fast: Use faster, less sensitive presets.
        slow: Use slower, more sensitive presets.
        vslow: Use very slow, very sensitive presets.
        max_indel: Do not look for indels longer than this.
        min_identity: Approximate minimum identity of reads to map (0.0 to 1.0).
        min_hits: Minimum number of seed hits required for candidate sites.
        local_alignment: Use local instead of global alignments.
        perfect_mode: Allow only perfect mappings.
        ambiguous_reads_mode: Behavior for ambiguously-mapped reads ('best', 'all', 'random', 'toss').
        sam_version: Set SAM version ('1.4' or '1.5').
        secondary_alignments: Allow secondary alignments for ambiguous reads.
        max_sites: Maximum number of sites to record per read.
        pairlen: Maximum insert size for paired-end reads.
        qtrim: Quality trimming mode ('l', 'r', 'lr', 'f' for left, right, both, or none).
        untrim: Untrim reads.
        trimq: Quality threshold for trimming.
        min_avg_quality: Reads with average quality below this will be discarded.
        nodisk: If True, do not write temporary files to disk.

    Returns:
        A dictionary containing the executed command, stdout, stderr, and a list of output files.
    """
    # --- Input Validation ---
    if not in_file.exists():
        raise FileNotFoundError(f"Input file not found: {in_file}")
    if not ref_file.exists():
        raise FileNotFoundError(f"Reference file not found: {ref_file}")
    if in2_file and not in2_file.exists():
        raise FileNotFoundError(f"Second input file not found: {in2_file}")

    if not (0.0 <= samplerate <= 1.0):
        raise ValueError("samplerate must be between 0.0 and 1.0")
    if not (0.0 <= min_identity <= 1.0):
        raise ValueError("min_identity must be between 0.0 and 1.0")

    valid_ambiguous_modes = {"best", "all", "random", "toss"}
    if ambiguous_reads_mode not in valid_ambiguous_modes:
        raise ValueError(f"ambiguous_reads_mode must be one of {valid_ambiguous_modes}")

    valid_sam_versions = {"1.4", "1.5"}
    if sam_version not in valid_sam_versions:
        raise ValueError(f"sam_version must be one of {valid_sam_versions}")
        
    valid_qtrim_modes = {"l", "r", "lr", "f"}
    if qtrim not in valid_qtrim_modes:
        raise ValueError(f"qtrim must be one of {valid_qtrim_modes}")

    # --- Command Construction ---
    cmd = ["bbmap.sh"]

    if java_memory:
        cmd.append(f"-Xmx{java_memory}")

    # Required parameters
    cmd.append(f"in={in_file}")
    cmd.append(f"ref={ref_file}")
    cmd.append(f"out={out_file}")

    # Optional file paths
    if in2_file:
        cmd.append(f"in2={in2_file}")
    if out_unmapped:
        cmd.append(f"outu={out_unmapped}")
    if bamscript:
        cmd.append(f"bamscript={bamscript}")
    if scafstats:
        cmd.append(f"scafstats={scafstats}")
    if covstats:
        cmd.append(f"covstats={covstats}")

    # Boolean flags (converted to 't'/'f')
    if overwrite:
        cmd.append("ow=t")
    if fast:
        cmd.append("fast=t")
    if slow:
        cmd.append("slow=t")
    if vslow:
        cmd.append("vslow=t")
    if local_alignment:
        cmd.append("local=t")
    if perfect_mode:
        cmd.append("perfectmode=t")
    if secondary_alignments:
        cmd.append("secondary=t")
    if untrim:
        cmd.append("untrim=t")
    if nodisk:
        cmd.append("nodisk=t")

    # String and numeric parameters
    if threads is not None:
        cmd.append(f"threads={threads}")
    if interleaved != "auto":
        cmd.append(f"interleaved={interleaved}")
    if reads != -1:
        cmd.append(f"reads={reads}")
    if samplerate != 1.0:
        cmd.append(f"samplerate={samplerate}")
    if max_indel != 16000:
        cmd.append(f"maxindel={max_indel}")
    if min_identity != 0.76:
        cmd.append(f"minid={min_identity}")
    if min_hits != 1:
        cmd.append(f"minhits={min_hits}")
    if ambiguous_reads_mode != "best":
        cmd.append(f"ambiguous={ambiguous_reads_mode}")
    if sam_version != "1.4":
        cmd.append(f"sam={sam_version}")
    if max_sites != 5:
        cmd.append(f"maxsites={max_sites}")
    if pairlen != 32000:
        cmd.append(f"pairlen={pairlen}")
    if qtrim != "f":
        cmd.append(f"qtrim={qtrim}")
    if trimq != 6:
        cmd.append(f"trimq={trimq}")
    if min_avg_quality != 0:
        cmd.append(f"maq={min_avg_quality}")

    # --- Subprocess Execution ---
    command_str = " ".join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
    except FileNotFoundError:
        return {
            "error": "bbmap.sh not found. Please ensure BBMap is in your PATH.",
            "command_executed": command_str,
            "stdout": "",
            "stderr": "Executable not found.",
        }
    except subprocess.CalledProcessError as e:
        return {
            "error": "BBMap execution failed.",
            "command_executed": command_str,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "return_code": e.returncode,
        }

    # --- Collect Output Files ---
    output_files = [str(out_file)]
    if out_unmapped:
        output_files.append(str(out_unmapped))
    if bamscript:
        output_files.append(str(bamscript))
    if scafstats:
        output_files.append(str(scafstats))
    if covstats:
        output_files.append(str(covstats))

    return {
        "command_executed": command_str,
        "stdout": result.stdout,
        "stderr": result.stderr,
        "output_files": output_files,
    }