Release v0.12.1

Author: markgeejw

openprotein/embeddings/future.py

@@ -96,8 +96,10 @@ def get_item(self, sequence: str | bytes) -> np.ndarray:
         return api.result_decode(data)
 
 
-Score = namedtuple("Score", ["name", "sequence", "score"])
-SingleSiteScore = namedtuple("SingleSiteScore", ["mut_code", "score"])
+Score = namedtuple("Score", ["name", "sequence", "score", "query_id"])
+Score.__new__.__defaults__ = (None,)
+SingleSiteScore = namedtuple("SingleSiteScore", ["mut_code", "score", "query_id"])
+SingleSiteScore.__new__.__defaults__ = (None,)
 S = TypeVar("S", bound=Union[Score, SingleSiteScore])
 
 
@@ -132,13 +134,18 @@ class EmbeddingsScoreFuture(BaseScoreFuture[Score]):
 
     def stream(self) -> Iterator[Score]:
         stream = api.request_get_score_result(session=self.session, job_id=self.id)
-        # name, sequence, ...
-        next(stream)  # ignore header
+        header = next(stream)
+        has_query_id = len(header) > 0 and header[0].strip().lower() == "query_id"
         for line in stream:
-            # combine scores into numpy array
-            scores = np.array([float(s) for s in line[2:]])
-            output = Score(name=line[0], sequence=line[1], score=scores)
-            yield output
+            if has_query_id:
+                query_id = line[0] if line[0] else None
+                name, sequence = line[1], line[2]
+                scores = np.array([float(s) for s in line[3:]])
+            else:
+                query_id = None
+                name, sequence = line[0], line[1]
+                scores = np.array([float(s) for s in line[2:]])
+            yield Score(name=name, sequence=sequence, score=scores, query_id=query_id)
 
 
 class EmbeddingsScoreSingleSiteFuture(BaseScoreFuture[SingleSiteScore]):
@@ -148,13 +155,18 @@ class EmbeddingsScoreSingleSiteFuture(BaseScoreFuture[SingleSiteScore]):
 
     def stream(self) -> Iterator[SingleSiteScore]:
         stream = api.request_get_score_result(session=self.session, job_id=self.id)
-        # name, sequence, ...
-        next(stream)  # ignore header
+        header = next(stream)
+        has_query_id = len(header) > 0 and header[0].strip().lower() == "query_id"
         for line in stream:
-            # combine scores into numpy array
-            scores = np.array([float(s) for s in line[1:]])
-            output = SingleSiteScore(mut_code=line[0], score=scores)
-            yield output
+            if has_query_id:
+                query_id = line[0] if line[0] else None
+                mut_code = line[1]
+                scores = np.array([float(s) for s in line[2:]])
+            else:
+                query_id = None
+                mut_code = line[0]
+                scores = np.array([float(s) for s in line[1:]])
+            yield SingleSiteScore(mut_code=mut_code, score=scores, query_id=query_id)
 
 
 class EmbeddingsGenerateFuture(BaseScoreFuture[Score]):
@@ -164,17 +176,18 @@ class EmbeddingsGenerateFuture(BaseScoreFuture[Score]):
 
     def stream(self) -> Iterator[Score]:
         stream = api.request_get_generate_result(session=self.session, job_id=self.id)
-        # name, sequence, ...
         header = next(stream)
-        has_query_id = (
-            len(header) > 2 and header[-1].strip().lower() == "query_id"
-        )
+        has_query_id = len(header) > 0 and header[0].strip().lower() == "query_id"
         for line in stream:
-            # combine scores into numpy array
-            score_values = line[2:-1] if has_query_id else line[2:]
-            scores = np.array([float(s) for s in score_values])
-            output = Score(name=line[0], sequence=line[1], score=scores)
-            yield output
+            if has_query_id:
+                query_id = line[0] if line[0] else None
+                name, sequence = line[1], line[2]
+                scores = np.array([float(s) for s in line[3:]])
+            else:
+                query_id = None
+                name, sequence = line[0], line[1]
+                scores = np.array([float(s) for s in line[2:]])
+            yield Score(name=name, sequence=sequence, score=scores, query_id=query_id)
 
     @property
     def sequences(self):

openprotein/errors.py

@@ -1,6 +1,10 @@
 from pydantic import BaseModel
 from requests import Response
 
+UPGRADE_MESSAGE = (
+    "If this issue persists, try upgrading the client: pip install --upgrade openprotein-python"
+)
+
 
 # Errors for OpenProtein
 class InvalidParameterError(Exception):
@@ -28,7 +32,7 @@ class APIError(Exception):
     """APIError"""
 
     def __init__(self, message: str):
-        self.message = message
+        self.message = f"{message}\n{UPGRADE_MESSAGE}"
         super().__init__(self.message)
 
 

openprotein/fold/__init__.py

@@ -17,7 +17,7 @@
 from .future import FoldResultFuture
 from .minifold import MiniFoldModel
 from .models import FoldModel
-from .protenix import ProtenixModel
+from .protenix import ProtenixConfidence, ProtenixModel
 from .rosettafold3 import RosettaFold3Model
 from .schemas import FoldJob, FoldMetadata
 
