openprotein-python 0.8.2__1-py3-none-any.whl

openprotein/align/schemas.py

@@ -0,0 +1,165 @@
+"""Schemas for OpenProtein align system."""
+
+from enum import Enum
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+from openprotein.jobs import Job, JobType
+
+
+class AlignType(str, Enum):
+    """
+    Enumeration of alignment types.
+
+    Attributes
+    ----------
+    INPUT : str
+        Raw input alignment.
+    MSA : str
+        Generated multiple sequence alignment.
+    PROMPT : str
+        Prompt-based alignment.
+    """
+
+    INPUT = "RAW"
+    MSA = "GENERATED"
+    PROMPT = "PROMPT"
+
+
+class MSASamplingMethod(str, Enum):
+    """
+    Enumeration of MSA sampling methods.
+
+    Attributes
+    ----------
+    RANDOM : str
+        Random sampling.
+    NEIGHBORS : str
+        Sampling based on neighbors.
+    NEIGHBORS_NO_LIMIT : str
+        Neighbor sampling without limit.
+    NEIGHBORS_NONGAP_NORM_NO_LIMIT : str
+        Neighbor sampling without gap normalization and without limit.
+    TOP : str
+        Top scoring sampling.
+    """
+
+    RANDOM = "RANDOM"
+    NEIGHBORS = "NEIGHBORS"
+    NEIGHBORS_NO_LIMIT = "NEIGHBORS_NO_LIMIT"
+    NEIGHBORS_NONGAP_NORM_NO_LIMIT = "NEIGHBORS_NONGAP_NORM_NO_LIMIT"
+    TOP = "TOP"
+
+
+class PromptPostParams(BaseModel):
+    """
+    Parameters for posting a prompt to generate an MSA.
+
+    Attributes
+    ----------
+    msa_id : str
+        Identifier for the MSA.
+    num_sequences : int or None, optional
+        Number of sequences to sample (default is None, must be >=0 and <100).
+    num_residues : int or None, optional
+        Number of residues to sample (default is None, must be >=0 and <24577).
+    method : MSASamplingMethod, optional
+        Sampling method to use (default is NEIGHBORS_NONGAP_NORM_NO_LIMIT).
+    homology_level : float, optional
+        Homology level threshold (default is 0.8, must be between 0 and 1).
+    max_similarity : float, optional
+        Maximum similarity threshold (default is 1.0, must be between 0 and 1).
+    min_similarity : float, optional
+        Minimum similarity threshold (default is 0.0, must be between 0 and 1).
+    always_include_seed_sequence : bool, optional
+        Whether to always include the seed sequence (default is False).
+    num_ensemble_prompts : int, optional
+        Number of ensemble prompts to generate (default is 1).
+    random_seed : int or None, optional
+        Random seed for reproducibility (default is None).
+    """
+
+    msa_id: str
+    num_sequences: int | None = Field(None, ge=0, lt=100)
+    num_residues: int | None = Field(None, ge=0, lt=24577)
+    method: MSASamplingMethod = MSASamplingMethod.NEIGHBORS_NONGAP_NORM_NO_LIMIT
+    homology_level: float = Field(0.8, ge=0, le=1)
+    max_similarity: float = Field(1.0, ge=0, le=1)
+    min_similarity: float = Field(0.0, ge=0, le=1)
+    always_include_seed_sequence: bool = False
+    num_ensemble_prompts: int = 1
+    random_seed: int | None = None
+
+
+class MSAJob(Job):
+    """
+    Base class for MSA-related jobs.
+
+    Attributes
+    ----------
+    job_type : Literal[JobType.align_align]
+        The type of job (must be JobType.align_align).
+    """
+
+    job_type: Literal[JobType.align_align]
+
+    @property
+    def msa_id(self):
+        """
+        Returns the MSA identifier for this job.
+
+        Returns
+        -------
+        str
+            The MSA identifier.
+        """
+        return self.msa_id
+
+
+class MafftJob(MSAJob, Job):
+    """
+    Job for running MAFFT alignment.
+
+    Attributes
+    ----------
+    job_type : Literal[JobType.mafft]
+        The type of job (must be JobType.mafft).
+    """
+
+    job_type: Literal[JobType.mafft]
+
+
+class ClustalOJob(MSAJob, Job):
+    """
+    Job for running Clustal Omega alignment.
+
+    Attributes
+    ----------
+    job_type : Literal[JobType.clustalo]
+        The type of job (must be JobType.clustalo).
+    """
+
+    job_type: Literal[JobType.clustalo]
+
+
+class AbNumberJob(MSAJob, Job):
+    """
+    Job for running AbNumber alignment.
+
+    Attributes
+    ----------
+    job_type : Literal[JobType.abnumber]
+        The type of job (must be JobType.abnumber).
+    """
+
+    job_type: Literal[JobType.abnumber]
+
+
+class AbNumberScheme(str, Enum):
+    """Antibody numbering scheme."""
+
+    IMGT = "imgt"
+    CHOTHIA = "chothia"
+    KABAT = "kabat"
+    AHO = "aho"

