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

adaptive_sdk/graphql_client/custom_queries.py

@@ -1,8 +1,8 @@
 from typing import Any, Dict, Optional
 from . import CompletionGroupBy
-from .custom_fields import AbcampaignFields, CompletionConnectionFields, CompletionFields, CompletionGroupDataConnectionFields, ComputePoolFields, CustomRecipeFields, DatasetFields, GraderFields, JobArtifactFields, JobConnectionFields, JobFields, JudgeFields, MetaObjectFields, MetricFields, ModelFields, PartitionFields, PrebuiltConfigDefinitionFields, PrebuiltCriteriaFields, RemoteConfigOutputFields, RemoteEnvFields, RoleFields, SystemPromptTemplateFields, TeamFields, ToolProviderFields, UsageAggregateItemFields, UsageAggregatePerUseCaseItemFields, UseCaseFields, UserFields
+from .custom_fields import AbcampaignFields, CompletionConnectionFields, CompletionFields, CompletionGroupDataConnectionFields, ComputePoolFields, ContractFields, CustomRecipeFields, DatasetFields, GlobalUsageFields, GraderFields, HarmonyGroupFields, JobArtifactFields, JobConnectionFields, JobFields, JudgeFields, MetaObjectFields, MetricFields, ModelFields, PrebuiltConfigDefinitionFields, PrebuiltCriteriaFields, RemoteConfigOutputFields, RemoteEnvFields, RoleFields, SystemPromptTemplateFields, TeamFields, ToolProviderFields, UsageAggregateItemFields, UsageAggregatePerUseCaseItemFields, UseCaseFields, UserFields
 from .custom_typing_fields import GraphQLField
-from .input_types import AbCampaignFilter, ArtifactFilter, CursorPageInput, CustomRecipeFilterInput, FeedbackFilterInput, ListCompletionsFilterInput, ListJobsFilterInput, ModelFilter, OrderPair, UsageFilterInput, UsagePerUseCaseFilterInput, UseCaseFilter
+from .input_types import AbCampaignFilter, ArtifactFilter, CursorPageInput, CustomRecipeFilterInput, FeedbackFilterInput, GlobalUsageFilterInput, ListCompletionsFilterInput, ListJobsFilterInput, ModelFilter, OrderPair, UsageFilterInput, UsagePerUseCaseFilterInput, UseCaseFilter
 
 class Query:
     """@private"""
@@ -19,6 +19,10 @@ class Query:
         cleared_arguments = {key: value for (key, value) in arguments.items() if value['value'] is not None}
         return AbcampaignFields(field_name='abCampaign', arguments=cleared_arguments)
 
+    @classmethod
+    def contract(cls) -> ContractFields:
+        return ContractFields(field_name='contract')
+
     @classmethod
     def custom_recipes(cls, use_case: str, filter: CustomRecipeFilterInput) -> CustomRecipeFields:
         arguments: Dict[str, Dict[str, Any]] = {'useCase': {'type': 'IdOrKey!', 'value': use_case}, 'filter': {'type': 'CustomRecipeFilterInput!', 'value': filter}}
@@ -79,6 +83,12 @@ class Query:
         cleared_arguments = {key: value for (key, value) in arguments.items() if value['value'] is not None}
         return UsageAggregatePerUseCaseItemFields(field_name='modelUsageByUseCase', arguments=cleared_arguments)
 
+    @classmethod
+    def global_usage(cls, filter: GlobalUsageFilterInput) -> GlobalUsageFields:
+        arguments: Dict[str, Dict[str, Any]] = {'filter': {'type': 'GlobalUsageFilterInput!', 'value': filter}}
+        cleared_arguments = {key: value for (key, value) in arguments.items() if value['value'] is not None}
+        return GlobalUsageFields(field_name='globalUsage', arguments=cleared_arguments)
+
     @classmethod
     def system_prompt_templates(cls) -> SystemPromptTemplateFields:
         return SystemPromptTemplateFields(field_name='systemPromptTemplates')
@@ -118,12 +128,12 @@ class Query:
         return ModelFields(field_name='model', arguments=cleared_arguments)
 
     @classmethod
-    def all_partitions(cls) -> PartitionFields:
-        return PartitionFields(field_name='allPartitions')
+    def all_harmony_groups(cls) -> HarmonyGroupFields:
+        return HarmonyGroupFields(field_name='allHarmonyGroups')
 
     @classmethod
-    def partitions(cls) -> PartitionFields:
-        return PartitionFields(field_name='partitions')
+    def harmony_groups(cls) -> HarmonyGroupFields:
+        return HarmonyGroupFields(field_name='harmonyGroups')
 
     @classmethod
     def compute_pools(cls) -> ComputePoolFields:

adaptive_sdk/graphql_client/custom_typing_fields.py

@@ -106,6 +106,13 @@ class BatchInferenceJobStageOutputGraphQLField(GraphQLField):
         self._alias = alias
         return self
 
+class BillingUsageGraphQLField(GraphQLField):
+    """@private"""
+
+    def alias(self, alias: str) -> 'BillingUsageGraphQLField':
+        self._alias = alias
+        return self
+
 class ChatMessageGraphQLField(GraphQLField):
     """@private"""
 
