adaptive-sdk 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

adaptive_sdk/graphql_client/enums.py

@@ -11,6 +11,10 @@ class AuthProviderKind(str, Enum):
     """@public"""
     OIDC = 'OIDC'
 
+class BillingCycle(str, Enum):
+    """@public"""
+    YEARLY = 'YEARLY'
+
 class CompletionGroupBy(str, Enum):
     """@public"""
     MODEL = 'MODEL'
@@ -24,12 +28,6 @@ class CompletionSource(str, Enum):
     DATASET = 'DATASET'
     EVALUATION = 'EVALUATION'
 
-class CompletionSourceOutput(str, Enum):
-    """@public"""
-    LIVE = 'LIVE'
-    OFFLINE = 'OFFLINE'
-    AUTOMATION = 'AUTOMATION'
-
 class ComputePoolCapability(str, Enum):
     """@public"""
     INFERENCE = 'INFERENCE'
@@ -47,9 +45,11 @@ class DatasetSource(str, Enum):
     UPLOADED = 'UPLOADED'
     GENERATED = 'GENERATED'
     GENERATING = 'GENERATING'
+    FILTERED = 'FILTERED'
 
 class DateBucketUnit(str, Enum):
     """@public"""
+    HOUR = 'HOUR'
     DAY = 'DAY'
     WEEK = 'WEEK'
     MONTH = 'MONTH'
@@ -57,21 +57,6 @@ class DateBucketUnit(str, Enum):
     YEAR = 'YEAR'
     NO_GROUP_BY_DATE = 'NO_GROUP_BY_DATE'
 
-class EvaluationJobStatus(str, Enum):
-    """@public"""
-    PENDING = 'PENDING'
-    RUNNING = 'RUNNING'
-    COMPLETED = 'COMPLETED'
-    FAILED = 'FAILED'
-    CANCELED = 'CANCELED'
-
-class EvaluationType(str, Enum):
-    """@public"""
-    AI_JUDGE = 'AI_JUDGE'
-    REMOTE_ENV = 'REMOTE_ENV'
-    GRADERS = 'GRADERS'
-    MULTI_AI_JUDGES = 'MULTI_AI_JUDGES'
-
 class ExternalModelProviderName(str, Enum):
     """@public"""
     AZURE = 'AZURE'
@@ -85,11 +70,6 @@ class FeedbackType(str, Enum):
     DIRECT = 'DIRECT'
     COMPARISON = 'COMPARISON'
 
-class FeedbackTypeOutput(str, Enum):
-    """@public"""
-    DIRECT = 'DIRECT'
-    PREFERENCE = 'PREFERENCE'
-
 class GraderTypeEnum(str, Enum):
     """@public"""
     JUDGE = 'JUDGE'
@@ -97,10 +77,16 @@ class GraderTypeEnum(str, Enum):
     REMOTE = 'REMOTE'
     CUSTOM = 'CUSTOM'
 
+class HarmonyStatus(str, Enum):
+    """@public"""
+    ONLINE = 'ONLINE'
+    OFFLINE = 'OFFLINE'
+
 class JobArtifactKind(str, Enum):
     """@public"""
     EVALUATION = 'EVALUATION'
     CUSTOM = 'CUSTOM'
+    DATASET = 'DATASET'
 
 class JobKind(str, Enum):
     """@public"""
@@ -186,11 +172,6 @@ class OpenAIModel(str, Enum):
     GPT5_NANO = 'GPT5_NANO'
     GPT41_MINI = 'GPT41_MINI'
 
-class PartitionStatus(str, Enum):
-    """@public"""
-    ONLINE = 'ONLINE'
-    OFFLINE = 'OFFLINE'
-
 class PrebuiltCriteriaKey(str, Enum):
     """@public"""
     FAITHFULNESS = 'FAITHFULNESS'
@@ -215,7 +196,7 @@ class RemoteEnvStatus(str, Enum):
     ONLINE = 'ONLINE'
     OFFLINE = 'OFFLINE'
 
