Source code for openprotein.app.services.fold

"""Application services for Fold."""

from openprotein.api import fold
from openprotein.app.models import (
    AlphaFold2Model,
    Boltz1Model,
    Boltz1xModel,
    Boltz2Model,
    ESMFoldModel,
    FoldModel,
    FoldResultFuture,
)
from openprotein.base import APISession




class FoldAPI:
    """
    This class defines a high level interface for accessing the fold API.
    """

    esmfold: ESMFoldModel
    alphafold2: AlphaFold2Model
    af2: AlphaFold2Model
    boltz_1: Boltz1Model
    boltz1: Boltz1Model
    boltz_1x: Boltz1xModel
    boltz1x: Boltz1xModel
    boltz_2: Boltz2Model
    boltz2: Boltz2Model



    def __init__(self, session: APISession):
        self.session = session
        self._load_models()


    def _load_models(self):
        # Dynamically add model instances as attributes - precludes any drift
        models = self.list_models()
        for model in models:
            model_name = model.id.replace("-", "_")  # hyphens out
            setattr(self, model_name, model)
        # Setup aliases
        self.af2 = self.alphafold2
        self.boltz1 = self.boltz_1
        self.boltz1x = self.boltz_1x
        self.boltz2 = self.boltz_2



    def list_models(self) -> list[FoldModel]:
        """list models available for creating folds of your sequences"""
        models = []
        for model_id in fold.fold_models_list_get(self.session):
            models.append(
                FoldModel.create(
                    session=self.session, model_id=model_id, default=FoldModel
                )
            )
        return models




    def get_model(self, model_id: str) -> FoldModel:
        """
        Get model by model_id. 

        FoldModel allows all the usual job manipulation: \
            e.g. making POST and GET requests for this model specifically. 


        Parameters
        ----------
        model_id : str
            the model identifier

        Returns
        -------
        FoldModel
            The model

        Raises
        ------
        HTTPError
            If the GET request does not succeed.
        """
        return FoldModel.create(
            session=self.session, model_id=model_id, default=FoldModel
        )




    def get_results(self, job) -> FoldResultFuture:
        """
        Retrieves the results of a fold job.

        Parameters
        ----------
        job : Job
            The fold job whose results are to be retrieved.

        Returns
        -------
        FoldResultFuture
            An instance of FoldResultFuture
        """
        return FoldResultFuture.create(job=job, session=self.session)