wandb 0.19.8__py3-none-any.whl → 0.19.10__py3-none-any.whl

wandb/apis/public/artifacts.py

@@ -3,49 +3,82 @@
 import json
 import re
 from copy import copy
-from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Sequence, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Iterable,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Sequence,
+    Type,
+    Union,
+)
 
+from typing_extensions import override
 from wandb_gql import Client, gql
 
 import wandb
 from wandb.apis import public
 from wandb.apis.normalize import normalize_exceptions
-from wandb.apis.paginator import Paginator
+from wandb.apis.paginator import Paginator, SizedPaginator
 from wandb.errors.term import termlog
-from wandb.sdk.artifacts._graphql_fragments import (
-    ARTIFACT_FILES_FRAGMENT,
-    ARTIFACTS_TYPES_FRAGMENT,
+from wandb.proto.wandb_deprecated import Deprecated
+from wandb.proto.wandb_internal_pb2 import ServerFeature
+from wandb.sdk.artifacts._generated import (
+    ARTIFACT_COLLECTION_MEMBERSHIP_FILES_GQL,
+    ARTIFACT_VERSION_FILES_GQL,
+    CREATE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL,
+    DELETE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL,
+    DELETE_ARTIFACT_PORTFOLIO_GQL,
+    DELETE_ARTIFACT_SEQUENCE_GQL,
+    MOVE_ARTIFACT_COLLECTION_GQL,
+    PROJECT_ARTIFACT_COLLECTION_GQL,
+    PROJECT_ARTIFACT_COLLECTIONS_GQL,
+    PROJECT_ARTIFACT_TYPE_GQL,
+    PROJECT_ARTIFACT_TYPES_GQL,
+    PROJECT_ARTIFACTS_GQL,
+    RUN_INPUT_ARTIFACTS_GQL,
+    RUN_OUTPUT_ARTIFACTS_GQL,
+    UPDATE_ARTIFACT_PORTFOLIO_GQL,
+    UPDATE_ARTIFACT_SEQUENCE_GQL,
+    ArtifactCollectionMembershipFiles,
+    ArtifactCollectionsFragment,
+    ArtifactsFragment,
+    ArtifactTypeFragment,
+    ArtifactTypesFragment,
+    ArtifactVersionFiles,
+    FilesFragment,
+    ProjectArtifactCollection,
+    ProjectArtifactCollections,
+    ProjectArtifacts,
+    ProjectArtifactType,
+    ProjectArtifactTypes,
+    RunInputArtifactsProjectRunInputArtifacts,
+    RunOutputArtifactsProjectRunOutputArtifacts,
 )
+from wandb.sdk.artifacts._graphql_fragments import omit_artifact_fields
+from wandb.sdk.internal.internal_api import Api as InternalApi
 from wandb.sdk.lib import deprecate
 
+from .utils import gql_compat
+
 if TYPE_CHECKING:
     from wandb.apis.public import RetryingClient, Run
 
 
-class ArtifactTypes(Paginator):
-    QUERY = gql(
-        """
-        query ProjectArtifacts(
-            $entityName: String!,
-            $projectName: String!,
-            $cursor: String,
-        ) {{
-            project(name: $projectName, entityName: $entityName) {{
-                artifactTypes(after: $cursor) {{
-                    ...ArtifactTypesFragment
-                }}
-            }}
-        }}
-        {}
-    """.format(ARTIFACTS_TYPES_FRAGMENT)
-    )
+class ArtifactTypes(Paginator["ArtifactType"]):
+    QUERY = gql(PROJECT_ARTIFACT_TYPES_GQL)
+
+    last_response: Optional[ArtifactTypesFragment]
 
     def __init__(
         self,
         client: Client,
         entity: str,
         project: str,
-        per_page: Optional[int] = 50,
+        per_page: int = 50,
     ):
         self.entity = entity
         self.project = project
@@ -54,41 +87,54 @@ class ArtifactTypes(Paginator):
             "entityName": entity,
             "projectName": project,
         }
-
         super().__init__(client, variable_values, per_page)
 
+    @override
+    def _update_response(self) -> None:
+        """Fetch and validate the response data for the current page."""
+        data = self.client.execute(self.QUERY, variable_values=self.variables)
+        result = ProjectArtifactTypes.model_validate(data)
+
+        # Extract the inner `*Connection` result for faster/easier access.
+        if not ((proj := result.project) and (conn := proj.artifact_types)):
+            raise ValueError(f"Unable to parse {type(self).__name__!r} response data")
+
+        self.last_response = ArtifactTypesFragment.model_validate(conn)
+
     @property
-    def length(self):
+    def length(self) -> None:
         # TODO
         return None
 
     @property
-    def more(self):
-        if self.last_response:
-            return self.last_response["project"]["artifactTypes"]["pageInfo"][
-                "hasNextPage"
-            ]
-        else:
+    def more(self) -> bool:
+        if self.last_response is None:
             return True
+        return self.last_response.page_info.has_next_page
 
     @property
