adaptive-sdk 0.1.8__tar.gz → 0.1.9__tar.gz

{adaptive_sdk-0.1.8 → adaptive_sdk-0.1.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: adaptive-sdk
-Version: 0.1.8
+Version: 0.1.9
 Summary: Python SDK for Adaptive Engine
 Author-email: Vincent Debergue <vincent@adaptive-ml.com>, Joao Moura <joao@adaptive-ml.com>, Yacine Bouraoui <yacine@adaptive-ml.com>
 Requires-Python: >=3.10

{adaptive_sdk-0.1.8 → adaptive_sdk-0.1.9}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "adaptive-sdk"
-version = "0.1.8"
+version = "0.1.9"
 description = "Python SDK for Adaptive Engine"
 authors = [
   { name = "Vincent Debergue", email = "vincent@adaptive-ml.com" },

{adaptive_sdk-0.1.8 → adaptive_sdk-0.1.9}/src/adaptive_sdk/graphql_client/__init__.py

@@ -12,6 +12,7 @@ from .cancel_job import CancelJob, CancelJobCancelJob
 from .client import GQLClient
 from .create_ab_campaign import CreateAbCampaign, CreateAbCampaignCreateAbCampaign
 from .create_custom_recipe import CreateCustomRecipe, CreateCustomRecipeCreateCustomRecipe
+from .create_dataset_from_multipart_upload import CreateDatasetFromMultipartUpload, CreateDatasetFromMultipartUploadCreateDatasetFromMultipartUpload
 from .create_grader import CreateGrader, CreateGraderCreateGrader
 from .create_job import CreateJob, CreateJobCreateJob
 from .create_judge import CreateJudge, CreateJudgeCreateJudge
@@ -21,6 +22,7 @@ from .create_role import CreateRole, CreateRoleCreateRole
 from .create_team import CreateTeam, CreateTeamCreateTeam
 from .create_use_case import CreateUseCase, CreateUseCaseCreateUseCase
 from .create_user import CreateUser, CreateUserCreateUser
+from .dataset_upload_processing_status import DatasetUploadProcessingStatus, DatasetUploadProcessingStatusDatasetUploadProcessingStatus
 from .delete_custom_recipe import DeleteCustomRecipe
 from .delete_dataset import DeleteDataset
 from .delete_grader import DeleteGrader, DeleteGraderDeleteGrader
@@ -78,4 +80,4 @@ from .update_judge import UpdateJudge, UpdateJudgeUpdateJudge
 from .update_model import UpdateModel, UpdateModelUpdateModelService
 from .update_model_compute_config import UpdateModelComputeConfig, UpdateModelComputeConfigUpdateModelComputeConfig
 from .update_user import UpdateUser, UpdateUserSetTeamMember, UpdateUserSetTeamMemberRole, UpdateUserSetTeamMemberTeam, UpdateUserSetTeamMemberUser
-__all__ = ['AbCampaignCreateData', 'AbCampaignDetailData', 'AbCampaignDetailDataMetric', 'AbCampaignDetailDataModels', 'AbCampaignDetailDataUseCase', 'AbCampaignFilter', 'AbCampaignReportData', 'AbCampaignReportDataVariants', 'AbCampaignReportDataVariantsComparisons', 'AbCampaignReportDataVariantsComparisonsVariant', 'AbCampaignReportDataVariantsInterval', 'AbCampaignReportDataVariantsVariant', 'AbcampaignCreate', 'AbcampaignStatus', 'AddExternalModel', 'AddExternalModelAddExternalModel', 'AddExternalModelInput', 'AddHFModel', 'AddHFModelImportHfModel', 'AddHFModelInput', 'AddModel', 'AddModelAddModel', 'AddModelAddModelBackbone', 'AddModelInput', 'AddRemoteEnv', 'AddRemoteEnvAddRemoteEnv', 'AnthropicProviderDataInput', 'ApiKeyCreate', 'ArtifactFilter', 'AsyncBaseClientOpenTelemetry', 'AsyncGQLClient', 'AttachModel', 'AttachModelToUseCase', 'AttachModelToUseCaseAttachModel', 'AuthProviderKind', 'AzureProviderDataInput', 'BaseClientOpenTelemetry', 'BaseModel', 'BillingCycle', 'CancelABCampaign', 'CancelAllocationInput', 'CancelJob', 'CancelJobCancelJob', 'CapabilityFilter', 'CompletionComparisonFeedbackData', 'CompletionComparisonFeedbackDataModel', 'CompletionComparisonFilterInput', 'CompletionData', 'CompletionDataChatMessages', 'CompletionDataComparisonFeedbacks', 'CompletionDataComparisonFeedbacksMetric', 'CompletionDataComparisonFeedbacksOtherCompletion', 'CompletionDataComparisonFeedbacksPreferedCompletion', 'CompletionDataComparisonFeedbacksUsecase', 'CompletionDataDirectFeedbacks', 'CompletionDataDirectFeedbacksMetric', 'CompletionDataLabels', 'CompletionDataMetadata', 'CompletionDataMetadataUsage', 'CompletionDataModel', 'CompletionFeedbackFilterInput', 'CompletionGroupBy', 'CompletionLabelValue', 'CompletionSource', 'CompletionsByFilters', 'CompletionsById', 'ComputePoolCapability', 'CreateAbCampaign', 'CreateAbCampaignCreateAbCampaign', 'CreateCustomRecipe', 'CreateCustomRecipeCreateCustomRecipe', 'CreateGrader', 'CreateGraderCreateGrader', 'CreateJob', 'CreateJobCreateJob', 'CreateJudge', 'CreateJudgeCreateJudge', 'CreateMetric', 'CreateMetricCreateMetric', 'CreatePrebuiltJudge', 'CreatePrebuiltJudgeCreatePrebuiltJudge', 'CreateRecipeInput', 'CreateRole', 'CreateRoleCreateRole', 'CreateTeam', 'CreateTeamCreateTeam', 'CreateToolProviderInput', 'CreateUseCase', 'CreateUseCaseCreateUseCase', 'CreateUser', 'CreateUserCreateUser', 'CursorPageInput', 'CustomConfigInput', 'CustomRecipeData', 'CustomRecipeDataCreatedBy', 'CustomRecipeDataLabels', 'CustomRecipeFilterInput', 'DatasetCompletionQuery', 'DatasetCreate', 'DatasetCreateFromFilters', 'DatasetCreateFromMultipartUpload', 'DatasetData', 'DatasetDataMetricsUsage', 'DatasetDataMetricsUsageMetric', 'DatasetKind', 'DatasetSource', 'DatasetStatus', 'DatasetUploadProcessingStatusInput', 'DateBucketUnit', 'DeleteCustomRecipe', 'DeleteDataset', 'DeleteGrader', 'DeleteGraderDeleteGrader', 'DeleteJudge', 'DeleteJudgeDeleteJudge', 'DeleteUser', 'DeleteUserDeleteUser', 'DeployModel', 'DescribeAbCampaign', 'DescribeAbCampaignAbCampaign', 'DescribeAbCampaignAbCampaignReport', 'DescribeDataset', 'DescribeDatasetDataset', 'DescribeInteraction', 'DescribeInteractionCompletion', 'DescribeJob', 'DescribeJobJob', 'DescribeMetric', 'DescribeMetricAdmin', 'DescribeMetricAdminMetric', 'DescribeMetricMetric', 'DescribeModel', 'DescribeModelAdmin', 'DescribeModelAdminModel', 'DescribeModelAdminModelBackbone', 'DescribeModelModel', 'DescribeModelModelBackbone', 'DescribeUseCase', 'DescribeUseCaseUseCase', 'EmojiInput', 'ExternalModelProviderName', 'FeedbackAddInput', 'FeedbackFilterInput', 'FeedbackType', 'FeedbackUpdateInput', 'FromGroupsQuery', 'GQLClient', 'GetCustomRecipe', 'GetCustomRecipeCustomRecipe', 'GetGrader', 'GetGraderGrader', 'GetJudge', 'GetJudgeJudge', 'GlobalUsageFilterInput', 'GoogleProviderDataInput', 'GraderConfigInput', 'GraderCreateInput', 'GraderData', 'GraderDataGraderConfigCustomConfigOutput', 'GraderDataGraderConfigJudgeConfigOutput', 'GraderDataGraderConfigJudgeConfigOutputExamples', 'GraderDataGraderConfigJudgeConfigOutputExamplesInput', 'GraderDataGraderConfigJudgeConfigOutputModel', 'GraderDataGraderConfigPrebuiltConfigOutput', 'GraderDataGraderConfigPrebuiltConfigOutputModel', 'GraderDataGraderConfigPrebuiltConfigOutputPrebuiltCriteria', 'GraderDataGraderConfigRemoteConfigOutput', 'GraderDataMetric', 'GraderDataUseCase', 'GraderTypeEnum', 'GraderUpdateInput', 'GraphQLClientError', 'GraphQLClientGraphQLError', 'GraphQLClientGraphQLMultiError', 'GraphQLClientHttpError', 'GraphQLClientInvalidResponseError', 'GroupSelection', 'GroupSelectionQuery', 'HarmonyGroupData', 'HarmonyGroupDataComputePool', 'HarmonyGroupDataOnlineModels', 'HarmonyStatus', 'JobArtifactFilter', 'JobArtifactKind', 'JobData', 'JobDataCreatedBy', 'JobDataDetails', 'JobDataDetailsArtifacts', 'JobDataDetailsArtifactsByproductsDatasetByproducts', 'JobDataDetailsArtifactsByproductsEvaluationByproducts', 'JobDataDetailsArtifactsByproductsEvaluationByproductsEvalResults', 'JobDataDetailsArtifactsByproductsEvaluationByproductsEvalResultsMetric', 'JobDataDetailsArtifactsByproductsEvaluationByproductsEvalResultsModelService', 'JobDataDetailsArtifactsByproductsModelByproducts', 'JobDataRecipe', 'JobDataStages', 'JobDataStagesInfoBatchInferenceJobStageOutput', 'JobDataStagesInfoEvalJobStageOutput', 'JobDataStagesInfoTrainingJobStageOutput', 'JobDataUseCase', 'JobInput', 'JobKind', 'JobStageOutputData', 'JobStageOutputDataInfoBatchInferenceJobStageOutput', 'JobStageOutputDataInfoEvalJobStageOutput', 'JobStageOutputDataInfoTrainingJobStageOutput', 'JobStatus', 'JobStatusOutput', 'JudgeCapability', 'JudgeConfigInput', 'JudgeCreate', 'JudgeData', 'JudgeDataExamples', 'JudgeDataExamplesInput', 'JudgeDataMetric', 'JudgeDataModel', 'JudgeExampleInput', 'JudgeExampleInputTurnEntry', 'JudgeUpdate', 'LabelFilter', 'LabelInput', 'LinkMetric', 'LinkMetricLinkMetric', 'ListAbCampaigns', 'ListAbCampaignsAbCampaigns', 'ListCompletionsFilterInput', 'ListComputePools', 'ListComputePoolsComputePools', 'ListComputePoolsComputePoolsHarmonyGroups', 'ListCustomRecipes', 'ListCustomRecipesCustomRecipes', 'ListDatasets', 'ListDatasetsDatasets', 'ListGraders', 'ListGradersGraders', 'ListGroupedInteractions', 'ListGroupedInteractionsCompletionsGrouped', 'ListGroupedInteractionsCompletionsGroupedNodes', 'ListGroupedInteractionsCompletionsGroupedNodesCompletions', 'ListGroupedInteractionsCompletionsGroupedNodesCompletionsNodes', 'ListGroupedInteractionsCompletionsGroupedNodesDirectFeedbacksStats', 'ListGroupedInteractionsCompletionsGroupedNodesDirectFeedbacksStatsMetric', 'ListGroupedInteractionsCompletionsGroupedPageInfo', 'ListHarmonyGroups', 'ListHarmonyGroupsHarmonyGroups', 'ListInteractions', 'ListInteractionsCompletions', 'ListInteractionsCompletionsNodes', 'ListInteractionsCompletionsPageInfo', 'ListJobs', 'ListJobsFilterInput', 'ListJobsJobs', 'ListJobsJobsNodes', 'ListJobsJobsPageInfo', 'ListJudgeVersions', 'ListJudgeVersionsJudgeVersions', 'ListJudges', 'ListJudgesJudges', 'ListMetrics', 'ListMetricsMetrics', 'ListModels', 'ListModelsModels', 'ListModelsModelsBackbone', 'ListPermissions', 'ListRemoteEnvs', 'ListRemoteEnvsRemoteEnvs', 'ListRoles', 'ListRolesRoles', 'ListTeams', 'ListTeamsTeams', 'ListUseCases', 'ListUseCasesUseCases', 'ListUsers', 'ListUsersUsers', 'LoadDataset', 'LoadDatasetCreateDataset', 'LockGrader', 'LockGraderLockGrader', 'Me', 'MeMe', 'MeMeApiKeys', 'MetricAggregation', 'MetricCreate', 'MetricData', 'MetricDataAdmin', 'MetricDataAdminUseCases', 'MetricGetOrCreate', 'MetricKind', 'MetricLink', 'MetricScoringType', 'MetricTrendInput', 'MetricUnlink', 'MetricWithContextData', 'ModelCapabilityFilter', 'ModelComputeConfigInput', 'ModelData', 'ModelDataAdmin', 'ModelDataAdminUseCases', 'ModelDataComputeConfig', 'ModelFilter', 'ModelOnline', 'ModelPlacementInput', 'ModelProviderDataInput', 'ModelServiceData', 'ModelServiceDataModel', 'ModelServiceDataModelBackbone', 'ModelServiceDisconnect', 'ModelServiceFilter', 'ModelserviceStatus', 'NvidiaProviderDataInput', 'OpenAIModel', 'OpenAIProviderDataInput', 'OrderPair', 'PrebuiltConfigInput', 'PrebuiltCriteriaKey', 'PrebuiltJudgeCreate', 'Protocol', 'ProviderName', 'RemoteConfigInput', 'RemoteEnvCreate', 'RemoteEnvData', 'RemoteEnvStatus', 'RemoveRemoteEnv', 'RemoveTeamMember', 'RemoveTeamMemberRemoveTeamMember', 'ResizePartitionInput', 'RoleCreate', 'SampleConfig', 'SelectionTypeInput', 'SessionStatus', 'ShareUseCase', 'ShareUseCaseShareUseCase', 'SortDirection', 'SystemPromptTemplateCreate', 'SystemPromptTemplateUpdate', 'TeamCreate', 'TeamMemberRemove', 'TeamMemberSet', 'TerminateModel', 'TestRemoteEnv', 'TestRemoteEnvTestRemoteEnvRemoteEnvTestOffline', 'TestRemoteEnvTestRemoteEnvRemoteEnvTestOnline', 'TimeRange', 'TimeseriesInput', 'TimeseriesInterval', 'UnitConfigInput', 'UnitPosition', 'UnlinkMetric', 'UpdateCompletion', 'UpdateCustomRecipe', 'UpdateCustomRecipeUpdateCustomRecipe', 'UpdateGrader', 'UpdateGraderUpdateGrader', 'UpdateJudge', 'UpdateJudgeUpdateJudge', 'UpdateModel', 'UpdateModelComputeConfig', 'UpdateModelComputeConfigUpdateModelComputeConfig', 'UpdateModelService', 'UpdateModelUpdateModelService', 'UpdateRecipeInput', 'UpdateToolProviderInput', 'UpdateUser', 'UpdateUserSetTeamMember', 'UpdateUserSetTeamMemberRole', 'UpdateUserSetTeamMemberTeam', 'UpdateUserSetTeamMemberUser', 'Upload', 'UsageFilterInput', 'UsagePerUseCaseFilterInput', 'UseCaseCreate', 'UseCaseData', 'UseCaseDataMetrics', 'UseCaseDataModelServices', 'UseCaseDataShares', 'UseCaseDataSharesRole', 'UseCaseDataSharesTeam', 'UseCaseFilter', 'UseCaseMetadataInput', 'UseCaseSettingsInput', 'UseCaseShareInput', 'UseCaseShares', 'UseCaseUpdate', 'UserCreate', 'UserCreateTeamWithRole', 'UserData', 'UserDataTeams', 'UserDataTeamsRole', 'UserDataTeamsTeam', 'WidgetInput']
+__all__ = ['AbCampaignCreateData', 'AbCampaignDetailData', 'AbCampaignDetailDataMetric', 'AbCampaignDetailDataModels', 'AbCampaignDetailDataUseCase', 'AbCampaignFilter', 'AbCampaignReportData', 'AbCampaignReportDataVariants', 'AbCampaignReportDataVariantsComparisons', 'AbCampaignReportDataVariantsComparisonsVariant', 'AbCampaignReportDataVariantsInterval', 'AbCampaignReportDataVariantsVariant', 'AbcampaignCreate', 'AbcampaignStatus', 'AddExternalModel', 'AddExternalModelAddExternalModel', 'AddExternalModelInput', 'AddHFModel', 'AddHFModelImportHfModel', 'AddHFModelInput', 'AddModel', 'AddModelAddModel', 'AddModelAddModelBackbone', 'AddModelInput', 'AddRemoteEnv', 'AddRemoteEnvAddRemoteEnv', 'AnthropicProviderDataInput', 'ApiKeyCreate', 'ArtifactFilter', 'AsyncBaseClientOpenTelemetry', 'AsyncGQLClient', 'AttachModel', 'AttachModelToUseCase', 'AttachModelToUseCaseAttachModel', 'AuthProviderKind', 'AzureProviderDataInput', 'BaseClientOpenTelemetry', 'BaseModel', 'BillingCycle', 'CancelABCampaign', 'CancelAllocationInput', 'CancelJob', 'CancelJobCancelJob', 'CapabilityFilter', 'CompletionComparisonFeedbackData', 'CompletionComparisonFeedbackDataModel', 'CompletionComparisonFilterInput', 'CompletionData', 'CompletionDataChatMessages', 'CompletionDataComparisonFeedbacks', 'CompletionDataComparisonFeedbacksMetric', 'CompletionDataComparisonFeedbacksOtherCompletion', 'CompletionDataComparisonFeedbacksPreferedCompletion', 'CompletionDataComparisonFeedbacksUsecase', 'CompletionDataDirectFeedbacks', 'CompletionDataDirectFeedbacksMetric', 'CompletionDataLabels', 'CompletionDataMetadata', 'CompletionDataMetadataUsage', 'CompletionDataModel', 'CompletionFeedbackFilterInput', 'CompletionGroupBy', 'CompletionLabelValue', 'CompletionSource', 'CompletionsByFilters', 'CompletionsById', 'ComputePoolCapability', 'CreateAbCampaign', 'CreateAbCampaignCreateAbCampaign', 'CreateCustomRecipe', 'CreateCustomRecipeCreateCustomRecipe', 'CreateDatasetFromMultipartUpload', 'CreateDatasetFromMultipartUploadCreateDatasetFromMultipartUpload', 'CreateGrader', 'CreateGraderCreateGrader', 'CreateJob', 'CreateJobCreateJob', 'CreateJudge', 'CreateJudgeCreateJudge', 'CreateMetric', 'CreateMetricCreateMetric', 'CreatePrebuiltJudge', 'CreatePrebuiltJudgeCreatePrebuiltJudge', 'CreateRecipeInput', 'CreateRole', 'CreateRoleCreateRole', 'CreateTeam', 'CreateTeamCreateTeam', 'CreateToolProviderInput', 'CreateUseCase', 'CreateUseCaseCreateUseCase', 'CreateUser', 'CreateUserCreateUser', 'CursorPageInput', 'CustomConfigInput', 'CustomRecipeData', 'CustomRecipeDataCreatedBy', 'CustomRecipeDataLabels', 'CustomRecipeFilterInput', 'DatasetCompletionQuery', 'DatasetCreate', 'DatasetCreateFromFilters', 'DatasetCreateFromMultipartUpload', 'DatasetData', 'DatasetDataMetricsUsage', 'DatasetDataMetricsUsageMetric', 'DatasetKind', 'DatasetSource', 'DatasetStatus', 'DatasetUploadProcessingStatus', 'DatasetUploadProcessingStatusDatasetUploadProcessingStatus', 'DatasetUploadProcessingStatusInput', 'DateBucketUnit', 'DeleteCustomRecipe', 'DeleteDataset', 'DeleteGrader', 'DeleteGraderDeleteGrader', 'DeleteJudge', 'DeleteJudgeDeleteJudge', 'DeleteUser', 'DeleteUserDeleteUser', 'DeployModel', 'DescribeAbCampaign', 'DescribeAbCampaignAbCampaign', 'DescribeAbCampaignAbCampaignReport', 'DescribeDataset', 'DescribeDatasetDataset', 'DescribeInteraction', 'DescribeInteractionCompletion', 'DescribeJob', 'DescribeJobJob', 'DescribeMetric', 'DescribeMetricAdmin', 'DescribeMetricAdminMetric', 'DescribeMetricMetric', 'DescribeModel', 'DescribeModelAdmin', 'DescribeModelAdminModel', 'DescribeModelAdminModelBackbone', 'DescribeModelModel', 'DescribeModelModelBackbone', 'DescribeUseCase', 'DescribeUseCaseUseCase', 'EmojiInput', 'ExternalModelProviderName', 'FeedbackAddInput', 'FeedbackFilterInput', 'FeedbackType', 'FeedbackUpdateInput', 'FromGroupsQuery', 'GQLClient', 'GetCustomRecipe', 'GetCustomRecipeCustomRecipe', 'GetGrader', 'GetGraderGrader', 'GetJudge', 'GetJudgeJudge', 'GlobalUsageFilterInput', 'GoogleProviderDataInput', 'GraderConfigInput', 'GraderCreateInput', 'GraderData', 'GraderDataGraderConfigCustomConfigOutput', 'GraderDataGraderConfigJudgeConfigOutput', 'GraderDataGraderConfigJudgeConfigOutputExamples', 'GraderDataGraderConfigJudgeConfigOutputExamplesInput', 'GraderDataGraderConfigJudgeConfigOutputModel', 'GraderDataGraderConfigPrebuiltConfigOutput', 'GraderDataGraderConfigPrebuiltConfigOutputModel', 'GraderDataGraderConfigPrebuiltConfigOutputPrebuiltCriteria', 'GraderDataGraderConfigRemoteConfigOutput', 'GraderDataMetric', 'GraderDataUseCase', 'GraderTypeEnum', 'GraderUpdateInput', 'GraphQLClientError', 'GraphQLClientGraphQLError', 'GraphQLClientGraphQLMultiError', 'GraphQLClientHttpError', 'GraphQLClientInvalidResponseError', 'GroupSelection', 'GroupSelectionQuery', 'HarmonyGroupData', 'HarmonyGroupDataComputePool', 'HarmonyGroupDataOnlineModels', 'HarmonyStatus', 'JobArtifactFilter', 'JobArtifactKind', 'JobData', 'JobDataCreatedBy', 'JobDataDetails', 'JobDataDetailsArtifacts', 'JobDataDetailsArtifactsByproductsDatasetByproducts', 'JobDataDetailsArtifactsByproductsEvaluationByproducts', 'JobDataDetailsArtifactsByproductsEvaluationByproductsEvalResults', 'JobDataDetailsArtifactsByproductsEvaluationByproductsEvalResultsMetric', 'JobDataDetailsArtifactsByproductsEvaluationByproductsEvalResultsModelService', 'JobDataDetailsArtifactsByproductsModelByproducts', 'JobDataRecipe', 'JobDataStages', 'JobDataStagesInfoBatchInferenceJobStageOutput', 'JobDataStagesInfoEvalJobStageOutput', 'JobDataStagesInfoTrainingJobStageOutput', 'JobDataUseCase', 'JobInput', 'JobKind', 'JobStageOutputData', 'JobStageOutputDataInfoBatchInferenceJobStageOutput', 'JobStageOutputDataInfoEvalJobStageOutput', 'JobStageOutputDataInfoTrainingJobStageOutput', 'JobStatus', 'JobStatusOutput', 'JudgeCapability', 'JudgeConfigInput', 'JudgeCreate', 'JudgeData', 'JudgeDataExamples', 'JudgeDataExamplesInput', 'JudgeDataMetric', 'JudgeDataModel', 'JudgeExampleInput', 'JudgeExampleInputTurnEntry', 'JudgeUpdate', 'LabelFilter', 'LabelInput', 'LinkMetric', 'LinkMetricLinkMetric', 'ListAbCampaigns', 'ListAbCampaignsAbCampaigns', 'ListCompletionsFilterInput', 'ListComputePools', 'ListComputePoolsComputePools', 'ListComputePoolsComputePoolsHarmonyGroups', 'ListCustomRecipes', 'ListCustomRecipesCustomRecipes', 'ListDatasets', 'ListDatasetsDatasets', 'ListGraders', 'ListGradersGraders', 'ListGroupedInteractions', 'ListGroupedInteractionsCompletionsGrouped', 'ListGroupedInteractionsCompletionsGroupedNodes', 'ListGroupedInteractionsCompletionsGroupedNodesCompletions', 'ListGroupedInteractionsCompletionsGroupedNodesCompletionsNodes', 'ListGroupedInteractionsCompletionsGroupedNodesDirectFeedbacksStats', 'ListGroupedInteractionsCompletionsGroupedNodesDirectFeedbacksStatsMetric', 'ListGroupedInteractionsCompletionsGroupedPageInfo', 'ListHarmonyGroups', 'ListHarmonyGroupsHarmonyGroups', 'ListInteractions', 'ListInteractionsCompletions', 'ListInteractionsCompletionsNodes', 'ListInteractionsCompletionsPageInfo', 'ListJobs', 'ListJobsFilterInput', 'ListJobsJobs', 'ListJobsJobsNodes', 'ListJobsJobsPageInfo', 'ListJudgeVersions', 'ListJudgeVersionsJudgeVersions', 'ListJudges', 'ListJudgesJudges', 'ListMetrics', 'ListMetricsMetrics', 'ListModels', 'ListModelsModels', 'ListModelsModelsBackbone', 'ListPermissions', 'ListRemoteEnvs', 'ListRemoteEnvsRemoteEnvs', 'ListRoles', 'ListRolesRoles', 'ListTeams', 'ListTeamsTeams', 'ListUseCases', 'ListUseCasesUseCases', 'ListUsers', 'ListUsersUsers', 'LoadDataset', 'LoadDatasetCreateDataset', 'LockGrader', 'LockGraderLockGrader', 'Me', 'MeMe', 'MeMeApiKeys', 'MetricAggregation', 'MetricCreate', 'MetricData', 'MetricDataAdmin', 'MetricDataAdminUseCases', 'MetricGetOrCreate', 'MetricKind', 'MetricLink', 'MetricScoringType', 'MetricTrendInput', 'MetricUnlink', 'MetricWithContextData', 'ModelCapabilityFilter', 'ModelComputeConfigInput', 'ModelData', 'ModelDataAdmin', 'ModelDataAdminUseCases', 'ModelDataComputeConfig', 'ModelFilter', 'ModelOnline', 'ModelPlacementInput', 'ModelProviderDataInput', 'ModelServiceData', 'ModelServiceDataModel', 'ModelServiceDataModelBackbone', 'ModelServiceDisconnect', 'ModelServiceFilter', 'ModelserviceStatus', 'NvidiaProviderDataInput', 'OpenAIModel', 'OpenAIProviderDataInput', 'OrderPair', 'PrebuiltConfigInput', 'PrebuiltCriteriaKey', 'PrebuiltJudgeCreate', 'Protocol', 'ProviderName', 'RemoteConfigInput', 'RemoteEnvCreate', 'RemoteEnvData', 'RemoteEnvStatus', 'RemoveRemoteEnv', 'RemoveTeamMember', 'RemoveTeamMemberRemoveTeamMember', 'ResizePartitionInput', 'RoleCreate', 'SampleConfig', 'SelectionTypeInput', 'SessionStatus', 'ShareUseCase', 'ShareUseCaseShareUseCase', 'SortDirection', 'SystemPromptTemplateCreate', 'SystemPromptTemplateUpdate', 'TeamCreate', 'TeamMemberRemove', 'TeamMemberSet', 'TerminateModel', 'TestRemoteEnv', 'TestRemoteEnvTestRemoteEnvRemoteEnvTestOffline', 'TestRemoteEnvTestRemoteEnvRemoteEnvTestOnline', 'TimeRange', 'TimeseriesInput', 'TimeseriesInterval', 'UnitConfigInput', 'UnitPosition', 'UnlinkMetric', 'UpdateCompletion', 'UpdateCustomRecipe', 'UpdateCustomRecipeUpdateCustomRecipe', 'UpdateGrader', 'UpdateGraderUpdateGrader', 'UpdateJudge', 'UpdateJudgeUpdateJudge', 'UpdateModel', 'UpdateModelComputeConfig', 'UpdateModelComputeConfigUpdateModelComputeConfig', 'UpdateModelService', 'UpdateModelUpdateModelService', 'UpdateRecipeInput', 'UpdateToolProviderInput', 'UpdateUser', 'UpdateUserSetTeamMember', 'UpdateUserSetTeamMemberRole', 'UpdateUserSetTeamMemberTeam', 'UpdateUserSetTeamMemberUser', 'Upload', 'UsageFilterInput', 'UsagePerUseCaseFilterInput', 'UseCaseCreate', 'UseCaseData', 'UseCaseDataMetrics', 'UseCaseDataModelServices', 'UseCaseDataShares', 'UseCaseDataSharesRole', 'UseCaseDataSharesTeam', 'UseCaseFilter', 'UseCaseMetadataInput', 'UseCaseSettingsInput', 'UseCaseShareInput', 'UseCaseShares', 'UseCaseUpdate', 'UserCreate', 'UserCreateTeamWithRole', 'UserData', 'UserDataTeams', 'UserDataTeamsRole', 'UserDataTeamsTeam', 'WidgetInput']

{adaptive_sdk-0.1.8 → adaptive_sdk-0.1.9}/src/adaptive_sdk/graphql_client/async_client.py

@@ -12,6 +12,7 @@ from .cancel_ab_campaign import CancelABCampaign
 from .cancel_job import CancelJob
 from .create_ab_campaign import CreateAbCampaign
 from .create_custom_recipe import CreateCustomRecipe
+from .create_dataset_from_multipart_upload import CreateDatasetFromMultipartUpload
 from .create_grader import CreateGrader
 from .create_job import CreateJob
 from .create_judge import CreateJudge
@@ -21,6 +22,7 @@ from .create_role import CreateRole
 from .create_team import CreateTeam
 from .create_use_case import CreateUseCase
 from .create_user import CreateUser
+from .dataset_upload_processing_status import DatasetUploadProcessingStatus
 from .delete_custom_recipe import DeleteCustomRecipe
 from .delete_dataset import DeleteDataset
 from .delete_grader import DeleteGrader
@@ -40,7 +42,7 @@ from .enums import CompletionGroupBy
 from .get_custom_recipe import GetCustomRecipe
 from .get_grader import GetGrader
 from .get_judge import GetJudge
-from .input_types import AbcampaignCreate, AbCampaignFilter, AddExternalModelInput, AddHFModelInput, AddModelInput, AttachModel, CreateRecipeInput, CursorPageInput, CustomRecipeFilterInput, DatasetCreate, GraderCreateInput, GraderUpdateInput, JobInput, JudgeCreate, JudgeUpdate, ListCompletionsFilterInput, ListJobsFilterInput, MetricCreate, MetricLink, MetricUnlink, ModelComputeConfigInput, ModelFilter, ModelPlacementInput, OrderPair, PrebuiltJudgeCreate, RemoteEnvCreate, RoleCreate, TeamCreate, TeamMemberRemove, TeamMemberSet, UpdateModelService, UpdateRecipeInput, UseCaseCreate, UseCaseShares, UserCreate
+from .input_types import AbcampaignCreate, AbCampaignFilter, AddExternalModelInput, AddHFModelInput, AddModelInput, AttachModel, CreateRecipeInput, CursorPageInput, CustomRecipeFilterInput, DatasetCreate, DatasetCreateFromMultipartUpload, DatasetUploadProcessingStatusInput, GraderCreateInput, GraderUpdateInput, JobInput, JudgeCreate, JudgeUpdate, ListCompletionsFilterInput, ListJobsFilterInput, MetricCreate, MetricLink, MetricUnlink, ModelComputeConfigInput, ModelFilter, ModelPlacementInput, OrderPair, PrebuiltJudgeCreate, RemoteEnvCreate, RoleCreate, TeamCreate, TeamMemberRemove, TeamMemberSet, UpdateModelService, UpdateRecipeInput, UseCaseCreate, UseCaseShares, UserCreate
 from .link_metric import LinkMetric
 from .list_ab_campaigns import ListAbCampaigns
 from .list_compute_pools import ListComputePools
@@ -356,6 +358,13 @@ class AsyncGQLClient(AsyncBaseClientOpenTelemetry):
         data = self.get_data(response)
         return CancelJob.model_validate(data)
 
+    async def create_dataset_from_multipart_upload(self, input: DatasetCreateFromMultipartUpload, **kwargs: Any) -> CreateDatasetFromMultipartUpload:
+        query = gql('\n            mutation CreateDatasetFromMultipartUpload($input: DatasetCreateFromMultipartUpload!) {\n              createDatasetFromMultipartUpload(input: $input) {\n                datasetId\n                status\n                totalParts\n                processedParts\n                progress\n                error\n              }\n            }\n            ')
+        variables: Dict[str, object] = {'input': input}
+        response = await self.execute(query=query, operation_name='CreateDatasetFromMultipartUpload', variables=variables, **kwargs)
+        data = self.get_data(response)
+        return CreateDatasetFromMultipartUpload.model_validate(data)
+
     async def list_datasets(self, input: Any, **kwargs: Any) -> ListDatasets:
         query = gql('\n            query ListDatasets($input: IdOrKey!) {\n              datasets(useCase: $input) {\n                ...DatasetData\n              }\n            }\n\n            fragment DatasetData on Dataset {\n              id\n              key\n              name\n              createdAt\n              kind\n              records\n              metricsUsage {\n                feedbackCount\n                comparisonCount\n                metric {\n                  ...MetricData\n                }\n              }\n            }\n\n            fragment MetricData on Metric {\n              id\n              key\n              name\n              kind\n              description\n              scoringType\n              createdAt\n              hasDirectFeedbacks\n              hasComparisonFeedbacks\n            }\n            ')
         variables: Dict[str, object] = {'input': input}
@@ -580,6 +589,13 @@ class AsyncGQLClient(AsyncBaseClientOpenTelemetry):
         data = self.get_data(response)
         return DescribeJob.model_validate(data)
 
+    async def dataset_upload_processing_status(self, input: DatasetUploadProcessingStatusInput, **kwargs: Any) -> DatasetUploadProcessingStatus:
+        query = gql('\n            query DatasetUploadProcessingStatus($input: DatasetUploadProcessingStatusInput!) {\n              datasetUploadProcessingStatus(input: $input) {\n                datasetId\n                status\n                totalParts\n                processedParts\n                progress\n                error\n              }\n            }\n            ')
+        variables: Dict[str, object] = {'input': input}
+        response = await self.execute(query=query, operation_name='DatasetUploadProcessingStatus', variables=variables, **kwargs)
+        data = self.get_data(response)
+        return DatasetUploadProcessingStatus.model_validate(data)
+
     async def execute_custom_operation(self, *fields: GraphQLField, operation_type: OperationType, operation_name: str) -> Dict[str, Any]:
         selections = self._build_selection_set(fields)
         combined_variables = self._combine_variables(fields)

{adaptive_sdk-0.1.8 → adaptive_sdk-0.1.9}/src/adaptive_sdk/graphql_client/client.py

@@ -12,6 +12,7 @@ from .cancel_ab_campaign import CancelABCampaign
 from .cancel_job import CancelJob
 from .create_ab_campaign import CreateAbCampaign
 from .create_custom_recipe import CreateCustomRecipe
+from .create_dataset_from_multipart_upload import CreateDatasetFromMultipartUpload
 from .create_grader import CreateGrader
 from .create_job import CreateJob
 from .create_judge import CreateJudge
@@ -21,6 +22,7 @@ from .create_role import CreateRole
 from .create_team import CreateTeam
 from .create_use_case import CreateUseCase
 from .create_user import CreateUser
+from .dataset_upload_processing_status import DatasetUploadProcessingStatus
 from .delete_custom_recipe import DeleteCustomRecipe
 from .delete_dataset import DeleteDataset
 from .delete_grader import DeleteGrader
@@ -40,7 +42,7 @@ from .enums import CompletionGroupBy
 from .get_custom_recipe import GetCustomRecipe
 from .get_grader import GetGrader
 from .get_judge import GetJudge
-from .input_types import AbcampaignCreate, AbCampaignFilter, AddExternalModelInput, AddHFModelInput, AddModelInput, AttachModel, CreateRecipeInput, CursorPageInput, CustomRecipeFilterInput, DatasetCreate, GraderCreateInput, GraderUpdateInput, JobInput, JudgeCreate, JudgeUpdate, ListCompletionsFilterInput, ListJobsFilterInput, MetricCreate, MetricLink, MetricUnlink, ModelComputeConfigInput, ModelFilter, ModelPlacementInput, OrderPair, PrebuiltJudgeCreate, RemoteEnvCreate, RoleCreate, TeamCreate, TeamMemberRemove, TeamMemberSet, UpdateModelService, UpdateRecipeInput, UseCaseCreate, UseCaseShares, UserCreate
+from .input_types import AbcampaignCreate, AbCampaignFilter, AddExternalModelInput, AddHFModelInput, AddModelInput, AttachModel, CreateRecipeInput, CursorPageInput, CustomRecipeFilterInput, DatasetCreate, DatasetCreateFromMultipartUpload, DatasetUploadProcessingStatusInput, GraderCreateInput, GraderUpdateInput, JobInput, JudgeCreate, JudgeUpdate, ListCompletionsFilterInput, ListJobsFilterInput, MetricCreate, MetricLink, MetricUnlink, ModelComputeConfigInput, ModelFilter, ModelPlacementInput, OrderPair, PrebuiltJudgeCreate, RemoteEnvCreate, RoleCreate, TeamCreate, TeamMemberRemove, TeamMemberSet, UpdateModelService, UpdateRecipeInput, UseCaseCreate, UseCaseShares, UserCreate
 from .link_metric import LinkMetric
 from .list_ab_campaigns import ListAbCampaigns
 from .list_compute_pools import ListComputePools
@@ -356,6 +358,13 @@ class GQLClient(BaseClientOpenTelemetry):
         data = self.get_data(response)
         return CancelJob.model_validate(data)
 
+    def create_dataset_from_multipart_upload(self, input: DatasetCreateFromMultipartUpload, **kwargs: Any) -> CreateDatasetFromMultipartUpload:
+        query = gql('\n            mutation CreateDatasetFromMultipartUpload($input: DatasetCreateFromMultipartUpload!) {\n              createDatasetFromMultipartUpload(input: $input) {\n                datasetId\n                status\n                totalParts\n                processedParts\n                progress\n                error\n              }\n            }\n            ')
+        variables: Dict[str, object] = {'input': input}
+        response = self.execute(query=query, operation_name='CreateDatasetFromMultipartUpload', variables=variables, **kwargs)
+        data = self.get_data(response)
+        return CreateDatasetFromMultipartUpload.model_validate(data)
+
     def list_datasets(self, input: str, **kwargs: Any) -> ListDatasets:
         query = gql('\n            query ListDatasets($input: IdOrKey!) {\n              datasets(useCase: $input) {\n                ...DatasetData\n              }\n            }\n\n            fragment DatasetData on Dataset {\n              id\n              key\n              name\n              createdAt\n              kind\n              records\n              metricsUsage {\n                feedbackCount\n                comparisonCount\n                metric {\n                  ...MetricData\n                }\n              }\n            }\n\n            fragment MetricData on Metric {\n              id\n              key\n              name\n              kind\n              description\n              scoringType\n              createdAt\n              hasDirectFeedbacks\n              hasComparisonFeedbacks\n            }\n            ')
         variables: Dict[str, object] = {'input': input}
@@ -580,6 +589,13 @@ class GQLClient(BaseClientOpenTelemetry):
         data = self.get_data(response)
         return DescribeJob.model_validate(data)
 
+    def dataset_upload_processing_status(self, input: DatasetUploadProcessingStatusInput, **kwargs: Any) -> DatasetUploadProcessingStatus:
+        query = gql('\n            query DatasetUploadProcessingStatus($input: DatasetUploadProcessingStatusInput!) {\n              datasetUploadProcessingStatus(input: $input) {\n                datasetId\n                status\n                totalParts\n                processedParts\n                progress\n                error\n              }\n            }\n            ')
+        variables: Dict[str, object] = {'input': input}
+        response = self.execute(query=query, operation_name='DatasetUploadProcessingStatus', variables=variables, **kwargs)
+        data = self.get_data(response)
+        return DatasetUploadProcessingStatus.model_validate(data)
+
     def execute_custom_operation(self, *fields: GraphQLField, operation_type: OperationType, operation_name: str) -> Dict[str, Any]:
         selections = self._build_selection_set(fields)
         combined_variables = self._combine_variables(fields)

adaptive_sdk-0.1.9/src/adaptive_sdk/graphql_client/create_dataset_from_multipart_upload.py

@@ -0,0 +1,18 @@
+from typing import Any, Optional
+from pydantic import Field
+from .base_model import BaseModel
+from .enums import SessionStatus
+
+class CreateDatasetFromMultipartUpload(BaseModel):
+    """@public"""
+    create_dataset_from_multipart_upload: 'CreateDatasetFromMultipartUploadCreateDatasetFromMultipartUpload' = Field(alias='createDatasetFromMultipartUpload')
+
+class CreateDatasetFromMultipartUploadCreateDatasetFromMultipartUpload(BaseModel):
+    """@public"""
+    dataset_id: Any = Field(alias='datasetId')
+    status: SessionStatus
+    total_parts: int = Field(alias='totalParts')
+    processed_parts: int = Field(alias='processedParts')
+    progress: float
+    error: Optional[str]
+CreateDatasetFromMultipartUpload.model_rebuild()

{adaptive_sdk-0.1.8 → adaptive_sdk-0.1.9}/src/adaptive_sdk/graphql_client/custom_queries.py

@@ -35,6 +35,12 @@ class Query:
         cleared_arguments = {key: value for key, value in arguments.items() if value['value'] is not None}
         return CustomRecipeFields(field_name='customRecipe', arguments=cleared_arguments)
 
+    @classmethod
+    def parse_recipe_schema(cls, recipe_content: str) -> GraphQLField:
+        arguments: Dict[str, Dict[str, Any]] = {'recipeContent': {'type': 'String!', 'value': recipe_content}}
+        cleared_arguments = {key: value for key, value in arguments.items() if value['value'] is not None}
+        return GraphQLField(field_name='parseRecipeSchema', arguments=cleared_arguments)
+
     @classmethod
     def datasets(cls, use_case: str) -> DatasetFields:
         arguments: Dict[str, Dict[str, Any]] = {'useCase': {'type': 'IdOrKey!', 'value': use_case}}

adaptive_sdk-0.1.9/src/adaptive_sdk/graphql_client/dataset_upload_processing_status.py

@@ -0,0 +1,18 @@
+from typing import Any, Optional
+from pydantic import Field
+from .base_model import BaseModel
+from .enums import SessionStatus
+
+class DatasetUploadProcessingStatus(BaseModel):
+    """@public"""
+    dataset_upload_processing_status: 'DatasetUploadProcessingStatusDatasetUploadProcessingStatus' = Field(alias='datasetUploadProcessingStatus')
+
+class DatasetUploadProcessingStatusDatasetUploadProcessingStatus(BaseModel):
+    """@public"""
+    dataset_id: Any = Field(alias='datasetId')
+    status: SessionStatus
+    total_parts: int = Field(alias='totalParts')
+    processed_parts: int = Field(alias='processedParts')
+    progress: float
+    error: Optional[str]
+DatasetUploadProcessingStatus.model_rebuild()

{adaptive_sdk-0.1.8 → adaptive_sdk-0.1.9}/src/adaptive_sdk/graphql_client/enums.py

@@ -49,6 +49,7 @@ class DatasetSource(str, Enum):
 
 class DatasetStatus(str, Enum):
     """@public"""
+    PENDING = 'PENDING'
     PROCESSING = 'PROCESSING'
     READY = 'READY'
     FAILED = 'FAILED'

{adaptive_sdk-0.1.8 → adaptive_sdk-0.1.9}/src/adaptive_sdk/resources/datasets.py

@@ -2,22 +2,20 @@ from __future__ import annotations
 import json
 import math
 import os
+import time
 from pathlib import Path
-from typing import List, Literal, TYPE_CHECKING
+from typing import List, TYPE_CHECKING
 
 from adaptive_sdk.graphql_client import (
     DatasetCreate,
     Upload,
-    LoadDatasetCreateDataset,
     ListDatasetsDatasets,
     DatasetData,
     DatasetCreateFromMultipartUpload,
+    DatasetUploadProcessingStatusInput,
+    SessionStatus,
 )
-from adaptive_sdk.graphql_client.custom_mutations import Mutation
-from adaptive_sdk.graphql_client.custom_fields import (
-    DatasetFields,
-    DatasetUploadProcessingStatusFields,
-)
+
 from adaptive_sdk.rest import rest_types
 from adaptive_sdk.error_handling import rest_error_handler
 
@@ -29,23 +27,12 @@ if TYPE_CHECKING:
 MIN_CHUNK_SIZE_BYTES = 5 * 1024 * 1024  # 5MB
 MAX_CHUNK_SIZE_BYTES = 100 * 1024 * 1024  # 100MB
 MAX_PARTS_COUNT = 10000
+INIT_CHUNKED_UPLOAD_ROUTE = "/upload/init"
+UPLOAD_PART_ROUTE = "/upload/part"
+ABORT_CHUNKED_UPLOAD_ROUTE = "/upload/abort"
 
 
 def _calculate_upload_parts(file_size: int) -> tuple[int, int]:
-    """
-    Calculate optimal number of parts and chunk size for multipart upload.
-
-    Strategy: Scale chunk size based on file size for optimal performance
-
-    Args:
-        file_size: Size of the file in bytes
-
-    Returns:
-        Tuple of (total_parts, chunk_size_bytes)
-
-    Raises:
-        ValueError: If file is too large to upload with the given constraints
-    """
     if file_size < MIN_CHUNK_SIZE_BYTES:
         raise ValueError(f"File size ({file_size:,} bytes) is too small for chunked upload")
 
@@ -97,19 +84,15 @@ class Datasets(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
         dataset_key: str,
         name: str | None = None,
         use_case: str | None = None,
-    ) -> LoadDatasetCreateDataset:
+    ) -> DatasetData:
         """
-         Upload a dataset from a file. File must be jsonl, where each line should match structure in example below.
+        Upload a dataset from a file. File must be jsonl, where each line should match supported structure.
 
         Args:
             file_path: Path to jsonl file.
             dataset_key: New dataset key.
             name: Optional name to render in UI; if `None`, defaults to same as `dataset_key`.
 
-        Example:
-        ```
-        {"messages": [{"role": "system", "content": "<optional system prompt>"}, {"role": "user", "content": "<user content>"}, {"role": "assistant", "content": "<assistant answer>"}], "completion": "hey"}
-        ```
         """
         file_size = os.path.getsize(file_path)
 
@@ -124,12 +107,8 @@ class Datasets(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
             )
             filename = Path(file_path).stem
             with open(file_path, "rb") as f:
-                file_upload = Upload(
-                    filename=filename, content=f, content_type="application/jsonl"
-                )
-                return self._gql_client.load_dataset(
-                    input=input, file=file_upload
-                ).create_dataset
+                file_upload = Upload(filename=filename, content=f, content_type="application/jsonl")
+                return self._gql_client.load_dataset(input=input, file=file_upload).create_dataset
 
     def _chunked_upload(
         self,
@@ -137,78 +116,73 @@ class Datasets(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
         dataset_key: str,
         name: str | None = None,
         use_case: str | None = None,
-    ) -> LoadDatasetCreateDataset:
+    ) -> DatasetData:
         """Upload large files using chunked upload via REST API."""
         file_size = os.path.getsize(file_path)
         total_parts, chunk_size = _calculate_upload_parts(file_size)
 
-        # Step 1: Initialize chunked upload session
         init_request = rest_types.InitChunkedUploadRequest(
             content_type="application/jsonl",
             metadata=None,
             total_parts_count=total_parts,
         )
-        response = self._rest_client.post(
-            "/upload/init", json=init_request.model_dump()
-        )
+        response = self._rest_client.post(INIT_CHUNKED_UPLOAD_ROUTE, json=init_request.model_dump())
         rest_error_handler(response)
-        init_response = rest_types.InitChunkedUploadResponse.model_validate(
-            response.json()
-        )
+        init_response = rest_types.InitChunkedUploadResponse.model_validate(response.json())
         session_id = init_response.session_id
 
         try:
-            # Step 2: Upload each part
             with open(file_path, "rb") as f:
                 for part_number in range(1, total_parts + 1):
                     chunk_data = f.read(chunk_size)
 
                     response = self._rest_client.post(
-                        "/upload/part",
+                        UPLOAD_PART_ROUTE,
                         params={"session_id": session_id, "part_number": part_number},
                         content=chunk_data,
                         headers={"Content-Type": "application/octet-stream"},
                     )
                     rest_error_handler(response)
 
-            # Step 3: Finalize upload by creating dataset from multipart upload
             input = DatasetCreateFromMultipartUpload(
                 useCase=self.use_case_key(use_case),
                 name=name if name else dataset_key,
                 key=dataset_key,
                 uploadSessionId=session_id,
             )
+            create_dataset_result = self._gql_client.create_dataset_from_multipart_upload(
+                input=input
+            ).create_dataset_from_multipart_upload
+
+            upload_done = False
+            while not upload_done:
+                check_progress_result = self._gql_client.dataset_upload_processing_status(
+                    input=DatasetUploadProcessingStatusInput(
+                        useCase=self.use_case_key(use_case), datasetId=create_dataset_result.dataset_id
+                    )
+                ).dataset_upload_processing_status
+                if check_progress_result.status == SessionStatus.DONE:
+                    upload_done = True
+                elif check_progress_result.status == SessionStatus.ERROR:
+                    raise Exception(f"Upload failed: {check_progress_result.error}")
+                else:
+                    time.sleep(2)
 
-            mutation_field = Mutation.create_dataset_from_multipart_upload(input=input)
-            mutation_field.fields(
-                DatasetUploadProcessingStatusFields.dataset_id,
-                DatasetUploadProcessingStatusFields.total_parts,
-                DatasetUploadProcessingStatusFields.processed_parts,
-                DatasetUploadProcessingStatusFields.progress,
-                DatasetUploadProcessingStatusFields.error,
-                DatasetUploadProcessingStatusFields.status,
-            )
-            result = self._gql_client.mutation(
-                mutation_field, operation_name="CreateDatasetFromMultipartUpload"
-            )
-            return LoadDatasetCreateDataset.model_validate(
-                result["createDatasetFromMultipartUpload"]
-            )
+            dataset_data = self.get(create_dataset_result.dataset_id, use_case=self.use_case_key(use_case))
+            assert dataset_data is not None
+
+            return dataset_data
 
-        except Exception as e:
-            # Abort the upload session on error
+        except Exception:
             try:
-                abort_request = rest_types.AbortChunkedUploadRequest(
-                    session_id=session_id
-                )
-                # delete with body works in runtime but not in type checking
-                self._rest_client.delete(  # type: ignore[call-arg]
-                    "/upload/abort",
-                    content=json.dumps(abort_request.model_dump()),
+                abort_request = rest_types.AbortChunkedUploadRequest(session_id=session_id)
+                self._rest_client.delete(
+                    ABORT_CHUNKED_UPLOAD_ROUTE,
+                    content=json.dumps(abort_request.model_dump()),  # type: ignore[call-arg]
                     headers={"Content-Type": "application/json"},
                 )
             except Exception:
-                pass  # Best effort cleanup
+                pass
             raise
 
     def list(
@@ -227,15 +201,11 @@ class Datasets(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
         Args:
             key: Dataset key.
         """
-        return self._gql_client.describe_dataset(
-            key, self.use_case_key(use_case)
-        ).dataset
+        return self._gql_client.describe_dataset(key, self.use_case_key(use_case)).dataset
 
     def delete(self, key: str, use_case: str | None = None) -> bool:
         """Delete dataset."""
-        return self._gql_client.delete_dataset(
-            id_or_key=key, use_case=self.use_case_key(use_case)
-        ).delete_dataset
+        return self._gql_client.delete_dataset(id_or_key=key, use_case=self.use_case_key(use_case)).delete_dataset
 
 
 class AsyncDatasets(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
@@ -249,7 +219,7 @@ class AsyncDatasets(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
         dataset_key: str,
         name: str | None = None,
         use_case: str | None = None,
-    ) -> LoadDatasetCreateDataset:
+    ) -> DatasetData:
         """
         Upload a dataset from a file. File must be jsonl, where each line should match structure in example below.
 
@@ -277,12 +247,8 @@ class AsyncDatasets(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
             )
             filename = Path(file_path).stem
             with open(file_path, "rb") as f:
-                file_upload = Upload(
-                    filename=filename, content=f, content_type="application/jsonl"
-                )
-                upload_result = await self._gql_client.load_dataset(
-                    input=input, file=file_upload
-                )
+                file_upload = Upload(filename=filename, content=f, content_type="application/jsonl")
+                upload_result = await self._gql_client.load_dataset(input=input, file=file_upload)
                 return upload_result.create_dataset
 
     async def _chunked_upload(
@@ -291,34 +257,28 @@ class AsyncDatasets(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
         dataset_key: str,
         name: str | None = None,
         use_case: str | None = None,
-    ) -> LoadDatasetCreateDataset:
+    ) -> DatasetData:
         """Upload large files using chunked upload via REST API."""
         file_size = os.path.getsize(file_path)
         total_parts, chunk_size = _calculate_upload_parts(file_size)
 
-        # Step 1: Initialize chunked upload session
         init_request = rest_types.InitChunkedUploadRequest(
             content_type="application/jsonl",
             metadata=None,
             total_parts_count=total_parts,
         )
-        response = await self._rest_client.post(
-            "/upload/init", json=init_request.model_dump()
-        )
+        response = await self._rest_client.post(INIT_CHUNKED_UPLOAD_ROUTE, json=init_request.model_dump())
         rest_error_handler(response)
-        init_response = rest_types.InitChunkedUploadResponse.model_validate(
-            response.json()
-        )
+        init_response = rest_types.InitChunkedUploadResponse.model_validate(response.json())
         session_id = init_response.session_id
 
         try:
-            # Step 2: Upload each part
             with open(file_path, "rb") as f:
                 for part_number in range(1, total_parts + 1):
                     chunk_data = f.read(chunk_size)
 
                     response = await self._rest_client.post(
-                        "/upload/part",
+                        UPLOAD_PART_ROUTE,
                         params={"session_id": session_id, "part_number": part_number},
                         content=chunk_data,
                         headers={"Content-Type": "application/octet-stream"},
@@ -331,33 +291,37 @@ class AsyncDatasets(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
                 key=dataset_key,
                 uploadSessionId=session_id,
             )
+            create_dataset_result = (
+                await self._gql_client.create_dataset_from_multipart_upload(input=input)
+            ).create_dataset_from_multipart_upload
+
+            upload_done = False
+            while not upload_done:
+                check_progress_result = (
+                    await self._gql_client.dataset_upload_processing_status(
+                        input=DatasetUploadProcessingStatusInput(
+                            useCase=self.use_case_key(use_case), datasetId=create_dataset_result.dataset_id
+                        )
+                    )
+                ).dataset_upload_processing_status
+                if check_progress_result.status == SessionStatus.DONE:
+                    upload_done = True
+                elif check_progress_result.status == SessionStatus.ERROR:
+                    raise Exception(f"Upload failed: {check_progress_result.error}")
+                else:
+                    time.sleep(2)
 
-            mutation_field = Mutation.create_dataset_from_multipart_upload(input=input)
-            mutation_field.fields(
-                DatasetUploadProcessingStatusFields.dataset_id,
-                DatasetUploadProcessingStatusFields.total_parts,
-                DatasetUploadProcessingStatusFields.processed_parts,
-                DatasetUploadProcessingStatusFields.progress,
-                DatasetUploadProcessingStatusFields.error,
-                DatasetUploadProcessingStatusFields.status,
-            )
-            result = await self._gql_client.mutation(
-                mutation_field, operation_name="CreateDatasetFromMultipartUpload"
-            )
-            return LoadDatasetCreateDataset.model_validate(
-                result["createDatasetFromMultipartUpload"]
-            )
+            dataset_data = await self.get(create_dataset_result.dataset_id, use_case=self.use_case_key(use_case))
+            assert dataset_data is not None
 
-        except Exception as e:
-            # Abort the upload session on error
+            return dataset_data
+
+        except Exception:
             try:
-                abort_request = rest_types.AbortChunkedUploadRequest(
-                    session_id=session_id
-                )
-                # delete with body works in runtime but not in type checking
-                await self._rest_client.delete(  # type: ignore[call-arg]
-                    "/upload/abort",
-                    content=json.dumps(abort_request.model_dump()),
+                abort_request = rest_types.AbortChunkedUploadRequest(session_id=session_id)
+                _ = await self._rest_client.delete(
+                    ABORT_CHUNKED_UPLOAD_ROUTE,
+                    content=json.dumps(abort_request.model_dump()),  # type: ignore[call-arg]
                     headers={"Content-Type": "application/json"},
                 )
             except Exception:
@@ -381,15 +345,11 @@ class AsyncDatasets(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
         Args:
             key: Dataset key.
         """
-        result = await self._gql_client.describe_dataset(
-            key, self.use_case_key(use_case)
-        )
+        result = await self._gql_client.describe_dataset(key, self.use_case_key(use_case))
         return result.dataset
 
     async def delete(self, key: str, use_case: str | None = None) -> bool:
         """Delete dataset."""
         return (
-            await self._gql_client.delete_dataset(
-                id_or_key=key, use_case=self.use_case_key(use_case)
-            )
+            await self._gql_client.delete_dataset(id_or_key=key, use_case=self.use_case_key(use_case))
         ).delete_dataset