Biomni/mcp_generated/mcp_bioawk/app/bioawk_server.py

import subprocess
from pathlib import Path
from typing import List, Optional, Dict, Any

# Assume 'mcp' is a globally available module in the execution environment.
# from some_mcp_library import mcp

@mcp.tool
def bioawk(
    program: Optional[str] = None,
    program_file: Optional[Path] = None,
    input_files: Optional[List[Path]] = None,
    output_file: Optional[Path] = None,
    format: Optional[str] = None,
    tab_separator: bool = False,
    include_header: bool = False,
    variables: Optional[List[str]] = None,
    field_separator: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Executes the bioawk command, a powerful stream editor for biological data formats.

    bioawk is an extension of awk that understands common biological data formats
    like FASTA, FASTQ, SAM, VCF, BED, and GFF. At least one of 'program' or
    'program_file' must be provided.

    Args:
        program: The AWK program string to execute. Mutually exclusive with 'program_file'.
        program_file: Path to a file containing the AWK program. Mutually exclusive with 'program'.
        input_files: A list of input files to process. If not provided, bioawk reads from standard input.
        output_file: Path to a file where the output (stdout) will be saved. If not provided, stdout is captured and returned.
        format: Input format (-c). Supported formats include 'fasta', 'fastq', 'sam', 'vcf', 'bed', 'gff'.
        tab_separator: Use tab as the input and output field separator (-t flag).
        include_header: Include header in the output for formats like VCF/SAM (-H flag).
        variables: A list of 'var=value' strings to define AWK variables (-v flag).
        field_separator: The input field separator string (-F flag).

    Returns:
        A dictionary containing the executed command, stdout, stderr, and a list of output files.
    """
    # --- Input Validation ---
    if not program and not program_file:
        raise ValueError("Either 'program' or 'program_file' must be provided.")
    if program and program_file:
        raise ValueError("'program' and 'program_file' are mutually exclusive and cannot be used together.")

    if program_file:
        if not program_file.is_file():
            raise FileNotFoundError(f"Program file not found: {program_file}")

    if input_files:
        for file_path in input_files:
            if not file_path.is_file():
                raise FileNotFoundError(f"Input file not found: {file_path}")

    VALID_FORMATS = {"fasta", "fastq", "sam", "vcf", "bed", "gff"}
    if format and format.lower() not in VALID_FORMATS:
        raise ValueError(f"Invalid format '{format}'. Must be one of {VALID_FORMATS}.")

    if variables:
        for var in variables:
            if "=" not in var:
                raise ValueError(f"Invalid variable assignment '{var}'. Must be in 'var=value' format.")

    # --- Command Construction ---
    cmd = ["bioawk"]

    if format:
        cmd.extend(["-c", format.lower()])
    if tab_separator:
        cmd.append("-t")
    if include_header:
        cmd.append("-H")
    if field_separator:
        cmd.extend(["-F", field_separator])
    if variables:
        for var in variables:
            cmd.extend(["-v", var])

    # The program or program_file argument must come after options but before input files.
    if program_file:
        cmd.extend(["-f", str(program_file)])
    elif program:
        cmd.append(program)

    if input_files:
        cmd.extend([str(p) for p in input_files])

    command_executed = " ".join(cmd)

    # --- Subprocess Execution ---
    try:
        if output_file:
            # Create parent directories if they don't exist
            output_file.parent.mkdir(parents=True, exist_ok=True)
            with open(output_file, "w") as f_out:
                result = subprocess.run(
                    cmd,
                    check=True,
                    stdout=f_out,
                    stderr=subprocess.PIPE,
                    text=True,
                )
            stdout_content = f"Output successfully written to {output_file}"
            output_files_list = [str(output_file)]
        else:
            result = subprocess.run(
                cmd,
                check=True,
                capture_output=True,
                text=True,
            )
            stdout_content = result.stdout
            output_files_list = []

        return {
            "command_executed": command_executed,
            "stdout": stdout_content,
            "stderr": result.stderr,
            "output_files": output_files_list,
        }
    except FileNotFoundError:
        raise RuntimeError("bioawk command not found. Please ensure it is installed and in your system's PATH.")
    except subprocess.CalledProcessError as e:
        # Return a structured error response if the command fails
        return {
            "command_executed": command_executed,
            "stdout": e.stdout or "",
            "stderr": e.stderr or f"bioawk exited with error code {e.returncode}",
            "output_files": [],
            "error": f"Command failed with exit code {e.returncode}"
        }