-    def cursor(self):
-        if self.last_response:
-            return self.last_response["project"]["artifactTypes"]["edges"][-1]["cursor"]
-        else:
+    def cursor(self) -> Optional[str]:
+        if self.last_response is None:
             return None
+        return self.last_response.edges[-1].cursor
 
-    def update_variables(self):
+    def update_variables(self) -> None:
         self.variables.update({"cursor": self.cursor})
 
-    def convert_objects(self):
-        if self.last_response["project"] is None:
+    def convert_objects(self) -> List["ArtifactType"]:
+        if self.last_response is None:
             return []
+
         return [
             ArtifactType(
-                self.client, self.entity, self.project, r["node"]["name"], r["node"]
+                client=self.client,
+                entity=self.entity,
+                project=self.project,
+                type_name=node.name,
+                attrs=node.model_dump(exclude_unset=True),
             )
-            for r in self.last_response["project"]["artifactTypes"]["edges"]
+            for edge in self.last_response.edges
+            if edge.node and (node := ArtifactTypeFragment.model_validate(edge.node))
         ]
 
 
@@ -110,39 +156,19 @@ class ArtifactType:
             self.load()
 
     def load(self):
-        query = gql(
-            """
-        query ProjectArtifactType(
-            $entityName: String!,
-            $projectName: String!,
-            $artifactTypeName: String!
-        ) {
-            project(name: $projectName, entityName: $entityName) {
-                artifactType(name: $artifactTypeName) {
-                    id
-                    name
-                    description
-                    createdAt
-                }
-            }
-        }
-        """
-        )
-        response: Optional[Mapping[str, Any]] = self.client.execute(
-            query,
+        data: Optional[Mapping[str, Any]] = self.client.execute(
+            gql(PROJECT_ARTIFACT_TYPE_GQL),
             variable_values={
                 "entityName": self.entity,
                 "projectName": self.project,
                 "artifactTypeName": self.type,
             },
         )
-        if (
-            response is None
-            or response.get("project") is None
-            or response["project"].get("artifactType") is None
-        ):
-            raise ValueError("Could not find artifact type {}".format(self.type))
-        self._attrs = response["project"]["artifactType"]
+        result = ProjectArtifactType.model_validate(data)
+        if not ((proj := result.project) and (artifact_type := proj.artifact_type)):
+            raise ValueError(f"Could not find artifact type {self.type}")
+
+        self._attrs = artifact_type.model_dump(exclude_unset=True)
         return self._attrs
 
     @property
@@ -167,14 +193,16 @@ class ArtifactType:
         return f"<ArtifactType {self.type}>"
 
 
-class ArtifactCollections(Paginator):
+class ArtifactCollections(SizedPaginator["ArtifactCollection"]):
+    last_response: Optional[ArtifactCollectionsFragment]
+
     def __init__(
         self,
         client: Client,
         entity: str,
         project: str,
         type_name: str,
-        per_page: Optional[int] = 50,
+        per_page: int = 50,
     ):
         self.entity = entity
         self.project = project
@@ -186,86 +214,65 @@ class ArtifactCollections(Paginator):
             "artifactTypeName": type_name,
         }
 
