Source code for tooluniverse.rcsb_advanced_search_tool

# rcsb_advanced_search_tool.py
"""
RCSB PDB Advanced Search Tool for ToolUniverse.

Provides attribute-based filtering of PDB structures using the RCSB Search API v2.
Supports filtering by organism, resolution, experimental method, molecular weight,
polymer description, and deposition date. Goes beyond simple text/sequence search
to enable complex multi-criterion structure discovery.

API: https://search.rcsb.org/
No authentication required. Free public access.
"""

import requests
from typing import Dict, Any
from .base_tool import BaseTool
from .tool_registry import register_tool

RCSB_SEARCH_URL = "https://search.rcsb.org/rcsbsearch/v2/query"

# Fix-R4B-2: `rows` was clamped to 50 and the paginate window start was
# hard-coded to 0, so results 51+ were structurally unreachable -- a search
# matching 81,865 entries could only ever expose its first 50, and a caller
# asking for rows=200 silently got 50 back with nothing in the response
# saying the request had been reduced. `start` makes the rest of the result
# set reachable; _paginate() reports the clamp instead of hiding it.
MAX_ROWS_PER_PAGE = 50




[docs]
@register_tool("RCSBAdvancedSearchTool")
class RCSBAdvancedSearchTool(BaseTool):
    """
    Advanced attribute-based search of the RCSB Protein Data Bank.

    Enables complex queries combining organism, resolution, experimental method,
    molecular weight, and more. Returns PDB IDs matching all criteria.

    No authentication required.
    """



[docs]
    def __init__(self, tool_config: Dict[str, Any]):
        super().__init__(tool_config)
        self.timeout = tool_config.get("timeout", 30)
        fields = tool_config.get("fields", {})
        self.endpoint = fields.get("endpoint", "advanced_search")





[docs]
    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Execute the RCSB advanced search."""
        try:
            return self._query(arguments)
        except requests.exceptions.Timeout:
            return {
                "status": "error",
                "error": f"RCSB Search API timed out after {self.timeout}s",
            }
        except requests.exceptions.ConnectionError:
            return {"status": "error", "error": "Failed to connect to RCSB Search API"}
        except requests.exceptions.HTTPError as e:
            msg = ""
            try:
                msg = e.response.json().get("message", "")[:200]
            except Exception:
                msg = str(e.response.status_code)
            return {"status": "error", "error": f"RCSB Search API error: {msg}"}
        except Exception as e:
            return {"status": "error", "error": f"Unexpected error: {str(e)}"}





[docs]
    @staticmethod
    def _paginate(arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Resolve the paginate window, reporting any clamp applied to `rows`.

        Returns the RCSB `paginate` block plus the bookkeeping needed to tell
        the caller how much of the result set they actually received.
        """
        requested_rows = (
            arguments.get("rows")
            or arguments.get("limit")
            or arguments.get("max_results")
            or 10
        )
        requested_rows = max(1, int(requested_rows))
        rows = min(requested_rows, MAX_ROWS_PER_PAGE)
        start = max(0, int(arguments.get("start") or arguments.get("offset") or 0))

        clamped = None
        if requested_rows > rows:
            clamped = (
                f"Requested rows={requested_rows} exceeds the per-page maximum of "
                f"{MAX_ROWS_PER_PAGE}; returned {rows}. Use 'start' to page through "
                "the remaining results."
            )
        return {"start": start, "rows": rows, "clamp_note": clamped}





[docs]
    @staticmethod
    def _pagination_metadata(window: Dict[str, Any], total: int, returned: int) -> Dict:
        """Build the metadata block describing this page of the result set."""
        meta: Dict[str, Any] = {
            "total_count": total,
            "returned": returned,
            "start": window["start"],
            "rows": window["rows"],
            "max_rows_per_page": MAX_ROWS_PER_PAGE,
        }

        notes = [n for n in (window["clamp_note"],) if n]
        next_start = window["start"] + returned
        if returned and next_start < total:
            notes.append(
                f"Showing results {window['start'] + 1}-{next_start} of {total}. "
                f"Re-run with start={next_start} for the next page."
            )
        if notes:
            meta["note"] = " ".join(notes)
        return meta





[docs]
    def _query(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Route to appropriate endpoint."""
        if self.endpoint == "advanced_search":
            return self._advanced_search(arguments)
        elif self.endpoint == "motif_search":
            return self._motif_search(arguments)
        else:
            return {"status": "error", "error": f"Unknown endpoint: {self.endpoint}"}