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

openprotein/fold/models.py

@@ -0,0 +1,139 @@
+"""Fold model representations which can be used directly for creating structure predictions."""
+
+from openprotein.base import APISession
+from openprotein.common import ModelMetadata
+
+from . import api
+from .future import FoldComplexResultFuture, FoldResultFuture
+
+
+class FoldModel:
+
+    # overridden by subclasses
+    # used to get correct fold model
+    model_id: list[str] | str
+
+    def __init__(
+        self,
+        session: APISession,
+        model_id: str,
+        metadata: ModelMetadata | None = None,
+    ):
+        self.session = session
+        self.id = model_id
+        self._metadata = metadata
+
+    def __str__(self) -> str:
+        return self.id
+
+    def __repr__(self) -> str:
+        return self.id
+
+    @classmethod
+    def get_model(cls):
+        """
+        Get the model_id(s) for this FoldModel subclass.
+
+        Returns
+        -------
+        list of str
+            List of model_id strings associated with this class.
+        """
+        if isinstance(cls.model_id, str):
+            return [cls.model_id]
+        return cls.model_id
+
+    @classmethod
+    def create(
+        cls,
+        session: APISession,
+        model_id: str,
+        default: type["FoldModel"] | None = None,
+        **kwargs,
+    ):
+        """
+        Create and return an instance of the appropriate FoldModel subclass based on the model_id.
+
+        Parameters
+        ----------
+        session : APISession
+            The API session to use.
+        model_id : str
+            The model identifier.
+        default : type[FoldModel] or None, optional
+            Default FoldModel subclass to use if no match is found.
+        **kwargs : dict, optional
+            Additional keyword arguments to pass to the model constructor.
+
+        Returns
+        -------
+        FoldModel
+            An instance of the appropriate FoldModel subclass.
+
+        Raises
+        ------
+        ValueError
+            If no suitable FoldModel subclass is found and no default is provided.
+        """
+        # Dynamically discover all subclasses of FoldModel
+        model_classes = FoldModel.__subclasses__()
+
+        # Find the FoldModel class that matches the model_id
+        for model_class in model_classes:
+            if model_id in model_class.get_model():
+                return model_class(session=session, model_id=model_id, **kwargs)
+        # default to FoldModel
+        if default is not None:
+            try:
+                return default(session=session, model_id=model_id, **kwargs)
+            except:
+                pass
+        raise ValueError(f"Unsupported model_id type: {model_id}")
+
+    @property
+    def metadata(self):
+        """
+        ModelMetadata : Model metadata for this model.
+        """
+        if self._metadata is None:
+            self._metadata = self.get_metadata()
+        return self._metadata
+
+    def get_metadata(self) -> ModelMetadata:
+        """
+        Get model metadata for this model.
+
+        Returns
+        -------
+        ModelMetadata
+            The metadata associated with this model.
+        """
+        return api.fold_model_get(self.session, self.id)
+
+    def fold(self, **kwargs) -> FoldResultFuture | FoldComplexResultFuture:
+        """
+        Fold a sequence using this model.
+
+        Parameters
+        ----------
+        **kwargs : dict, optional
+            Additional keyword arguments to pass to the underlying
+            `fold` request.
+
+        Returns
+        -------
+        FoldResultFuture or FoldComplexResultFuture
+            Future object representing the fold result.
+        """
+        return FoldResultFuture.create(
+            session=self.session,
+            job=api.fold_models_post(
+                session=self.session,
+                model_id=(
+                    model_id
+                    if isinstance(model_id := self.model_id, str)
+                    else model_id[0]
+                ),
+                **kwargs,
+            ),
+        )

openprotein/fold/schemas.py

@@ -0,0 +1,39 @@
+"""Schema for OpenProtein fold system."""
+
+from typing import Literal
+
+from pydantic import BaseModel
+
+from openprotein.jobs import BatchJob, Job, JobType
+
+
+class FoldMetadata(BaseModel):
+    """
+    Metadata for a folding job.
+
+    Attributes
+    ----------
+    job_id : str
+        Unique identifier for the job.
+    model_id : str
+        Identifier for the model used in the job.
+    args : dict or None, optional
+        Additional arguments for the job.
+    """
+
+    job_id: str
+    model_id: str
+    args: dict | None = None
+
+
+class FoldJob(Job, BatchJob):
+    """
+    Folding job class.
+
+    Attributes
+    ----------
+    job_type : Literal[JobType.embeddings_fold]
+        The type of job, set to embeddings_fold.
+    """
+
+    job_type: Literal[JobType.embeddings_fold]

openprotein/jobs/__init__.py

@@ -0,0 +1,9 @@
+"""
+Jobs module for OpenProtein.
+
+isort:skip_file
+"""
+
+from .schemas import BatchJob, Job, JobStatus, JobType
+from .futures import Future, MappedFuture, PagedFuture, StreamingFuture
+from .jobs import JobsAPI

openprotein/jobs/api.py

@@ -0,0 +1,71 @@
+"""Jobs and job-centric flows."""
+
+from typing import List
+
+from pydantic import TypeAdapter
+
+from openprotein.base import APISession
+
+from .schemas import Job
+
+
+def job_args_get(session: APISession, job_id: str) -> dict:
+    """Get job."""
+    endpoint = f"v1/jobs/{job_id}/args"
+    response = session.get(endpoint)
+    return dict(**response.json())
+
+
+def job_get(session: APISession, job_id: str) -> Job:
+    """Get job."""
+    endpoint = f"v1/jobs/{job_id}"
+    response = session.get(endpoint)
+    return TypeAdapter(Job).validate_python(response.json())
+
+
+def jobs_list(
+    session: APISession,
+    status: str | None = None,
+    job_type: str | None = None,
+    assay_id: str | None = None,
+    more_recent_than: str | None = None,
+    limit: int | None = None,
+) -> List[Job]:
+    """
+    Retrieve a list of jobs filtered by specific criteria.
+
+    Parameters
+    ----------
+    session : APISession
+        The current API session for communication with the server.
+    status : str, optional
+        Filter by job status. If None, jobs of all statuses are retrieved. Default is None.
+    job_type : str, optional
+        Filter by Filter. If None, jobs of all types are retrieved. Default is None.
+    assay_id : str, optional
+        Filter by assay. If None, jobs for all assays are retrieved. Default is None.
+    more_recent_than : str, optional
+        Retrieve jobs that are more recent than a specified date. If None, no date filtering is applied. Default is None.
+
+    Returns
+    -------
+    List[Job]
+        A list of Job instances that match the specified criteria.
+    """
+    endpoint = "v1/jobs"
+
+    params = {}
+    if status is not None:
+        params["status"] = status
+    if job_type is not None:
+        params["job_type"] = job_type
+    if assay_id is not None:
+        params["assay_id"] = assay_id
+    if more_recent_than is not None:
+        params["more_recent_than"] = more_recent_than
+    if limit is not None:
+        params["limit"] = limit
+
+    response = session.get(endpoint, params=params)
+    # return jobs, not futures
+    return TypeAdapter(List[Job]).validate_python(response.json())