-        self.QUERY = gql(
-            """
-            query ProjectArtifactCollections(
-                $entityName: String!,
-                $projectName: String!,
-                $artifactTypeName: String!
-                $cursor: String,
-            ) {{
-                project(name: $projectName, entityName: $entityName) {{
-                    artifactType(name: $artifactTypeName) {{
-                        artifactCollections: {}(after: $cursor) {{
-                            pageInfo {{
-                                endCursor
-                                hasNextPage
-                            }}
-                            totalCount
-                            edges {{
-                                node {{
-                                    id
-                                    name
-                                    description
-                                    createdAt
-                                }}
-                                cursor
-                            }}
-                        }}
-                    }}
-                }}
-            }}
-        """.format(
-                artifact_collection_plural_edge_name(
-                    server_supports_artifact_collections_gql_edges(client)
-                )
-            )
+        if server_supports_artifact_collections_gql_edges(client):
+            rename_fields = None
+        else:
+            rename_fields = {"artifactCollections": "artifactSequences"}
+
+        self.QUERY = gql_compat(
+            PROJECT_ARTIFACT_COLLECTIONS_GQL, rename_fields=rename_fields
         )
 
         super().__init__(client, variable_values, per_page)
 
+    @override
+    def _update_response(self) -> None:
+        """Fetch and validate the response data for the current page."""
+        data = self.client.execute(self.QUERY, variable_values=self.variables)
+        result = ProjectArtifactCollections.model_validate(data)
+
+        # Extract the inner `*Connection` result for faster/easier access.
+        if not (
+            (proj := result.project)
+            and (type_ := proj.artifact_type)
+            and (conn := type_.artifact_collections)
+        ):
+            raise ValueError(f"Unable to parse {type(self).__name__!r} response data")
+
+        self.last_response = ArtifactCollectionsFragment.model_validate(conn)
+
     @property
     def length(self):
-        if self.last_response:
-            return self.last_response["project"]["artifactType"]["artifactCollections"][
-                "totalCount"
-            ]
-        else:
+        if self.last_response is None:
             return None
+        return self.last_response.total_count
 
     @property
     def more(self):
-        if self.last_response:
-            return self.last_response["project"]["artifactType"]["artifactCollections"][
-                "pageInfo"
-            ]["hasNextPage"]
-        else:
+        if self.last_response is None:
             return True
+        return self.last_response.page_info.has_next_page
 
     @property
     def cursor(self):
-        if self.last_response:
-            return self.last_response["project"]["artifactType"]["artifactCollections"][
-                "edges"
-            ][-1]["cursor"]
-        else:
+        if self.last_response is None:
             return None
+        return self.last_response.edges[-1].cursor
 
-    def update_variables(self):
+    def update_variables(self) -> None:
         self.variables.update({"cursor": self.cursor})
 
-    def convert_objects(self):
+    def convert_objects(self) -> List["ArtifactCollection"]:
         return [
             ArtifactCollection(
-                self.client,
-                self.entity,
-                self.project,
-                r["node"]["name"],
-                self.type_name,
+                client=self.client,
+                entity=self.entity,
+                project=self.project,
+                name=node.name,
+                type=self.type_name,
             )
-            for r in self.last_response["project"]["artifactType"][
-                "artifactCollections"
-            ]["edges"]
+            for edge in self.last_response.edges
+            if (node := edge.node)
         ]
 
 
@@ -298,83 +305,38 @@ class ArtifactCollection:
         self.organization = organization
 
     @property
-    def id(self):
+    def id(self) -> str:
         return self._attrs["id"]
 
     @normalize_exceptions
-    def artifacts(self, per_page=50):
+    def artifacts(self, per_page: int = 50) -> "Artifacts":
         """Artifacts."""
         return Artifacts(
-            self.client,
-            self.entity,
-            self.project,
-            self._saved_name,
-            self._saved_type,
+            client=self.client,
+            entity=self.entity,
+            project=self.project,
+            collection_name=self._saved_name,
+            type=self._saved_type,
             per_page=per_page,
         )
 
     @property
-    def aliases(self):
+    def aliases(self) -> List[str]:
         """Artifact Collection Aliases."""
         return self._aliases
 
     @property
-    def created_at(self):
+    def created_at(self) -> str:
         return self._created_at
 
     def load(self):
-        query = gql(
-            """
-        query ArtifactCollection(
-            $entityName: String!,
-            $projectName: String!,
-            $artifactTypeName: String!,
-            $artifactCollectionName: String!,
-            $cursor: String,
-            $perPage: Int = 1000
-        ) {{
-            project(name: $projectName, entityName: $entityName) {{
-                artifactType(name: $artifactTypeName) {{
-                    artifactCollection: {}(name: $artifactCollectionName) {{
-                        id
-                        name
-                        description
-                        createdAt
-                        tags {{
-                            edges {{
-                                node {{
-                                    id
-                                    name
-                                }}
-                            }}
-                        }}
-                        aliases(after: $cursor, first: $perPage){{
-                            edges {{
-                                node {{
-                                    alias
-                                }}
-                                cursor
-                            }}
-                            pageInfo {{
-                                endCursor
-                                hasNextPage
-                            }}
-                        }}
-                    }}
-                    artifactSequence(name: $artifactCollectionName) {{
-                        __typename
-                    }}
-                }}
-            }}
-        }}
-        """.format(
-                artifact_collection_edge_name(
-                    server_supports_artifact_collections_gql_edges(self.client)
-                )
-            )
-        )
+        if server_supports_artifact_collections_gql_edges(self.client):
+            rename_fields = None
+        else:
+            rename_fields = {"artifactCollection": "artifactSequence"}
+
         response = self.client.execute(
-            query,
+            gql_compat(PROJECT_ARTIFACT_COLLECTION_GQL, rename_fields=rename_fields),
             variable_values={
                 "entityName": self.entity,
                 "projectName": self.project,
@@ -382,26 +344,30 @@ class ArtifactCollection:
                 "artifactCollectionName": self._saved_name,
             },
         )
-        if (
-            response is None
-            or response.get("project") is None
-            or response["project"].get("artifactType") is None
-            or response["project"]["artifactType"].get("artifactCollection") is None
+
+        result = ProjectArtifactCollection.model_validate(response)
+
+        if not (
+            result.project
+            and (proj := result.project)
+            and (type_ := proj.artifact_type)
+            and (collection := type_.artifact_collection)
         ):
-            raise ValueError("Could not find artifact type {}".format(self._saved_type))
-        sequence = response["project"]["artifactType"]["artifactSequence"]
+            raise ValueError(f"Could not find artifact type {self._saved_type}")
+
+        sequence = type_.artifact_sequence
         self._is_sequence = (
-            sequence is not None and sequence["__typename"] == "ArtifactSequence"
-        )
+            sequence is not None
+        ) and sequence.typename__ == "ArtifactSequence"
 
         if self._attrs is None:
-            self._attrs = response["project"]["artifactType"]["artifactCollection"]
+            self._attrs = collection.model_dump(exclude_unset=True)
         return self._attrs
 
     def change_type(self, new_type: str) -> None:
         """Deprecated, change type directly with `save` instead."""
         deprecate.deprecate(
-            field_name=deprecate.Deprecated.artifact_collection__change_type,
+            field_name=Deprecated.artifact_collection__change_type,
             warning_message="ArtifactCollection.change_type(type) is deprecated, use ArtifactCollection.save() instead.",
         )
 
@@ -410,32 +376,13 @@ class ArtifactCollection:
         termlog(
             f"Changing artifact collection type of {self._saved_type} to {new_type}"
         )
-        template = """
-            mutation MoveArtifactCollection(
-                $artifactSequenceID: ID!
-                $destinationArtifactTypeName: String!
-            ) {
-                moveArtifactSequence(
-                input: {
-                    artifactSequenceID: $artifactSequenceID
-                    destinationArtifactTypeName: $destinationArtifactTypeName
-                }
-                ) {
-                artifactCollection {
-                    id
-                    name
-                    description
-                    __typename
-                }
-                }
-            }
-            """
-        variable_values = {
-            "artifactSequenceID": self.id,
-            "destinationArtifactTypeName": new_type,
-        }
-        mutation = gql(template)
-        self.client.execute(mutation, variable_values=variable_values)
+        self.client.execute(
+            gql(MOVE_ARTIFACT_COLLECTION_GQL),
+            variable_values={
+                "artifactSequenceID": self.id,
+                "destinationArtifactTypeName": new_type,
+            },
+        )
         self._saved_type = new_type
         self._type = new_type
 
@@ -444,40 +391,19 @@ class ArtifactCollection:
         return self._is_sequence
 
     @normalize_exceptions
-    def delete(self):
+    def delete(self) -> None:
         """Delete the entire artifact collection."""
-        if self.is_sequence():
-            mutation = gql(
-                """
-                mutation deleteArtifactSequence($id: ID!) {
-                    deleteArtifactSequence(input: {
-                        artifactSequenceID: $id
-                    }) {
-                        artifactCollection {
-                            state
-                        }
-                    }
-                }
-                """
-            )
-        else:
-            mutation = gql(
-                """
-                mutation deleteArtifactPortfolio($id: ID!) {
-                    deleteArtifactPortfolio(input: {
-                        artifactPortfolioID: $id
-                    }) {
-                        artifactCollection {
-                            state
-                        }
-                    }
-                }
-                """
-            )
-        self.client.execute(mutation, variable_values={"id": self.id})
+        self.client.execute(
+            gql(
+                DELETE_ARTIFACT_SEQUENCE_GQL
+                if self.is_sequence()
+                else DELETE_ARTIFACT_PORTFOLIO_GQL
+            ),
+            variable_values={"id": self.id},
+        )
 
     @property
-    def description(self):
+    def description(self) -> str:
         """A description of the artifact collection."""
         return self._description
 
@@ -486,7 +412,7 @@ class ArtifactCollection:
         self._description = description
 
     @property
-    def tags(self):
+    def tags(self) -> List[str]:
         """The tags associated with the artifact collection."""
         return self._tags
 
@@ -499,7 +425,7 @@ class ArtifactCollection:
         self._tags = tags
 
     @property
-    def name(self):
+    def name(self) -> str:
         """The name of the artifact collection."""
         return self._name
 
@@ -520,202 +446,80 @@ class ArtifactCollection:
             )
         self._type = type
 
