openprotein-python 0.14.2__tar.gz → 0.16.0__tar.gz

{openprotein_python-0.14.2 → openprotein_python-0.16.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openprotein-python
-Version: 0.14.2
+Version: 0.16.0
 Summary: OpenProtein Python interface.
 Author-email: Mark Gee <markgee@ne47.bio>, "Timothy Truong Jr." <ttruong@ne47.bio>, Tristan Bepler <tbepler@ne47.bio>
 License-Expression: MIT
@@ -8,6 +8,7 @@ License-File: LICENSE.txt
 Classifier: Development Status :: 4 - Beta
 Classifier: Programming Language :: Python :: 3
 Requires-Python: <3.15,>=3.10
+Requires-Dist: filelock<4,>=3.13
 Requires-Dist: gemmi<0.8,>=0.6.0
 Requires-Dist: numpy<3,>=1.9
 Requires-Dist: pandas<3,>=2.2.2

openprotein_python-0.16.0/openprotein/auth_store.py

@@ -0,0 +1,79 @@
+import json
+import os
+import tempfile
+from pathlib import Path
+
+from filelock import FileLock
+from pydantic import BaseModel, ValidationError
+
+
+class TokenRecord(BaseModel):
+    """A persisted auth record for one identity."""
+
+    access_token: str
+    refresh_token: str | None = None
+    expires_at: float = 0.0
+
+
+class TokenStore:
+    """Persist the rotating refresh token across processes.
+
+    The on-disk file is a JSON object keyed by identity
+    ``"{backend}|{username}"`` -> ``TokenRecord``. The read-modify-write
+    critical section is guarded by a cross-process ``filelock`` so two
+    processes never refresh the same rotating token concurrently.
+    """
+
+    def __init__(self, dir: Path | None = None):
+        self.dir = Path(dir) if dir is not None else Path.home() / ".openprotein"
+        self.dir.mkdir(parents=True, exist_ok=True)
+        self.path = self.dir / "token.json"
+        # One FileLock instance per store -> reentrant within this process.
+        self._lock = FileLock(str(self.dir / "token.lock"))
+
+    @staticmethod
+    def identity(backend: str, username: str) -> str:
+        return f"{backend}|{username}"
+
+    def lock(self):
+        """Context manager for the cross-process critical section."""
+        return self._lock
+
+    def _read_all(self) -> dict:
+        # Caller must hold self._lock: this is the read half of the
+        # read-modify-write critical section in load()/save().
+        try:
+            return json.loads(self.path.read_text())
+        except (OSError, json.JSONDecodeError):
+            # Missing or corrupt file -> behave as if there were no records.
+            return {}
+
+    def load(self, identity: str) -> TokenRecord | None:
+        with self._lock:
+            data = self._read_all()
+            entry = data.get(identity)
+            if entry is None:
+                return None
+            try:
+                return TokenRecord.model_validate(entry)
+            except ValidationError:
+                return None
+
+    def save(self, identity: str, record: TokenRecord) -> None:
+        with self._lock:
+            data = self._read_all()
+            data[identity] = record.model_dump()
+            self._atomic_write(data)
+
+    def _atomic_write(self, data: dict) -> None:
+        fd, tmp = tempfile.mkstemp(dir=str(self.dir), prefix="token.", suffix=".tmp")
+        try:
+            with os.fdopen(fd, "w") as f:
+                json.dump(data, f)
+            os.chmod(tmp, 0o600)
+            os.replace(tmp, self.path)
+        except BaseException:
+            # Don't leave a temp file behind on failure.
+            if os.path.exists(tmp):
+                os.unlink(tmp)
+            raise

openprotein_python-0.16.0/openprotein/base.py

@@ -0,0 +1,372 @@
+import sys
+import threading
+import time
+import warnings
+from typing import Mapping, Sequence
+from urllib.parse import urljoin
+
+import requests
+import requests.auth
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry  # type: ignore
+
+from openprotein._version import __version__
+from openprotein.auth_store import TokenRecord, TokenStore
+from openprotein.errors import AuthError, HTTPError
+
+# Refresh the access token this many seconds before it actually expires,
+# to absorb clock skew and request latency. 30s is a conservative buffer for
+# typical clock drift plus a round-trip against a ~15-minute access token.
+SKEW_SECONDS = 30
+
+
+class BearerAuth(requests.auth.AuthBase):
+    """
+    See https://stackoverflow.com/a/58055668
+    """
+
+    def __init__(self, username, token):
+        self._username = username
+        self._token = token
+
+    @property
+    def token(self):
+        warnings.warn(
+            "DeprecationWarning: Accessing session.auth.token is deprecated and will throw an error in the future. Use session.auth._token if you need it."
+        )
+        return self._token
+
+    def __call__(self, r):
+        r.headers["Authorization"] = "Bearer " + self._token
+        return r
+
+
+class APISession(requests.Session):
+    """
+    A class to handle API sessions. This class provides a connection session to the OpenProtein API.
+    """
+
+    def __init__(
+        self,
+        username: str,
+        password: str,
+        backend: str,
+        timeout: int = 180,
+    ):
+        if not username or not password:
+            raise AuthError(
+                "Expected username and password. Or use environment variables `OPENPROTEIN_USERNAME` and `OPENPROTEIN_PASSWORD`. Or provide these variables (`username` and `password`) in ~/.openprotein/config.toml."
+            )
+        super().__init__()
+        self.backend = backend
+        self.verify = True
+        self.timeout = timeout
+
+        # Custom retry strategies
+        # auto retry for pesky connection reset errors and others
+        # 503 will catch if BE is refreshing
+        retry = Retry(
+            total=4,
+            backoff_factor=3,  # 0,1,4,13s
+            status_forcelist=[500, 502, 503, 504, 101, 104],
+        )
+        adapter = HTTPAdapter(max_retries=retry)
+        self.mount("https://", adapter)
+        self.headers["User-Agent"] = f"openprotein-python/{__version__}"
+
+        # Auth / refresh state. Initialised before any network call so that
+        # request() can read _refreshing / _legacy_mode safely.
+        self._username = username
+        self._password = password
+        self._token_store = TokenStore()
+        self._identity = self._token_store.identity(backend, username)
+        # _refresh_lock serializes concurrent refreshes across threads.
+        # _refreshing (see the _refreshing property) is a thread-local recursion
+        # guard: the refresh POST flows back through request(), so the refreshing
+        # thread must skip the refresh logic, but OTHER threads sharing this
+        # session (e.g. via stream_parallel) must not be suppressed by it.
+        self._refresh_lock = threading.Lock()
+        self._tl = threading.local()
+        self._refresh_token: str | None = None
+        self._expires_at: float = 0.0
+        self._legacy_mode: bool = False
+
+        self._init_auth()
+
+    @property
+    def _refreshing(self) -> bool:
+        return getattr(self._tl, "refreshing", False)
+
+    @_refreshing.setter
+    def _refreshing(self, value: bool) -> None:
+        self._tl.refreshing = value
+
+    def _init_auth(self):
+        """Adopt a valid stored record if present, otherwise log in fresh."""
+        record = self._token_store.load(self._identity)
+        if record is not None and record.refresh_token:
+            self._adopt(record)
+        else:
+            self.login(self._username, self._password)
+
+    def post(self, url, data=None, json=None, **kwargs):
+        r"""Sends a POST request. Returns :class:`Response` object.
+
+        :param url: URL for the new :class:`Request` object.
+        :param data: (optional) Dictionary, list of tuples, bytes, or file-like
+            object to send in the body of the :class:`Request`.
+        :param json: (optional) json to send in the body of the :class:`Request`.
+        :param \*\*kwargs: Optional arguments that ``request`` takes.
+        :rtype: requests.Response
+        """
+        timeout = self.timeout
+        if "timeout" in kwargs:
+            timeout = kwargs.pop("timeout")
+
+        return self.request(
+            "POST", url, data=data, json=json, timeout=timeout, **kwargs
+        )
+
+    def login(self, username: str, password: str):
+        """
+        Authenticate connection to OpenProtein with your credentials.
+
+        Parameters
+        -----------
+        username: str
+            username
+        password: str
+            password
+        """
+        # unset the auth first
+        self.auth = None
+        self._get_auth_token(username, password)
+
+    def _auth_post(self, url: str, **kwargs):
+        """POST for the auth endpoints, with the refresh logic suppressed.
+
+        Sets ``_refreshing`` so the proactive/reactive refresh in request()
+        does not recurse while we are logging in or refreshing.
+        """
+        self._refreshing = True
+        try:
+            return self.post(url, **kwargs)
+        finally:
+            self._refreshing = False
+
+    def _get_auth_token(self, username: str, password: str):
+        url = urljoin(self.backend, "v1/auth/login")
+        try:
+            response = self._auth_post(
+                url,
+                data={"username": username, "password": password},
+                headers={"X-Token-Delivery": "body"},
+                timeout=3,
+            )
+        except HTTPError as e:
+            if e.status_code == 404:
+                # Backend not yet upgraded -> use the legacy endpoint.
+                return self._legacy_login(username, password)
+            raise AuthError(
+                "Authentication failed. Please check your credentials and connection."
+            ) from e
+        self._apply_token_response(response.json(), username)
+
+    def _legacy_login(self, username: str, password: str):
+        url = urljoin(self.backend, "v1/login/access-token")
+        try:
+            response = self._auth_post(
+                url,
+                data={"username": username, "password": password},
+                timeout=3,
+            )
+        except HTTPError as e:
+            raise AuthError(
+                "Authentication failed. Please check your credentials and connection."
+            ) from e
+        token = response.json().get("access_token")
+        if token is None:
+            raise AuthError("Unable to authenticate with given credentials.")
+        self._legacy_mode = True
+        self._refresh_token = None
+        self._expires_at = 0.0
+        self._set_token(username, token)
+
+    def _apply_token_response(self, result: dict, username: str):
+        """Parse a login/refresh response, update memory + store, set the bearer."""
+        access_token = result.get("access_token")
+        if access_token is None:
+            raise AuthError("Unable to authenticate with given credentials.")
+        refresh_token = result.get("refresh_token")
+        if refresh_token is None:
+            # Cookie-mode / pre-rotation backend -> no refresh available.
+            self._legacy_mode = True
+            self._refresh_token = None
+            self._expires_at = 0.0
+        else:
+            self._legacy_mode = False
+            self._refresh_token = refresh_token
+            # A missing/zero expires_in is treated as "expired now", which makes
+            # the proactive check refresh on the next request rather than trust a
+            # token of unknown lifetime.
+            self._expires_at = time.time() + (result.get("expires_in") or 0)
+            self._token_store.save(
+                self._identity,
+                TokenRecord(
+                    access_token=access_token,
+                    refresh_token=refresh_token,
+                    expires_at=self._expires_at,
+                ),
+            )
+        self._set_token(username, access_token)
+
+    def _set_token(self, username: str, access_token: str):
+        if isinstance(self.auth, BearerAuth):
+            self.auth._token = access_token
+        else:
+            self.auth = BearerAuth(username=username, token=access_token)
+
+    def _adopt(self, record: TokenRecord):
+        """Adopt an existing token record (from disk or a fresher in-store copy).
+
+        Callers must only pass a record whose ``refresh_token`` is set — a
+        record without one is not adoptable (it would leave the session in
+        non-legacy mode with nothing to refresh).
+        """
+        self._legacy_mode = False
+        self._refresh_token = record.refresh_token
+        self._expires_at = record.expires_at
+        self._set_token(self._username, record.access_token)
+
+    def _refresh(self):
+        """Refresh the access token, serialized across threads and processes.
+
+        Re-reads the store under the lock first: if another thread/process
+        already rotated the token, adopt that instead of refreshing again
+        (replaying a superseded refresh token would revoke the session).
+        """
+        with self._refresh_lock, self._token_store.lock():
+            record = self._token_store.load(self._identity)
+            if (
+                record is not None
+                and record.refresh_token
+                and record.expires_at > self._expires_at
+            ):
+                self._adopt(record)
+                return
+
+            if self._refresh_token is None:
+                self._relogin()
+                return
+
+            url = urljoin(self.backend, "v1/auth/refresh")
+            try:
+                response = self._auth_post(
+                    url,
+                    json={"refresh_token": self._refresh_token},
+                    headers={"X-Token-Delivery": "body"},
+                    timeout=10,
+                )
+            except HTTPError as e:
+                # Recover by re-logging-in when the refresh token is rejected
+                # (401/403) or when the endpoint is absent (404 from a backend
+                # that issued a refresh token but doesn't expose /auth/refresh,
+                # e.g. mid-rollback) — re-login itself drops to legacy mode if
+                # the backend has no refresh support. Other statuses are real
+                # errors and propagate.
+                if e.status_code in (401, 403, 404):
+                    self._relogin()
+                    return
+                raise
+            self._apply_token_response(response.json(), self._username)
+
+    def _relogin(self):
+        """Re-establish a session from the in-memory credentials."""
+        try:
+            self.login(self._username, self._password)
+        except AuthError:
+            raise
+        except Exception as e:
+            raise AuthError("Session expired and re-login failed.") from e
+
+    def _should_refresh_proactively(self) -> bool:
+        return (
+            not self._refreshing
+            and not self._legacy_mode
+            and self._refresh_token is not None
+            and time.time() >= self._expires_at - SKEW_SECONDS
+        )
+
+    def request(self, method: str, url: str, *args, **kwargs):
+        if self._should_refresh_proactively():
+            self._refresh()
+
+        full_url = urljoin(self.backend, url)
+        response = super().request(method, full_url, *args, **kwargs)
+
+        # Reactive refresh: a 401 means the access token expired between the
+        # proactive check and the call (or skew). Refresh once and retry once.
+        if (
+            response.status_code == 401
+            and not self._refreshing
+            and not self._legacy_mode
+        ):
+            self._refresh()
+            response = super().request(method, full_url, *args, **kwargs)
+
+        if (js := kwargs.get("json")) and js is not None:
+            if _total_size(js) > 1e6:
+                warnings.warn(
+                    "The requested payload is >1MB. There might be some delays or issues in processing. If the request fails, please try again with smaller sizes."
+                )
+
+        # intercept CloudFront errors
+        if "cloudfront" in response.headers.get("Server", "").lower():
+            if response.status_code in (502, 503):
+                raise CloudFrontError(
+                    f"We're experiencing a temporary backend issue via CloudFront. Please try again later. Error {response.status_code}."
+                )
+            elif response.status_code == 504:
+                raise TimeoutError(
+                    "Your request took too long to process likely due to it's size. Please try breaking it up into smaller requests if possible."
+                )
+        elif not response.ok:
+            # raise custom exception that prints better error message than requests.HTTPError
+            raise HTTPError(response)
+        return response
+
+
+def _total_size(o: Sequence | Mapping | object, seen=None):
+    """Recursively finds size of objects including contents."""
+    if seen is None:
+        seen = set()
+    obj_id = id(o)
+    if obj_id in seen:
+        return 0
+    seen.add(obj_id)
+    size = sys.getsizeof(o)
+    if isinstance(o, dict):
+        size += sum((_total_size(k, seen) + _total_size(v, seen)) for k, v in o.items())
+    elif isinstance(o, (list, tuple, set, frozenset)):
+        size += sum(_total_size(i, seen) for i in o)
+    return size
+
+
+class TimeoutError(requests.exceptions.HTTPError):
+    """
+    An Exception raised due to timeout, possibly from overly large
+    requests.
+    """
+
+
+class CloudFrontError(requests.exceptions.HTTPError):
+    """
+    An Exception raised due to CloudFront.
+
+    This is usually due to the strict timeout from CloudFront.
+    AWS CloudFront limits responses to return within 2 minutes.
+    This can be a bit prohibitive for our system that tends to
+    deal with large data. It is usually safe to just ignore/retry upon
+    hitting this error. Our system will scale up and still handle
+    the job.
+    """

{openprotein_python-0.14.2 → openprotein_python-0.16.0}/openprotein/embeddings/models.py

@@ -17,6 +17,7 @@ from . import api
 from .future import EmbeddingsResultFuture
 
 if TYPE_CHECKING:
+    from openprotein.predictor.schemas import KernelType
     from openprotein.predictor import PredictorModel
     from openprotein.svd import SVDModel
     from openprotein.umap import UMAPModel
@@ -352,6 +353,9 @@ class EmbeddingModel:
         reduction: ReductionType,
         name: str | None = None,
         description: str | None = None,
+        kernel: "KernelType | str" = "rbf",
+        period: float | None = None,
+        alpha: float | None = None,
         **kwargs,
     ) -> "PredictorModel":
         """
@@ -363,6 +367,13 @@ class EmbeddingModel:
             Assay to fit GP on.
         properties : list of str
             Properties in the assay to fit the GP on.
+        kernel : KernelType or str, optional
+            Kernel for the GP, e.g. "rbf", "matern52", "periodic",
+            "rational_quadratic". Defaults to "rbf".
+        period : float or None, optional
+            Period length; only valid for the "periodic" kernel (must be > 0).
+        alpha : float or None, optional
+            Scale-mixture parameter; only valid for "rational_quadratic" (must be > 0).
         reduction : ReductionType
             Type of embedding reduction to use for computing features. PLM must use reduction.
         name : str or None, optional
@@ -398,6 +409,9 @@ class EmbeddingModel:
             reduction=reduction,
             name=name,
             description=description,
+            kernel=kernel,
+            period=period,
+            alpha=alpha,
             **kwargs,
         )
 

{openprotein_python-0.14.2 → openprotein_python-0.16.0}/openprotein/embeddings/poet.py

@@ -17,6 +17,7 @@ from .future import (
 from .models import EmbeddingModel
 
 if TYPE_CHECKING:
+    from openprotein.predictor.schemas import KernelType
     from openprotein.predictor import PredictorModel
     from openprotein.svd import SVDModel
     from openprotein.umap import UMAPModel
@@ -432,6 +433,9 @@ class PoETModel(EmbeddingModel):
         name: str | None = None,
         description: str | None = None,
         prompt: str | Prompt | None = None,
+        kernel: "KernelType | str" = "rbf",
+        period: float | None = None,
+        alpha: float | None = None,
         **kwargs,
     ) -> "PredictorModel":
         """