-class SelectionTypeOutput(str, Enum):
+class SelectionTypeInput(str, Enum):
     """@public"""
     ALL = 'ALL'
     RANDOM = 'RANDOM'
@@ -234,26 +215,7 @@ class TimeseriesInterval(str, Enum):
     MONTH = 'MONTH'
     QUARTER = 'QUARTER'
     YEAR = 'YEAR'
-
-class TrainingJobStatus(str, Enum):
-    """@public"""
-    PENDING = 'PENDING'
-    RUNNING = 'RUNNING'
-    COMPLETED = 'COMPLETED'
-    FAILED = 'FAILED'
-    CANCELED = 'CANCELED'
-
-class TrainingMetadataOutputAlignmentMethod(str, Enum):
-    """@public"""
-    DPO = 'DPO'
-    PPO = 'PPO'
-    SFT = 'SFT'
-    GRPO = 'GRPO'
-
-class TrainingMetadataOutputTrainingType(str, Enum):
-    """@public"""
-    FULL_WEIGHTS = 'FULL_WEIGHTS'
-    PARAMETER_EFFICIENT = 'PARAMETER_EFFICIENT'
+    MILLENNIUM = 'MILLENNIUM'
 
 class UnitPosition(str, Enum):
     """@public"""

adaptive_sdk/graphql_client/fragments.py

@@ -1,7 +1,7 @@
 from typing import Annotated, Any, List, Literal, Optional, Union
 from pydantic import Field
 from .base_model import BaseModel
-from .enums import AbcampaignStatus, CompletionSource, CompletionSourceOutput, DatasetKind, FeedbackType, GraderTypeEnum, JobArtifactKind, JobKind, JobStatus, JobStatusOutput, JudgeCapability, MetricKind, MetricScoringType, ModelKindFilter, ModelOnline, PartitionStatus, PrebuiltCriteriaKey, ProviderName, RemoteEnvStatus
+from .enums import AbcampaignStatus, CompletionSource, DatasetKind, FeedbackType, GraderTypeEnum, HarmonyStatus, JobArtifactKind, JobKind, JobStatus, JobStatusOutput, JudgeCapability, MetricKind, MetricScoringType, ModelKindFilter, ModelOnline, PrebuiltCriteriaKey, ProviderName, RemoteEnvStatus
 
 class AbCampaignCreateData(BaseModel):
     """@public"""
@@ -330,6 +330,27 @@ class GraderDataMetric(MetricData):
     """@public"""
     pass
 
+class HarmonyGroupData(BaseModel):
+    """@public"""
+    id: Any
+    key: str
+    compute_pool: Optional['HarmonyGroupDataComputePool'] = Field(alias='computePool')
+    status: HarmonyStatus
+    url: str
+    world_size: int = Field(alias='worldSize')
+    gpu_types: str = Field(alias='gpuTypes')
+    created_at: int = Field(alias='createdAt')
+    online_models: List['HarmonyGroupDataOnlineModels'] = Field(alias='onlineModels')
+
+class HarmonyGroupDataComputePool(BaseModel):
+    """@public"""
+    key: str
+    name: str
+
+class HarmonyGroupDataOnlineModels(ModelData):
+    """@public"""
+    pass
+
 class JobData(BaseModel):
     """@public"""
     id: Any
@@ -472,30 +493,6 @@ class JudgeDataMetric(MetricData):
     """@public"""
     pass
 