-    def _update_collection(self):
-        mutation = gql("""
-            mutation UpdateArtifactCollection(
-                $artifactSequenceID: ID!
-                $name: String
-                $description: String
-            ) {
-                updateArtifactSequence(
-                input: {
-                    artifactSequenceID: $artifactSequenceID
-                    name: $name
-                    description: $description
-                }
-                ) {
-                artifactCollection {
-                    id
-                    name
-                    description
-                }
-                }
-            }
-        """)
-
-        variable_values = {
-            "artifactSequenceID": self.id,
-            "name": self._name,
-            "description": self.description,
-        }
-        self.client.execute(mutation, variable_values=variable_values)
+    def _update_collection(self) -> None:
+        self.client.execute(
+            gql(
+                UPDATE_ARTIFACT_SEQUENCE_GQL
+                if self.is_sequence()
+                else UPDATE_ARTIFACT_PORTFOLIO_GQL
+            ),
+            variable_values={
+                "id": self.id,
+                "name": self.name,
+                "description": self.description,
+            },
+        )
         self._saved_name = self._name
 
-    def _update_collection_type(self):
-        type_mutation = gql("""
-            mutation MoveArtifactCollection(
-                $artifactSequenceID: ID!
-                $destinationArtifactTypeName: String!
-            ) {
-                moveArtifactSequence(
-                input: {
-                    artifactSequenceID: $artifactSequenceID
-                    destinationArtifactTypeName: $destinationArtifactTypeName
-                }
-                ) {
-                artifactCollection {
-                    id
-                    name
-                    description
-                    __typename
-                }
-                }
-            }
-            """)
-
-        variable_values = {
-            "artifactSequenceID": self.id,
-            "destinationArtifactTypeName": self._type,
-        }
-        self.client.execute(type_mutation, variable_values=variable_values)
+    def _update_collection_type(self) -> None:
+        self.client.execute(
+            gql(MOVE_ARTIFACT_COLLECTION_GQL),
+            variable_values={
+                "artifactSequenceID": self.id,
+                "destinationArtifactTypeName": self.type,
+            },
+        )
         self._saved_type = self._type
 
