adaptive-sdk 0.1.13__py3-none-any.whl → 0.1.14__py3-none-any.whl

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, CompletionGroupBy, CompletionSource, DatasetKind, DatasetSource, DateBucketUnit, ExternalModelProviderName, FeedbackType, GraderTypeEnum, JobArtifactKind, JobKind, JobStatus, MetricAggregation, MetricKind, MetricScoringType, ModelCapabilityFilter, ModelOnline, OpenAIModel, PrebuiltCriteriaKey, Protocol, SelectionTypeInput, SortDirection, TimeseriesInterval, UnitPosition
+from .enums import AbcampaignStatus, CompletionGroupBy, CompletionSource, DatasetKind, DatasetSource, DateBucketUnit, ExternalModelProviderName, FeedbackType, GraderTypeEnum, JobArtifactKind, JobKind, JobStatus, MetricAggregation, MetricKind, MetricScoringType, ModelCapabilityFilter, ModelOnline, PrebuiltCriteriaKey, Protocol, SelectionTypeInput, SortDirection, TimeseriesInterval, UnitPosition
 
 class AbCampaignFilter(BaseModel):
     """@private"""
@@ -42,10 +42,16 @@ class AddModelInput(BaseModel):
     name: str
     key: Optional[str] = None
 
+class AddModelToUseCaseInput(BaseModel):
+    """@private"""
+    use_case: str = Field(alias='useCase')
+    model: str
+
 class AnthropicProviderDataInput(BaseModel):
     """@private"""
     api_key: str = Field(alias='apiKey')
     external_model_id: str = Field(alias='externalModelId')
+    endpoint: Optional[str] = None
 
 class ApiKeyCreate(BaseModel):
     """@private"""
@@ -58,15 +64,6 @@ class ArtifactFilter(BaseModel):
     kinds: Optional[List[JobArtifactKind]] = None
     job_id: Optional[Any] = Field(alias='jobId', default=None)
 
-class AttachModel(BaseModel):
-    """@private"""
-    use_case: str = Field(alias='useCase')
-    model: str
-    attached: bool = True
-    placement: Optional['ModelPlacementInput'] = None
-    wait: bool = False
-    'Wait for the model to be deployed or not'
-
 class CancelAllocationInput(BaseModel):
     """@private"""
     harmony_group: str = Field(alias='harmonyGroup')
@@ -93,6 +90,38 @@ class CompletionFeedbackFilterInput(BaseModel):
     reasons: Optional[List[str]] = None
     user: Optional[Any] = None
 
+class CompletionFilterExpression(BaseModel):
+    """@private
+Advanced filter expression supporting AND/OR/NOT logic"""
+    and_: Optional[List['CompletionFilterExpression']] = Field(alias='and', default=None)
+    'Combine multiple conditions with AND (all must match)'
+    or_: Optional[List['CompletionFilterExpression']] = Field(alias='or', default=None)
+    'Combine multiple conditions with OR (at least one must match)'
+    not_: Optional['CompletionFilterExpression'] = Field(alias='not', default=None)
+    'Negate a condition'
+    timerange: Optional['TimeRange'] = None
+    'Filter by time'
+    model: Optional['IdOrKeyCondition'] = None
+    'Filter by model'
+    label: Optional['LabelCondition'] = None
+    'Filter by label key-value pairs'
+    feedbacks: Optional['FeedbackCondition'] = None
+    'Filter by feedback/metric values'
+    source: Optional[CompletionSource] = None
+    'Filter by completion source'
+    prompt_hash: Optional['StringCondition'] = Field(alias='promptHash', default=None)
+    'Filter by prompt hash'
+    session_id: Optional[Any] = Field(alias='sessionId', default=None)
+    'Filter by session ID'
+    user_id: Optional[Any] = Field(alias='userId', default=None)
+    'Filter by user ID'
+    completion_id: Optional[Any] = Field(alias='completionId', default=None)
+    'Filter by completion ID'
+    completion: Optional['TextCondition'] = None
+    'Filter by completion content'
+    prompt: Optional['TextCondition'] = None
+    'Filter by prompt content'
+
 class CompletionLabelValue(BaseModel):
     """@private"""
     key: str
@@ -174,6 +203,18 @@ class DatasetUploadProcessingStatusInput(BaseModel):
     use_case: str = Field(alias='useCase')
     dataset_id: Any = Field(alias='datasetId')
 