@@ -211,6 +218,13 @@ class ComputePoolGraphQLField(GraphQLField):
         self._alias = alias
         return self
 
+class ContractGraphQLField(GraphQLField):
+    """@private"""
+
+    def alias(self, alias: str) -> 'ContractGraphQLField':
+        self._alias = alias
+        return self
+
 class CustomConfigOutputGraphQLField(GraphQLField):
     """@private"""
 
@@ -348,6 +362,13 @@ class EvaluationResultGraphQLField(GraphQLField):
         self._alias = alias
         return self
 
+class GlobalUsageGraphQLField(GraphQLField):
+    """@private"""
+
+    def alias(self, alias: str) -> 'GlobalUsageGraphQLField':
+        self._alias = alias
+        return self
+
 class GpuAllocationGraphQLField(GraphQLField):
     """@private"""
 
@@ -394,6 +415,13 @@ class GuidelinesTrainingParamsOutputGraphQLField(GraphQLField):
         self._alias = alias
         return self
 
+class HarmonyGroupGraphQLField(GraphQLField):
+    """@private"""
+
+    def alias(self, alias: str) -> 'HarmonyGroupGraphQLField':
+        self._alias = alias
+        return self
+
 class InteractionOutputGraphQLField(GraphQLField):
     """@private"""
 
@@ -605,13 +633,6 @@ class PageInfoGraphQLField(GraphQLField):
         self._alias = alias
         return self
 
-class PartitionGraphQLField(GraphQLField):
-    """@private"""
-
-    def alias(self, alias: str) -> 'PartitionGraphQLField':
-        self._alias = alias
-        return self
-
 class PpotrainingParamsOutputGraphQLField(GraphQLField):
     """@private"""
 
@@ -908,6 +929,20 @@ class UsageAggregatePerUseCaseItemGraphQLField(GraphQLField):
         self._alias = alias
         return self
 
+class UsageStatsGraphQLField(GraphQLField):
+    """@private"""
+
+    def alias(self, alias: str) -> 'UsageStatsGraphQLField':
+        self._alias = alias
+        return self
+
+class UsageStatsByModelGraphQLField(GraphQLField):
+    """@private"""
+
+    def alias(self, alias: str) -> 'UsageStatsByModelGraphQLField':
+        self._alias = alias
+        return self
+
 class UseCaseGraphQLField(GraphQLField):
     """@private"""
 

adaptive_sdk/graphql_client/describe_job.py

@@ -0,0 +1,12 @@
+from typing import Optional
+from .base_model import BaseModel
+from .fragments import JobData
+
+class DescribeJob(BaseModel):
+    """@public"""
+    job: Optional['DescribeJobJob']
+
+class DescribeJobJob(JobData):
+    """@public"""
+    pass
+DescribeJob.model_rebuild()

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'
@@ -50,6 +54,7 @@ class DatasetSource(str, Enum):
 
 class DateBucketUnit(str, Enum):
     """@public"""
+    HOUR = 'HOUR'
     DAY = 'DAY'
     WEEK = 'WEEK'
     MONTH = 'MONTH'
@@ -97,6 +102,11 @@ class GraderTypeEnum(str, Enum):
     REMOTE = 'REMOTE'
     CUSTOM = 'CUSTOM'
 
+class HarmonyStatus(str, Enum):
+    """@public"""
+    ONLINE = 'ONLINE'
+    OFFLINE = 'OFFLINE'
+
 class JobArtifactKind(str, Enum):
     """@public"""
     EVALUATION = 'EVALUATION'
@@ -186,11 +196,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'
@@ -234,6 +239,7 @@ class TimeseriesInterval(str, Enum):
     MONTH = 'MONTH'
     QUARTER = 'QUARTER'
     YEAR = 'YEAR'
+    MILLENNIUM = 'MILLENNIUM'
 
 class TrainingJobStatus(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, CompletionSourceOutput, 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
@@ -566,27 +587,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,6 +677,7 @@ CustomRecipeData.model_rebuild()
 DatasetData.model_rebuild()
 ModelData.model_rebuild()
 GraderData.model_rebuild()
+HarmonyGroupData.model_rebuild()
 JobData.model_rebuild()
 JobStageOutputData.model_rebuild()
 JudgeData.model_rebuild()
@@ -685,7 +686,6 @@ 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

@@ -160,6 +160,12 @@ class FeedbackUpdateInput(BaseModel):
     value: Optional[Any] = None
     details: Optional[str] = None
 
+class GlobalUsageFilterInput(BaseModel):
+    """@private"""
+    timerange: Optional['TimeRange'] = None
+    interval: DateBucketUnit
+    timezone: Optional[str] = None
+
 class GoogleProviderDataInput(BaseModel):
     """@private"""
     api_key: str = Field(alias='apiKey')
@@ -464,13 +470,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"""
@@ -537,6 +544,7 @@ CustomRecipeFilterInput.model_rebuild()
 DatasetGenerate.model_rebuild()
 DatasetGenerationConfig.model_rebuild()
 FeedbackFilterInput.model_rebuild()
+GlobalUsageFilterInput.model_rebuild()
 GraderConfigInput.model_rebuild()
 GraderCreateInput.model_rebuild()
 GraderUpdateInput.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