-    def _update_portfolio(self):
-        mutation = gql("""
-            mutation UpdateArtifactPortfolio(
-                $artifactPortfolioID: ID!
-                $name: String
-                $description: String
-            ) {
-                updateArtifactPortfolio(
-                input: {
-                    artifactPortfolioID: $artifactPortfolioID
-                    name: $name
-                    description: $description
-                }
-                ) {
-                artifactCollection {
-                    id
-                    name
-                    description
-                }
-                }
-            }
-        """)
-        variable_values = {
-            "artifactPortfolioID": self.id,
-            "name": self._name,
-            "description": self.description,
-        }
-        self.client.execute(mutation, variable_values=variable_values)
-        self._saved_name = self._name
-
-    def _add_tags(self, tags_to_add):
-        add_mutation = gql(
-            """
-            mutation CreateArtifactCollectionTagAssignments(
-                $entityName: String!
-                $projectName: String!
-                $artifactCollectionName: String!
-                $tags: [TagInput!]!
-            ) {
-                createArtifactCollectionTagAssignments(
-                input: {
-                    entityName: $entityName
-                    projectName: $projectName
-                    artifactCollectionName: $artifactCollectionName
-                    tags: $tags
-                }
-                ) {
-                tags {
-                    id
-                    name
-                    tagCategoryName
-                }
-                }
-            }
-            """
-        )
+    def _add_tags(self, tags_to_add: Iterable[str]) -> None:
         self.client.execute(
-            add_mutation,
+            gql(CREATE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL),
             variable_values={
                 "entityName": self.entity,
                 "projectName": self.project,
                 "artifactCollectionName": self._saved_name,
-                "tags": [
-                    {
-                        "tagName": tag,
-                    }
-                    for tag in tags_to_add
-                ],
+                "tags": [{"tagName": tag} for tag in tags_to_add],
             },
         )
 
-    def _delete_tags(self, tags_to_delete):
-        delete_mutation = gql(
-            """
-            mutation DeleteArtifactCollectionTagAssignments(
-                $entityName: String!
-                $projectName: String!
-                $artifactCollectionName: String!
-                $tags: [TagInput!]!
-            ) {
-                deleteArtifactCollectionTagAssignments(
-                input: {
-                    entityName: $entityName
-                    projectName: $projectName
-                    artifactCollectionName: $artifactCollectionName
-                    tags: $tags
-                }
-                ) {
-                success
-                }
-            }
-            """
-        )
+    def _delete_tags(self, tags_to_delete: Iterable[str]) -> None:
         self.client.execute(
-            delete_mutation,
+            gql(DELETE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL),
             variable_values={
                 "entityName": self.entity,
                 "projectName": self.project,
                 "artifactCollectionName": self._saved_name,
-                "tags": [
-                    {
-                        "tagName": tag,
-                    }
-                    for tag in tags_to_delete
-                ],
+                "tags": [{"tagName": tag} for tag in tags_to_delete],
             },
         )
 
     def save(self) -> None:
         """Persist any changes made to the artifact collection."""
-        if self.is_sequence():
-            self._update_collection()
+        self._update_collection()
 
-            if self._saved_type != self._type:
-                self._update_collection_type()
-        else:
-            self._update_portfolio()
+        if self.is_sequence() and (self._saved_type != self._type):
+            self._update_collection_type()
 
-        tags_to_add = set(self._tags) - set(self._saved_tags)
-        tags_to_delete = set(self._saved_tags) - set(self._tags)
-        if len(tags_to_add) > 0:
+        current_tags = set(self._tags)
+        saved_tags = set(self._saved_tags)
+        if tags_to_add := (current_tags - saved_tags):
             self._add_tags(tags_to_add)
-        if len(tags_to_delete) > 0:
+        if tags_to_delete := (saved_tags - current_tags):
             self._delete_tags(tags_to_delete)
         self._saved_tags = copy(self._tags)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"<ArtifactCollection {self._name} ({self._type})>"
 
 
-class Artifacts(Paginator):
+class Artifacts(SizedPaginator["wandb.Artifact"]):
     """An iterable collection of artifact versions associated with a project and optional filter.
 
     This is generally used indirectly via the `Api`.artifact_versions method.
     """
 
+    last_response: Optional[ArtifactsFragment]
+
     def __init__(
         self,
         client: Client,
@@ -728,8 +532,6 @@ class Artifacts(Paginator):
         per_page: int = 50,
         tags: Optional[Union[str, List[str]]] = None,
     ):
-        from wandb.sdk.artifacts.artifact import _gql_artifact_fragment
-
         self.entity = entity
         self.collection_name = collection_name
         self.type = type
