Release v0.8.6

- Fix boltz confidence retrieval
- Properly handle when model metadata input tokens is null, e.g. msa for AF2
- Some documentation fixes
- Add metadata for fold

Author: markgeejw

anaconda_build/meta.yaml

@@ -1,6 +1,6 @@
 package:
   name: openprotein-python
-  version: "0.8.5"
+  version: "0.8.6"
 
 source:
   path: ../

openprotein/common/model_metadata.py

@@ -28,6 +28,6 @@ class ModelMetadata(BaseModel):
     max_sequence_length: int | None = None
     dimension: int
     output_types: list[str]
-    input_tokens: list[str]
+    input_tokens: list[str] | None
     output_tokens: list[str] | None = None
     token_descriptions: list[list[TokenInfo]]

openprotein/fold/alphafold2.py

@@ -47,10 +47,8 @@ def fold(
             number of times to recycle models
         num_models : int
             number of models to train - best model will be used
-        max_msa : Union[str, int]
-            maximum number of sequences in the msa to use.
-        relax_max_iterations : int
-            maximum number of iterations
+        num_relax : int
+            maximum number of iterations for relax
 
         Returns
         -------
@@ -61,6 +59,7 @@ def fold(
                 "Inputs to AlphaFold 2 have been updated. 'msa' should be supplied as 'proteins' argument. Support will be dropped in the future."
             )
             proteins = kwargs["msa"]
+            assert isinstance(proteins, MSAFuture), "Expected msa to be an MSAFuture"
         if "ligands" in kwargs or "dnas" in kwargs or "rnas" in kwargs:
             with warnings.catch_warnings():
                 warnings.simplefilter("always")  # Force warning to always show
@@ -73,6 +72,10 @@ def fold(
             msa_to_seed: dict[str, Counter] = dict()
             for protein in proteins:
                 if (msa := protein.msa) is not None:
+                    if isinstance(msa, Protein.NullMSA):
+                        raise ValueError(
+                            "AlphaFold 2 expects MSA and does not support single sequence mode"
+                        )
                     msa_id = msa.id if isinstance(msa, MSAFuture) else msa
                     if msa_id in msa_to_seed:
                         seeds = msa_to_seed[msa_id]

openprotein/fold/future.py

@@ -9,11 +9,11 @@
 from openprotein import config
 from openprotein.base import APISession
 from openprotein.chains import DNA, RNA, Ligand
-from openprotein.jobs import Future, MappedFuture
+from openprotein.jobs import Future, JobsAPI, MappedFuture
 from openprotein.protein import Protein
 
 from . import api
-from .schemas import FoldJob
+from .schemas import FoldJob, FoldMetadata
 
 if TYPE_CHECKING:
     from .boltz import BoltzAffinity, BoltzConfidence
@@ -34,34 +34,39 @@ class FoldResultFuture(MappedFuture, Future):
     def __init__(
         self,
         session: APISession,
-        job: FoldJob,
+        job: FoldJob | None = None,
+        metadata: FoldMetadata | None = None,
         sequences: list[bytes] | None = None,
         max_workers: int = config.MAX_CONCURRENT_WORKERS,
     ):
         """
         Initialize a FoldResultFuture instance.
 
-        Parameters
-        ----------
-        session : APISession
-            The API session to use for requests.
-        job : FoldJob
-            The fold job associated with this future.
-        sequences : list[bytes], optional
-            List of sequences submitted for the fold request. If None, sequences will be fetched.
-        max_workers : int, optional
-            Maximum number of concurrent workers. Default is config.MAX_CONCURRENT_WORKERS.
+        Takes in either a fold job, or the fold job metadata.
+
+        :meta private:
         """
-        super().__init__(session, job, max_workers)
+        # initialize the fold job metadata
+        if metadata is None:
+            if job is None or job.job_id is None:
+                raise ValueError("Expected fold metadata or job")
+            metadata = api.fold_get(session, job.job_id)
+        self._metadata = metadata
+        if job is None:
+            jobs_api = getattr(session, "jobs", None)
+            assert isinstance(jobs_api, JobsAPI)
+            job = FoldJob.create(jobs_api.get_job(job_id=metadata.job_id))
         if sequences is None:
             sequences = api.fold_get_sequences(self.session, job_id=job.job_id)
         self._sequences = sequences
+        super().__init__(session, job, max_workers)
 
     @classmethod
     def create(
         cls: type[Self],
         session: APISession,
-        job: FoldJob,
+        job: FoldJob | None = None,
+        metadata: FoldMetadata | None = None,
         **kwargs,
     ) -> "Self | FoldComplexResultFuture":
         """
@@ -81,7 +86,13 @@ def create(
         FoldResultFuture or FoldComplexResultFuture
             An instance of FoldResultFuture or FoldComplexResultFuture depending on the model.
         """
-        model_id = api.fold_get(session=session, job_id=job.job_id).model_id
+        if job is not None:
+            job_id = job.job_id
+        elif metadata is not None:
+            job_id = metadata.job_id
+        else:
+            raise ValueError("Expected fold metadata or job")
+        model_id = api.fold_get(session=session, job_id=job_id).model_id
         if model_id.startswith("boltz") or model_id.startswith("alphafold"):
             return FoldComplexResultFuture(session=session, job=job, **kwargs)
         else:
@@ -101,22 +112,6 @@ def sequences(self) -> list[bytes]:
             self._sequences = api.fold_get_sequences(self.session, self.job.job_id)
         return self._sequences
 
-    @property
-    def model_id(self) -> str:
-        """
-        Get the model ID used for the fold request.
-
-        Returns
-        -------
-        str
-            Model ID.
-        """
-        if self._model_id is None:
-            self._model_id = api.fold_get(
-                session=self.session, job_id=self.job.job_id
-            ).model_id
-        return self._model_id
-
     @property
     def id(self):
         """
@@ -129,6 +124,17 @@ def id(self):
         """
         return self.job.job_id
 
+
+    @property
+    def metadata(self) -> FoldMetadata:
+        """The fold metadata."""
+        return self._metadata
+
+    @property
+    def model_id(self) -> str:
+        """The fold model used."""
+        return self._metadata.model_id
+
     def __keys__(self):
         """
         Get the list of sequences submitted for the fold request.
@@ -189,7 +195,8 @@ class FoldComplexResultFuture(Future):
     def __init__(
         self,
         session: APISession,
-        job: FoldJob,
+        job: FoldJob | None = None,
+        metadata: FoldMetadata | None = None,
         model_id: str | None = None,
         proteins: list[Protein] | None = None,
         ligands: list[Ligand] | None = None,
@@ -216,6 +223,16 @@ def __init__(
         rnas : list[RNA], optional
             List of RNAs submitted for fold request.
         """
+        # initialize the fold job metadata
+        if metadata is None:
+            if job is None or job.job_id is None:
+                raise ValueError("Expected fold metadata or job")
+            metadata = api.fold_get(session, job.job_id)
+        self._metadata = metadata
+        if job is None:
+            jobs_api = getattr(session, "jobs", None)
+            assert isinstance(jobs_api, JobsAPI)
+            job = FoldJob.create(jobs_api.get_job(job_id=metadata.job_id))
         super().__init__(session, job)
         self._model_id = model_id
         self._proteins = proteins
@@ -229,6 +246,11 @@ def __init__(
         self._confidence: list["BoltzConfidence"] | None = None
         self._affinity: "BoltzAffinity | None" = None
 
+    @property
+    def metadata(self) -> FoldMetadata:
+        """The fold metadata."""
+        return self._metadata
+
     @property
     def model_id(self) -> str:
         """
@@ -433,6 +455,8 @@ def confidence(self) -> list["BoltzConfidence"]:
         AttributeError
             If confidence is not supported for the model.
         """
+        from .boltz import BoltzConfidence
+
         if self.model_id not in {"boltz-1", "boltz-1x", "boltz-2"}:
             raise AttributeError("confidence not supported for non-Boltz model")
         if self._confidence is None: