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