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

openprotein/prompt/models.py

@@ -0,0 +1,142 @@
+from openprotein.base import APISession
+from openprotein.jobs import Future, JobsAPI
+from openprotein.protein import Protein
+
+from . import api
+from .schemas import PromptJob, PromptMetadata, QueryMetadata
+
+
+class Prompt(Future):
+    """Prompt which contains a set of sequences and/or structures used to condition the PoET models."""
+
+    metadata: PromptMetadata
+    session: APISession
+    job: PromptJob | None
+
+    def __init__(
+        self,
+        session: APISession,
+        job: PromptJob | None = None,
+        metadata: PromptMetadata | None = None,
+        num_replicates: int | None = None,
+    ):
+        """
+        Initialize a new Prompt instance.
+
+        Parameters
+        ----------
+        session : APISession
+            An APISession object used for interacting with the API.
+        job: PromptJob | None
+            A PromptJob containing information about the optional prompt job.
+        metadata : PromptMetadata
+            A PromptMetadata object containing metadata for the prompt.
+        """
+        """Initializes with either job get or svd metadata get."""
+        if metadata is None:
+            # use job to fetch metadata
+            if job is None:
+                raise ValueError("Expected prompt metadata or job")
+            # if no num_replicates, we need an api call to get the info
+            if num_replicates is None:
+                metadata = api.get_prompt_metadata(
+                    session=session, prompt_id=job.job_id
+                )
+            # else we can just build the metadata from the job
+            else:
+                metadata = PromptMetadata(
+                    id=job.job_id,
+                    name=job.job_id,
+                    description=None,
+                    created_date=job.created_date,
+                    num_replicates=num_replicates,
+                    job_id=job.job_id,
+                    status=job.status,
+                )
+        self.metadata = metadata
+        self.session = session
+        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)
+
+    def __str__(self) -> str:
+        return str(self.metadata)
+
+    def __repr__(self) -> str:
+        return repr(self.metadata)
+
+    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):
+        if self.job is None:
+            return None
+        return super()._wait_job(**kwargs)
+
+    @property
+    def id(self):
+        return self.metadata.id
+
+    @property
+    def name(self):
+        return self.metadata.name
+
+    @property
+    def description(self):
+        return self.metadata.description
+
+    @property
+    def created_date(self):
+        return self.metadata.created_date
+
+    @property
+    def num_replicates(self):
+        return self.metadata.num_replicates
+
+    @property
+    def status(self):
+        if self.job is not None:
+            return super().status
+        return self.metadata.status
+
+
+class Query:
+    """Query containing a sequence/structure used to query the PoET-2 model which opens up new workflows."""
+
+    metadata: QueryMetadata
+    session: APISession
+
+    def __init__(self, session: APISession, metadata: QueryMetadata):
+        """
+        Initialize a new Query instance.
+
+        Parameters
+        ----------
+        session : APISession
+            An APISession object used for interacting with the API.
+        metadata : QueryMetadata
+            A QueryMetadata object containing metadata for the query.
+        """
+        self.metadata = metadata
+        self.session = session
+
+    def __str__(self) -> str:
+        return str(self.metadata)
+
+    def __repr__(self) -> str:
+        return repr(self.metadata)
+
+    def get(self) -> Protein:
+        query = api.get_query(session=self.session, query_id=str(self.id))
+        return query
+
+    @property
+    def id(self):
+        return self.metadata.id
+
+    @property
+    def created_date(self):
+        return self.metadata.created_date

openprotein/prompt/prompt.py