-class ListCompletionsFilterOutputData(BaseModel):
-    """@public"""
-    use_case: str = Field(alias='useCase')
-    models: Optional[List[str]]
-    timerange: Optional['ListCompletionsFilterOutputDataTimerange']
-    feedbacks: Optional[List['ListCompletionsFilterOutputDataFeedbacks']]
-    labels: Optional[List['ListCompletionsFilterOutputDataLabels']]
-    tags: Optional[List[str]]
-    source: Optional[List[CompletionSourceOutput]]
-
-class ListCompletionsFilterOutputDataTimerange(BaseModel):
-    """@public"""
-    from_: Any = Field(alias='from')
-    to: Any
-
-class ListCompletionsFilterOutputDataFeedbacks(BaseModel):
-    """@public"""
-    metric: str
-
-class ListCompletionsFilterOutputDataLabels(BaseModel):
-    """@public"""
-    key: str
-    value: Optional[List[str]]
-
 class MetricDataAdmin(BaseModel):
     """@public"""
     id: Any
@@ -566,27 +563,6 @@ class ModelServiceDataModelBackbone(ModelData):
     """@public"""
     pass
 
-class PartitionData(BaseModel):
-    """@public"""
-    id: Any
-    key: str
-    compute_pool: Optional['PartitionDataComputePool'] = Field(alias='computePool')
-    status: PartitionStatus
-    url: str
-    world_size: int = Field(alias='worldSize')
-    gpu_types: str = Field(alias='gpuTypes')
-    created_at: int = Field(alias='createdAt')
-    online_models: List['PartitionDataOnlineModels'] = Field(alias='onlineModels')
-
-class PartitionDataComputePool(BaseModel):
-    """@public"""
-    key: str
-    name: str
-
-class PartitionDataOnlineModels(ModelData):
-    """@public"""
-    pass
-
 class RemoteEnvData(BaseModel):
     """@public"""
     id: Any
@@ -677,15 +653,14 @@ CustomRecipeData.model_rebuild()
 DatasetData.model_rebuild()
 ModelData.model_rebuild()
 GraderData.model_rebuild()
+HarmonyGroupData.model_rebuild()
 JobData.model_rebuild()
 JobStageOutputData.model_rebuild()
 JudgeData.model_rebuild()
-ListCompletionsFilterOutputData.model_rebuild()
 MetricDataAdmin.model_rebuild()
 MetricWithContextData.model_rebuild()
 ModelDataAdmin.model_rebuild()
 ModelServiceData.model_rebuild()
-PartitionData.model_rebuild()
 RemoteEnvData.model_rebuild()
 UseCaseData.model_rebuild()
 UserData.model_rebuild()

adaptive_sdk/graphql_client/get_custom_recipe.py

@@ -1,10 +1,11 @@
+from typing import Optional
 from pydantic import Field
 from .base_model import BaseModel
 from .fragments import CustomRecipeData
 
 class GetCustomRecipe(BaseModel):
     """@public"""
-    custom_recipe: 'GetCustomRecipeCustomRecipe' = Field(alias='customRecipe')
+    custom_recipe: Optional['GetCustomRecipeCustomRecipe'] = Field(alias='customRecipe')
 
 class GetCustomRecipeCustomRecipe(CustomRecipeData):
     """@public"""

adaptive_sdk/graphql_client/input_types.py

@@ -1,7 +1,7 @@
 from typing import Any, List, Optional
 from pydantic import Field
 from .base_model import BaseModel
-from .enums import AbcampaignStatus, CompletionSource, DatasetSource, DateBucketUnit, ExternalModelProviderName, FeedbackType, GraderTypeEnum, JobArtifactKind, JobKind, JobStatus, MetricAggregation, MetricKind, MetricScoringType, ModelKindFilter, ModelOnline, OpenAIModel, PrebuiltCriteriaKey, Protocol, SortDirection, TimeseriesInterval, UnitPosition
+from .enums import AbcampaignStatus, CompletionGroupBy, CompletionSource, DatasetKind, DatasetSource, DateBucketUnit, ExternalModelProviderName, FeedbackType, GraderTypeEnum, JobArtifactKind, JobKind, JobStatus, MetricAggregation, MetricKind, MetricScoringType, ModelKindFilter, ModelOnline, OpenAIModel, PrebuiltCriteriaKey, Protocol, SelectionTypeInput, SortDirection, TimeseriesInterval, UnitPosition
 
 class AbCampaignFilter(BaseModel):
     """@private"""