openprotein/base.py

@@ -0,0 +1,181 @@
+import os
+import sys
+import warnings
+from collections.abc import Container, Mapping
+from typing import Union
+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
+
+import openprotein.config as config
+from openprotein.errors import APIError, AuthError, HTTPError
+
+USERNAME = os.getenv("OPENPROTEIN_USERNAME")
+PASSWORD = os.getenv("OPENPROTEIN_PASSWORD")
+BACKEND = os.getenv("OPENPROTEIN_API_BACKEND", "https://api.openprotein.ai/api/")
+
+
+class BearerAuth(requests.auth.AuthBase):
+    """
+    See https://stackoverflow.com/a/58055668
+    """
+
+    def __init__(self, token):
+        self.token = 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.
+
+    Parameters
+    ----------
+    username : str
+        The username of the user.
+    password : str
+        The password of the user.
+
+    Examples
+    --------
+    >>> session = APISession("username", "password")
+    """
+
+    def __init__(
+        self,
+        username: str | None = USERNAME,
+        password: str | None = PASSWORD,
+        backend: str = BACKEND,
+        timeout: int = 180,
+    ):
+        if not username or not password:
+            raise AuthError(
+                "Expected username and password. Or use environment variables `OPENPROTEIN_USERNAME` and `OPENPROTEIN_PASSWORD`"
+            )
+        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.login(username, 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.auth = self._get_auth_token(username, password)
+
+    def _get_auth_token(self, username: str, password: str):
+        endpoint = "v1/login/access-token"
+        url = urljoin(self.backend, endpoint)
+        try:
+            response = self.post(
+                url, data={"username": username, "password": password}, timeout=3
+            )
+        except HTTPError as e:
+            # if an error occured during auth, we raise an AuthError with reference to the HTTPError
+            raise AuthError(
+                f"Authentication failed. Please check your credentials and connection."
+            ) from e
+
+        result = response.json()
+        token = result.get("access_token")
+        if token is None:
+            raise AuthError("Unable to authenticate with given credentials.")
+        return BearerAuth(token)
+
+    def request(self, method: str, url: str, *args, **kwargs):
+        full_url = urljoin(self.backend, url)
+        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, 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 RestEndpoint:
+    pass
+
+
+class TimeoutError(requests.exceptions.HTTPError):
+    pass
+
+
+class CloudFrontError(requests.exceptions.HTTPError):
+    pass

openprotein/chains.py

@@ -0,0 +1,88 @@
+"""Additional chains that can be used with OpenProtein."""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class DNA:
+    """
+    Represents a DNA sequence.
+
+    Attributes:
+        sequence (str): The nucleotide sequence of the DNA.
+    """
+
+    sequence: str
+    chain_id: str | list[str] | None = None
+    cyclic: bool = False
+
+    def __init__(
+        self,
+        sequence: str,
+        chain_id: str | list[str] | None = None,
+        cyclic: bool = False,
+    ):
+        # validate the sequence matches DNA
+        if not all(nt in set("ACGT") for nt in sequence.upper()):
+            raise ValueError("Sequence contains invalid DNA nucleotides.")
+        self.sequence = sequence
+        self.chain_id = chain_id
+        self.cyclic = cyclic
+
+
+@dataclass
+class RNA:
+    """
+    Represents an RNA sequence.
+
+    Attributes:
+        sequence (str): The nucleotide sequence of the RNA.
+    """
+
+    sequence: str
+    chain_id: str | list[str] | None = None
+    cyclic: bool = False
+
+    def __init__(
+        self,
+        sequence: str,
+        chain_id: str | list[str] | None = None,
+        cyclic: bool = False,
+    ):
+        # validate the sequence matches RNA
+        if not all(nt in set("ACGU") for nt in sequence.upper()):
+            raise ValueError("Sequence contains invalid RNA nucleotides.")
+        self.sequence = sequence
+        self.chain_id = chain_id
+        self.cyclic = cyclic
+
+
+@dataclass
+class Ligand:
+    """
+    Represents a ligand with optional Chemical Component Dictionary (CCD) identifier and SMILES string.
+
+    Requires either a CCD identifier or SMILES string.
+
+    Attributes:
+        ccd (str | None): The CCD identifier for the ligand.
+        smiles (str | None): The SMILES representation of the ligand.
+    """
+
+    chain_id: str | list[str] | None = None
+    ccd: str | None = None
+    smiles: str | None = None
+
+    def __init__(
+        self,
+        *,
+        chain_id: str | list[str] | None = None,
+        ccd: str | None = None,
+        smiles: str | None = None,
+    ):
+        self.chain_id = chain_id
+        if (ccd is None and smiles is None) or (ccd is not None and smiles is not None):
+            raise ValueError("Exactly one of 'ccd' or 'smiles' must be provided.")
+        # TODO add validation
+        self.ccd = ccd
+        self.smiles = smiles

openprotein/common/__init__.py

@@ -0,0 +1,5 @@
+"""Common classes and utilities for OpenProtein."""
+
+from .features import FeatureType
+from .model_metadata import ModelDescription, ModelMetadata, TokenInfo
+from .reduction import ReductionType

openprotein/common/features.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class FeatureType(str, Enum):
+
+    PLM = "PLM"
+    SVD = "SVD"

openprotein/common/model_metadata.py

@@ -0,0 +1,33 @@
+"""Model metadata for OpenProtein models."""
+
+from pydantic import BaseModel, Field
+
+
+class ModelDescription(BaseModel):
+    """Description of available protein embedding models."""
+
+    citation_title: str | None = None
+    doi: str | None = None
+    summary: str = "Protein language model for embeddings"
+
+
+class TokenInfo(BaseModel):
+    """Information about the tokens used in the embedding model."""
+
+    id: int
+    token: str
+    primary: bool
+    description: str
+
+
+class ModelMetadata(BaseModel):
+    """Metadata about available protein embedding models."""
+
+    id: str = Field(..., alias="model_id")
+    description: ModelDescription
+    max_sequence_length: int | None = None
+    dimension: int
+    output_types: list[str]
+    input_tokens: list[str]
+    output_tokens: list[str] | None = None
+    token_descriptions: list[list[TokenInfo]]

openprotein/common/reduction.py

@@ -0,0 +1,8 @@
+"""Reduction types used in OpenProtein."""
+
+from enum import Enum
+
+
+class ReductionType(str, Enum):
+    MEAN = "MEAN"
+    SUM = "SUM"

openprotein/config.py

@@ -0,0 +1,9 @@
+POLLING_INTERVAL = 5
+POLLING_TIMEOUT = 900
+
+MAX_CONCURRENT_WORKERS = 10
+
+BASE_PAGE_SIZE = 10000
+POET_PAGE_SIZE = 50000
+POET_MAX_PAGE_SIZE = 50000
+EMBEDDING_PAGE_SIZE = 16

openprotein/csv.py

@@ -0,0 +1,31 @@
+import csv
+from typing import Iterator
+
+
+def ensure_str_lines(
+    lines: Iterator[str] | Iterator[bytes], encoding="utf-8"
+) -> Iterator[str]:
+    for line in lines:
+        if isinstance(line, bytes):
+            yield line.decode(encoding)
+        else:
+            yield line
+
+
+def parse_stream(lines: Iterator[str] | Iterator[bytes]) -> Iterator[list[str]]:
+    """
+    Returns a CSV reader from a requests.Response object.
+
+    Parameters
+    ----------
+    response : requests.Response
+        The response object to parse.
+
+    Returns
+    -------
+    csv.reader
+        A csv reader object for the response.
+    """
+    reader = csv.reader(ensure_str_lines(lines))
+    for row in reader:
+        yield row

openprotein/data/__init__.py

@@ -0,0 +1,9 @@
+"""
+Dataset management module for OpenProtein.
+
+isort:skip_file
+"""
+
+from .schemas import AssayDataPage, AssayDataRow, AssayMetadata
+from .assaydataset import AssayDataset
+from .data import DataAPI