@@ -745,154 +547,107 @@ class Artifacts(Paginator):
             "collection": self.collection_name,
             "filters": json.dumps(self.filters),
         }
-        self.QUERY = gql(
-            """
-            query Artifacts($project: String!, $entity: String!, $type: String!, $collection: String!, $cursor: String, $perPage: Int = 50, $order: String, $filters: JSONString) {{
-                project(name: $project, entityName: $entity) {{
-                    artifactType(name: $type) {{
-                        artifactCollection: {}(name: $collection) {{
-                            name
-                            artifacts(filters: $filters, after: $cursor, first: $perPage, order: $order) {{
-                                totalCount
-                                edges {{
-                                    node {{
-                                        ...ArtifactFragment
-                                    }}
-                                    version
-                                    cursor
-                                }}
-                                pageInfo {{
-                                    endCursor
-                                    hasNextPage
-                                }}
-                            }}
-                        }}
-                    }}
-                }}
-            }}
-            {}
-            """.format(
-                artifact_collection_edge_name(
-                    server_supports_artifact_collections_gql_edges(client)
-                ),
-                _gql_artifact_fragment(),
-            )
+
+        if server_supports_artifact_collections_gql_edges(client):
+            rename_fields = None
+        else:
+            rename_fields = {"artifactCollection": "artifactSequence"}
+
+        self.QUERY = gql_compat(
+            PROJECT_ARTIFACTS_GQL,
+            omit_fields=omit_artifact_fields(api=InternalApi()),
+            rename_fields=rename_fields,
         )
+
         super().__init__(client, variables, per_page)
 
+    @override
+    def _update_response(self) -> None:
+        data = self.client.execute(self.QUERY, variable_values=self.variables)
+        result = ProjectArtifacts.model_validate(data)
+
+        # Extract the inner `*Connection` result for faster/easier access.
+        if not (
+            (proj := result.project)
+            and (type_ := proj.artifact_type)
+            and (collection := type_.artifact_collection)
+            and (conn := collection.artifacts)
+        ):
+            raise ValueError(f"Unable to parse {type(self).__name__!r} response data")
+
+        self.last_response = ArtifactsFragment.model_validate(conn)
+
     @property
-    def length(self):
-        if self.last_response:
-            return self.last_response["project"]["artifactType"]["artifactCollection"][
-                "artifacts"
-            ]["totalCount"]
-        else:
+    def length(self) -> Optional[int]:
+        if self.last_response is None:
             return None
+        return self.last_response.total_count
 
     @property
-    def more(self):
-        if self.last_response:
-            return self.last_response["project"]["artifactType"]["artifactCollection"][
-                "artifacts"
-            ]["pageInfo"]["hasNextPage"]
-        else:
+    def more(self) -> bool:
+        if self.last_response is None:
             return True
+        return self.last_response.page_info.has_next_page
 
     @property
-    def cursor(self):
-        if self.last_response:
-            return self.last_response["project"]["artifactType"]["artifactCollection"][
-                "artifacts"
-            ]["edges"][-1]["cursor"]
-        else:
+    def cursor(self) -> Optional[str]:
+        if self.last_response is None:
             return None
+        return self.last_response.edges[-1].cursor
 
-    def convert_objects(self):
-        collection = self.last_response["project"]["artifactType"]["artifactCollection"]
-        artifact_edges = collection.get("artifacts", {}).get("edges", [])
+    def convert_objects(self) -> List["wandb.Artifact"]:
+        artifact_edges = (edge for edge in self.last_response.edges if edge.node)
         artifacts = (
             wandb.Artifact._from_attrs(
-                self.entity,
-                self.project,
-                self.collection_name + ":" + a["version"],
-                a["node"],
-                self.client,
+                entity=self.entity,
+                project=self.project,
+                name=f"{self.collection_name}:{edge.version}",
+                attrs=edge.node.model_dump(exclude_unset=True),
+                client=self.client,
             )
-            for a in artifact_edges
+            for edge in artifact_edges
         )
         required_tags = set(self.tags or [])
-        return [
-            artifact for artifact in artifacts if required_tags.issubset(artifact.tags)
-        ]
+        return [art for art in artifacts if required_tags.issubset(art.tags)]
 
 
-class RunArtifacts(Paginator):
-    def __init__(
-        self, client: Client, run: "Run", mode="logged", per_page: Optional[int] = 50
-    ):
-        from wandb.sdk.artifacts.artifact import _gql_artifact_fragment
-
-        output_query = gql(
-            """
-            query RunOutputArtifacts(
-                $entity: String!, $project: String!, $runName: String!, $cursor: String, $perPage: Int,
-            ) {
-                project(name: $project, entityName: $entity) {
-                    run(name: $runName) {
-                        outputArtifacts(after: $cursor, first: $perPage) {
-                            totalCount
-                            edges {
-                                node {
-                                    ...ArtifactFragment
-                                }
-                                cursor
-                            }
-                            pageInfo {
-                                endCursor
-                                hasNextPage
-                            }
-                        }
-                    }
-                }
-            }
-            """
-            + _gql_artifact_fragment()
-        )
+class RunArtifacts(SizedPaginator["wandb.Artifact"]):
+    last_response: Union[
+        RunOutputArtifactsProjectRunOutputArtifacts,
+        RunInputArtifactsProjectRunInputArtifacts,
+    ]
 
-        input_query = gql(
-            """
-            query RunInputArtifacts(
-                $entity: String!, $project: String!, $runName: String!, $cursor: String, $perPage: Int,
-            ) {
-                project(name: $project, entityName: $entity) {
-                    run(name: $runName) {
-                        inputArtifacts(after: $cursor, first: $perPage) {
-                            totalCount
-                            edges {
-                                node {
-                                    ...ArtifactFragment
-                                }
-                                cursor
-                            }
-                            pageInfo {
-                                endCursor
-                                hasNextPage
-                            }
-                        }
-                    }
-                }
-            }
-            """
-            + _gql_artifact_fragment()
-        )
+    #: The pydantic model used to parse the (inner part of the) raw response.
+    _response_cls: Type[
+        Union[
+            RunOutputArtifactsProjectRunOutputArtifacts,
+            RunInputArtifactsProjectRunInputArtifacts,
+        ]
+    ]
 
+    def __init__(
+        self,
+        client: Client,
+        run: "Run",
+        mode: Literal["logged", "used"] = "logged",
+        per_page: int = 50,
+    ):
         self.run = run
+
         if mode == "logged":
             self.run_key = "outputArtifacts"
-            self.QUERY = output_query
+            self.QUERY = gql_compat(
+                RUN_OUTPUT_ARTIFACTS_GQL,
+                omit_fields=omit_artifact_fields(api=InternalApi()),
+            )
+            self._response_cls = RunOutputArtifactsProjectRunOutputArtifacts
         elif mode == "used":
             self.run_key = "inputArtifacts"
-            self.QUERY = input_query
+            self.QUERY = gql_compat(
+                RUN_INPUT_ARTIFACTS_GQL,
+                omit_fields=omit_artifact_fields(api=InternalApi()),
+            )
+            self._response_cls = RunInputArtifactsProjectRunInputArtifacts
         else:
             raise ValueError("mode must be logged or used")
 
@@ -901,72 +656,50 @@ class RunArtifacts(Paginator):
             "project": run.project,
             "runName": run.id,
         }
