Biomanus/biomni/graph/schema_extractor.py

from __future__ import annotations

import re
from collections import Counter

from biomni.graph.workflow_priors import WORKFLOW_PRIORS, rank_workflow_priors


class ToolSchemaExtractor:
    """Infer lightweight graph semantics from MCP/tool metadata."""

    CAPABILITY_PATTERNS = {
        "normalization": {"normalize", "normalization", "normalise", "normalisation", "size factor", "vst", "rlog"},
        "quantification": {"quant", "quantification", "salmon", "kallisto", "abundance", "transcript"},
        "read_counting": {"featurecounts", "htseq", "count reads", "read count", "summarizeoverlaps"},
        "differential_expression": {
            "deseq2",
            "differential expression",
            "differentially expressed",
            "deg",
            "edger",
            "limma",
            "log2fc",
            "fold change",
        },
        "pathway_enrichment": {
            "gsea",
            "fgsea",
            "clusterprofiler",
            "enrichment",
            "gene set",
            "msigdb",
            "kegg",
            "reactome",
            "gene ontology",
        },
        "data_inspection": {"inspect", "summary", "schema", "preview", "describe", "table"},
        "exact_sequence_analysis": {"exact short-sequence primitive", "exact sequence primitive"},
        "format_conversion": {"convert", "transform", "parse", "merge", "filter", "export"},
        "single_cell_analysis": {"single-cell", "singlecell", "scrna", "scanpy", "seurat", "clustering", "umap"},
        "alignment": {"align", "alignment", "bwa", "star", "map"},
        "variant_calling": {"variant", "vcf", "gatk", "snp", "mutation"},
        "variant_filtering": {"filter variant", "bcftools filter", "vcf filter", "normalize variant"},
        "taxonomic_classification": {"taxonomy", "taxonomic", "kraken", "metaphlan", "centrifuge", "bracken"},
        "genome_annotation": {"prokka", "annotate genome", "genome annotation"},
        "orthology_clustering": {"orthofinder", "ortholog", "orthogroup", "cluster annotation"},
        "annotation": {"annotate", "annotation", "label", "classify"},
        "visualization": {"plot", "visual", "umap", "tsne", "heatmap", "volcano"},
        "database_search": {
            "database",
            "lookup",
            "search",
            "pubmed",
            "ensembl",
            "clinvar",
            "uniprot",
            "biomart",
            "msigdb",
            "mirdb",
            "gtrd",
            "gene set",
        },
        "quality_control": {"qc", "quality", "filter", "trim", "clean"},
        "statistics": {"stat", "qvalue", "pvalue", "regression", "test", "significance"},
    }

    DB_RESOURCE_HINTS = {
        "ensembl",
        "clinvar",
        "uniprot",
        "pubmed",
        "biomart",
        "msigdb",
        "mirdb",
        "gtrd",
        "gene set",
        "ontology",
        "disease ontology",
        "cell atlas",
        "annotation database",
        "metadata query",
        "association table",
    }

    PATHWAY_HINTS = {
        "gsea",
        "fgsea",
        "clusterprofiler",
        "enrichment",
        "gene set",
        "pathway",
        "kegg",
        "reactome",
        "gene ontology",
        "msigdb",
    }

    SINGLE_CELL_HINTS = {
        "single-cell",
        "singlecell",
        "scrna",
        "scanpy",
        "seurat",
        "singlecellexperiment",
        "cell atlas",
        "cell type",
        "umap",
    }

    DIFFERENTIAL_EXPRESSION_HINTS = {
        "differential expression",
        "deseq2",
        "edger",
        "limma",
        "deg",
        "fold change",
        "log2fc",
        "padj",
    }

    SERVER_CAPABILITY_ALLOWLISTS = {
        "bwa": {"alignment"},
        "hisat2": {"alignment"},
        "star": {"alignment"},
        "kallisto": {"quantification"},
        "salmon": {"quantification"},
        "bioconductor-tximport": {"quantification", "format_conversion", "read_counting"},
        "bioconductor-deseq2": {"normalization", "differential_expression", "statistics", "visualization"},
        "cnvkit": {"variant_calling", "variant_filtering", "annotation", "visualization"},
        "samtools": {"alignment", "format_conversion", "data_inspection"},
        "bcftools": {"variant_calling", "variant_filtering", "annotation", "format_conversion"},
        "gatk": {"variant_calling", "variant_filtering"},
        "gatk4": {"variant_calling", "variant_filtering"},
        "freebayes": {"variant_calling"},
        "ucsc-cell-browser": {"single_cell_analysis", "visualization", "data_inspection"},
        "vpt": {"statistics", "annotation", "format_conversion"},
        "r-dwls": {"statistics", "single_cell_analysis"},
        "seacells": {"single_cell_analysis"},
        "harpy": {"single_cell_analysis", "statistics"},
    }

    DATATYPE_PATTERNS = {
        "anndata": {"anndata", "h5ad"},
        "annotated_variant_table": {"annotated variant", "variant annotation table", "annotated vcf"},
        "fastq": {"fastq", "fasta", "reads"},
        "metagenomic_reads": {"metagenomic reads", "metagenomics reads", "microbiome reads"},
        "bam": {"bam", "sam"},
        "vcf": {"vcf", "variant"},
        "reference_genome": {"reference genome", "genome index", "fasta", "genomic fna"},
        "annotation_gff": {"gff", "gtf", "genomic annotation"},
        "quant_table": {"quant.sf", "abundance table", "transcript abundance"},
        "csv": {"csv", "tsv", "table"},
        "json": {"json"},
        "rds": {"rds", "seurat", "spatial object"},
        "count_matrix": {"count", "counts", "matrix", "raw counts"},
        "raw_count_matrix": {
            "raw count",
            "raw count matrix",
            "raw count matrices",
            "raw counts",
            "integer counts",
            "count matrix",
            "count matrices",
        },
        "normalized_count_matrix": {"normalized count", "normalised count", "vst", "rlog", "size factor"},
        "sample_metadata": {"sample metadata", "metadata table", "condition", "phenotype", "group labels"},
        "de_table": {"differential expression", "deg", "log2fc", "padj", "de table"},
        "differential_expression_table": {
            "differential expression table",
            "differential expression",
            "deg",
            "log2fc",
            "padj",
            "de table",
        },
        "significant_gene_list": {"significant genes", "up-regulated genes", "down-regulated genes", "de genes"},
        "ranked_gene_list": {"ranked list", "ranked genes", "prerank", "gene ranking"},
        "gene_list": {"gene list", "genes"},
        "gene_symbol_list": {"gene symbol", "gene symbols", "symbol list"},
        "entrez_gene_list": {"entrez", "entrez id", "entrez gene"},
        "pathway_table": {"pathway", "enrichment", "gsea result"},
        "kegg_enrichment_table": {"kegg enrichment", "enrichkegg", "kegg pathway"},
        "go_enrichment_table": {"go enrichment", "gene ontology enrichment"},
        "merged_pathway_table": {"merged pathway", "shared pathway", "comparative pathway"},
        "taxonomic_profile": {"taxonomic profile", "relative abundance", "taxonomy table"},
        "differential_abundance_table": {"differential abundance", "abundance difference"},
        "genome_assembly": {"assembly", "genome assembly", "contig", "scaffold"},
        "cluster_table": {"cluster table", "cluster annotation", "orthogroup"},
        "marker_gene_table": {"marker gene", "cluster marker", "cell type marker"},
        "database_record": {"database record", "knowledge source", "annotation record", "metadata entry"},
        "sequence_feature": {"motif", "primer", "binding site", "sequence feature"},
        "variant_summary": {"variant summary", "variant report"},
        "report_table": {"final table", "output table", "result table"},
        "image": {"png", "pdf", "plot", "figure", "image"},
        "text": {"txt", "text", "markdown", "report", "summary"},
    }

    OPERATION_SPECS = {
        "database_lookup": {
            "patterns": {"database", "lookup", "search", "disgenet", "omim", "clinvar", "ensembl", "uniprot"},
            "accepts": ["text", "gene_symbol_list"],
            "produces": ["database_record"],
            "capabilities": ["database_search", "annotation"],
            "stage": "input_acquisition",
        },
        "sequence_feature_extraction": {
            "patterns": {"sequence", "fasta", "motif", "primer", "binding site", "reverse complement"},
            "accepts": ["text", "fastq", "reference_genome"],
            "produces": ["sequence_feature"],
            "capabilities": ["alignment", "annotation", "format_conversion"],
            "stage": "analysis",
        },
        "base_counting": {
            "patterns": {"base count", "count bases", "nucleotide count", "composition", "alphabet frequency"},
            "accepts": ["text", "sequence_feature"],
            "produces": ["sequence_feature", "report_table"],
            "capabilities": ["statistics", "data_inspection", "exact_sequence_analysis"],
            "stage": "analysis",
        },
        "gc_content_calculation": {
            "patterns": {"gc content", "percent gc", "gc percentage", "base composition"},
            "accepts": ["text", "sequence_feature"],
            "produces": ["report_table"],
            "capabilities": ["statistics", "data_inspection", "exact_sequence_analysis"],
            "stage": "analysis",
        },
        "reverse_complement": {
            "patterns": {"reverse complement", "revcomp", "reverse-complement", "rc sequence"},
            "accepts": ["text", "sequence_feature"],
            "produces": ["sequence_feature"],
            "capabilities": ["format_conversion", "sequence_feature_extraction", "exact_sequence_analysis"],
            "stage": "analysis",
        },
        "orf_detection": {
            "patterns": {"orf", "open reading frame", "translate", "translation", "coding sequence"},
            "accepts": ["text", "sequence_feature"],
            "produces": ["sequence_feature", "report_table"],
            "capabilities": ["annotation", "sequence_feature_extraction", "exact_sequence_analysis"],
            "stage": "analysis",
        },
        "restriction_digest": {
            "patterns": {"restriction digest", "digest", "restriction enzyme", "cut site", "fragments"},
            "accepts": ["text", "sequence_feature"],
            "produces": ["sequence_feature", "report_table"],
            "capabilities": ["sequence_feature_extraction", "data_inspection", "exact_sequence_analysis"],
            "stage": "analysis",
        },
        "primer_validation": {
            "patterns": {"primer", "pcr", "gibson", "amplicon", "primer pair", "tm", "linearize", "linearized"},
            "accepts": ["text", "sequence_feature"],
            "produces": ["sequence_feature", "report_table"],
            "capabilities": ["sequence_feature_extraction", "alignment", "exact_sequence_analysis"],
            "stage": "analysis",
        },
        "sequence_alignment": {
            "patterns": {"align sequence", "sequence alignment", "local alignment", "exact match", "pairwise alignment"},
            "accepts": ["text", "sequence_feature"],
            "produces": ["sequence_feature"],
            "capabilities": ["alignment"],
            "stage": "analysis",
        },
        "kozak_scoring": {
            "patterns": {"kozak", "translation efficiency", "start codon context", "human cell"},
            "accepts": ["text", "sequence_feature"],
            "produces": ["report_table"],
            "capabilities": ["sequence_feature_extraction", "statistics", "exact_sequence_analysis"],
            "stage": "analysis",
        },
        "quality_control": {
            "patterns": {"fastqc", "multiqc", "quality control", "qc"},
            "accepts": ["fastq"],
            "produces": ["text", "image"],
            "capabilities": ["quality_control"],
            "stage": "preprocessing",
        },
        "quantification": {
            "patterns": {"quantify", "quantification", "salmon", "kallisto", "abundance", "quant.sf"},
            "accepts": ["fastq", "bam", "reference_genome"],
            "produces": ["quant_table"],
            "capabilities": ["quantification"],
            "stage": "analysis",
        },
        "alignment": {
            "patterns": {"align reads", "alignment", "hisat2", "star", "bwa", "bowtie2", "map reads"},
            "accepts": ["fastq", "reference_genome"],
            "produces": ["bam"],
            "capabilities": ["alignment"],
            "stage": "preprocessing",
        },
        "read_counting": {
            "patterns": {"featurecounts", "htseq", "count reads", "read counts", "gene count"},
            "accepts": ["bam", "annotation_gff"],
            "produces": ["raw_count_matrix"],
            "capabilities": ["read_counting"],
            "stage": "analysis",
        },
        "count_normalization": {
            "patterns": {"normalize counts", "normalise counts", "normalization", "size factor", "vst", "rlog"},
            "accepts": ["raw_count_matrix", "sample_metadata"],
            "produces": ["normalized_count_matrix"],
            "capabilities": ["normalization"],
            "stage": "analysis",
        },
        "differential_expression": {
            "patterns": {"differential expression", "deseq2", "deg", "edger", "limma", "log2fc", "padj"},
            "accepts": ["raw_count_matrix", "normalized_count_matrix", "sample_metadata"],
            "produces": ["differential_expression_table"],
            "capabilities": ["differential_expression", "statistics"],
            "constraints": ["requires_replicates", "bulk_rna_seq"],
            "stage": "analysis",
        },
        "gene_filtering": {
            "patterns": {"filter genes", "significant genes", "up-regulated", "down-regulated", "padj"},
            "accepts": ["differential_expression_table"],
            "produces": ["significant_gene_list", "gene_symbol_list"],
            "capabilities": ["format_conversion", "statistics"],
            "stage": "downstream",
        },
        "gene_id_conversion": {
            "patterns": {"gene id", "convert gene", "entrez", "symbol", "ortholog", "id conversion"},
            "accepts": ["gene_symbol_list", "significant_gene_list"],
            "produces": ["entrez_gene_list"],
            "capabilities": ["annotation", "format_conversion", "genome_annotation"],
            "stage": "downstream",
        },
        "kegg_enrichment": {
            "patterns": {"kegg", "pathway enrichment", "enrichkegg", "clusterprofiler"},
            "accepts": ["entrez_gene_list", "ranked_gene_list", "significant_gene_list"],
            "produces": ["kegg_enrichment_table"],
            "capabilities": ["pathway_enrichment"],
            "constraints": ["requires_entrez_id"],
            "stage": "downstream",
        },
        "go_enrichment": {
            "patterns": {"go enrichment", "gene ontology", "enrichgo", "ontology"},
            "accepts": ["entrez_gene_list", "ranked_gene_list", "significant_gene_list"],
            "produces": ["go_enrichment_table"],
            "capabilities": ["pathway_enrichment"],
            "stage": "downstream",
        },
        "pathway_merge": {
            "patterns": {"compare shared pathways", "shared pathways", "merge pathways", "compare pathways"},
            "accepts": ["kegg_enrichment_table", "go_enrichment_table", "pathway_table"],
            "produces": ["merged_pathway_table"],
            "capabilities": ["format_conversion", "data_inspection"],
            "stage": "reporting",
        },
        "csv_export": {
            "patterns": {"csv", "output table", "save table", "export"},
            "accepts": [
                "merged_pathway_table",
                "kegg_enrichment_table",
                "go_enrichment_table",
                "pathway_table",
                "differential_expression_table",
                "variant_summary",
                "taxonomic_profile",
            ],
            "produces": ["report_table"],
            "capabilities": ["format_conversion"],
            "stage": "reporting",
        },
        "variant_annotation": {
            "patterns": {"variant annotation", "vep", "snpeff", "annotate variant", "clinvar"},
            "accepts": ["vcf"],
            "produces": ["annotated_variant_table"],
            "capabilities": ["annotation", "variant_filtering"],
            "stage": "analysis",
        },
        "variant_summary": {
            "patterns": {"variant summary", "summarize variants", "variant report"},
            "accepts": ["annotated_variant_table", "vcf"],
            "produces": ["variant_summary"],
            "capabilities": ["data_inspection", "format_conversion"],
            "stage": "reporting",
        },
        "single_cell_clustering": {
            "patterns": {"single-cell", "singlecell", "scrna", "scanpy", "seurat", "clustering", "umap"},
            "accepts": ["anndata", "rds"],
            "produces": ["cluster_table"],
            "capabilities": ["single_cell_analysis"],
            "stage": "analysis",
        },
        "marker_gene_detection": {
            "patterns": {"marker genes", "cluster markers", "cell type markers"},
            "accepts": ["anndata", "cluster_table"],
            "produces": ["marker_gene_table", "gene_symbol_list"],
            "capabilities": ["single_cell_analysis", "statistics"],
            "stage": "downstream",
        },
        "taxonomic_classification": {
            "patterns": {"taxonomy", "taxonomic", "kraken", "metaphlan", "centrifuge", "bracken"},
            "accepts": ["fastq", "metagenomic_reads"],
            "produces": ["taxonomic_profile"],
            "capabilities": ["taxonomic_classification"],
            "stage": "analysis",
        },
        "differential_abundance": {
            "patterns": {"differential abundance", "abundance difference", "compare abundance"},
            "accepts": ["taxonomic_profile", "sample_metadata"],
            "produces": ["differential_abundance_table"],
            "capabilities": ["statistics"],
            "stage": "downstream",
        },
    }

    STAGE_RULES = [
        ("input_acquisition", {"database_search"}),
        (
            "preprocessing",
            {"data_inspection", "format_conversion", "quality_control", "alignment", "normalization"},
        ),
        (
            "analysis",
            {
                "quantification",
                "read_counting",
                "differential_expression",
                "single_cell_analysis",
                "variant_calling",
                "variant_filtering",
                "taxonomic_classification",
                "orthology_clustering",
                "statistics",
            },
        ),
        ("downstream", {"pathway_enrichment", "annotation", "genome_annotation"}),
        ("reporting", {"visualization"}),
    ]

    STOPWORDS = {
        "the",
        "and",
        "for",
        "with",
        "from",
        "that",
        "this",
        "into",
        "using",
        "data",
        "tool",
        "file",
        "files",
        "analysis",
        "result",
        "results",
        "input",
        "output",
    }

    CURATED_SERVER_SEMANTICS = {
        "salmon": {
            "capabilities": ["quantification"],
            "consumes": ["fastq", "reference_genome"],
            "produces": ["quant_table"],
            "operations": ["quantification"],
            "stage": "analysis",
        },
        "bioconductor-tximport": {
            "capabilities": ["quantification", "format_conversion"],
            "consumes": ["quant_table", "annotation_gff"],
            "produces": ["count_matrix"],
            "operations": ["read_counting"],
            "stage": "analysis",
        },
        "bioconductor-deseq2": {
            "capabilities": ["normalization", "differential_expression", "statistics", "visualization"],
            "consumes": ["raw_count_matrix", "sample_metadata", "count_matrix"],
            "produces": ["differential_expression_table", "de_table", "image"],
            "operations": ["count_normalization", "differential_expression", "gene_filtering"],
            "stage": "analysis",
        },
        "hisat2": {
            "capabilities": ["alignment"],
            "consumes": ["fastq", "reference_genome", "annotation_gff"],
            "produces": ["bam"],
            "operations": ["alignment"],
            "stage": "preprocessing",
        },
        "star": {
            "capabilities": ["alignment"],
            "consumes": ["fastq", "reference_genome", "annotation_gff"],
            "produces": ["bam"],
            "operations": ["alignment"],
            "stage": "preprocessing",
        },
        "htseq": {
            "capabilities": ["read_counting"],
            "consumes": ["bam", "annotation_gff"],
            "produces": ["raw_count_matrix", "count_matrix"],
            "operations": ["read_counting"],
            "stage": "analysis",
        },
        "subread": {
            "capabilities": ["alignment", "read_counting"],
            "consumes": ["fastq", "bam", "reference_genome", "annotation_gff"],
            "produces": ["bam", "raw_count_matrix", "count_matrix"],
            "operations": ["alignment", "read_counting"],
            "stage": "analysis",
        },
        "fastqc": {
            "capabilities": ["quality_control"],
            "consumes": ["fastq"],
            "produces": ["text"],
            "operations": ["quality_control"],
            "stage": "preprocessing",
        },
        "multiqc": {
            "capabilities": ["quality_control", "visualization"],
            "consumes": ["text"],
            "produces": ["text", "image"],
            "operations": ["quality_control", "csv_export"],
            "stage": "reporting",
        },
        "gffutils": {
            "capabilities": ["genome_annotation", "format_conversion", "data_inspection"],
            "consumes": ["annotation_gff", "reference_genome"],
            "produces": ["annotation_gff", "gene_list", "gene_symbol_list", "csv"],
            "operations": ["gene_id_conversion", "database_lookup"],
            "stage": "preprocessing",
        },
        "bcftools": {
            "capabilities": ["variant_calling", "variant_filtering", "annotation"],
            "consumes": ["bam", "vcf", "reference_genome"],
            "produces": ["vcf", "annotated_variant_table", "variant_summary"],
            "operations": ["variant_annotation", "variant_summary"],
            "stage": "analysis",
        },
        "seqfu": {
            "capabilities": ["data_inspection", "format_conversion", "statistics"],
            "consumes": ["text", "fastq", "sequence_feature"],
            "produces": ["text", "sequence_feature", "report_table"],
            "operations": ["base_counting", "gc_content_calculation", "reverse_complement", "sequence_feature_extraction"],
            "stage": "analysis",
        },
        "bioawk": {
            "capabilities": ["data_inspection", "format_conversion", "statistics"],
            "consumes": ["text", "fastq", "sequence_feature"],
            "produces": ["text", "sequence_feature", "report_table"],
            "operations": ["base_counting", "gc_content_calculation", "sequence_feature_extraction"],
            "stage": "analysis",
        },
        "primer3": {
            "capabilities": ["sequence_feature_extraction", "alignment", "statistics"],
            "consumes": ["text", "sequence_feature"],
            "produces": ["sequence_feature", "report_table"],
            "operations": ["primer_validation", "sequence_feature_extraction"],
            "stage": "analysis",
        },
        "bioconductor-decipher": {
            "capabilities": ["alignment", "annotation", "sequence_feature_extraction"],
            "consumes": ["text", "sequence_feature", "fastq", "reference_genome"],
            "produces": ["sequence_feature", "report_table"],
            "operations": ["primer_validation", "sequence_alignment", "sequence_feature_extraction"],
            "stage": "analysis",
        },
        "pyrodigal": {
            "capabilities": ["annotation", "sequence_feature_extraction"],
            "consumes": ["text", "sequence_feature", "reference_genome"],
            "produces": ["sequence_feature", "report_table"],
            "operations": ["orf_detection", "sequence_feature_extraction"],
            "stage": "analysis",
        },
        "mummer": {
            "capabilities": ["alignment", "data_inspection"],
            "consumes": ["text", "sequence_feature", "reference_genome"],
            "produces": ["sequence_feature", "report_table"],
            "operations": ["sequence_alignment", "primer_validation"],
            "stage": "analysis",
        },
        "parasail-python": {
            "capabilities": ["alignment", "statistics"],
            "consumes": ["text", "sequence_feature"],
            "produces": ["sequence_feature", "report_table"],
            "operations": ["sequence_alignment", "primer_validation"],
            "stage": "analysis",
        },
        "viennarna": {
            "capabilities": ["sequence_feature_extraction", "statistics"],
            "consumes": ["text", "sequence_feature"],
            "produces": ["sequence_feature", "report_table"],
            "operations": ["kozak_scoring", "sequence_feature_extraction"],
            "stage": "analysis",
        },
        "sequence-operations": {
            "capabilities": ["sequence_feature_extraction", "statistics", "data_inspection", "alignment"],
            "consumes": ["text", "sequence_feature"],
            "produces": ["sequence_feature", "report_table"],
            "operations": ["sequence_feature_extraction"],
            "stage": "analysis",
        },
        "local-bio-cache": {
            "capabilities": ["database_search", "alignment", "data_inspection"],
            "consumes": ["text", "sequence_feature", "reference_genome"],
            "produces": ["database_record", "sequence_feature", "report_table"],
            "operations": ["database_lookup", "sequence_alignment", "sequence_feature_extraction"],
            "stage": "analysis",
        },
        "gatk": {
            "capabilities": ["variant_calling", "variant_filtering"],
            "consumes": ["bam", "reference_genome"],
            "produces": ["vcf"],
            "operations": ["variant_annotation"],
            "stage": "analysis",
        },
        "freebayes": {
            "capabilities": ["variant_calling"],
            "consumes": ["bam", "reference_genome"],
            "produces": ["vcf"],
            "operations": ["variant_annotation"],
            "stage": "analysis",
        },
        "kraken2": {
            "capabilities": ["taxonomic_classification"],
            "consumes": ["fastq", "metagenomic_reads"],
            "produces": ["taxonomic_profile"],
            "operations": ["taxonomic_classification"],
            "stage": "analysis",
        },
        "metaphlan": {
            "capabilities": ["taxonomic_classification"],
            "consumes": ["fastq", "metagenomic_reads"],
            "produces": ["taxonomic_profile"],
            "operations": ["taxonomic_classification"],
            "stage": "analysis",
        },
        "prokka": {
            "capabilities": ["genome_annotation"],
            "consumes": ["genome_assembly"],
            "produces": ["annotation_gff", "gene_list"],
            "operations": ["gene_id_conversion"],
            "stage": "analysis",
        },
        "orthofinder": {
            "capabilities": ["orthology_clustering"],
            "consumes": ["gene_list"],
            "produces": ["cluster_table"],
            "operations": ["gene_id_conversion"],
            "stage": "analysis",
        },
    }

    CURATED_TOOL_SEMANTICS = {
        ("sequence-operations", "gc_content"): {
            "capabilities": ["statistics", "data_inspection"],
            "consumes": ["text", "sequence_feature"],
            "produces": ["report_table"],
            "operations": ["base_counting", "gc_content_calculation"],
            "stage": "analysis",
            "override": True,
        },
        ("sequence-operations", "reverse_complement"): {
            "capabilities": ["format_conversion", "sequence_feature_extraction", "exact_sequence_analysis"],
            "consumes": ["text", "sequence_feature"],
            "produces": ["sequence_feature"],
            "operations": ["reverse_complement"],
            "stage": "analysis",
            "override": True,
        },
        ("sequence-operations", "find_orfs"): {
            "capabilities": ["annotation", "sequence_feature_extraction", "exact_sequence_analysis"],
            "consumes": ["text", "sequence_feature"],
            "produces": ["sequence_feature", "report_table"],
            "operations": ["orf_detection"],
            "stage": "analysis",
            "override": True,
        },
        ("sequence-operations", "amino_acid_at_position_in_longest_orf"): {
            "capabilities": ["annotation", "sequence_feature_extraction", "exact_sequence_analysis"],
            "consumes": ["text", "sequence_feature"],
            "produces": ["sequence_feature", "report_table"],
            "operations": ["orf_detection"],
            "stage": "analysis",
            "override": True,
        },
        ("sequence-operations", "restriction_digest"): {
            "capabilities": ["sequence_feature_extraction", "data_inspection", "exact_sequence_analysis"],
            "consumes": ["text", "sequence_feature"],
            "produces": ["sequence_feature", "report_table"],
            "operations": ["restriction_digest"],
            "stage": "analysis",
            "override": True,
        },
        ("sequence-operations", "pcr_primer_analysis"): {
            "capabilities": ["sequence_feature_extraction", "alignment", "exact_sequence_analysis"],
            "consumes": ["text", "sequence_feature"],
            "produces": ["sequence_feature", "report_table"],
            "operations": ["primer_validation", "sequence_alignment"],
            "stage": "analysis",
            "override": True,
        },
        ("sequence-operations", "kozak_score"): {
            "capabilities": ["sequence_feature_extraction", "statistics", "exact_sequence_analysis"],
            "consumes": ["text", "sequence_feature"],
            "produces": ["report_table"],
            "operations": ["kozak_scoring"],
            "stage": "analysis",
            "override": True,
        },
    }

    def tokenize(self, text: str) -> set[str]:
        return {
            token
            for token in re.findall(r"[a-z0-9_\-\.]{2,}", (text or "").lower())
            if token not in self.STOPWORDS
        }

    def extract_tool_semantics(self, tool: dict, server_name: str | None = None) -> dict:
        name = str(tool.get("name", ""))
        description = str(tool.get("description", ""))
        params = tool.get("inputSchema", {}).get("properties", {}) or tool.get("parameters", {}) or {}
        param_text = " ".join(
            f"{param_name} {spec.get('type', '')} {spec.get('description', '')}" for param_name, spec in params.items()
        )
        combined = " ".join(part for part in [server_name or "", name, description, param_text] if part)
        tokens = self.tokenize(combined)

        capabilities = self._match_labels(combined, self.CAPABILITY_PATTERNS)
        consumes = self._infer_consumed_types(combined, params)
        produces = self._infer_produced_types(combined, params, capabilities)
        operations = self._infer_tool_operations(combined, capabilities)
        constraints = self._infer_constraints(combined)
        stage = self._infer_stage(capabilities, name, description)

        curated = self._curated_semantics(server_name, name)
        if curated:
            if curated.get("override"):
                capabilities = list(curated.get("capabilities", capabilities))
                consumes = list(curated.get("consumes", consumes))
                produces = list(curated.get("produces", produces))
                operations = list(curated.get("operations", operations))
                constraints = list(curated.get("constraints", constraints))
            else:
                capabilities = self._merge_labels(capabilities, curated.get("capabilities", []))
                consumes = self._merge_labels(consumes, curated.get("consumes", []))
                produces = self._merge_labels(produces, curated.get("produces", []))
                operations = self._merge_labels(operations, curated.get("operations", []))
                constraints = self._merge_labels(constraints, curated.get("constraints", []))
            stage = curated.get("stage") or stage

        operation_semantics = self._merge_operation_semantics(operations, consumes, produces, constraints)
        consumes = operation_semantics["consumes"]
        produces = operation_semantics["produces"]
        constraints = operation_semantics["constraints"]

        return {
            "name": name,
            "description": description,
            "keywords": sorted(tokens),
            "capabilities": capabilities,
            "consumes": consumes,
            "produces": produces,
            "operations": operations,
            "constraints": constraints,
            "stage": stage,
            "text": combined,
        }

    def extract_server_semantics(self, server_entry: dict) -> dict:
        tool_semantics = [
            self.extract_tool_semantics(tool, server_name=server_entry.get("name")) for tool in server_entry.get("tools", [])
        ]
        capability_counts = Counter()
        datatype_counts = Counter()
        operation_counts = Counter()
        constraint_counts = Counter()
        stage_counts = Counter()
        keywords: set[str] = self.tokenize(
            " ".join(
                [
                    server_entry.get("name", ""),
                    server_entry.get("summary", ""),
                    " ".join(server_entry.get("keywords", [])),
                ]
            )
        )
        for tool in tool_semantics:
            capability_counts.update(tool["capabilities"])
            datatype_counts.update(tool["consumes"])
            datatype_counts.update(tool["produces"])
            operation_counts.update(tool.get("operations", []))
            constraint_counts.update(tool.get("constraints", []))
            if tool["stage"]:
                stage_counts.update([tool["stage"]])
            keywords.update(tool["keywords"])

        capabilities = [item for item, _ in capability_counts.most_common()]
        capabilities = self._sanitize_server_capabilities(server_entry, capabilities, tool_semantics)
        stage_counts = Counter(
            tool["stage"]
            for tool in tool_semantics
            if tool.get("stage") and any(cap in capabilities for cap in tool.get("capabilities", []))
        )

        return {
            "capabilities": capabilities,
            "datatypes": [item for item, _ in datatype_counts.most_common()],
            "operations": [item for item, _ in operation_counts.most_common()],
            "constraints": [item for item, _ in constraint_counts.most_common()],
            "stages": [item for item, _ in stage_counts.most_common()],
            "keywords": sorted(keywords),
            "tool_semantics": tool_semantics,
        }

    def _curated_semantics(self, server_name: str | None, tool_name: str | None = None) -> dict:
        if not server_name:
            return {}
        server_key = server_name.lower()
        tool_key = (tool_name or "").lower()
        if tool_key:
            curated_tool = self.CURATED_TOOL_SEMANTICS.get((server_key, tool_key), {})
            if curated_tool:
                return curated_tool
        curated = self.CURATED_SERVER_SEMANTICS.get(server_key, {})
        if not curated and server_key.startswith("bioconductor-"):
            curated = self.CURATED_SERVER_SEMANTICS.get(server_key.replace("_", "-"), {})
        return curated

    def _merge_labels(self, primary: list[str], additions: list[str]) -> list[str]:
        merged = []
        for item in list(primary or []) + list(additions or []):
            if item and item not in merged:
                merged.append(item)
        return merged

    def _sanitize_server_capabilities(
        self,
        server_entry: dict,
        capabilities: list[str],
        tool_semantics: list[dict],
    ) -> list[str]:
        server_name = str(server_entry.get("name", "")).lower()
        category = str(server_entry.get("category", "")).lower()
        combined_text = " ".join(
            [
                server_name,
                str(server_entry.get("summary", "")),
                " ".join(server_entry.get("keywords", [])),
                " ".join(tool.get("text", "") for tool in tool_semantics),
            ]
        ).lower()

        sanitized = list(capabilities)

        if "database_search" in sanitized and not any(hint in combined_text for hint in self.DB_RESOURCE_HINTS):
            sanitized = [cap for cap in sanitized if cap != "database_search"]

        if "pathway_enrichment" in sanitized and not any(hint in combined_text for hint in self.PATHWAY_HINTS):
            sanitized = [cap for cap in sanitized if cap != "pathway_enrichment"]

        if "single_cell_analysis" in sanitized and not any(hint in combined_text for hint in self.SINGLE_CELL_HINTS):
            sanitized = [cap for cap in sanitized if cap != "single_cell_analysis"]

        if "differential_expression" in sanitized and not any(
            hint in combined_text for hint in self.DIFFERENTIAL_EXPRESSION_HINTS
        ):
            sanitized = [cap for cap in sanitized if cap != "differential_expression"]

        if category in {"single_cell", "transcriptomics"} and "database_search" in sanitized:
            if not any(hint in combined_text for hint in {"cell atlas", "metadata query", "annotation database"}):
                sanitized = [cap for cap in sanitized if cap != "database_search"]

        allowed = self.SERVER_CAPABILITY_ALLOWLISTS.get(server_name)
        if allowed:
            sanitized = [cap for cap in sanitized if cap in allowed]

        return sanitized

    def infer_query_semantics(self, query_context: dict) -> dict:
        text_parts = [
            query_context.get("question", ""),
            query_context.get("original_query", ""),
            query_context.get("rewritten_query", ""),
            query_context.get("retrieval_query", ""),
            " ".join(query_context.get("subtasks", [])),
            " ".join(query_context.get("keywords", [])),
            " ".join(query_context.get("categories", [])),
        ]
        text = " ".join(part for part in text_parts if part)
        capabilities = self._match_labels(text, self.CAPABILITY_PATTERNS)
        datatypes = self._match_labels(text, self.DATATYPE_PATTERNS)
        operations = self._infer_operations(text)
        for operation in operations:
            capabilities = self._merge_labels(capabilities, self.OPERATION_SPECS.get(operation, {}).get("capabilities", []))
        stages = self._infer_query_stages(query_context, capabilities)
        operation_tokens = {token for operation in operations for token in operation.split("_")}
        constraints = self._infer_constraints(text)
        task_spec = self.infer_task_spec(query_context, text, datatypes, operations, constraints)
        operation_anchors = self._merge_labels(operations, task_spec.get("required_operations", []))
        return {
            "keywords": sorted(self.tokenize(text) | operation_tokens),
            "capabilities": capabilities,
            "datatypes": datatypes,
            "operations": operations,
            "operation_anchors": operation_anchors,
            "stages": stages,
            "organism": self._infer_organism(text),
            "branch_labels": self._infer_branch_labels(text),
            "workflow_type": self._infer_workflow_type(text, operations),
            "output_schema": self._infer_output_schema(text),
            "constraints": constraints,
            "task_spec": task_spec,
        }

    def infer_task_spec(
        self,
        query_context: dict,
        text: str,
        datatypes: list[str],
        operations: list[str],
        constraints: list[str],
    ) -> dict:
        categories = {str(category).lower() for category in query_context.get("categories", [])}
        capability_set = set(self._match_labels(text, self.CAPABILITY_PATTERNS))
        datatype_set = set(datatypes)
        operation_set = set(operations)
        ranked_priors = rank_workflow_priors(
            WORKFLOW_PRIORS,
            text=text,
            categories=categories,
            capabilities=capability_set,
            datatypes=datatype_set,
            operations=operation_set,
        )

        input_types = list(datatypes)
        target_output = {"semantic_type": "report_table"}
        required_operations = list(operations)
        assay = None

        # These priors are graph-structured soft constraints. They expose
        # operation anchors and preferred type signatures, but the concrete
        # operation path is still selected later by constrained subgraph search.
        for prior, _score in ranked_priors:
            input_types = self._merge_labels(input_types, list(prior.input_types))
            required_operations = self._merge_labels(required_operations, list(prior.operation_hints))
            assay = assay or prior.assay

        if ranked_priors:
            best_prior, best_score = ranked_priors[0]
            target_output = {"semantic_type": best_prior.target_output}
            for prior, score in ranked_priors[1:]:
                if prior.target_output == "report_table" and score >= best_score - 1.0:
                    target_output = {"semantic_type": "report_table"}
                    break

        organism = self._infer_organism(text)
        constraint_payload = {
            "labels": constraints,
            "organism": organism,
            "assay": assay,
        }
        if organism:
            constraint_payload["labels"] = self._merge_labels(constraint_payload["labels"], [f"organism_{organism}"])

        return {
            "input_types": input_types or ["text"],
            "target_output": target_output,
            "constraints": constraint_payload,
            "required_operations": required_operations,
            "workflow_type": self._infer_workflow_type(text, required_operations),
            "workflow_priors": [
                {"name": prior.name, "score": score, "target_output": prior.target_output}
                for prior, score in ranked_priors
            ],
        }

    def _match_labels(self, text: str, pattern_map: dict[str, set[str]]) -> list[str]:
        lowered = (text or "").lower()
        matched = []
        for label, hints in pattern_map.items():
            if any(hint in lowered for hint in hints):
                matched.append(label)
        return matched

    def _infer_operations(self, text: str) -> list[str]:
        lowered = (text or "").lower()
        operations = []
        for operation, spec in self.OPERATION_SPECS.items():
            hints = spec.get("patterns", set())
            if any(hint in lowered for hint in hints):
                operations.append(operation)
        return operations

    def _infer_tool_operations(self, text: str, capabilities: list[str]) -> list[str]:
        operations = set(self._infer_operations(text))
        capability_set = set(capabilities)
        capability_defaults = {
            "quality_control": ["quality_control"],
            "alignment": ["alignment"],
            "quantification": ["quantification"],
            "read_counting": ["read_counting"],
            "normalization": ["count_normalization"],
            "differential_expression": ["differential_expression"],
            "database_search": ["database_lookup"],
            "taxonomic_classification": ["taxonomic_classification"],
            "single_cell_analysis": ["single_cell_clustering"],
            "variant_filtering": ["variant_annotation"],
        }
        for capability, defaults in capability_defaults.items():
            if capability in capability_set:
                operations.update(defaults)
        return sorted(operations)

    def _infer_constraints(self, text: str) -> list[str]:
        lowered = (text or "").lower()
        constraints = []
        if any(token in lowered for token in ("mouse", "murine", "mus musculus")):
            constraints.append("organism_mouse")
        if any(token in lowered for token in ("human", "homo sapiens")):
            constraints.append("organism_human")
        if any(token in lowered for token in ("bulk rna", "bulk rna-seq", "deseq2", "raw counts")):
            constraints.append("bulk_rna_seq")
        if any(token in lowered for token in ("integer count", "raw count", "raw counts")):
            constraints.append("requires_raw_integer_counts")
        if "entrez" in lowered:
            constraints.append("requires_entrez_id")
        if any(token in lowered for token in ("replicate", "replicates")):
            constraints.append("requires_replicates")
        if any(token in lowered for token in ("paired-end", "paired end", "read1", "read2")):
            constraints.append("paired_end_reads")
        return constraints

    def _merge_operation_semantics(
        self,
        operations: list[str],
        consumes: list[str],
        produces: list[str],
        constraints: list[str],
    ) -> dict:
        merged_consumes = list(consumes)
        merged_produces = list(produces)
        merged_constraints = list(constraints)
        for operation in operations:
            spec = self.OPERATION_SPECS.get(operation, {})
            merged_consumes = self._merge_labels(merged_consumes, spec.get("accepts", []))
            merged_produces = self._merge_labels(merged_produces, spec.get("produces", []))
            merged_constraints = self._merge_labels(merged_constraints, spec.get("constraints", []))
        return {
            "consumes": merged_consumes,
            "produces": merged_produces,
            "constraints": merged_constraints,
        }

    def _infer_organism(self, text: str) -> str | None:
        lowered = (text or "").lower()
        if any(token in lowered for token in ("mouse", "murine", "mus musculus")):
            return "mouse"
        if any(token in lowered for token in ("human", "homo sapiens")):
            return "human"
        return None

    def _infer_branch_labels(self, text: str) -> list[str]:
        labels = []
        for pattern in (r"\b5xfad\b", r"\b3xtg[-_ ]?ad\b", r"\bps?[30o][0o]1s\b"):
            for match in re.findall(pattern, text or "", flags=re.IGNORECASE):
                normalized = match.replace(" ", "-")
                if normalized.lower() == "3xtg-ad":
                    normalized = "3xTG-AD"
                elif normalized.lower() == "5xfad":
                    normalized = "5xFAD"
                elif normalized.lower() in {"ps301s", "ps3o1s", "p301s", "p3o1s"}:
                    normalized = "PS301S"
                if normalized not in labels:
                    labels.append(normalized)
        return labels

    def _infer_workflow_type(self, text: str, operations: list[str]) -> str:
        lowered = (text or "").lower()
        has_comparison = any(token in lowered for token in ("compare", "comparative", "shared", "across"))
        if has_comparison and len(self._infer_branch_labels(text)) >= 2:
            return "fork_join_comparative_analysis"
        if len(operations) > 1:
            return "linear_typed_transformation"
        return "single_step"

    def _infer_output_schema(self, text: str) -> dict:
        lowered = (text or "").lower()
        if "pathway" in lowered and any(token in lowered for token in ("5xfad", "3xtg", "ps301s", "p301s")):
            return {
                "columns": ["pathway", "5xFAD_pvalue", "3xTG_AD_pvalue", "PS3O1S_pvalue"],
                "join_key": "pathway",
            }
        if "csv" in lowered:
            return {"format": "csv"}
        return {}

    def _infer_consumed_types(self, text: str, params: dict) -> list[str]:
        lowered = (text or "").lower()
        matched = set(self._match_labels(lowered, self.DATATYPE_PATTERNS))
        for param_name, spec in params.items():
            blob = f"{param_name} {spec.get('description', '')}".lower()
            matched.update(self._match_labels(blob, self.DATATYPE_PATTERNS))
            if any(keyword in param_name.lower() for keyword in ("input", "file", "path", "filename")):
                matched.update(self._match_labels(blob, self.DATATYPE_PATTERNS))
        return sorted(matched)

    def _infer_produced_types(self, text: str, params: dict, capabilities: list[str]) -> list[str]:
        lowered = (text or "").lower()
        produced = set()
        for label, hints in self.DATATYPE_PATTERNS.items():
            if any(hint in lowered for hint in hints) and any(
                marker in lowered for marker in ("save", "write", "export", "output", "generate", "produce")
            ):
                produced.add(label)
        for param_name, spec in params.items():
            blob = f"{param_name} {spec.get('description', '')}".lower()
            if any(keyword in param_name.lower() for keyword in ("output", "save", "dest", "dir")):
                produced.update(self._match_labels(blob, self.DATATYPE_PATTERNS))
        if "pathway_enrichment" in capabilities:
            produced.add("pathway_table")
        if "differential_expression" in capabilities:
            produced.add("de_table")
        if "visualization" in capabilities:
            produced.add("image")
        return sorted(produced)

    def _infer_stage(self, capabilities: list[str], name: str, description: str) -> str:
        lowered = f"{name} {description}".lower()
        if any(token in lowered for token in ("plot", "figure", "visual")):
            return "reporting"
        # Prefer later analytical stages over generic preprocessing hints.
        priority_rules = [
            ("downstream", {"pathway_enrichment", "annotation", "genome_annotation"}),
            (
                "analysis",
                {
                    "quantification",
                    "read_counting",
                    "differential_expression",
                    "single_cell_analysis",
                    "variant_calling",
                    "variant_filtering",
                    "taxonomic_classification",
                    "orthology_clustering",
                    "statistics",
                },
            ),
            ("preprocessing", {"data_inspection", "format_conversion", "quality_control", "alignment", "normalization"}),
            ("input_acquisition", {"database_search"}),
        ]
        for stage, stage_capabilities in priority_rules:
            if any(capability in stage_capabilities for capability in capabilities):
                return stage
        for stage, stage_capabilities in self.STAGE_RULES:
            if any(capability in stage_capabilities for capability in capabilities):
                return stage
        return "analysis"

    def _infer_query_stages(self, query_context: dict, capabilities: list[str]) -> list[str]:
        text = " ".join(query_context.get("subtasks", [])) + " " + query_context.get("rewritten_query", "")
        lowered = text.lower()
        stages = []
        if any(token in lowered for token in ("inspect", "load", "schema", "preview", "normalize", "normalise")):
            stages.append("preprocessing")
        if any(
            token in lowered
            for token in (
                "analyze",
                "differential",
                "compare",
                "cluster",
                "quantify",
                "count",
                "variant",
                "taxonomy",
            )
        ):
            stages.append("analysis")
        if any(token in lowered for token in ("gsea", "enrichment", "pathway")):
            stages.append("downstream")
        if any(token in lowered for token in ("summarize", "report", "save plot", "visualize")):
            stages.append("reporting")
        if not stages:
            for stage, stage_capabilities in self.STAGE_RULES:
                if any(capability in stage_capabilities for capability in capabilities):
                    stages.append(stage)
        return stages