+class DeleteModelInput(BaseModel):
+    """@private"""
+    model: str
+
+class DeployModelInput(BaseModel):
+    """@private"""
+    use_case: str = Field(alias='useCase')
+    model: str
+    placement: Optional['ModelPlacementInput'] = None
+    wait: bool = False
+    'Wait for the model to be deployed or not'
+
 class EmojiInput(BaseModel):
     """@private"""
     native: str
@@ -185,6 +226,18 @@ class FeedbackAddInput(BaseModel):
     reason: Optional[str] = None
     user_id: Optional[Any] = Field(alias='userId', default=None)
 
+class FeedbackCondition(BaseModel):
+    """@private
+Feedback/metric filter condition with numeric comparisons"""
+    metric: str
+    'Metric to filter by'
+    value: Optional['FloatNumericCondition'] = None
+    'Numeric value condition'
+    reasons: Optional[List[str]] = None
+    'Filter by feedback reasons'
+    user: Optional[Any] = None
+    'Filter by user who gave the feedback'
+
 class FeedbackFilterInput(BaseModel):
     """@private"""
     labels: Optional[List['LabelFilter']] = None
@@ -194,6 +247,22 @@ class FeedbackUpdateInput(BaseModel):
     value: Optional[Any] = None
     details: Optional[str] = None
 
+class FloatNumericCondition(BaseModel):
+    """@private
+Numeric matching condition for filter expressions, parameterized by the numeric type"""
+    eq: Optional[float] = None
+    'Equal to value'
+    neq: Optional[float] = None
+    'Not Equal to value'
+    gt: Optional[float] = None
+    'Greater than value'
+    gte: Optional[float] = None
+    'Greater than or equal to value'
+    lt: Optional[float] = None
+    'Less than value'
+    lte: Optional[float] = None
+    'Less than or equal to value'
+
 class FromGroupsQuery(BaseModel):
     """@private"""
     filters: 'ListCompletionsFilterInput'
@@ -211,6 +280,7 @@ class GoogleProviderDataInput(BaseModel):
     """@private"""
     api_key: str = Field(alias='apiKey')
     external_model_id: str = Field(alias='externalModelId')
+    endpoint: Optional[str] = None
 
 class GraderConfigInput(BaseModel):
     """@private"""
@@ -243,6 +313,32 @@ class GroupSelectionQuery(BaseModel):
     group_id: str = Field(alias='groupId')
     selection: 'GroupSelection'
 
+class IdOrKeyCondition(BaseModel):
+    """@private
+String matching condition for filter expressions"""
+    eq: Optional[str] = None
+    'Exact match'
+    in_: Optional[List[str]] = Field(alias='in', default=None)
+    'Match any of the provided values (OR)'
+    neq: Optional[str] = None
+    'Does not equal'
+
+class IntegerNumericCondition(BaseModel):
+    """@private
+Numeric matching condition for filter expressions, parameterized by the numeric type"""
+    eq: Optional[int] = None
+    'Equal to value'
+    neq: Optional[int] = None
+    'Not Equal to value'
+    gt: Optional[int] = None
+    'Greater than value'
+    gte: Optional[int] = None
+    'Greater than or equal to value'
+    lt: Optional[int] = None
+    'Less than value'
+    lte: Optional[int] = None
+    'Less than or equal to value'
+
 class JobArtifactFilter(BaseModel):
     """@private"""
     kinds: List[JobArtifactKind]
@@ -293,6 +389,14 @@ class JudgeUpdate(BaseModel):
     examples: Optional[List['JudgeExampleInput']] = None
     model: Optional[str] = None
 
+class LabelCondition(BaseModel):
+    """@private
+Label-specific filter condition"""
+    key: str
+    'Label key'
+    value: Optional['StringCondition'] = None
+    'Label value condition (optional - if not set, just checks for key existence)'
+
 class LabelFilter(BaseModel):
     """@private"""
     key: str
@@ -316,6 +420,10 @@ class ListCompletionsFilterInput(BaseModel):
     prompt_hash: Optional[str] = Field(alias='promptHash', default=None)
     completion_id: Optional[Any] = Field(alias='completionId', default=None)
     source: Optional[List[CompletionSource]] = None
+    completion: Optional[str] = None
+    prompt: Optional[str] = None
+    advanced_filter: Optional['CompletionFilterExpression'] = Field(alias='advancedFilter', default=None)
+    'Advanced filter supporting AND/OR/NOT logic\nWhen set, this takes precedence over the simple filter fields above\n(except use_case which is always required)'
 
 class ListJobsFilterInput(BaseModel):
     """@private"""