@@ -91,6 +91,15 @@ class CompletionLabelValue(BaseModel):
     key: str
     value: str
 
+class CompletionsByFilters(BaseModel):
+    """@private"""
+    filters: 'ListCompletionsFilterInput'
+    exclude: List[Any]
+
+class CompletionsById(BaseModel):
+    """@private"""
+    include: List[Any]
+
 class CreateRecipeInput(BaseModel):
     """@private"""
     name: str
@@ -120,6 +129,12 @@ class CustomRecipeFilterInput(BaseModel):
     """@private"""
     labels: Optional[List['LabelFilter']] = None
 
+class DatasetCompletionQuery(BaseModel):
+    """@private"""
+    from_selection: Optional['CompletionsById'] = Field(alias='fromSelection', default=None)
+    from_filters: Optional['CompletionsByFilters'] = Field(alias='fromFilters', default=None)
+    from_groups: Optional['FromGroupsQuery'] = Field(alias='fromGroups', default=None)
+
 class DatasetCreate(BaseModel):
     """@private"""
     use_case: str = Field(alias='useCase')
@@ -127,18 +142,16 @@ class DatasetCreate(BaseModel):
     key: Optional[str] = None
     source: Optional[DatasetSource] = None
 
-class DatasetGenerate(BaseModel):
+class DatasetCreateFromFilters(BaseModel):
     """@private"""
     use_case: str = Field(alias='useCase')
     name: str
     key: Optional[str] = None
-    compute_pool: Optional[str] = Field(alias='computePool', default=None)
-    config: 'DatasetGenerationConfig'
-    num_gpus: int = Field(alias='numGpus', default=1)
-
-class DatasetGenerationConfig(BaseModel):
-    """@private"""
-    rag: Optional['RagdataGenerationConfig'] = None
+    completion_query: 'DatasetCompletionQuery' = Field(alias='completionQuery')
+    sample_config: Optional['SampleConfig'] = Field(alias='sampleConfig', default=None)
+    feedback_filters: Optional['FeedbackFilterInput'] = Field(alias='feedbackFilters', default=None)
+    kind: DatasetKind
+    metrics: Optional[List[str]] = None
 
 class EmojiInput(BaseModel):
     """@private"""
@@ -160,6 +173,18 @@ class FeedbackUpdateInput(BaseModel):
     value: Optional[Any] = None
     details: Optional[str] = None
 
+class FromGroupsQuery(BaseModel):
+    """@private"""
+    filters: 'ListCompletionsFilterInput'
+    grouping: CompletionGroupBy
+    groups: List['GroupSelectionQuery']
+
+class GlobalUsageFilterInput(BaseModel):
+    """@private"""
+    timerange: Optional['TimeRange'] = None
+    interval: DateBucketUnit
+    timezone: Optional[str] = None
+
 class GoogleProviderDataInput(BaseModel):
     """@private"""
     api_key: str = Field(alias='apiKey')
@@ -186,6 +211,16 @@ class GraderUpdateInput(BaseModel):
     grader_type: Optional[GraderTypeEnum] = Field(alias='graderType', default=None)
     grader_config: Optional['GraderConfigInput'] = Field(alias='graderConfig', default=None)
 
+class GroupSelection(BaseModel):
+    """@private"""
+    exclude: Optional[List[Any]] = None
+    select_only: Optional[List[Any]] = Field(alias='selectOnly', default=None)
+
+class GroupSelectionQuery(BaseModel):
+    """@private"""
+    group_id: str = Field(alias='groupId')
+    selection: 'GroupSelection'
+
 class JobArtifactFilter(BaseModel):
     """@private"""
     kinds: List[JobArtifactKind]
@@ -274,6 +309,7 @@ class MetricCreate(BaseModel):
     kind: MetricKind
     scoring_type: MetricScoringType = Field(alias='scoringType', default=MetricScoringType.HIGHER_IS_BETTER)
     description: Optional[str] = None