@@ -443,6 +447,13 @@ class PoETModel(EmbeddingModel):
             Assay to fit GP on.
         properties : list of str
             Properties in the assay to fit the GP on.
+        kernel : KernelType or str, optional
+            Kernel for the GP, e.g. "rbf", "matern52", "periodic",
+            "rational_quadratic". Defaults to "rbf".
+        period : float or None, optional
+            Period length; only valid for the "periodic" kernel (must be > 0).
+        alpha : float or None, optional
+            Scale-mixture parameter; only valid for "rational_quadratic" (must be > 0).
         prompt : str or Prompt or None, optional
             Prompt from an align workflow to condition the PoET model.
         **kwargs
@@ -464,5 +475,8 @@ class PoETModel(EmbeddingModel):
             name=name,
             description=description,
             prompt_id=prompt_id,
+            kernel=kernel,
+            period=period,
+            alpha=alpha,
             **kwargs,
         )

{openprotein_python-0.14.2 → openprotein_python-0.16.0}/openprotein/embeddings/poet2.py

@@ -21,6 +21,7 @@ from .models import EmbeddingModel
 from .poet import PoETModel
 
 if TYPE_CHECKING:
+    from openprotein.predictor.schemas import KernelType
     from openprotein.models.structure_generation import StructureGenerationFuture
     from openprotein.predictor import PredictorModel
     from openprotein.svd import SVDModel