@@ -369,6 +477,8 @@ class ModelFilter(BaseModel):
     capabilities: Optional['CapabilityFilter'] = None
     view_all: Optional[bool] = Field(alias='viewAll', default=None)
     online: Optional[List[ModelOnline]] = None
+    published: Optional[bool] = None
+    size: Optional['IntegerNumericCondition'] = None
 
 class ModelPlacementInput(BaseModel):
     """@private"""
@@ -378,23 +488,22 @@ class ModelPlacementInput(BaseModel):
 class ModelProviderDataInput(BaseModel):
     """@private"""
     open_ai: Optional['OpenAIProviderDataInput'] = Field(alias='openAI', default=None)
+    legacy_open_ai: Optional['OpenAIProviderDataInput'] = Field(alias='legacyOpenAI', default=None)
     google: Optional['GoogleProviderDataInput'] = None
     anthropic: Optional['AnthropicProviderDataInput'] = None
 
-class ModelServiceDisconnect(BaseModel):
-    """@private"""
-    use_case: str = Field(alias='useCase')
-    model_service: str = Field(alias='modelService')
-
 class ModelServiceFilter(BaseModel):
     """@private"""
     model: Optional[str] = None
     capabilities: Optional['CapabilityFilter'] = None
+    active_only: bool = Field(alias='activeOnly', default=True)
+    'If true (default), only return model services whose model has a binding.\nIf false, return all model services regardless of binding status.'
 
 class OpenAIProviderDataInput(BaseModel):
     """@private"""
     api_key: str = Field(alias='apiKey')
-    external_model_id: OpenAIModel = Field(alias='externalModelId')
+    external_model_id: str = Field(alias='externalModelId')
+    endpoint: Optional[str] = None
 
 class OrderPair(BaseModel):
     """@private"""
@@ -424,6 +533,11 @@ class RemoteEnvCreate(BaseModel):
     name: Optional[str] = None
     description: Optional[str] = None
 
+class RemoveModelFromUseCaseInput(BaseModel):
+    """@private"""
+    use_case: str = Field(alias='useCase')
+    model: str
+
 class ResizePartitionInput(BaseModel):
     """@private"""
     harmony_group: str = Field(alias='harmonyGroup')
@@ -444,6 +558,16 @@ class SearchInput(BaseModel):
     """@private"""
     query: str
 
+class StringCondition(BaseModel):
+    """@private
+String matching condition for filter expressions"""
+    eq: Optional[str] = None
+    'Exact match'
+    in_: Optional[List[str]] = Field(alias='in', default=None)
+    'Match any of the provided values (OR)'
+    neq: Optional[str] = None
+    'Does not equal'
+
 class SystemPromptTemplateCreate(BaseModel):
     """@private"""
     name: str
@@ -472,6 +596,13 @@ class TeamMemberSet(BaseModel):
     team: str
     role: str
 
+class TextCondition(BaseModel):
+    """@private"""
+    eq: Optional[str] = None
+    'Exact match'
+    contains: Optional[str] = None
+    'Text contains this substring (case insensitive)'
+
 class TimeRange(BaseModel):
     """@private"""
     from_: int | str = Field(alias='from')
@@ -502,12 +633,17 @@ class UpdateCompletion(BaseModel):
     metadata: Optional[Any] = None
     'set metadata associated with this prompt for use with external reward servers'
 
+class UpdateModelInput(BaseModel):
+    """@private"""
+    model: str
+    published: Optional[bool] = None
+    stable: Optional[bool] = None
+
 class UpdateModelService(BaseModel):
     """@private"""
     use_case: str = Field(alias='useCase')
     model_service: str = Field(alias='modelService')
     is_default: Optional[bool] = Field(alias='isDefault', default=None)
-    attached: Optional[bool] = None
     desired_online: Optional[bool] = Field(alias='desiredOnline', default=None)
     name: Optional[str] = None
     system_prompt_template: Optional[Any] = Field(alias='systemPromptTemplate', default=None)
@@ -597,12 +733,14 @@ class WidgetInput(BaseModel):
     aggregation: MetricAggregation
     unit: 'UnitConfigInput'
 AddExternalModelInput.model_rebuild()
-AttachModel.model_rebuild()
+CompletionFilterExpression.model_rebuild()
 CompletionsByFilters.model_rebuild()
 CreateRecipeInput.model_rebuild()
 CustomRecipeFilterInput.model_rebuild()
 DatasetCompletionQuery.model_rebuild()
 DatasetCreateFromFilters.model_rebuild()