-
         super().__init__(client, variable_values, per_page)
 
+    @override
+    def _update_response(self) -> None:
+        data = self.client.execute(self.QUERY, variable_values=self.variables)
+
+        # Extract the inner `*Connection` result for faster/easier access.
+        inner_data = data["project"]["run"][self.run_key]
+        self.last_response = self._response_cls.model_validate(inner_data)
+
     @property
-    def length(self):
-        if self.last_response:
-            return self.last_response["project"]["run"][self.run_key]["totalCount"]
-        else:
+    def length(self) -> Optional[int]:
+        if self.last_response is None:
             return None
+        return self.last_response.total_count
 
     @property
-    def more(self):
-        if self.last_response:
-            return self.last_response["project"]["run"][self.run_key]["pageInfo"][
-                "hasNextPage"
-            ]
-        else:
+    def more(self) -> bool:
+        if self.last_response is None:
             return True
+        return self.last_response.page_info.has_next_page
 
     @property
-    def cursor(self):
-        if self.last_response:
-            return self.last_response["project"]["run"][self.run_key]["edges"][-1][
-                "cursor"
-            ]
-        else:
+    def cursor(self) -> Optional[str]:
+        if self.last_response is None:
             return None
+        return self.last_response.edges[-1].cursor
 
-    def convert_objects(self):
+    def convert_objects(self) -> List["wandb.Artifact"]:
         return [
             wandb.Artifact._from_attrs(
-                r["node"]["artifactSequence"]["project"]["entityName"],
-                r["node"]["artifactSequence"]["project"]["name"],
-                "{}:v{}".format(
-                    r["node"]["artifactSequence"]["name"], r["node"]["versionIndex"]
-                ),
-                r["node"],
-                self.client,
+                entity=node.artifact_sequence.project.entity_name,
+                project=node.artifact_sequence.project.name,
+                name=f"{node.artifact_sequence.name}:v{node.version_index}",
+                attrs=node.model_dump(exclude_unset=True),
+                client=self.client,
             )
-            for r in self.last_response["project"]["run"][self.run_key]["edges"]
+            for r in self.last_response.edges
+            if (node := r.node)
         ]
 
 
-class ArtifactFiles(Paginator):
-    QUERY = gql(
-        """
-        query ArtifactFiles(
-            $entityName: String!,
-            $projectName: String!,
-            $artifactTypeName: String!,
-            $artifactName: String!
-            $fileNames: [String!],
-            $fileCursor: String,
-            $fileLimit: Int = 50
-        ) {{
-            project(name: $projectName, entityName: $entityName) {{
-                artifactType(name: $artifactTypeName) {{
-                    artifact(name: $artifactName) {{
-                        ...ArtifactFilesFragment
-                    }}
-                }}
-            }}
-        }}
-        {}
-    """.format(ARTIFACT_FILES_FRAGMENT)
-    )
+class ArtifactFiles(SizedPaginator["public.File"]):
+    last_response: Optional[FilesFragment]
 
     def __init__(
         self,
@@ -975,59 +708,92 @@ class ArtifactFiles(Paginator):
         names: Optional[Sequence[str]] = None,
         per_page: int = 50,
     ):