@@ -533,6 +534,9 @@ class PoET2Model(PoETModel, EmbeddingModel):
         query: str | bytes | Protein | Complex | Query | None = None,
         use_query_structure_in_decoder: bool = True,
         decoder_type: Literal["mlm", "clm"] | None = None,
+        kernel: "KernelType | str" = "rbf",
+        period: float | None = None,
+        alpha: float | None = None,
         **kwargs,
     ) -> "PredictorModel":
         """
@@ -544,6 +548,13 @@ class PoET2Model(PoETModel, EmbeddingModel):
             Assay to fit GP on.
         properties : list of str
             Properties in the assay to fit the GP on.
+        kernel : KernelType or str, optional
+            Kernel for the GP, e.g. "rbf", "matern52", "periodic",
+            "rational_quadratic". Defaults to "rbf".
+        period : float or None, optional
+            Period length; only valid for the "periodic" kernel (must be > 0).
+        alpha : float or None, optional
+            Scale-mixture parameter; only valid for "rational_quadratic" (must be > 0).
         prompt : str or Prompt or None, optional
             Prompt from an align workflow to condition PoET model.
         query : str or bytes or Protein or Complex or Query or None, optional
@@ -570,5 +581,8 @@ class PoET2Model(PoETModel, EmbeddingModel):
             query_id=query_id,
             use_query_structure_in_decoder=use_query_structure_in_decoder,
             decoder_type=decoder_type,
+            kernel=kernel,
+            period=period,
+            alpha=alpha,
             **kwargs,
         )

{openprotein_python-0.14.2 → openprotein_python-0.16.0}/openprotein/fold/common.py

@@ -13,40 +13,28 @@ from openprotein.prompt import PromptAPI
 def normalize_inputs(
     proteins: Sequence[Complex | Protein | str | bytes],
 ):
-    # collate the id's used
-    used_ids = []
     normalized_complexes: list[Complex] = []
-    remaining_proteins: list[Protein] = []
-    remaining_protein_strings: list[str] = []
     for protein in proteins:
-        if isinstance(protein, Protein):
-            # collate these to init with id_gen
-            remaining_proteins.append(protein)
-        elif isinstance(protein, Complex):
-            used_ids.extend(list(protein.get_chains().keys()))
+        if isinstance(protein, Complex):
             normalized_complexes.append(protein)
+        elif isinstance(protein, Protein):
+            # a lone protein becomes a single-chain complex
+            id_gen = id_generator()
+            normalized_complexes.append(Complex(chains={next(id_gen): protein}))
         else:
             if isinstance(protein, bytes):
                 protein = protein.decode()
-            # handle ':'-delimited
+            # a ':'-delimited string is one complex, one chain per split sequence.
+            # chain ids only need to be unique within the complex, so use a fresh
+            # generator per complex.
+            id_gen = id_generator()
+            chains: dict[str, Protein] = {}
             for seq in protein.split(":"):
-                # collate these to init with id_gen
-                remaining_protein_strings.append(seq)
-
-    # auto generate the chain ids
-    id_gen = id_generator(used_ids)
-
-    # add the remaining proteins with id gen
-    for protein in remaining_proteins:
-        id = next(id_gen)
-        normalized_complexes.append(Complex(chains={id: protein}))
-
-    for protein_str in remaining_protein_strings:
-        id = next(id_gen)
-        protein = Protein(sequence=protein_str)
-        # protein strings default to null msa
-        protein.msa = Protein.NullMSA
-        normalized_complexes.append(Complex(chains={id: protein}))
+                chain = Protein(sequence=seq)
+                # protein strings default to null msa
+                chain.msa = Protein.NullMSA
+                chains[next(id_gen)] = chain
+            normalized_complexes.append(Complex(chains=chains))
 
     return normalized_complexes
 

{openprotein_python-0.14.2 → openprotein_python-0.16.0}/openprotein/models/__init__.py

@@ -19,6 +19,7 @@ from .foundation.boltzgen import BoltzGenFuture, BoltzGenModel
 from .foundation.esmif1 import ESMIF1Model
 from .foundation.proteinmpnn import ProteinMPNNModel
 from .foundation.rfdiffusion import RFdiffusionFuture, RFdiffusionModel
+from .foundation.solublempnn import SolubleMPNNModel
 from .models import ModelsAPI
 from .structure_generation import (
     StructureGenerationFuture,

openprotein_python-0.16.0/openprotein/models/foundation/solublempnn.py

@@ -0,0 +1,25 @@
+"""SolubleMPNN model: ProteinMPNN variant tuned for soluble protein targets."""
+
+from .proteinmpnn import ProteinMPNNModel
+
+
+class SolubleMPNNModel(ProteinMPNNModel):
+    """
+    Class for SolubleMPNN model.
+
+    Soluble-protein variant of ProteinMPNN.
+
+    Examples
+    --------
+    View specific model details (including supported tokens) with the `?` operator.
+
+    Examples
+    --------
+    .. code-block:: ipython3
+
+        >>> import openprotein
+        >>> session = openprotein.connect(username="user", password="password")
+        >>> session.models.solublempnn?
+    """
+
+    model_id = "solublempnn"