+DeployModelInput.model_rebuild()
+FeedbackCondition.model_rebuild()
 FeedbackFilterInput.model_rebuild()
 FromGroupsQuery.model_rebuild()
 GlobalUsageFilterInput.model_rebuild()
@@ -614,6 +752,7 @@ JudgeConfigInput.model_rebuild()
 JudgeCreate.model_rebuild()
 JudgeExampleInput.model_rebuild()
 JudgeUpdate.model_rebuild()
+LabelCondition.model_rebuild()
 ListCompletionsFilterInput.model_rebuild()
 ListJobsFilterInput.model_rebuild()
 MetricGetOrCreate.model_rebuild()

adaptive_sdk/graphql_client/remove_model_from_use_case.py

@@ -0,0 +1,6 @@
+from pydantic import Field
+from .base_model import BaseModel
+
+class RemoveModelFromUseCase(BaseModel):
+    """@public"""
+    remove_model_from_use_case: bool = Field(alias='removeModelFromUseCase')

adaptive_sdk/input_types/typed_dicts.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
-from typing import TypedDict, List, Any, TypeAlias, Literal
-from typing_extensions import Required, NotRequired
+
+from typing import Any, List, Literal, TypeAlias, TypedDict
+
+from typing_extensions import NotRequired, Required
 
 
 class ChatMessage(TypedDict, total=True):
@@ -24,28 +26,26 @@ class CompletionComparisonFilterInput(TypedDict, total=True):
     metric: Required[str]
 
 
+class NumericCondition(TypedDict, total=False):
+    eq: NotRequired[float]
+    neq: NotRequired[float]
+    gt: NotRequired[float]
+    gte: NotRequired[float]
+    lt: NotRequired[float]
+    lte: NotRequired[float]
+
+
 class CompletionFeedbackFilterInput(TypedDict, total=False):
     """
     Filter for completion metric feedbacks.
 
     Args:
         metric: Feedback key logged against.
-        gt: >
-        gte: >=
-        eq: ==
-        neq: !=
-        lt: <
-        lte: <=
         user: Feedbacks logged by `user` id.
     """
 
     metric: Required[str]
-    gt: NotRequired[float]
-    gte: NotRequired[float]
-    eq: NotRequired[float]
-    neq: NotRequired[float]
-    lt: NotRequired[float]
-    lte: NotRequired[float]
+    value: NotRequired[NumericCondition]
     reasons: NotRequired[List[str]]
     user: NotRequired[Any]
 
@@ -135,7 +135,6 @@ class ListCompletionsFilterInput(TypedDict, total=False):
     labels: NotRequired[List["CompletionLabelFilter"]]
     prompt_hash: NotRequired[str]
     completion_id: NotRequired[Any]
-    tags: NotRequired[List[str]]
     source: NotRequired[List[CompletionSource]]
 
 

adaptive_sdk/resources/__init__.py

@@ -1,4 +1,5 @@
 from .abtests import ABTests, AsyncABTests
+from .artifacts import Artifacts, AsyncArtifacts
 from .chat import Chat, AsyncChat
 from .compute_pools import ComputePools, AsyncComputePools  # type: ignore[attr-defined]
 from .recipes import Recipes, AsyncRecipes