+    unit: Optional[str] = None
 
 class MetricGetOrCreate(BaseModel):
     """@private"""
@@ -360,12 +396,6 @@ class PrebuiltJudgeCreate(BaseModel):
     model: str
     prebuilt_criteria_key: PrebuiltCriteriaKey = Field(alias='prebuiltCriteriaKey')
 
-class RagdataGenerationConfig(BaseModel):
-    """@private"""
-    chunks_per_question: int = Field(alias='chunksPerQuestion')
-    model: str
-    system_prompt: Optional[str] = Field(alias='systemPrompt', default=None)
-
 class RemoteConfigInput(BaseModel):
     """@private"""
     url: str
@@ -377,12 +407,22 @@ class RemoteEnvCreate(BaseModel):
     name: Optional[str] = None
     description: Optional[str] = None
 
+class ResizePartitionInput(BaseModel):
+    """@private"""
+    harmony_group: str = Field(alias='harmonyGroup')
+    size: int
+
 class RoleCreate(BaseModel):
     """@private"""
     key: Optional[str] = None
     name: str
     permissions: List[str]
 
+class SampleConfig(BaseModel):
+    """@private"""
+    selection_type: SelectionTypeInput = Field(alias='selectionType')
+    sample_size: Optional[int] = Field(alias='sampleSize', default=None)
+
 class SystemPromptTemplateCreate(BaseModel):
     """@private"""
     name: str
@@ -464,13 +504,14 @@ class UpdateToolProviderInput(BaseModel):
 class UsageFilterInput(BaseModel):
     """@private"""
     model_id: Any = Field(alias='modelId')
-    timerange: 'TimeRange'
+    timerange: Optional['TimeRange'] = None
     unit: DateBucketUnit
+    timezone: Optional[str] = None
 
 class UsagePerUseCaseFilterInput(BaseModel):
     """@private"""
     model_id: Any = Field(alias='modelId')
-    timerange: 'TimeRange'
+    timerange: Optional['TimeRange'] = None
 
 class UseCaseCreate(BaseModel):
     """@private"""
@@ -532,14 +573,18 @@ class WidgetInput(BaseModel):
     unit: 'UnitConfigInput'
 AddExternalModelInput.model_rebuild()
 AttachModel.model_rebuild()
+CompletionsByFilters.model_rebuild()
 CreateRecipeInput.model_rebuild()
 CustomRecipeFilterInput.model_rebuild()
-DatasetGenerate.model_rebuild()
-DatasetGenerationConfig.model_rebuild()
+DatasetCompletionQuery.model_rebuild()
+DatasetCreateFromFilters.model_rebuild()
 FeedbackFilterInput.model_rebuild()
+FromGroupsQuery.model_rebuild()
+GlobalUsageFilterInput.model_rebuild()
 GraderConfigInput.model_rebuild()
 GraderCreateInput.model_rebuild()
 GraderUpdateInput.model_rebuild()
+GroupSelectionQuery.model_rebuild()
 JudgeConfigInput.model_rebuild()
 JudgeCreate.model_rebuild()
 JudgeExampleInput.model_rebuild()

adaptive_sdk/graphql_client/list_compute_pools.py

@@ -2,7 +2,7 @@ from typing import Any, List
 from pydantic import Field
 from .base_model import BaseModel
 from .enums import ComputePoolCapability
-from .fragments import PartitionData
+from .fragments import HarmonyGroupData
 
 class ListComputePools(BaseModel):
     """@public"""
@@ -15,9 +15,9 @@ class ListComputePoolsComputePools(BaseModel):
     name: str
     created_at: int = Field(alias='createdAt')
     capabilities: List[ComputePoolCapability]
-    partitions: List['ListComputePoolsComputePoolsPartitions']
+    harmony_groups: List['ListComputePoolsComputePoolsHarmonyGroups'] = Field(alias='harmonyGroups')
 