@@ -28,6 +28,7 @@
     "ESMFoldModel",
     "MiniFoldModel",
     "AlphaFold2Model",
+    "ProtenixConfidence",
     "ProtenixModel",
     "Boltz1Model",
     "Boltz1xModel",

openprotein/fold/future.py

@@ -20,9 +20,10 @@
 
 if TYPE_CHECKING:
     from .boltz import BoltzAffinity, BoltzConfidence
+    from .protenix import ProtenixConfidence
 
 FoldResult: typing.TypeAlias = (
-    "Structure | np.ndarray | pd.DataFrame | BoltzAffinity | list[BoltzConfidence]"
+    "Structure | np.ndarray | pd.DataFrame | BoltzAffinity | list[BoltzConfidence] | list[ProtenixConfidence]"
 )
 
 
@@ -307,9 +308,14 @@ def get_item(
 
                 data = TypeAdapter(BoltzAffinity).validate_python(data)
             elif key == "confidence":
-                from .boltz import BoltzConfidence
+                if self.model_id == "protenix":
+                    from .protenix import ProtenixConfidence
 
-                data = TypeAdapter(list[BoltzConfidence]).validate_python(data)
+                    data = TypeAdapter(list[ProtenixConfidence]).validate_python(data)
+                else:
+                    from .boltz import BoltzConfidence
+
+                    data = TypeAdapter(list[BoltzConfidence]).validate_python(data)
             return data  # ty: ignore[invalid-return-type]
 
     @typing.overload
@@ -588,7 +594,9 @@ def get_metrics(self) -> list[pd.DataFrame]:
             self._metrics = metrics
         return copy.deepcopy(self._metrics)
 
-    def get_confidence(self) -> list[list["BoltzConfidence"]]:
+    def get_confidence(
+        self,
+    ) -> "list[list[BoltzConfidence]] | list[list[ProtenixConfidence]]":
         """
         Retrieve the confidences of the structure prediction.
 
@@ -598,8 +606,8 @@ def get_confidence(self) -> list[list["BoltzConfidence"]]:
 
         Returns
         -------
-        list[list[BoltzConfidence]]
-            List of list of BoltzConfidence objects.
+        list[list[BoltzConfidence]] | list[list[ProtenixConfidence]]
+            List of list of confidence objects (model-specific schema).
 
         Raises
         ------

openprotein/fold/protenix.py

@@ -2,6 +2,8 @@
 
 from collections.abc import Sequence
 
+from pydantic import BaseModel
+
 from openprotein.align import MSAFuture
 from openprotein.base import APISession
 from openprotein.common import ModelMetadata
@@ -19,6 +21,64 @@
 from .models import FoldModel
 
 
+class ProtenixConfidence(BaseModel):
+    """
+    Per-sample confidence scores from a Protenix structure prediction.
+
+    Attributes
+    ----------
+    ranking_score : float
+        Composite ranking metric: ``0.8 * iptm + 0.2 * ptm - 100 * has_clash``.
+    ptm : float
+        Predicted TM-score for the full complex.
+    iptm : float
+        Interface pTM aggregated over inter-chain residue pairs.
+    plddt : float
+        Mean per-atom pLDDT in [0, 100].
+    gpde : float
+        Global PDE weighted by contact probabilities.
+    has_clash : float
+        Binary clash flag (1.0 if atomic clashes detected, else 0.0).
+    num_recycles : int
+        Number of recycling iterations used.
+    disorder : float
+        Disorder score (currently always 0.0).
+    chain_ptm : list[float]
+        Per-chain pTM scores, indexed by chain.
+    chain_iptm : list[float]
+        Per-chain ipTM scores.
+    chain_plddt : list[float]
+        Per-chain mean pLDDT scores.
+    chain_gpde : list[float]
+        Per-chain global PDE scores.
+    chain_pair_iptm : list[list[float]]
+        Chain-pair ipTM matrix.
+    chain_pair_iptm_global : list[list[float]]
+        Chain-pair ipTM matrix with ligand-aware weighting.
+    chain_pair_gpde : list[list[float]]
+        Chain-pair global PDE matrix.
+    """
+
+    ranking_score: float
+    ptm: float
+    iptm: float
+    plddt: float
+    gpde: float
+    has_clash: float
+    num_recycles: int
+    disorder: float
+    chain_ptm: list[float]
+    chain_iptm: list[float]
+    chain_plddt: list[float]
+    chain_gpde: list[float]
+    chain_pair_iptm: list[list[float]]
+    chain_pair_iptm_global: list[list[float]]
+    chain_pair_gpde: list[list[float]]
+
+    class Config:
+        extra = "allow"
+
+
 class ProtenixModel(FoldModel):
     """
     Class providing inference endpoints for Protenix structure prediction.

openprotein/jobs/futures.py

@@ -25,7 +25,7 @@
 
 from openprotein import config
 from openprotein.base import APISession
-from openprotein.errors import TimeoutException
+from openprotein.errors import APIError, TimeoutException
 from openprotein.jobs.schemas import Job, JobStatus
 
 from . import api
@@ -389,7 +389,12 @@ def wait(
         time.sleep(1)  # buffer for BE to register job
         job = self._wait_job(interval=interval, timeout=timeout, verbose=verbose)
         self.job = job
-        return self.get()
+        try:
+            return self.get()
+        except APIError:
+            raise
+        except Exception as e:
+            raise APIError(f"Failed to retrieve results: {e}") from e
 
 
 class StreamingFuture(Future[list[V]], ABC, Generic[V]):
@@ -433,15 +438,20 @@ def get(self, verbose: bool = False, **kwargs) -> list[V]:
             A list containing all results from the job.
 
         """
-        generator = self.stream(**kwargs)
-        if verbose:
-            total = None
-            if hasattr(self, "__len__"):
-                total = len(self)  # type: ignore - static type checker doesnt know
-            generator = tqdm.tqdm(
-                generator, desc="Retrieving", total=total, position=0, mininterval=1.0
-            )
-        return [entry for entry in generator]
+        try:
+            generator = self.stream(**kwargs)
+            if verbose:
+                total = None
+                if hasattr(self, "__len__"):
+                    total = len(self)  # type: ignore - static type checker doesnt know
+                generator = tqdm.tqdm(
+                    generator, desc="Retrieving", total=total, position=0, mininterval=1.0
+                )
+            return [entry for entry in generator]
+        except APIError:
+            raise
+        except Exception as e:
+            raise APIError(f"Failed to parse results: {e}") from e
 
     def wait(
         self,

openprotein/prompt/models.py

@@ -1,3 +1,4 @@
+from openprotein import config
 from openprotein.base import APISession
 from openprotein.jobs import Future, JobsAPI
 from openprotein.molecules import Complex, Protein
@@ -55,11 +56,12 @@ def __init__(
                 )
         self.metadata = metadata
         self.session = session
+        self.job = None  # default for uploaded
         if self.metadata.job_id is not None:
             jobs_api = getattr(session, "jobs", None)
             assert isinstance(jobs_api, JobsAPI)
             job = PromptJob.create(jobs_api.get_job(job_id=self.metadata.job_id))
-            super().__init__(session, job)
+            self.job = job
 
     def __str__(self) -> str:
         return str(self.metadata)
@@ -77,10 +79,15 @@ def get(self) -> list[list[Protein]]:
         context = api.get_prompt(session=self.session, prompt_id=str(self.id))
         return context
 
-    def _wait_job(self, **kwargs):
+    def _wait_job(
+        self,
+        interval: float = config.POLLING_INTERVAL,
+        timeout: int | None = None,
+        verbose: bool = False,
+    ):
         if self.job is None:
             return None
-        return super()._wait_job(**kwargs)
+        return super()._wait_job(interval, timeout, verbose)
 
     @property
     def id(self):

openprotein/utils/chain_id.py

@@ -1,16 +1,14 @@
 import re
 import string
 
-valid_id_pattern = re.compile(r"^[A-Z]{1,5}$|^\d{1,5}$")
+valid_id_pattern = re.compile(r"^[A-Za-z0-9]{1,5}$")
 
 
 def is_valid_id(id_str: str) -> bool:
     """
-    Check if the id_str matches the valid pattern for IDs (1-5 uppercase or 1-5 digits).
+    Check if the id_str matches the valid pattern for IDs (1-5 uppercase or digits).
     """
-    if not id_str or len(id_str) > 5:
-        return False
-    return bool(valid_id_pattern.fullmatch(id_str))
+    return bool(id_str and valid_id_pattern.fullmatch(id_str))
 
 
 def id_generator(

pyproject.toml

@@ -89,7 +89,7 @@ markers = [
   "slow: Slow-running tests",
   "integration: Integration tests",
 ]
-addopts = "-v --strict-markers --tb=short --disable-warnings"
+addopts = "-v --strict-markers --tb=short --disable-warnings -m 'not e2e'"
 timeout = 1200
 testpaths = ["tests"]
 log_cli = false

tests/e2e/test_embeddings_e2e.py

@@ -393,6 +393,24 @@ def test_e2e_poet2_generate_with_query_fanout(session: OpenProtein):
     )
 
 
+@pytest.mark.e2e
+def test_e2e_poet2_generate_with_prompt(session: OpenProtein):
+    """Validate PoET2 generate with a prompt that has already reached SUCCESS."""
+    n_sequences = 2
+    context = ["ACDEFGHIKLMNPQRSTVWY", "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQ"]
+    prompt = session.prompt.create_prompt(context)
+    assert prompt.wait_until_done(timeout=TIMEOUT)
+
+    future = session.embedding.poet2.generate(
+        prompt=prompt,
+        num_samples=n_sequences,
+        temperature=1.0,
+    )
+    assert future.wait_until_done(timeout=GENERATE_TIMEOUT)
+    results = future.get()
+    _assert_generated_sequences(results=results, expected_count=n_sequences)
+
+
 @pytest.mark.e2e
 def test_e2e_proteinmpnn_score_not_implemented(session: OpenProtein):
     with pytest.raises(NotImplementedError, match="Score not yet implemented"):