@@ -17,6 +18,7 @@ from .users import Users, AsyncUsers
 
 __all__ = [
     "ABTests",
+    "Artifacts",
     "Chat",
     "ComputePools",
     "Recipes",
@@ -33,6 +35,7 @@ __all__ = [
     "UseCase",
     "Users",
     "AsyncABTests",
+    "AsyncArtifacts",
     "AsyncChat",
     "AsyncComputePools",
     "AsyncRecipes",

adaptive_sdk/resources/artifacts.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+from .base_resource import SyncAPIResource, AsyncAPIResource
+
+if TYPE_CHECKING:
+    from adaptive_sdk.client import Adaptive, AsyncAdaptive
+
+
+class Artifacts(SyncAPIResource):  # type: ignore[misc]
+    """
+    Resource to interact with job artifacts.
+    """
+
+    def __init__(self, client: Adaptive) -> None:
+        SyncAPIResource.__init__(self, client)
+
+    def download(self, artifact_id: str, destination_path: str) -> None:
+        """
+        Download an artifact file to a local path.
+
+        Args:
+            artifact_id: The UUID of the artifact to download.
+            destination_path: Local file path where the artifact will be saved.
+
+        Raises:
+            HTTPError: If the download fails or the artifact is not found.
+        """
+        download_url = f"/artifacts/{artifact_id}/download"
+        response = self._rest_client.get(download_url)
+        response.raise_for_status()
+
+        with open(destination_path, "wb") as f:
+            f.write(response.content)
+
+
+class AsyncArtifacts(AsyncAPIResource):  # type: ignore[misc]
+    """
+    Async resource to interact with job artifacts.
+    """
+
+    def __init__(self, client: AsyncAdaptive) -> None:
+        AsyncAPIResource.__init__(self, client)
+
+    async def download(self, artifact_id: str, destination_path: str) -> None:
+        """
+        Download an artifact file to a local path.
+
+        Args:
+            artifact_id: The UUID of the artifact to download.
+            destination_path: Local file path where the artifact will be saved.
+
+        Raises:
+            HTTPError: If the download fails or the artifact is not found.
+        """
+        download_url = f"/artifacts/{artifact_id}/download"
+        response = await self._rest_client.get(download_url)
+        response.raise_for_status()
+
+        with open(destination_path, "wb") as f:
+            f.write(response.content)

adaptive_sdk/resources/chat.py

@@ -48,6 +48,7 @@ class Chat(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
         ab_campaign: str | None = None,
         n: int | None = None,
         labels: Dict[str, str] | None = None,
+        store: bool | None = None,
     ) -> rest_types.ChatResponse: ...
 
     @overload
@@ -67,6 +68,7 @@ class Chat(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
         ab_campaign: str | None = None,
         n: int | None = None,
         labels: Dict[str, str] | None = None,
+        store: bool | None = None,
     ) -> Generator[rest_types.ChatResponseChunk, None, None]: ...
 
     def create(
@@ -85,6 +87,7 @@ class Chat(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
         ab_campaign: str | None = None,
         n: int | None = None,
         labels: Dict[str, str] | None = None,
+        store: bool | None = None,
     ) -> rest_types.ChatResponse | Generator[rest_types.ChatResponseChunk, None, None]:
         """
         Create a chat completion.
@@ -135,6 +138,7 @@ class Chat(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
             ab_campaign=ab_campaign,
             n=n,
             labels=labels,
+            store=store,
         )
         if input.stream:
             return self._stream(input)
@@ -142,9 +146,7 @@ class Chat(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
         rest_error_handler(r)
         return rest_types.ChatResponse.model_validate(r.json())
 
-    def _stream(
-        self, input: rest_types.ChatInput
-    ) -> Generator[rest_types.ChatResponseChunk, None, None]:
+    def _stream(self, input: rest_types.ChatInput) -> Generator[rest_types.ChatResponseChunk, None, None]:
         import httpx_sse
 
         with httpx_sse.connect_sse(
@@ -183,6 +185,7 @@ class AsyncChat(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
         ab_campaign: str | None = None,
         n: int | None = None,
         labels: Dict[str, str] | None = None,
+        store: bool | None = None,
     ) -> rest_types.ChatResponse: ...
 
     @overload  # type: ignore[no-redef, misc]
@@ -202,6 +205,7 @@ class AsyncChat(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
         ab_campaign: str | None = None,
         n: int | None = None,
         labels: Dict[str, str] | None = None,
+        store: bool | None = None,
     ) -> AsyncGenerator[rest_types.ChatResponseChunk, None]: ...
 
     async def create(  # type: ignore
@@ -220,6 +224,7 @@ class AsyncChat(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
         ab_campaign: str | None = None,
         n: int | None = None,
         labels: Dict[str, str] | None = None,
+        store: bool | None = None,
     ) -> rest_types.ChatResponse | AsyncGenerator[rest_types.ChatResponseChunk, None]:
         """
         Create a chat completion.
@@ -270,18 +275,15 @@ class AsyncChat(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
             ab_campaign=ab_campaign,
             n=n,
             labels=labels,
+            store=store,
         )
         if input.stream:
             return self._stream(input)
-        r = await self._rest_client.post(
-            ROUTE, json=input.model_dump(exclude_none=True)
-        )
+        r = await self._rest_client.post(ROUTE, json=input.model_dump(exclude_none=True))
         rest_error_handler(r)
         return rest_types.ChatResponse.model_validate(r.json())
 
-    async def _stream(
-        self, input: rest_types.ChatInput
-    ) -> AsyncGenerator[rest_types.ChatResponseChunk, None]:
+    async def _stream(self, input: rest_types.ChatInput) -> AsyncGenerator[rest_types.ChatResponseChunk, None]:
         import httpx_sse
 
         async with httpx_sse.aconnect_sse(