@@ -0,0 +1,130 @@
+"""Prompt API providing the interface to create prompts for use with PoET models."""
+
+from typing import List, Sequence
+
+from openprotein.base import APISession
+from openprotein.protein import Protein
+
+from . import api
+from .models import Prompt, Query
+from .schemas import Context
+
+
+class PromptAPI:
+    """Prompt API providing the interface to create prompts for use with PoET models."""
+
+    def __init__(self, session: APISession):
+        self.session = session
+
+    def create_prompt(
+        self,
+        context: Context | Sequence[Context],
+        name: str | None = None,
+        description: str | None = None,
+    ) -> Prompt:
+        """
+        Create a prompt.
+
+        Parameters
+        ----------
+        context : Context | Sequence[Context]
+            context or list of contexts, where each context is a Sequence of str,
+            bytes, and/or Protein
+        query : Optional[bytes | str | Protein]
+            Optional query provided as sequence/structure
+        name : str
+            Name of the prompt.
+        description : Optional[str]
+            Description of the prompt.
+
+        Returns
+        -------
+        PromptMetadata
+            Metadata of the created prompt.
+        """
+        return Prompt(
+            session=self.session,
+            metadata=api.create_prompt(
+                session=self.session,
+                context=context,
+                name=name,
+                description=description,
+            ),
+        )
+
+    def get_prompt(self, prompt_id: str) -> Prompt:
+        """
+        Get the prompt for a given prompt ID.
+
+        Parameters
+        ----------
+        prompt_id : str
+            The prompt ID.
+
+        Returns
+        -------
+        BinaryIO
+            The prompt data in binary format.
+        """
+        return Prompt(
+            session=self.session,
+            metadata=api.get_prompt_metadata(session=self.session, prompt_id=prompt_id),
+        )
+
+    def list_prompts(self) -> List[Prompt]:
+        """
+        List all prompts.
+
+        Returns
+        -------
+        List[PromptMetadata]
+            List of prompt metadata.
+        """
+        return [
+            Prompt(session=self.session, metadata=p)
+            for p in api.list_prompts(session=self.session)
+        ]
+
+    def create_query(
+        self,
+        query: str | bytes | Protein,
+    ) -> Query:
+        """
+        Create a query.
+
+        Parameters
+        ----------
+        query : Optional[bytes | str | Protein]
+            Optional query provided as sequence/structure
+
+        Returns
+        -------
+        QueryMetadata
+            Metadata of the created query.
+        """
+        return Query(
+            session=self.session,
+            metadata=api.create_query(
+                session=self.session,
+                query=query,
+            ),
+        )
+
+    def get_query(self, query_id: str) -> Query:
+        """
+        Get the query for a given query ID.
+
+        Parameters
+        ----------
+        query_id : str
+            The query ID.
+
+        Returns
+        -------
+        BinaryIO
+            The query data in binary format.
+        """
+        return Query(
+            session=self.session,
+            metadata=api.get_query_metadata(session=self.session, query_id=query_id),
+        )

openprotein/prompt/schemas.py

@@ -0,0 +1,49 @@
+from datetime import datetime
+from typing import Literal, Sequence
+
+from pydantic import BaseModel, Field
+
+from openprotein.jobs import Job, JobStatus, JobType
+from openprotein.protein import Protein
+
+Context = Sequence[bytes | str | Protein]
+
+
+class PromptJob(Job):
+    """A representation of a prompt job."""
+
+    job_type: Literal[JobType.align_prompt]
+
+    @property
+    def msa_id(self):
+        """ID of the underlying MSA."""
+        return self.msa_id
+
+    @property
+    def prompt_id(self):
+        """Prompt ID."""
+        return self.job_id
+
+
+class PromptMetadata(BaseModel):
+    """Metadata about a prompt."""
+
+    id: str = Field(description="Prompt unique identifier.")
+    name: str = Field(description="Name of the prompt")
+    description: str | None = Field(
+        None,
+        description="Description of the prompt",
+    )
+    created_date: datetime = Field(description="The date the prompt was created.")
+    num_replicates: int = Field(description="Number of replicates provided as context.")
+    job_id: str | None = Field(
+        None, description="The job_id of the sampling job, if it exists."
+    )
+    status: JobStatus = Field(description="The status of the prompt.")
+
+
+class QueryMetadata(BaseModel):
+    """Metadata about a query."""
+
+    id: str = Field(description="Query unique identifier.")
+    created_date: datetime = Field(description="The date the query was created.")