cradle-sdk 0.1.1__py3-none-any.whl

cradle/sdk/types/task.py

@@ -0,0 +1,985 @@
+from abc import ABC
+from enum import StrEnum
+from typing import Annotated, ClassVar, Literal
+
+from pydantic import BaseModel, Field
+
+from .assembly import AssemblyFeature
+from .common import (
+    ArchivableResourceMixin,
+    BaseListResponse,
+    Context,
+    ErrorResponseMixin,
+    ResourceResponse,
+)
+
+
+class ParametersBase(BaseModel):
+    pass
+
+
+class TableSourceTable(BaseModel):
+    kind: Literal["TABLE"] = Field(default="TABLE")
+    table: str
+
+
+class TableSourceQuery(BaseModel):
+    kind: Literal["QUERY"] = Field(default="QUERY")
+    query: str
+
+
+class TableSourceTaskResult(BaseModel):
+    kind: Literal["TASK_RESULT"] = Field(default="TASK_RESULT")
+    table: str
+    task_id: int | list[int]
+
+
+class FeatureAnnotation(StrEnum):
+    """Annotation to be configured per residue in a domain:
+
+    UNRELIABLE: The evolutionary context is not reliable, eg in CDRs of antibodies.
+    IGNORE: A part of the sequence that should not be used for evolutionary search and won't be mutated, this can be tags, linkers or other artificial subsequences.
+    """
+
+    UNRELIABLE = "UNRELIABLE"
+    IGNORE = "IGNORE"
+
+
+class DomainFeatureItem(BaseModel):
+    annotation: FeatureAnnotation = Field(
+        description="Annotation for the residues. `UNRELIABLE`: do not look at the MSA to infer conservation-based features (e.g. CDRs) for the protein `IGNORE`: do not change this position at all (e.g. for linkers) and do not use to infer any conservation-based features."
+    )
+    ranges: list[tuple[int, int]] = Field(
+        description="List of ranges where the annotation applies. Each range is a tuple of (start, end). Indices are 0-based and inclusive for the start and exclusive for the end. For example, (0, 5) blocks positions 0 to 4."
+    )
+
+
+class ProteinModelType(StrEnum):
+    DEFAULT = "DEFAULT"
+    ANTIBODY = "ANTIBODY"
+
+
+class ScaleType(StrEnum):
+    MULTIPLICATIVE = "MULTIPLICATIVE"
+    ADDITIVE = "ADDITIVE"
+    RANK = "RANK"
+
+
+class Assay(BaseModel):
+    """Metadata of an assay."""
+
+    assay_id: str = Field(description="ID of the assay")
+    name: str = Field(description="Human readable name of the assay")
+    scale_type: ScaleType = Field(
+        description="""The scale type defines how the assay behaves.
+
+* **Additive**: Assay values are comparable in magnitude over batches, a delta of eg 5 in one batch is equivalent to 5 in another batch.
+* **Multiplicative**: Assay values are comparable over batches by applying a multiplication factor, eg fold improvement: `score_variantB_batch1 = factor(typically starting sequence score) * score_variantB_batch`.
+* **Rank**: Assay values are not comparable over batches, we can only assume the ranking within a batch is correct."""
+    )
+    unit: str | None = Field(default=None, description="Unit of the assay - used for display purposes only")
+
+
+class OptimizationDirection(StrEnum):
+    MAXIMIZE = "MAXIMIZE"
+    MINIMIZE = "MINIMIZE"
+
+
+class PrimaryObjective(BaseModel):
+    assay_id: str = Field(description="The ID of the assay to be optimized.")
+    direction: OptimizationDirection = Field(description="Whether the maximize or minimize the assay value.")
+
+
+class TrainParameters(ParametersBase):
+    task_type: ClassVar[str] = "train/v1"
+
+    homologs: "TableInput" = Field(
+        description="The table of homologous sequences to be used for training the base model."
+    )
+    domain_features: dict[
+        str,
+        Annotated[
+            list[DomainFeatureItem],
+            Field(
+                description="List of domain annotations. A residue may be assigned zero or one annotations. Ranges which are not annotated are considered reliable, i.e. the MSA in those ranges is used to infer conservation-based features."
+            ),
+        ],
+    ] = Field(
+        default_factory=dict,
+        description="Mapping of protein sequences to their respective domain features. Each range in the domain features must be a valid index into the corresponding protein sequence.",
+    )
+    protein_model_type: ProteinModelType = Field(
+        default=ProteinModelType.DEFAULT,
+        description="Specifies the type of base model to use for training: use `DEFAULT` for single-chain proteins and `ANTIBODY` for multi-chain antibodies.",
+    )
+    dataset: "TableInput" = Field(description="The name of the input table to be used for model training.")
+    assays: list[Assay] = Field(default_factory=list, description="List of assay metadata entries")
+    primary_objective: PrimaryObjective = Field(
+        description="The primary objective on which the sampler model is conditioned. Should correspond to an assay in `dataset`."
+    )
+
+
+class AbsoluteConstraint(BaseModel):
+    type_: Literal["ABSOLUTE_CONSTRAINT"] = Field(
+        default="ABSOLUTE_CONSTRAINT", alias="type", validation_alias="type", serialization_alias="type"
+    )
+    assay_id: str = Field(description="The ID of the constrained assay.")
+    direction: OptimizationDirection = Field(
+        description="Whether the assay value should be above (MAXIMIZE) or below (MINIMIZE) the threshold."
+    )
+    threshold: float = Field(description="The threshold assay value.")
+
+
+class RelativeConstraint(BaseModel):
+    type_: Literal["RELATIVE_CONSTRAINT"] = Field(
+        default="RELATIVE_CONSTRAINT", alias="type", validation_alias="type", serialization_alias="type"
+    )
+    assay_id: str = Field(description="The ID of the constrained assay.")
+    direction: OptimizationDirection = Field(
+        description="Whether the assay value of new sequences should be above (MAXIMIZE) or below (MINIMIZE) the assay value of the `relative_to` sequence."
+    )
+    relative_to: str = Field(description="The sequence relative to which new sequences are constrained.")
+
+
+class ArtifactParam(BaseModel):
+    artifact_id: int
+
+
+class BlockedAAItem(BaseModel):
+    blocked_aas: str = Field(
+        description="One or more amino acid letter code to be blocked from being mutated to. Use '*' to indicate that a position is fully blocked from being mutated."
+    )
+    ranges: list[tuple[int, int]] = Field(
+        description="List of ranges where the mutations are blocked. Each range is a tuple of (start, end). Indices are 0-based and inclusive for the start and exclusive for the end. For example, (0, 5) blocks positions 0 to 4."
+    )
+
+
+class BlockedMotifItem(BaseModel):
+    blocked_motif: str = Field(description="Amino acid motif to be blocked from appearing in the final sequence.")
+    ranges: list[tuple[int, int]] = Field(
+        description="List of ranges where the motif is blocked. Each range is a tuple of (start, end). Indices are 0-based and inclusive for the start and exclusive for the end. For example, (0, 5) blocks positions 0 to 4. The entire motif must fit into the specified positions to be blocked."
+    )
+
+
+class TemplateInputs(BaseModel):
+    """Configuration for template-based sequence generation."""
+
+    template_id: str = Field(
+        description="Identifier for the template used to identify the source of generated sequences in the output."
+    )
+    sequence: str = Field(description="The template amino acid sequence")
+    num_results: int = Field(
+        description="Number of sequences generated from this template to be added to the final plate."
+    )
+    min_mutations: int = Field(default=1, description="Minimum number of mutations per sequence.")
+    max_mutations: int = Field(default=8, description="Maximum number of mutations per sequence.")
+    hit_redundancy: int = Field(
+        default=10,
+        description="Number of top candidates for which to optimize the average score. A higher number means less diversity.",
+    )
+    blocked_aas: list[BlockedAAItem] = Field(
+        default_factory=list,
+        description="List of blocked amino acids which may not be mutated to, with their respective ranges.",
+    )
+    blocked_motifs: list[BlockedMotifItem] = Field(
+        default_factory=list, description="List of blocked amino acid motifs with their respective ranges."
+    )
+    blocked_regexps: list[str] | None = Field(
+        default_factory=list, description="Regular expressions that the sequence cannot match"
+    )
+    must_match_regexps: list[str] | None = Field(
+        default_factory=list,
+        description="Generated sequences must match all of the given regular expressions. Can be used to enforce certain motifs.",
+    )
+
+
+class EngineerParameters(ParametersBase):
+    task_type: ClassVar[str] = "engineer/v1"
+
+    dataset: "TableInput" = Field(description="The table of assayed data, to be used for exploitation.")
+    assays: list[Assay] = Field(default_factory=list, description="List of assay metadata entries")
+    primary_objective: PrimaryObjective = Field(description="The primary objective the samplers are trained for.")
+    constraints: "Constraints" = Field(default_factory=list, description="List of assay constraints")
+    samplers: list[ArtifactParam] = Field(description="The sampling models used to generate candidate sequences.")
+    scorer: ArtifactParam = Field(description="Scoring model used to predict assay values for candidate sequences.")
+    template_sequences: list[TemplateInputs] = Field(description="Initial sequences used as templates for generation")
+
+
+class MonomerAssembly(BaseModel):
+    monomer: str = Field(description="The monomer sequence")
+
+
+class MonomerSamplingTemplate(BaseModel):
+    kind: Literal["MONOMER"] = Field(default="MONOMER")
+    assembly: MonomerAssembly = Field(description="The template sequence to use for generation.")
+    sampler: ArtifactParam = Field(description="The trained sampler artifact used to generate candidate sequences.")
+
+
+class VhVlAssembly(BaseModel):
+    vh: str | None = Field(description="The heavy chain sequence")
+    vl: str | None = Field(description="The light chain sequence")
+
+
+class VhVlSamplingTemplate(BaseModel):
+    kind: Literal["VHVL"] = Field(default="VHVL")
+    assembly: VhVlAssembly = Field(description="The template sequence to use for generation.")
+    sampler: ArtifactParam = Field(description="The trained sampler artifact used to generate candidate sequences.")
+
+
+class RangeAnnotation(BaseModel):
+    ranges: dict[str, list[tuple[int, int]]] = Field(
+        description="A mapping from polymer names to annotation ranges. Each range is a tuple of (start, end). Indices are 0-based and inclusive for the start and exclusive for the end. For example, (0, 5) refers to positions 0 to 4."
+    )
+
+
+class Requirement(BaseModel):
+    """Assembly generation requirements specifying required and forbidden features on the generated assemblies.
+
+    Each requirement defines constraints for a specific region of the assembly.
+    Multiple requirements can target different or overlapping regions.
+    """
+
+    apply_to: RangeAnnotation
+    motifs: list[Annotated[str, Field(description="A motif or regular expression of amino acids.")]]
+    operator: Literal["none", "all", "any"] = Field(
+        description="The choice of how to interpret the `motifs` together as a requirement."
+    )
+
+
+class ModelBasedGenerator(BaseModel):
+    """Uses a trained sampler model to generate assemblies."""
+
+    type: Literal["MODEL_BASED"] = Field(default="MODEL_BASED")
+    template: MonomerSamplingTemplate | VhVlSamplingTemplate = Field(
+        description="The template sequence to use for generation.", discriminator="kind"
+    )
+    requirements: list[Requirement] = Field(
+        default_factory=list,
+        description="Assembly based generation requirements to satisfy. All requirements must be met for a valid assembly.",
+    )
+    discourage_mutations: RangeAnnotation | None = Field(
+        default=None, description="Expression defining regions where mutations should be discouraged."
+    )
+    min_mutations: int = Field(description="Minimum number of mutations from the template assembly.")
+    max_mutations: int = Field(description="Maximum number of mutations from the template assembly.")
+
+
+class MonomerSource(BaseModel):
+    kind: Literal["MONOMER"] = Field(default="MONOMER")
+    source: "TableInput"
+
+
+class VhVlSource(BaseModel):
+    kind: Literal["VHVL"] = Field(default="VHVL")
+    source: "TableInput"
+
+
+class Explicit(BaseModel):
+    """A pre-defined table of assemblies to use directly for generation without modification.
+
+    Assemblies will be evaluated in the order provided.
+    """
+
+    type: Literal["EXPLICIT"] = Field(default="EXPLICIT")
+    source: MonomerSource | VhVlSource = Field(
+        description="A pre-defined table of assemblies to use directly for generation without modification. Assemblies will be evaluated in the order provided.",
+        discriminator="kind",
+    )
+
+
+class AllowedMutation(BaseModel):
+    in_range: RangeAnnotation = Field(
+        description="The annotation expression defining the positions at which the amino acids listed in `amino_acids` can be used."
+    )
+    amino_acids: list[Annotated[str, Field(description="A single amino acid.")]]
+
+
+class CombinatorialGenerator(BaseModel):
+    """Combinatorial assembly generation from explicit mutation.
+
+    Systematically generates assemblies by combining allowed mutations up to a
+    specified maximum number.
+    """
+
+    type: Literal["COMBINATORIAL"] = Field(default="COMBINATORIAL")
+    template: MonomerAssembly | VhVlAssembly = Field(
+        description="The assembly that serves as template for generation, including any annotations to be used in `requirements`."
+    )
+    max_mutations: int = Field(
+        description="Maximum number of mutation to combine in generation. Lower values will result in fewer total candidate sequences evaluated for this generator."
+    )
+    requirements: list[Requirement] = Field(
+        default_factory=list,
+        description="Assembly based generation requirements to satisfy. All requirements must be met for a valid sequence.",
+    )
+    allowed_mutations: list[AllowedMutation] = Field(
+        description="The complete set of mutations that can be used in generation."
+    )
+
+
+class Parallel(BaseModel):
+    """An ensemble of generators to use in parallel.
+
+    All generators in the ensemble run independently and each contribute an equal
+    number of candidate assemblies for evaluation. This does not guarantee that each
+    generator will contribute an equal number of assemblies to `selected_assemblies`,
+    as assemblies from some generators in `ensemble` may perform better than others.
+    """
+
+    type: Literal["PARALLEL"] = Field(default="PARALLEL")
+    ensemble: list["Generator"] = Field(
+        default_factory=list, description="The list of generators to run in parallel within this ensemble."
+    )
+
+
+class Sequential(BaseModel):
+    """An ensemble of generators to evaluate sequentially.
+
+    Each generator will have the opportunity to evaluate an equal number of candidate
+    assemblies. However, if earlier generators are able to satisfy the requirements of
+    the `ranker`, later generators may never be used. In this way this method of
+    ensembling expresses a strict preference among the generators in `ensemble`.
+    """
+
+    type: Literal["SEQUENTIAL"] = Field(default="SEQUENTIAL")
+    ensemble: list["Generator"] = Field(
+        default_factory=list, description="The list of generators to run in sequence within this ensemble."
+    )
+
+
+class Measure(BaseModel):
+    assay_id: str = Field(description="The ID of the constrained assay.")
+    direction: OptimizationDirection = Field(
+        description="Whether the assay value should be above (MAXIMIZE) or below (MINIMIZE) the threshold."
+    )
+
+
+class RelativeTo(BaseModel):
+    reference: MonomerAssembly | VhVlAssembly = Field(description="The reference assembly.")
+    margin: float | None = Field(
+        default=None,
+        description="The margin (additive assay) or desired fold improvement (multiplicative assay) relative to the reference assembly.",
+    )
+
+
+class PrimaryObjectiveV2(BaseModel):
+    measure: Measure
+    reference: float | RelativeTo = Field(
+        description="The reference relative to which the primary objective is defined."
+    )
+
+
+class ConstraintV2(BaseModel):
+    measure: Measure
+    threshold: float | RelativeTo = Field(
+        description="The threshold value for the assay, expressed as an absolute value or relative to another sequence. Specify relative constraints where possible. Absolute values cannot be used for assays with a RANK scale type, and must be positive for assay with a MULTIPLICATIVE scale type."
+    )
+
+
+class MonomerScorerConfig(BaseModel):
+    kind: Literal["MONOMER_SCORER"] = Field(default="MONOMER_SCORER")
+    scorer: ArtifactParam = Field(description="The trained scoring model artifact used to score sequences.")
+    controls: "TableInput | None" = Field(
+        default=None,
+        description="Optional control sequences to include in the final output for experimental validation. These controls should generally be included in the assayed plate, as `selected_sequences` will be intentionally diverse from these sequences. However these sequences will *not* be included in the `selected_sequences` output table, so -- if used -- take care to include reinclude them. ",
+    )
+
+
+class VhVlScorerConfig(BaseModel):
+    kind: Literal["VHVL_SCORER"] = Field(default="VHVL_SCORER")
+    scorer: ArtifactParam = Field(description="The trained scoring model artifact used to score sequences.")
+    controls: "TableInput | None" = Field(
+        default=None,
+        description="Optional control sequences to include in the final output for experimental validation. These controls should generally be included in the assayed plate, as `selected_sequences` will be intentionally diverse from these sequences. However these sequences will *not* be included in the `selected_sequences` output table, so -- if used -- take care to include reinclude them. ",
+    )
+
+
+class ScoreBasedRanker(BaseModel):
+    """Configuration for ranking and selecting generated assemblies."""
+
+    type: Literal["SCORE_BASED"] = Field(default="SCORE_BASED")
+    primary_objective: PrimaryObjectiveV2 = Field(description="The primary optimization objective of `engineer`.")
+    constraints: "ConstraintsV2" = Field(description="List of constraints `engineer` will attempt to satisfy.")
+    scorer_config: MonomerScorerConfig | VhVlScorerConfig = Field(
+        description="The trained scoring model artifact used to score sequences, along with any optional control sequences to include in the final output for experimental validation.",
+        discriminator="kind",
+    )
+    hit_redundancy: int = Field(
+        default=3,
+        description="""The maximum number of similar variants within the proposed set.
+
+We say two variants are "similar" when their predicted performances are highly correlated.
+At lower values this will mean on average fewer hits, but greater diversity within the set,
+while at maximum value (`hit_redundancy=num_selected`) the selection is purely greedy, with
+no requirement on diversity within the set.""",
+    )
+
+
+class EngineerParametersV2(ParametersBase):
+    task_type: ClassVar[str] = "engineer/v2"
+
+    generator: "Generator" = Field(
+        description="The generator configuration or configurations that define requirements and strategies for generating candidate assemblies.",
+        discriminator="type",
+    )
+    ranker: ScoreBasedRanker = Field(
+        description="The configuration defining how candidate assemblies are evaluated for inclusion in `selected_assemblies`."
+    )
+    num_assemblies: int = Field(
+        description="The target number of selected assemblies to output. In rare cases of very difficult `constraints` this task may produce a smaller number."
+    )
+
+
+class DiversifyTemplateInputs(BaseModel):
+    """Configuration for template-based sequence generation."""
+
+    template_id: str = Field(
+        description="Identifier for the template used to identify the source of generated sequences in the output."
+    )
+    sequence: str = Field(description="The template amino acid sequence")
+    num_results: int = Field(
+        description="Number of sequences to generate for the final plate from the template sequence."
+    )
+    min_mutations: int = Field(default=1, description="Minimum number of mutations per sequence")
+    max_mutations: int = Field(default=4, description="Maximum number of mutations per sequence")
+    blocked_aas: list[BlockedAAItem] = Field(
+        default_factory=list,
+        description="List of blocked amino acids which may not be mutated to, with their respective ranges.",
+    )
+    blocked_motifs: list[BlockedMotifItem] = Field(
+        default_factory=list, description="List of blocked amino acid motifs with their respective ranges."
+    )
+
+
+class DiversifyParameters(ParametersBase):
+    """Parameters for the diversify task."""
+
+    task_type: ClassVar[str] = "diversify/v1"
+
+    template_sequence: DiversifyTemplateInputs = Field(description="The template amino acid sequence to diversify.")
+    homologs: "TableInput" = Field(
+        description="Table of homologs to use for generator training. In the simplest case, this is just the result of a multiple sequence alignment (MSA) against the sequence to optimize."
+    )
+    domain_features: dict[
+        str,
+        Annotated[
+            list[DomainFeatureItem],
+            Field(
+                description="List of domain annotations. A residue may be assigned zero or one annotations. Ranges which are not annotated are considered reliable, i.e. the MSA in those ranges is used to infer conservation-based features."
+            ),
+        ],
+    ] = Field(
+        default_factory=dict,
+        description="Properties for each domain. A domain is one or more subsequences of the protein sequence, for example the active site of an enzyme, CDRs of an antibody, heavy and light chains of an scFv, etc. ",
+    )
+    protein_model_type: ProteinModelType = Field(
+        default=ProteinModelType.DEFAULT,
+        description="Specifies the type of base model to use for training: use `DEFAULT` for single-chain proteins and `ANTIBODY` for multi-chain antibodies.",
+    )
+
+
+class MonomerSamplerConfig(BaseModel):
+    kind: Literal["MONOMER_SAMPLER"] = Field(default="MONOMER_SAMPLER")
+    sampler: ArtifactParam = Field(description="The trained sampler artifact used to compute sequence likelihoods.")
+    reference_assembly: MonomerAssembly = Field(
+        description="The reference assembly used as the denominator when computing sequence likelihood ratios."
+    )
+    controls: "TableInput | None" = Field(
+        default=None,
+        description="Optional control assemblies to include in the final output for experimental validation. These controls should generally be included in the assayed plate, as `selected_assemblies` will be intentionally diverse from these assemblies. However these assemblies will *not* be included in the `selected_assemblies` output table, so -- if used -- take care to include reinclude them. ",
+    )
+
+
+class VhVlSamplerConfig(BaseModel):
+    kind: Literal["VHVL_SAMPLER"] = Field(default="VHVL_SAMPLER")
+    sampler: ArtifactParam = Field(description="The trained sampler artifact used to compute assembly likelihoods.")
+    reference_assembly: VhVlAssembly = Field(
+        description="The reference assembly used as the denominator when computing assembly likelihood ratios."
+    )
+    controls: "TableInput | None" = Field(
+        default=None,
+        description="Optional control assemblies to include in the final output for experimental validation. These controls should generally be included in the assayed plate, as `selected_assemblies` will be intentionally diverse from these assemblies. However these assemblies will *not* be included in the `selected_assemblies` output table, so -- if used -- take care to include reinclude them. ",
+    )
+
+
+class LikelihoodBasedRanker(BaseModel):
+    """Configuration for ranking and selecting generated assemblies based on likelihoods."""
+
+    type: Literal["LIKELIHOOD_BASED"] = Field(default="LIKELIHOOD_BASED")
+    sampler_config: MonomerSamplerConfig | VhVlSamplerConfig = Field(
+        description="The trained sampler artifact used to compute assembly likelihoods, along with any optional control assemblies to include in the final output for experimental validation.",
+        discriminator="kind",
+    )
+    hit_redundancy: int = Field(
+        default=3,
+        description="""The maximum number of similar variants within the proposed set.
+
+We say two variants are "similar" when their predicted performances are highly correlated.
+At lower values this will mean on average fewer hits, but greater diversity within the set,
+while at maximum value (`hit_redundancy=num_selected`) the selection is purely greedy, with
+no requirement on diversity within the set.""",
+    )
+
+
+class DiversifyParametersV2(ParametersBase):
+    task_type: ClassVar[str] = "diversify/v2"
+
+    generator: "Generator" = Field(
+        description="The generator configuration or configurations that define requirements and strategies for generating candidate assemblies.",
+        discriminator="type",
+    )
+    ranker: LikelihoodBasedRanker = Field(
+        description="The configuration defining how candidate assemblies are evaluated for inclusion in `selected_assemblies`."
+    )
+    num_assemblies: int = Field(
+        description="The target number of `selected_assemblies` to output. In rare cases, the actual number of output assemblies may be slightly lower."
+    )
+
+
+class AnalyzeDataParameters(ParametersBase):
+    task_type: ClassVar[str] = "analyze.data/v1"
+
+    reference_sequence: str = Field(description="The sequence to consider as a reference for computing mutations.")
+    dataset: "TableInput" = Field(description="The assay data to be used for model training.")
+    assays: list[Assay] = Field(default_factory=list, description="List of assay metadata entries")
+
+
+class AnalyzeDataParametersV2(ParametersBase):
+    task_type: ClassVar[str] = "analyze.data/v2"
+
+    dataset: "TableInput" = Field(description="The data to analyze before training.")
+    reference: MonomerAssembly | VhVlAssembly = Field(
+        description="The reference assembly used for computing mutations."
+    )
+    assays: list[Assay] = Field(default_factory=list, description="The assays present in the dataset.")
+
+
+class TemplateMetadata(BaseModel):
+    template_id: str = Field(
+        description="Identifier for the template used to identify the source of generated sequences."
+    )
+    sequence: str = Field(description="Amino acid template sequence used for generation.")
+    blocked_aas: list[BlockedAAItem] = Field(
+        default_factory=list,
+        description="List of blocked amino acids which may not be mutated to, with their respective ranges.",
+    )
+    min_mutations: int
+    max_mutations: int
+    structure: ArtifactParam | None = Field(default=None)
+
+
+class AnalyzeDiversifyParameters(ParametersBase):
+    task_type: ClassVar[str] = "analyze.diversify/v1"
+
+    selected_sequences: "TableInput" = Field(description="The name of the input table with the selected sequences.")
+    generated_sequences: "TableInput" = Field(
+        description="The name of the input view/table with the generated sequences."
+    )
+    template_metadata: TemplateMetadata = Field(description="Metadata about templates.")
+
+
+class AnalyzeTrainParameters(ParametersBase):
+    task_type: ClassVar[str] = "analyze.train/v1"
+
+    prediction_data: "TableInput" = Field(
+        description="Table of predicted assay values for the input dataset under the scorer."
+    )
+    generator_prediction_data: "TableInput" = Field(
+        description="Pseudo log-likelihoods of the sequences in the dataset under the sampler."
+    )
+    assays: list[Assay] = Field(default_factory=list, description="List of assay metadata entries")
+    primary_objective: PrimaryObjective = Field(
+        description="The primary objective on which the sampler model is conditioned."
+    )
+    constraints: "Constraints" = Field(default_factory=list, description="The constraints applied to training.")
+
+
+class AnalyzeEngineerParameters(ParametersBase):
+    task_type: ClassVar[str] = "analyze.engineer/v1"
+
+    selected_sequences: "TableInput" = Field(description="The name of the input table with the selected sequences.")
+    engineered_sequences: "TableInput" = Field(
+        description="The name of the input view/table with the generated sequences."
+    )
+    assays: list[Assay] = Field(
+        default_factory=list, description="The list of objectives that were optimized during training."
+    )
+    primary_objective: PrimaryObjective = Field(description="The primary objective to be optimized.")
+    constraints: "Constraints" = Field(default_factory=list, description="The constraints applied to engineering.")
+    template_metadata: list[TemplateMetadata] = Field(description="Metadata about templates.")
+    new_mutation_ratio: float = Field(description="Ratio of new mutations in the generated sequences.")
+    scorer: ArtifactParam = Field(description="Scoring model used to predict assay values for candidate sequences.")
+
+
+class Homologs(BaseModel):
+    type_: Literal["HOMOLOGS"] = Field(
+        default="HOMOLOGS", alias="type", validation_alias="type", serialization_alias="type"
+    )
+    homologs: "TableInput" = Field(description="Precomputed homologs table.")
+
+
+class ProteinDatabase(StrEnum):
+    UNIREF_30 = "UNIREF_30"
+    OAS_90 = "OAS_90"
+
+
+class DatabaseSearch(BaseModel):
+    type_: Literal["DATABASE_SEARCH"] = Field(
+        default="DATABASE_SEARCH", alias="type", validation_alias="type", serialization_alias="type"
+    )
+    seed_domains: list[str] = Field(description="A list of domain sequences to search for.")
+    database: ProteinDatabase = Field(
+        default=ProteinDatabase.UNIREF_30, description="Type of the protein database to use."
+    )
+
+
+class SearchParameters(ParametersBase):
+    task_type: ClassVar[str] = "search/v1"
+
+    source: Homologs | DatabaseSearch = Field(
+        description="Precomputed homologs table or seed sequences to search for. If a table of precomputed sequence homologs is provided, the entries in the `seed_domain` column of this table may be used forsequence feature computation. For example, if the sequence to be optimized is an scFv, the `seed_domain` column would contain the heavy and light chain scaffold sequences of the scFv. These 2 scaffold sequences will be used to search for homologs and infer features such as the CDRs.",
+        discriminator="type_",
+    )
+    antibody_sequence_features: bool = Field(
+        default=False, description="Whether to compute sequence features. Only supported for antibodies."
+    )
+
+
+class UnirefDatabaseSearch(BaseModel):
+    type_: Literal["UNIREF_SEARCH"] = Field(
+        default="UNIREF_SEARCH", alias="type", validation_alias="type", serialization_alias="type"
+    )
+    seed_assemblies: list[MonomerAssembly] = Field(description="A list of monomeric assemblies to search for.")
+
+
+class OASDatabaseSearch(BaseModel):
+    type_: Literal["OAS_SEARCH"] = Field(
+        default="OAS_SEARCH", alias="type", validation_alias="type", serialization_alias="type"
+    )
+    seed_assemblies: list[VhVlAssembly] = Field(description="The list of light/heavy chain assemblies to search for.")
+    search_unpaired: bool = Field(default=True, description="Whether to search for unpaired heavy/light chains.")
+    search_paired: bool = Field(default=True, description="Whether to search for paired heavy/light chains.")
+    annotate_assembly_features: bool = Field(
+        default=True, description="Whether to infer CDR features based on OAS annotations."
+    )
+
+
+class SearchParametersV2(ParametersBase):
+    task_type: ClassVar[str] = "search/v2"
+
+    source: UnirefDatabaseSearch | OASDatabaseSearch = Field(
+        description="The source database to search against, along with the assemblies to search for, and database-specific configuration.",
+        discriminator="type_",
+    )
+
+
+class PreRanked(BaseModel):
+    """Configuration for selecting from pre-ranked assemblies.
+
+    With this option, the assemblies passed for selection are assumed to be ranked in descending
+    order of desirability. Selection is based on this order and subject to diversity constraints that
+    can be controlled via the `hit_redundancy` parameter.
+    """
+
+    type: Literal["PRE_RANKED"] = Field(default="PRE_RANKED")
+    hit_redundancy: int = Field(
+        default=3,
+        description="""The maximum number of similar variants within the proposed set.
+
+We say two variants are "similar" when their predicted performances are highly correlated.
+At lower values this will mean on average fewer hits, but greater diversity within the set,
+while at maximum value (`hit_redundancy=num_assemblies`) the selection is purely greedy, with
+no requirement on diversity within the set.""",
+    )
+
+
+class SelectParametersV2(ParametersBase):
+    task_type: ClassVar[str] = "select/v2"
+
+    assemblies: MonomerSource | VhVlSource = Field(
+        description="The source of assemblies to select from.", discriminator="kind"
+    )
+    ranker: (
+        Annotated[
+            LikelihoodBasedRanker | ScoreBasedRanker,
+            Field(description="The ranking strategy for selecting candidate assemblies.", discriminator="type"),
+        ]
+        | PreRanked
+    ) = Field(
+        description="The configuration defining how candidate assemblies are evaluated for inclusion in `selected_assemblies`."
+    )
+    min_assemblies_to_select: int = Field(
+        description="The minimum number of assemblies to select. The number of selected candidates is dynamically determined and may be greater than this number. If an exact number of assemblies is desired, set `min_assemblies_to_select` and `max_assemblies_to_select` to the same value."
+    )
+    max_assemblies_to_select: int | None = Field(
+        default=None,
+        description="An optional (inclusive) upper limit on the number of assemblies to select. If provided, the task will select between `min_assemblies_to_select` and `max_assemblies_to_select` assemblies.",
+    )
+
+
+class MonomerData(BaseModel):
+    kind: Literal["MONOMER"] = Field(default="MONOMER")
+    homologs: "TableInput" = Field(description="The homologous assemblies used for training.")
+    dataset: "TableInput | None" = Field(
+        description="The assayed assemblies used for training. If None, predictor finetuning and generator conditioning will be skipped."
+    )
+
+
+class VhVlData(BaseModel):
+    kind: Literal["VHVL"] = Field(default="VHVL")
+    homologs: "TableInput" = Field(description="The homologous assemblies used for training.")
+    dataset: "TableInput | None" = Field(
+        description="The assayed assemblies used for training. If None, predictor finetuning and generator conditioning will be skipped."
+    )
+
+
+class TrainParametersV2(ParametersBase):
+    task_type: ClassVar[str] = "train/v2"
+
+    data: MonomerData | VhVlData = Field(
+        description="The training data, including homologs and assayed sequences.", discriminator="kind"
+    )
+    assays: list[Assay] = Field(
+        default_factory=list,
+        description="The assays used for sampler conditioning (on the `primary_objective`) and scorer finetuning and evaluation. Must correspond to assays present in `data.dataset`. If `data.dataset` is None, `assays` must be empty. Otherwise, it must contain at least one assay.",
+    )
+    primary_objective: Measure | None = Field(
+        default=None,
+        description="The assay on which the sampler model is conditioned. Should correspond to an assay in `dataset`. If `data.dataset` is None, `primary_objective` must not be set. Otherwise, it must be set.",
+    )
+
+
+class ResultBase(BaseModel):
+    pass
+
+
+class AbstractTableResult(BaseModel):
+    table: str = Field(description="The reference to the table the result was written to.")
+
+
+class TrainResult(ResultBase):
+    base_sampler: ArtifactParam = Field(
+        description="A sampler which attempts to produce plausible sequences as understood relative to the passed `homologs`. "
+    )
+    conditioned_sampler: ArtifactParam | None = Field(
+        description="A sampler which attempts to produce high quality sequences as understood relative to the passed `primary_objective`."
+    )
+    scorer: ArtifactParam = Field(description="The scoring model, trained on the assay data provided.")
+    prediction_data: AbstractTableResult = Field(
+        description="Predicted assay values on all folds of the input dataset."
+    )
+    generator_prediction_data: AbstractTableResult = Field(
+        description="Pseudo log-likelihoods of the sequences in the dataset under the generator(s)."
+    )
+    report: ArtifactParam = Field(description="A report on the result of the task.")
+
+
+class MonomerModels(BaseModel):
+    kind: Literal["MONOMER"] = Field(default="MONOMER")
+    sampler: ArtifactParam = Field(description="The monomer sampler model.")
+    scorer: ArtifactParam | None = Field(description="The monomer scorer model.")
+    conditioned_sampler: ArtifactParam | None = Field(
+        description="The monomer sampler model conditioned on the primary objective."
+    )
+
+
+class VhVlModels(BaseModel):
+    kind: Literal["VH_VL"] = Field(default="VH_VL")
+    sampler: ArtifactParam = Field(description="The VhVl sampler model.")
+    scorer: ArtifactParam | None = Field(description="The VhVl scorer model.")
+    conditioned_sampler: ArtifactParam | None = Field(
+        description="The VhVl sampler model conditioned on the primary objective."
+    )
+
+
+class TrainResultV2(ResultBase):
+    models: MonomerModels | VhVlModels = Field(description="The trained models.", discriminator="kind")
+    report: ArtifactParam | None = Field(description="A report on the result of the task.")
+
+
+class EngineerResult(ResultBase):
+    selected_sequences: AbstractTableResult = Field(description="Table containing the selected sequences.")
+    engineered_sequences: AbstractTableResult = Field(description="Table containing the generated sequences.")
+    num_evaluated_seqs: int = Field(
+        description="Number of sequences evaluated in silico (only the ones sent to the predictor)."
+    )
+    num_generated_seqs: int = Field(
+        description="Number of sequences generated in silico (this is > num_evaluated_seqs, as it includes generated sequences that were not sent to the predictor, due to lower likelihood)."
+    )
+    new_mutation_ratio: float = Field(
+        description="Number of mutations in `selected_candidates` that do not appear in train/test datasets."
+    )
+    template_metadata: list[TemplateMetadata] = Field(description="Metadata about templates.")
+    report: ArtifactParam = Field(description="A report on the result of the task.")
+
+
+class EngineerResultV2(ResultBase):
+    selected_assemblies: AbstractTableResult = Field(
+        description="Table containing the selected assemblies output by `engineer`."
+    )
+    report: ArtifactParam | None = Field(description="A report on the result of the task.")
+
+
+class AnalyzeResult(ResultBase):
+    report: ArtifactParam = Field(description="The report data")
+
+
+class AnalyzeResultV2(ResultBase):
+    report: ArtifactParam = Field(description="The analysis report")
+
+
+class SearchResult(ResultBase):
+    domain_features: dict[
+        str,
+        Annotated[
+            list[DomainFeatureItem],
+            Field(
+                description="List of domain annotations. A residue may be assigned zero or one annotations. Ranges which are not annotated are considered reliable, i.e. the MSA in those ranges is used to infer conservation-based features."
+            ),
+        ],
+    ] = Field(
+        default_factory=dict,
+        description="Mapping of protein sequences to their respective domain features. Each range in the domain features must be a valid index into the corresponding protein sequence.",
+    )
+    homologs: AbstractTableResult
+
+
+class DiversifyResult(ResultBase):
+    """Result of the diversify task."""
+
+    selected_sequences: AbstractTableResult = Field(description="Table containing the selected sequences.")
+    generated_sequences: AbstractTableResult = Field(description="Table containing the generated sequences.")
+    template_metadata: TemplateMetadata = Field(description="Metadata about templates.")
+    report: ArtifactParam = Field(description="A report on the result of the task.")
+
+
+class DiversifyResultV2(ResultBase):
+    selected_assemblies: AbstractTableResult = Field(
+        description="Table containing the selected assemblies output by `diversify`."
+    )
+    report: ArtifactParam | None = Field(description="A report on the result of the task.")
+
+
+class SearchResultV2(ResultBase):
+    assembly_features: list[dict[str, AssemblyFeature]] = Field(
+        default_factory=list,
+        description="Parallel array (wrt to `seed_assemblies`) of features, keyed by feature name.",
+    )
+    homologs: AbstractTableResult
+    report: ArtifactParam | None = Field(default=None, description="A report on the result of the task.")
+
+
+class SelectResultV2(ResultBase):
+    selected_assemblies: AbstractTableResult = Field(
+        description="Table containing the selected assemblies output by `select` sorted in descending order of preference."
+    )
+
+
+class TaskState(StrEnum):
+    INIT = "INIT"
+    PREPARING_INPUTS = "PREPARING_INPUTS"
+    LAUNCHING = "LAUNCHING"
+    EXECUTING = "EXECUTING"
+    LOADING_RESULTS = "LOADING_RESULTS"
+    CANCELLING = "CANCELLING"
+    RECOVERING = "RECOVERING"
+    FAILED = "FAILED"
+    COMPLETED = "COMPLETED"
+    CANCELLED = "CANCELLED"
+
+
+class TaskResponse(ResourceResponse, ArchivableResourceMixin, ErrorResponseMixin):
+    type: str
+    context: "Context" = Field(
+        description="The context in which this task runs. Tasks started in a project or round context will only 'see' table data that belongs to the associated project.",
+        discriminator="kind",
+    )
+    name: str | None = Field(description="An optional name of the task. It must be unique within the task's context.")
+    description: str | None = Field(description="A description of the task.")
+    data_version_id: int | None = Field(description="The data version at which table inputs are observed.")
+    data_load_id: int | None = Field(description="ID of the data load for the task's results.")
+    parameters: "TaskParameters" = Field(description="The parameters of the task.")
+    result: (
+        TrainResult
+        | TrainResultV2
+        | EngineerResult
+        | EngineerResultV2
+        | AnalyzeResult
+        | AnalyzeResultV2
+        | SearchResult
+        | DiversifyResult
+        | DiversifyResultV2
+        | SearchResultV2
+        | SelectResultV2
+        | None
+    ) = Field(description="The result of the completed task.")
+    state: TaskState = Field(description="The current state of the task's execution.")
+
+
+class TaskUpdateBase(BaseModel, ABC):
+    name: str | None = Field(
+        default=None,
+        description="Optional name for the task. If provided it will ensure that only one task with that name exists within the task's context.",
+    )
+    description: str | None = Field(
+        default=None, description="Optional description to provide additional notes about the executed task."
+    )
+
+
+class TaskUpdate(TaskUpdateBase):
+    pass
+
+
+class ListTaskResponse(BaseListResponse):
+    items: list[TaskResponse]
+
+
+class TaskCreate(TaskUpdateBase):
+    parameters: "TaskParameters" = Field(description="The parameters for the task.")
+    context: "Context" = Field(description="The context in which the task is executed.", discriminator="kind")
+    data_version_id: int | None = Field(
+        default=None,
+        description="For tasks that have parameters referencing data tables, this specifies the version of the table data to use. If no version is specified, it defaults to the most recent data version.",
+    )
+
+
+TaskParameters = (
+    TrainParameters
+    | EngineerParameters
+    | EngineerParametersV2
+    | DiversifyParameters
+    | DiversifyParametersV2
+    | AnalyzeDataParameters
+    | AnalyzeDataParametersV2
+    | AnalyzeDiversifyParameters
+    | AnalyzeTrainParameters
+    | AnalyzeEngineerParameters
+    | SearchParameters
+    | SearchParametersV2
+    | SelectParametersV2
+    | TrainParametersV2
+)
+TaskResult = (
+    TrainResult
+    | TrainResultV2
+    | EngineerResult
+    | EngineerResultV2
+    | AnalyzeResult
+    | AnalyzeResultV2
+    | SearchResult
+    | DiversifyResult
+    | DiversifyResultV2
+    | SearchResultV2
+    | SelectResultV2
+)
+TableInput = list[dict[str, str | int | float | None]] | TableSourceTable | TableSourceQuery | TableSourceTaskResult
+Constraint = Annotated[AbsoluteConstraint | RelativeConstraint, Field(discriminator="type_")]
+Constraints = Annotated[
+    list[Annotated[AbsoluteConstraint | RelativeConstraint, Field(discriminator="type_")]],
+    Field(default_factory=list, description="List of assay constraints"),
+]
+Generator = Annotated[
+    Sequential | Parallel | ModelBasedGenerator | Explicit | CombinatorialGenerator,
+    Field(description="The generation strategy for creating candidate sequences.", discriminator="type"),
+]
+ConstraintsV2 = Annotated[list[ConstraintV2], Field(default_factory=list, description="List of assay constraints")]