feat: pull protocols.io AIND workspace, create ProtocolModel (#253)

* feat: pull protocols.io AIND workspace, create ProtocolModel

* tests: coverage on protoocls

* chore: lint

* chore: typo

* refactor: title -> name

* tests: fix test

* chore: docstrings

* fix: API pagination doesn't work

* fix: properly deal with versions by suffixing them to title

* feat: from_doi/from_url functions

* refactor: use proper regex for URLs

* chore: lint

* chore: docstring

* fix: type errors

* feat: added version

* refactor: remove authors, add version

* chore: unused import

Author: dbirman

scripts/get_protocols.py

@@ -0,0 +1,92 @@
+"""Script to grab protocol information from protocols.IO using the official API.
+
+Requires you set PROTOCOLS_CLIENT_TOKEN environment variable with your API token, from
+protocols.io/developers
+"""
+
+import re
+import requests
+import pandas as pd
+import os
+
+
+OUTPUT_CSV = "src/aind_data_schema_models/_generators/models/protocols.csv"
+WORKSPACE_URI = "allen-institute-for-neural-dynamics"
+API_BASE = "https://www.protocols.io/api"
+
+
+def get_access_token():
+    """Get access token from environment variable."""
+    token = os.environ.get("PROTOCOLS_CLIENT_TOKEN")
+    if not token:
+        raise RuntimeError("PROTOCOLS_CLIENT_TOKEN environment variable not set.")
+    return token
+
+
+def get_workspace_protocols(token, workspace_uri):
+    """Get protocols from a given workspace."""
+    headers = {"Authorization": f"Bearer {token}"}
+    all_items = []
+    page = 1
+    while True:
+        url = f"{API_BASE}/v3/workspaces/{workspace_uri}/protocols?page={page}&page_size=10000"
+        resp = requests.get(url, headers=headers)
+        resp.raise_for_status()
+        data = resp.json()
+        items = data.get("items", [])
+        if not items:
+            break
+        all_items.extend(items)
+        # Check pagination info
+        pagination = data.get("pagination", {})
+        total_pages = pagination.get("total_pages", 1)
+        if page >= total_pages:
+            break
+        page += 1
+    return all_items
+
+
+def parse_protocol(protocol: dict) -> list[dict]:
+    """Parse protocol information to extract title and DOI."""
+    title = protocol.get("title", "")
+
+    versions = protocol.get("versions", [])
+
+    protocols = []
+
+    for version in versions:
+        version_doi = version.get("doi", "")
+        version_doi = re.sub(r"^(https?://)?(dx\.)?doi\.org/", "", version_doi)
+        version_id = int(version.get("version_id", "")) + 1  # version_id is 0-based
+
+        protocols.append({"title": title, "DOI": version_doi, "version": version_id})
+
+    return protocols
+
+
+def main():
+    """Main function to fetch protocols and save to CSV."""
+    token = get_access_token()
+    print("Fetching protocols from workspace...")
+    protocols = get_workspace_protocols(token, WORKSPACE_URI)
+    print(f"Found {len(protocols)} protocols.")
+    # Deduplicate by DOI
+    results_by_doi = {}
+    for p in protocols:
+        print(f"Processing protocol: {p.get('title', '')}")
+        parsed_protocols = parse_protocol(p)
+        for proto in parsed_protocols:
+            doi = proto["DOI"]
+            if doi and doi not in results_by_doi:
+                results_by_doi[doi] = proto
+    results = list(results_by_doi.values())
+    # Ensure output directory exists
+    os.makedirs(os.path.dirname(OUTPUT_CSV), exist_ok=True)
+    # Write CSV using pandas
+    df = pd.DataFrame(results)
+    df.to_csv(OUTPUT_CSV, index=False)
+    print(f"Saved {len(results)} protocols to {OUTPUT_CSV}")
+
+
+if __name__ == "__main__":
+    main()

src/aind_data_schema_models/_generators/generator.py

@@ -34,6 +34,16 @@ def load_data(data_type: str, root_path: str) -> pd.DataFrame:
     return data
 
 
+def regex_search(value, pattern):
+    """Perform regex search on a value and return matched groups."""
+    import re
+
+    match = re.search(pattern, value)
+    if match:
+        return match.groups()
+    return []
+
+
 def generate_code(data_type: str, root_path: str, isort: bool = True, black: bool = True):
     """Generate code from the template type
 
@@ -62,6 +72,8 @@ def generate_code(data_type: str, root_path: str, isort: bool = True, black: boo
     env.filters["to_class_name"] = to_class_name
     env.filters["to_class_name_underscored"] = to_class_name_underscored
     env.filters["unique_rows"] = lambda data, key: data.drop_duplicates(subset=key)
+
+    env.filters["regex_search"] = regex_search
     rendered_template = env.from_string(template)
 
     # Render template with data

src/aind_data_schema_models/_generators/models/protocols.csv

@@ -0,0 +1,68 @@
+title,DOI,version
+Solenoid Valve Calibration for Behavior Rigs Utilizing Water Reward,10.17504/protocols.io.261gerq7dl47/v1,1
+Running a Dynamic Foraging Behavior Task in Mice,10.17504/protocols.io.5jyl8p4m6g2w/v1,1
+Whole Brain Embedding for SmartSPIM - EasyIndex with 2% Agarose,10.17504/protocols.io.3byl4jpn8lo5/v1,1
+Stereotaxic Injection by Nanoject Protocol,10.17504/protocols.io.bd8qi9vw,1
+Stereotaxic Injection by Nanoject Protocol,10.17504/protocols.io.besfjebn,2
+Stereotaxic Injection by Nanoject Protocol,10.17504/protocols.io.bgpujvnw,3
+Stereotaxic Injection by Nanoject Protocol,10.17504/protocols.io.bp2l6nr7kgqe/v4,4
+Stereotaxic Injection by Nanoject Protocol,10.17504/protocols.io.bp2l6nr7kgqe/v5,5
+Stereotaxic Injection by Nanoject Protocol,10.17504/protocols.io.bp2l6nr7kgqe/v6,6
+Stereotaxic Injection by Nanoject Protocol,10.17504/protocols.io.bp2l6nr7kgqe/v7,7
+Stereotaxic Injection by Iontophoresis,10.17504/protocols.io.bd8ti9wn,1
+Stereotaxic Injection by Iontophoresis,10.17504/protocols.io.besgjebw,2
+Stereotaxic Injection by Iontophoresis,10.17504/protocols.io.bgpvjvn6,3
+Stereotaxic Injection by Iontophoresis,10.17504/protocols.io.14egn8ewzg5d/v4,4
+Stereotaxic Injection by Iontophoresis,10.17504/protocols.io.14egn8ewzg5d/v5,5
+Stereotaxic Injection by Iontophoresis,10.17504/protocols.io.14egn8ewzg5d/v6,6
+Stereotaxic Injection by Iontophoresis,10.17504/protocols.io.14egn8ewzg5d/v7,7
+Mouse Habituation - Head Fixation into Tube,10.17504/protocols.io.rm7vzxd74gx1/v1,1
+Mouse Water Restriction,10.17504/protocols.io.x54v9pn34g3e/v1,1
+Modified Frame-projected  Independent Fiber  Photometry (FIP) System_Hardware,10.17504/protocols.io.261ge39edl47/v1,1
+Modified Frame-projected  Independent Fiber  Photometry (FIP) System_Hardware,10.17504/protocols.io.261ge39edl47/v2,2
+Making Agarose for use in acute in vivo Electrophysiology Experiments ,10.17504/protocols.io.5jyl8py89g2w/v1,1
+Plug Removal for acute in vivo Electrophysiology Experiments,10.17504/protocols.io.eq2lywz8qvx9/v1,1
+Intraperitoneal Injection in an Adult Mouse,10.17504/protocols.io.bfzgjp3w,1
+Intraperitoneal Injection in an Adult Mouse,10.17504/protocols.io.5qpvo5w8dl4o/v2,2
+Multiplexed RNA FISH on Expanded Mouse Brain Slices,10.17504/protocols.io.dm6gpzj28lzp/v1,1
+General Setup and Takedown Procedures for Rodent Neurosurgery,10.17504/protocols.io.kqdg392o7g25/v1,1
+General Setup and Takedown Procedures for Rodent Neurosurgery,10.17504/protocols.io.kqdg392o7g25/v2,2
+Dual Hemisphere Craniotomy for Electrophysiology,10.17504/protocols.io.rm7vzjoe2lx1/v1,1
+Aqueous (SBiP) Delipidation for Whole Mouse Brain After morphoFISH Perfusion ,10.17504/protocols.io.rm7vzjz54lx1/v1,1
+Mouse Habituation - Head Fixation on Disk,10.17504/protocols.io.j8nlkojmxv5r/v1,1
+Mouse Habituation - Head Fixation on Disk,10.17504/protocols.io.j8nlkojmxv5r/v2,2
+Preparation of Lipopolysaccharide for Intraperitoneal Injection,10.17504/protocols.io.14egn9y1ml5d/v1,1
+Temporal Assessment of Immune Response,10.17504/protocols.io.5jyl8dqx6g2w/v1,1
+Mouse VAB Catheter Maintenance,10.17504/protocols.io.8epv52o5dv1b/v1,1
+Processing Blood Intended for Olink Assay,10.17504/protocols.io.261ger81jl47/v1,1
+BARseq 2.5,10.17504/protocols.io.kqdg3ke9qv25/v1,1
+Aqueous (SBiP) Delipidation of a Whole Mouse Brain,10.17504/protocols.io.n2bvj81mwgk5/v1,1
+Aqueous (SBiP) Delipidation of a Whole Mouse Brain,10.17504/protocols.io.n2bvj81mwgk5/v2,2
+Tetrahydrofuran and Dichloromethane Delipidation of a Whole Mouse Brain,10.17504/protocols.io.36wgqj1kxvk5/v1,1
+Tetrahydrofuran and Dichloromethane Delipidation of a Whole Mouse Brain,10.17504/protocols.io.36wgqj1kxvk5/v2,2
+Whole Mouse Brain Delipidation - Dichloromethane,10.17504/protocols.io.dm6gpj7n5gzp/v1,1
+"Whole Mouse Brain Delipidation, Immunolabeling, and Expansion Microscopy",10.17504/protocols.io.n92ldpwjxl5b/v1,1
+Immunolabeling of a Whole Mouse Brain,10.17504/protocols.io.ewov1okwylr2/v1,1
+Structural MRI Using the University of Washington 14T Vertical Bore Bruker MRI,10.17504/protocols.io.3byl4j8p2lo5/v1,1
+Duragel Application for Acute Electrophysiology Recordings,10.17504/protocols.io.14egn2dwqg5d/v1,1
+SmartSPIM setup and alignment,10.17504/protocols.io.5jyl8jyb7g2w/v1,1
+Refractive Index Matching - EasyIndex,10.17504/protocols.io.kxygx965kg8j/v1,1
+Preparing a 3D Printed Implant for Acute In Vivo Electrophysiology,10.17504/protocols.io.6qpvr4jmogmk/v1,1
+Imaging cleared mouse brains on SmartSPIM,10.17504/protocols.io.3byl4jo1rlo5/v1,1
+Refractive Index Matching - Ethyl Cinnamate,10.17504/protocols.io.n2bvj8k4bgk5/v1,1
+DAPI Staining Mouse Brain Sections,10.17504/protocols.io.3byl4jm6rlo5/v1,1
+Modified Frame-projected Independent Fiber Photometry (FIP) System_Triggering system,10.17504/protocols.io.kxygx3e6wg8j/v1,1
+Multi-Site Optic Fiber Implants,10.17504/protocols.io.6qpvr3dqovmk/v1,1
+Immunohistochemistry (IHC) Staining Mouse Brain Sections,10.17504/protocols.io.5qpvo3b7bv4o/v1,1
+Sectioning Mouse Brain with Sliding Microtome,10.17504/protocols.io.5jyl8p97rg2w/v1,1
+Mounting and Coverslipping Mouse Brain Sections,10.17504/protocols.io.n92ldmpy7l5b/v1,1
+Stereotactic Injections with Headframe Implant,10.17504/protocols.io.eq2lyj72elx9/v1,1
+"Protocol Collection: Perfusing, Sectioning, IHC, Mounting and Coverslipping Mouse Brain Specimens",10.17504/protocols.io.kxygx3yxkg8j/v1,1
+Mouse Cardiac Perfusion Fixation and Brain Collection,10.17504/protocols.io.bd8vi9w6,1
+Mouse Cardiac Perfusion Fixation and Brain Collection,10.17504/protocols.io.besijece,2
+Mouse Cardiac Perfusion Fixation and Brain Collection,10.17504/protocols.io.beudjes6,3
+Mouse Cardiac Perfusion Fixation and Brain Collection,10.17504/protocols.io.be2djga6,4
+Mouse Cardiac Perfusion Fixation and Brain Collection,10.17504/protocols.io.bg5vjy66,5
+Mouse Cardiac Perfusion Fixation and Brain Collection,10.17504/protocols.io.8epv51bejl1b/v6,6
+Mouse Cardiac Perfusion Fixation and Brain Collection,10.17504/protocols.io.8epv51bejl1b/v7,7
+Mouse Cardiac Perfusion Fixation and Brain Collection,10.17504/protocols.io.8epv51bejl1b/v8,8

src/aind_data_schema_models/_generators/models/registries.csv

@@ -1,5 +1,6 @@
 name,abbreviation
 Addgene,ADDGENE
+Digital Object Identifier,DOI
 Edinburgh Mouse Atlas Project,EMAPA
 Mouse Genome Informatics,MGI
 National Center for Biotechnology Information,NCBI

src/aind_data_schema_models/_generators/templates/protocols.txt

@@ -0,0 +1,52 @@
+"""Protocols"""
+{% raw -%}
+from pydantic import Field, ConfigDict
+from typing import Union
+from typing_extensions import Annotated
+from aind_data_schema_models.pid_names import BaseName
+from aind_data_schema_models.registries import Registry
+import re
+{% endraw %}
+
+class ProtocolModel(BaseName):
+    """Base model for protocol"""
+    model_config = ConfigDict(frozen=True)
+    name: str
+    version: int
+    registry: Registry
+    registry_identifier: str
+
+{% for _, row in data.iterrows() %}
+class {{ row['title'] | to_class_name_underscored }}_V{{ row['version'] }}(ProtocolModel):
+    """Model {{ row['title'] }}"""
+    name: str = "{{ row['title'] }}"
+    version: int = {{ row['version'] }}
+    registry: Registry = Registry.DOI
+    registry_identifier: str = "{{ row['DOI'] }}"
+
+
+{% endfor %}
+
+class Protocols:
+    """Protocols"""
+{% for _, row in data.iterrows() %}
+    {{ row['title'] | to_class_name | upper }}_V{{ row['version'] }} = {{ row['title'] | to_class_name_underscored }}_V{{ row['version'] }}()
+{% endfor %}
+
+    ALL = tuple(ProtocolModel.__subclasses__())
+
+    ONE_OF = Annotated[Union[tuple(ProtocolModel.__subclasses__())], Field(discriminator="title")]
+
+    doi_map = {m().registry_identifier: m() for m in ALL if getattr(m(), "registry_identifier", None)}
+
+    @classmethod
+    def from_doi(cls, doi: str) -> ProtocolModel:
+        """Return protocol model by DOI."""
+        return cls.doi_map.get(doi, None)
+
+    @classmethod
+    def from_url(cls, url: str) -> ProtocolModel:
+        """Return protocol model by DOI, stripping URL prefixes."""
+        # Remove any leading protocol/domain up to the DOI
+        doi = re.sub(r'^(https?://)?(dx\.)?doi\.org/', '', url)
+        return cls.from_doi(doi)