+        self.query_via_membership = InternalApi()._check_server_feature_with_fallback(
+            ServerFeature.ARTIFACT_COLLECTION_MEMBERSHIP_FILES
+        )
         self.artifact = artifact
-        variables = {
-            "entityName": artifact.source_entity,
-            "projectName": artifact.source_project,
-            "artifactTypeName": artifact.type,
-            "artifactName": artifact.source_name,
-            "fileNames": names,
-        }
+
+        if self.query_via_membership:
+            query_str = ARTIFACT_COLLECTION_MEMBERSHIP_FILES_GQL
+            variables = {
+                "entityName": artifact.entity,
+                "projectName": artifact.project,
+                "artifactName": artifact.name.split(":")[0],
+                "artifactVersionIndex": artifact.version,
+                "fileNames": names,
+            }
+        else:
+            query_str = ARTIFACT_VERSION_FILES_GQL
+            variables = {
+                "entityName": artifact.source_entity,
+                "projectName": artifact.source_project,
+                "artifactName": artifact.source_name,
+                "artifactTypeName": artifact.type,
+                "fileNames": names,
+            }
+
         # The server must advertise at least SDK 0.12.21
         # to get storagePath
         if not client.version_supported("0.12.21"):
-            self.QUERY = gql(self.QUERY.loc.source.body.replace("storagePath\n", ""))
+            self.QUERY = gql_compat(query_str, omit_fields={"storagePath"})
+        else:
+            self.QUERY = gql(query_str)
+
         super().__init__(client, variables, per_page)
 
+    @override
+    def _update_response(self) -> None:
+        data = self.client.execute(self.QUERY, variable_values=self.variables)
+
+        # Extract the inner `*Connection` result for faster/easier access.
+        if self.query_via_membership:
+            result = ArtifactCollectionMembershipFiles.model_validate(data)
+            conn = result.project.artifact_collection.artifact_membership.files
+        else:
+            result = ArtifactVersionFiles.model_validate(data)
+            conn = result.project.artifact_type.artifact.files
+
+        if conn is None:
+            raise ValueError(f"Unable to parse {type(self).__name__!r} response data")
+
+        self.last_response = FilesFragment.model_validate(conn)
+
     @property
-    def path(self):
+    def path(self) -> List[str]:
         return [self.artifact.entity, self.artifact.project, self.artifact.name]
 
     @property
-    def length(self):
+    def length(self) -> int:
         return self.artifact.file_count
 
     @property
-    def more(self):
-        if self.last_response:
-            return self.last_response["project"]["artifactType"]["artifact"]["files"][
-                "pageInfo"
-            ]["hasNextPage"]
-        else:
+    def more(self) -> bool:
+        if self.last_response is None:
             return True
+        return self.last_response.page_info.has_next_page
 
     @property
-    def cursor(self):
-        if self.last_response:
-            return self.last_response["project"]["artifactType"]["artifact"]["files"][
-                "edges"
-            ][-1]["cursor"]
-        else:
+    def cursor(self) -> Optional[str]:
+        if self.last_response is None:
             return None
+        return self.last_response.edges[-1].cursor
 
-    def update_variables(self):
+    def update_variables(self) -> None:
         self.variables.update({"fileLimit": self.per_page, "fileCursor": self.cursor})
 
-    def convert_objects(self):
+    def convert_objects(self) -> List["public.File"]:
         return [
-            public.File(self.client, r["node"])
-            for r in self.last_response["project"]["artifactType"]["artifact"]["files"][
-                "edges"
-            ]
+            public.File(
+                client=self.client,
+                attrs=node.model_dump(exclude_unset=True),
+            )
+            for edge in self.last_response.edges
+            if (node := edge.node)
         ]
 
-    def __repr__(self):
-        return "<ArtifactFiles {} ({})>".format("/".join(self.path), len(self))
+    def __repr__(self) -> str:
+        path_str = "/".join(self.path)
+        return f"<ArtifactFiles {path_str} ({len(self)})>"
 
 
 def server_supports_artifact_collections_gql_edges(
@@ -1043,21 +809,3 @@ def server_supports_artifact_collections_gql_edges(
             "W&B Local Server version does not support ArtifactCollection gql edges; falling back to using legacy ArtifactSequence. Please update server to at least version 0.9.50."
         )
     return supported
-
-
-def artifact_collection_edge_name(server_supports_artifact_collections: bool) -> str:
-    return (
-        "artifactCollection"
-        if server_supports_artifact_collections
-        else "artifactSequence"
-    )
-
-
-def artifact_collection_plural_edge_name(
-    server_supports_artifact_collections: bool,
-) -> str:
-    return (
-        "artifactCollections"
-        if server_supports_artifact_collections
-        else "artifactSequences"
-    )