-class ListComputePoolsComputePoolsPartitions(PartitionData):
+class ListComputePoolsComputePoolsHarmonyGroups(HarmonyGroupData):
     """@public"""
     pass
 ListComputePools.model_rebuild()

adaptive_sdk/graphql_client/list_harmony_groups.py

@@ -0,0 +1,13 @@
+from typing import List
+from pydantic import Field
+from .base_model import BaseModel
+from .fragments import HarmonyGroupData
+
+class ListHarmonyGroups(BaseModel):
+    """@public"""
+    harmony_groups: List['ListHarmonyGroupsHarmonyGroups'] = Field(alias='harmonyGroups')
+
+class ListHarmonyGroupsHarmonyGroups(HarmonyGroupData):
+    """@public"""
+    pass
+ListHarmonyGroups.model_rebuild()

adaptive_sdk/resources/jobs.py

@@ -26,6 +26,9 @@ class Jobs(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
         SyncAPIResource.__init__(self, client)
         UseCaseResource.__init__(self, client)
 
+    def get(self, job_id: str) -> JobData | None:
+        return self._gql_client.describe_job(id=job_id).job
+
     def list(
         self,
         first: int | None = 100,
@@ -52,13 +55,11 @@ class Jobs(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
             useCase=use_case,
             kind=[JobKind(k) for k in kind] if kind else None,
         )
-        return self._gql_client.list_jobs(
-            page=page, filter=validated_filter, order=[]
-        ).jobs
+        return self._gql_client.list_jobs(page=page, filter=validated_filter, order=[]).jobs
 
-    def create(
+    def run(
         self,
-        recipe: str,
+        recipe_key: str,
         args: dict[str, Any],
         num_gpus: int,
         use_case: str | None = None,
@@ -66,7 +67,7 @@ class Jobs(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
     ) -> JobData:
         return self._gql_client.create_job(
             input=JobInput(
-                recipe=recipe,
+                recipe=recipe_key,
                 args=args,
                 useCase=self.use_case_key(use_case),
                 computePool=compute_pool,
@@ -74,7 +75,7 @@ class Jobs(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
             )
         ).create_job
 
-    def cancel_custom(self, job_id: str) -> JobData:
+    def cancel(self, job_id: str) -> JobData:
         return self._gql_client.cancel_job(job_id=job_id).cancel_job
 
 
@@ -87,6 +88,9 @@ class AsyncJobs(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
         AsyncAPIResource.__init__(self, client)
         UseCaseResource.__init__(self, client)
 
+    async def get(self, job_id: str) -> JobData | None:
+        return (await self._gql_client.describe_job(id=job_id)).job
+
     async def list(
         self,
         first: int | None = 100,
@@ -112,13 +116,11 @@ class AsyncJobs(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
             useCase=self.use_case_key(use_case),
             kind=[JobKind[k] for k in kind] if kind else None,
         )
-        return (
-            await self._gql_client.list_jobs(page=page, filter=validated_filter)
-        ).jobs
+        return (await self._gql_client.list_jobs(page=page, filter=validated_filter)).jobs
 
-    async def create_custom(
+    async def run(
         self,
-        recipe: str,
+        recipe_key: str,
         args: dict[str, Any],
         num_gpus: int,
         use_case: str | None = None,
@@ -127,7 +129,7 @@ class AsyncJobs(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
         return (
             await self._gql_client.create_job(
                 input=JobInput(
-                    recipe=recipe,
+                    recipe=recipe_key,
                     args=args,
                     useCase=self.use_case_key(use_case),
                     computePool=compute_pool,
@@ -136,5 +138,5 @@ class AsyncJobs(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
             )
         ).create_job
 
-    async def cancel_custom(self, job_id: str) -> JobData:
+    async def cancel(self, job_id: str) -> JobData:
         return (await self._gql_client.cancel_job(job_id=job_id)).cancel_job