Source code for tooluniverse.nextstrain_tool
# nextstrain_tool.py
"""
Nextstrain REST API tool for ToolUniverse.
Nextstrain is an open-source project to harness the scientific and
public health potential of pathogen genome data. It provides real-time
tracking of evolving pathogens through phylogenetic analysis.
API: https://nextstrain.org/charon
No authentication required. Free for all use.
"""
import requests
from typing import Dict, Any
from .base_tool import BaseTool
from .tool_registry import register_tool
NEXTSTRAIN_BASE_URL = "https://nextstrain.org/charon"
# Default number of dataset paths listed per pathogen. Callers can raise this
# (or pass 0 for "no cap") via the ``datasets_per_pathogen`` argument.
DEFAULT_DATASETS_PER_PATHOGEN = 10
[docs]
@register_tool("NextstrainTool")
class NextstrainTool(BaseTool):
"""
Tool for querying Nextstrain, the pathogen evolution tracker.
Provides access to phylogenetic datasets for various pathogens
including influenza, SARS-CoV-2, Zika, Ebola, Dengue, and more.
Returns metadata and phylogenetic tree data.
No authentication required.
"""
[docs]
def __init__(self, tool_config: Dict[str, Any]):
super().__init__(tool_config)
self.timeout = tool_config.get("timeout", 45)
self.endpoint_type = tool_config.get("fields", {}).get(
"endpoint_type", "list_datasets"
)
[docs]
def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Execute the Nextstrain API call."""
try:
return self._dispatch(arguments)
except requests.exceptions.Timeout:
return {
"status": "error",
"error": f"Nextstrain API request timed out after {self.timeout} seconds",
}
except requests.exceptions.ConnectionError:
return {
"status": "error",
"error": "Failed to connect to Nextstrain API. Check network connectivity.",
}
except requests.exceptions.HTTPError as e:
return {
"status": "error",
"error": f"Nextstrain API HTTP error: {e.response.status_code}",
}
except Exception as e:
return {
"status": "error",
"error": f"Unexpected error querying Nextstrain: {str(e)}",
}
[docs]
def _dispatch(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Route to appropriate endpoint based on config."""
if self.endpoint_type == "list_datasets":
return self._list_datasets(arguments)
elif self.endpoint_type == "get_dataset":
return self._get_dataset(arguments)
else:
return {
"status": "error",
"error": f"Unknown endpoint_type: {self.endpoint_type}",
}
[docs]
def _list_datasets(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""List available Nextstrain pathogen datasets."""
pathogen_filter = (arguments.get("pathogen") or "").lower()
# dataset_limit is the paging spelling of the same per-pathogen cap.
raw_limit = arguments.get("datasets_per_pathogen")
if raw_limit is None or raw_limit == "":
raw_limit = arguments.get("dataset_limit")
if raw_limit is None or raw_limit == "":
per_pathogen_limit = DEFAULT_DATASETS_PER_PATHOGEN
else:
try:
per_pathogen_limit = int(raw_limit)
except (TypeError, ValueError):
return {
"status": "error",
"error": (
"datasets_per_pathogen must be a non-negative integer "
"(0 means return every dataset)."
),
}
if per_pathogen_limit < 0:
return {
"status": "error",
"error": (
"datasets_per_pathogen must be a non-negative integer "
"(0 means return every dataset)."
),
}
raw_offset = arguments.get("dataset_offset")
try:
dataset_offset = 0 if raw_offset in (None, "") else int(raw_offset)
except (TypeError, ValueError):
dataset_offset = -1
if dataset_offset < 0:
return {
"status": "error",
"error": "dataset_offset must be a non-negative integer.",
}
url = f"{NEXTSTRAIN_BASE_URL}/getAvailable"
response = requests.get(url, timeout=self.timeout)
response.raise_for_status()
raw = response.json()
all_datasets = raw.get("datasets", [])
# Group by pathogen (first segment of the request path)
pathogen_groups = {}
for ds in all_datasets:
request_path = ds.get("request", "")
if not request_path:
continue
pathogen = request_path.split("/")[0]
pathogen_groups.setdefault(pathogen, []).append(request_path)
# Filter by pathogen if specified
if pathogen_filter:
filtered = {}
for p, paths in pathogen_groups.items():
if pathogen_filter in p.lower():
filtered[p] = paths
pathogen_groups = filtered
# Build response
results = []
for pathogen, paths in sorted(pathogen_groups.items()):
ordered = sorted(paths)
remaining = ordered[dataset_offset:]
listed = (
remaining if per_pathogen_limit == 0 else remaining[:per_pathogen_limit]
)
next_offset = dataset_offset + len(listed)
has_more = next_offset < len(ordered)
entry = {
"pathogen": pathogen,
"dataset_count": len(ordered),
"datasets": listed,
"datasets_listed": len(listed),
"datasets_returned": len(listed),
"dataset_offset": dataset_offset,
"has_more": has_more,
"next_dataset_offset": next_offset if has_more else None,
}
if len(listed) < len(ordered):
entry["datasets_truncated"] = True
results.append(entry)
# 'total_datasets' is the true catalogue size (sum of every pathogen's
# dataset_count), never the size of the possibly-truncated listings.
total_datasets = sum(r["dataset_count"] for r in results)
listed_datasets = sum(r["datasets_listed"] for r in results)
truncated = listed_datasets < total_datasets
response = {
"status": "success",
"data": results,
"truncated": truncated,
"metadata": {
"source": "Nextstrain",
"total_pathogens": len(results),
"total_datasets": total_datasets,
"listed_datasets": listed_datasets,
"returned_datasets": listed_datasets,
"datasets_per_pathogen": per_pathogen_limit,
"dataset_offset": dataset_offset,
"groups_with_more": sum(1 for r in results if r["has_more"]),
"truncated": truncated,
"filter": pathogen_filter or "(none)",
"endpoint": "list_datasets",
},
}
if truncated:
largest = max(r["dataset_count"] for r in results)
if per_pathogen_limit:
scope = f"capped at {per_pathogen_limit} entries"
else:
scope = f"starts at dataset_offset={dataset_offset}"
note = (
f"Listed {listed_datasets} of {total_datasets} datasets: each pathogen's "
f"'datasets' array is {scope}, while its "
f"'dataset_count' reports the true number available. To retrieve the rest, "
f"re-run with datasets_per_pathogen=0 (no cap) or a higher cap "
f"(datasets_per_pathogen={largest} covers the largest pathogen), page "
f"with dataset_offset set to each pathogen's 'next_dataset_offset' "
f"while 'has_more' is true, and/or narrow the query with the "
f"'pathogen' filter."
)
response["truncation_note"] = note
response["metadata"]["truncation_note"] = note
return response
[docs]
def _get_dataset(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Get metadata and tree summary for a Nextstrain dataset."""
dataset = arguments.get("dataset", "")
if not dataset:
return {
"status": "error",
"error": "dataset parameter is required (e.g., 'zika', 'ebola', 'flu/seasonal/h3n2/ha/2y')",
}
url = f"{NEXTSTRAIN_BASE_URL}/getDataset"
params = {"prefix": dataset}
response = requests.get(url, params=params, timeout=self.timeout)
response.raise_for_status()
raw = response.json()
meta = raw.get("meta", {})
tree = raw.get("tree", {})
# Count sequences (leaves in tree)
def count_leaves(node):
if not isinstance(node, dict):
return 0
children = node.get("children", [])
if not children:
return 1
return sum(count_leaves(c) for c in children)
num_sequences = count_leaves(tree)
# Extract tree root attributes
root_attrs = tree.get("node_attrs", {})
root_info = {}
for key, val in root_attrs.items():
if isinstance(val, dict) and "value" in val:
root_info[key] = val["value"]
elif not isinstance(val, dict):
root_info[key] = val
# Data provenance
provenance = meta.get("data_provenance", [])
prov_names = []
for p in provenance:
if isinstance(p, dict):
prov_names.append(p.get("name", ""))
# Maintainers
maintainers = []
for m in meta.get("maintainers", []):
if isinstance(m, dict):
maintainers.append(m.get("name", ""))
result = {
"dataset": dataset,
"title": meta.get("title", ""),
"updated": meta.get("updated", ""),
"build_url": meta.get("build_url", ""),
"num_sequences": num_sequences,
"data_provenance": prov_names,
"maintainers": maintainers,
"root_attributes": root_info,
}
# Color-by options (complete list — these are short keys, so nothing is
# dropped; the count is reported alongside so callers can verify.)
colorings = meta.get("colorings", [])
if colorings:
coloring_keys = [
c.get("key", "") for c in colorings if isinstance(c, dict)
]
result["available_colorings"] = coloring_keys
result["available_colorings_count"] = len(coloring_keys)
result["available_colorings_complete"] = True
return {
"status": "success",
"data": result,
"metadata": {
"source": "Nextstrain",
"query": dataset,
"version": raw.get("version", ""),
"endpoint": "get_dataset",
},
}