made small fixes

Author: vaibhavatlan

pyatlan/model/assets/core/a_i_application.py

@@ -4,13 +4,17 @@
 
 from __future__ import annotations
 
-from typing import ClassVar, List, Optional, overload
+from typing import ClassVar, List, Optional, Set, overload
 
 from pydantic.v1 import Field, validator
 
-from pyatlan.model.enums import AIApplicationDevelopmentStage, CertificateStatus
+from pyatlan.model.enums import (
+    AIApplicationDevelopmentStage,
+    AtlanConnectorType,
+    CertificateStatus,
+)
 from pyatlan.model.fields.atlan_fields import KeywordField, RelationField
-from pyatlan.utils import init_guid, validate_required_fields, to_camel_case
+from pyatlan.utils import init_guid, to_camel_case, validate_required_fields
 
 from .a_i import AI
 
@@ -36,9 +40,9 @@ def creator(
         name: str,
         ai_application_version: str,
         ai_application_development_stage: AIApplicationDevelopmentStage,
-        owner_groups: set[str],
-        owner_users: set[str],
-        models: list[AIModel],
+        owner_groups: Set[str],
+        owner_users: Set[str],
+        models: List[AIModel],
     ) -> AIApplication: ...
 
     @classmethod
@@ -49,9 +53,9 @@ def creator(
         name: str,
         ai_application_version: str,
         ai_application_development_stage: AIApplicationDevelopmentStage,
-        owner_groups: Optional[set[str]] = set(),
-        owner_users: Optional[set[str]] = set(),
-        models: Optional[list[AIModel]] = [],
+        owner_groups: Optional[Set[str]] = set(),
+        owner_users: Optional[Set[str]] = set(),
+        models: Optional[List[AIModel]] = [],
     ) -> AIApplication:
         validate_required_fields(
             ["name", "ai_application_version", "ai_application_development_stage"],
@@ -163,9 +167,9 @@ def creator(
             name: str,
             ai_application_version: str,
             ai_application_development_stage: AIApplicationDevelopmentStage,
-            owner_groups: Optional[set[str]] = set(),
-            owner_users: Optional[set[str]] = set(),
-            models: Optional[list[AIModel]] = [],
+            owner_groups: Optional[Set[str]] = set(),
+            owner_users: Optional[Set[str]] = set(),
+            models: Optional[List[AIModel]] = [],
         ) -> AIApplication.Attributes:
             validate_required_fields(
                 ["name", "ai_application_version", "ai_application_development_stage"],
@@ -175,7 +179,7 @@ def creator(
             return AIApplication.Attributes(
                 name=name,
                 qualified_name=f"default/ai/aiapplication/{name_camel_case}",
-                connector_name="ai",
+                connector_name=AtlanConnectorType.AI.value,
                 ai_application_version=ai_application_version,
                 ai_application_development_stage=ai_application_development_stage,
                 owner_groups=owner_groups,

pyatlan/model/assets/core/a_i_model.py

@@ -4,11 +4,11 @@
 
 from __future__ import annotations
 
-from typing import ClassVar, List, Optional, overload
+from typing import ClassVar, List, Optional, Set, overload
 
 from pydantic.v1 import Field, validator
 
-from pyatlan.model.enums import AIDatasetType, AIModelStatus
+from pyatlan.model.enums import AIDatasetType, AIModelStatus, AtlanConnectorType
 from pyatlan.model.fields.atlan_fields import KeywordField, RelationField, TextField
 from pyatlan.utils import (
     get_epoch_timestamp,
@@ -40,8 +40,8 @@ def creator(
         *,
         name: str,
         ai_model_status: AIModelStatus,
-        owner_groups: set[str],
-        owner_users: set[str],
+        owner_groups: Set[str],
+        owner_users: Set[str],
         ai_model_version: str,
     ) -> AIModel: ...
 
@@ -52,8 +52,8 @@ def creator(
         *,
         name: str,
         ai_model_status: AIModelStatus,
-        owner_groups: Optional[set[str]] = set(),
-        owner_users: Optional[set[str]] = set(),
+        owner_groups: Optional[Set[str]] = set(),
+        owner_users: Optional[Set[str]] = set(),
         ai_model_version: Optional[str] = None,
     ) -> AIModel:
         validate_required_fields(
@@ -221,8 +221,8 @@ def creator(
             *,
             name: str,
             ai_model_status: AIModelStatus,
-            owner_groups: Optional[set[str]] = set(),
-            owner_users: Optional[set[str]] = set(),
+            owner_groups: Optional[Set[str]] = set(),
+            owner_users: Optional[Set[str]] = set(),
             ai_model_version: Optional[str] = None,
         ) -> AIModel.Attributes:
             validate_required_fields(
@@ -233,7 +233,7 @@ def creator(
             return AIModel.Attributes(
                 name=name,
                 qualified_name=f"default/ai/aiapplication/{name_camel_case}",
-                connector_name="ai",
+                connector_name=AtlanConnectorType.AI.value,
                 ai_model_status=ai_model_status,
                 ai_model_version=ai_model_version,
                 owner_groups=owner_groups,

pyatlan/model/enums.py

@@ -138,6 +138,7 @@ class AtlanConnectionCategory(str, Enum):
     SCHEMA_REGISTRY = "schema-registry"
     APP = "app"
     CUSTOM = "custom"
+    AI = "ai"
 
 
 class AtlanConnectorType(str, Enum, metaclass=utils.ExtendableEnumMeta):
@@ -381,6 +382,7 @@ def get_connector_name(
     CRATEDB = ("cratedb", AtlanConnectionCategory.DATABASE)
     KX = ("kx", AtlanConnectionCategory.DATABASE)
     DOCUMENTDB = ("documentdb", AtlanConnectionCategory.DATABASE)
+    AI = ("ai", AtlanConnectionCategory.AI)
 
 
 class AtlanCustomAttributePrimitiveType(str, Enum):