Skip to content

Commit 837bf82

Browse files
vaibhavatlanAryamanz29
vaibhavatlan
authored and
Aryamanz29
committed
Added a_i_model creator()
1 parent 873cf59 commit 837bf82

File tree

2 files changed

+305
-1
lines changed

2 files changed

+305
-1
lines changed

‎pyatlan/model/assets/core/a_i_model.py

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,65 @@
44

55
from __future__ import annotations
66

7-
from typing import ClassVar, List, Optional
7+
from typing import ClassVar, List, Optional, overload
88

99
from pydantic.v1 import Field, validator
1010

1111
from pyatlan.model.enums import AIModelStatus
1212
from pyatlan.model.fields.atlan_fields import KeywordField, RelationField, TextField
13+
from pyatlan.utils import init_guid, to_camel_case, validate_required_fields
1314

1415
from .a_i import AI
1516

1617

1718
class AIModel(AI):
1819
"""Description"""
1920

21+
@overload
22+
@classmethod
23+
def creator(
24+
cls,
25+
*,
26+
name: str,
27+
ai_model_status: AIModelStatus,
28+
) -> AIModel: ...
29+
30+
@overload
31+
@classmethod
32+
def creator(
33+
cls,
34+
*,
35+
name: str,
36+
ai_model_status: AIModelStatus,
37+
owner_groups: set[str],
38+
owner_users: set[str],
39+
ai_model_version: str,
40+
) -> AIModel: ...
41+
42+
@classmethod
43+
@init_guid
44+
def creator(
45+
cls,
46+
*,
47+
name: str,
48+
ai_model_status: AIModelStatus,
49+
owner_groups: Optional[set[str]] = set(),
50+
owner_users: Optional[set[str]] = set(),
51+
ai_model_version: Optional[str] = None,
52+
) -> AIModel:
53+
validate_required_fields(
54+
["name", "ai_model_status"],
55+
[name, ai_model_status],
56+
)
57+
attributes = AIModel.Attributes.creator(
58+
name=name,
59+
ai_model_status=ai_model_status,
60+
owner_groups=owner_groups,
61+
owner_users=owner_users,
62+
ai_model_version=ai_model_version,
63+
)
64+
return cls(attributes=attributes)
65+
2066
type_name: str = Field(default="AIModel", allow_mutation=False)
2167

2268
@validator("type_name")
@@ -129,6 +175,33 @@ class Attributes(AI.Attributes):
129175
default=None, description=""
130176
) # relationship
131177

178+
@classmethod
179+
@init_guid
180+
def creator(
181+
cls,
182+
*,
183+
name: str,
184+
ai_model_status: AIModelStatus,
185+
owner_groups: Optional[set[str]] = set(),
186+
owner_users: Optional[set[str]] = set(),
187+
ai_model_version: Optional[str] = None,
188+
) -> AIModel.Attributes:
189+
validate_required_fields(
190+
["name", "ai_model_status"],
191+
[name, ai_model_status],
192+
)
193+
name_camel_case = to_camel_case(name)
194+
return AIModel.Attributes(
195+
name=name,
196+
qualified_name=f"default/ai/aiapplication/{name_camel_case}",
197+
connector_name="ai",
198+
ai_model_status=ai_model_status,
199+
ai_model_version=ai_model_version,
200+
owner_groups=owner_groups,
201+
owner_users=owner_users,
202+
asset_cover_image="/assets/default-product-cover-DeQonY47.webp",
203+
)
204+
132205
attributes: AIModel.Attributes = Field(
133206
default_factory=lambda: AIModel.Attributes(),
134207
description=(

‎test_model.json

Lines changed: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
1+
{
2+
"entities": [
3+
{
4+
"attributes": {
5+
"name": "final-model-test",
6+
"aiModelStatus": "ACTIVE",
7+
"assetCoverImage": "/assets/default-product-cover-DeQonY47.webp",
8+
"ownerGroups": [],
9+
"ownerUsers": [],
10+
"qualifiedName": "default/ai/aiapplication/finalModelTest"
11+
},
12+
"typeName": "AIModel"
13+
}
14+
]
15+
}
16+
17+
18+
{
19+
"entities": [
20+
{
21+
"typeName": "Process",
22+
"attributes": {
23+
"name": "psdk_my_db_a7ljF -> final-model-test",
24+
"qualifiedName": "default/ai/dataset/937c19beebb3f99fb02e59f14702d86d47623cd519c3f8aec44136cfc574dd18",
25+
"aiDatasetType": "TRAINING",
26+
"inputs": [
27+
{
28+
"typeName": "Database",
29+
"guid": "f9503834-afdf-4cb1-a332-c5306355d134"
30+
}
31+
],
32+
"outputs": [
33+
{
34+
"typeName": "AIModel",
35+
"guid": "d1bb9565-bc37-4bdf-8133-7aad425c60bb"
36+
}
37+
],
38+
"__state": "ACTIVE"
39+
}
40+
},
41+
{
42+
"typeName": "Process",
43+
"attributes": {
44+
"name": "/CITIES/_TEST_3 -> final-model-test",
45+
"qualifiedName": "default/ai/dataset/4578fd36d3956f38b3123555bc5fac951e76ea7b3517e7b988e71b8191f36f3e",
46+
"aiDatasetType": "TRAINING",
47+
"inputs": [
48+
{
49+
"typeName": "Table",
50+
"guid": "57b97883-06c5-4565-9f0f-2cc8bcfb773e"
51+
}
52+
],
53+
"outputs": [
54+
{
55+
"typeName": "AIModel",
56+
"guid": "d1bb9565-bc37-4bdf-8133-7aad425c60bb"
57+
}
58+
],
59+
"__state": "ACTIVE"
60+
}
61+
},
62+
{
63+
"typeName": "Process",
64+
"attributes": {
65+
"name": "psdk_my_db_a7ljF -> final-model-test",
66+
"qualifiedName": "default/ai/dataset/9b6c7210d288bb9bf84924260a0ca0d0646ab5ecfbb88bafe57ecf8c1d943848",
67+
"aiDatasetType": "TESTING",
68+
"inputs": [
69+
{
70+
"typeName": "Database",
71+
"guid": "f9503834-afdf-4cb1-a332-c5306355d134"
72+
}
73+
],
74+
"outputs": [
75+
{
76+
"typeName": "AIModel",
77+
"guid": "d1bb9565-bc37-4bdf-8133-7aad425c60bb"
78+
}
79+
],
80+
"__state": "ACTIVE"
81+
}
82+
},
83+
{
84+
"typeName": "Process",
85+
"attributes": {
86+
"name": "/CITIES/_TEST_3 -> final-model-test",
87+
"qualifiedName": "default/ai/dataset/f02c189b43104c1a0af09b523a411f061d0ea5fbfdb0ccb992e0b00f450faf40",
88+
"aiDatasetType": "TESTING",
89+
"inputs": [
90+
{
91+
"typeName": "Table",
92+
"guid": "57b97883-06c5-4565-9f0f-2cc8bcfb773e"
93+
}
94+
],
95+
"outputs": [
96+
{
97+
"typeName": "AIModel",
98+
"guid": "d1bb9565-bc37-4bdf-8133-7aad425c60bb"
99+
}
100+
],
101+
"__state": "ACTIVE"
102+
}
103+
},
104+
{
105+
"typeName": "Process",
106+
"attributes": {
107+
"name": "/cities_test -> final-model-test",
108+
"qualifiedName": "default/ai/dataset/31a0fd6150afa6cb27cacf25a9edfc0b33776c78b42daf9138c2a5b79543a93e",
109+
"aiDatasetType": "TESTING",
110+
"inputs": [
111+
{
112+
"typeName": "View",
113+
"guid": "be0d1b1e-2b6d-40d7-85f0-56a0f785dce8"
114+
}
115+
],
116+
"outputs": [
117+
{
118+
"typeName": "AIModel",
119+
"guid": "d1bb9565-bc37-4bdf-8133-7aad425c60bb"
120+
}
121+
],
122+
"__state": "ACTIVE"
123+
}
124+
},
125+
{
126+
"typeName": "Process",
127+
"attributes": {
128+
"name": "/CITIES/_TEST_3 -> final-model-test",
129+
"qualifiedName": "default/ai/dataset/41db7c162d30f5a2d38f0cc99ddbcf54c64bd235e8fa88e75bec025a1c8f72c0",
130+
"aiDatasetType": "INFERENCE",
131+
"inputs": [
132+
{
133+
"typeName": "Table",
134+
"guid": "57b97883-06c5-4565-9f0f-2cc8bcfb773e"
135+
}
136+
],
137+
"outputs": [
138+
{
139+
"typeName": "AIModel",
140+
"guid": "d1bb9565-bc37-4bdf-8133-7aad425c60bb"
141+
}
142+
],
143+
"__state": "ACTIVE"
144+
}
145+
},
146+
{
147+
"typeName": "Process",
148+
"attributes": {
149+
"name": "/cities_test -> final-model-test",
150+
"qualifiedName": "default/ai/dataset/8e0a0cfe752ada96758bd2702568966a5fde153f50760b5444fc4a1029fed098",
151+
"aiDatasetType": "INFERENCE",
152+
"inputs": [
153+
{
154+
"typeName": "View",
155+
"guid": "be0d1b1e-2b6d-40d7-85f0-56a0f785dce8"
156+
}
157+
],
158+
"outputs": [
159+
{
160+
"typeName": "AIModel",
161+
"guid": "d1bb9565-bc37-4bdf-8133-7aad425c60bb"
162+
}
163+
],
164+
"__state": "ACTIVE"
165+
}
166+
},
167+
{
168+
"typeName": "Process",
169+
"attributes": {
170+
"name": "psdk_my_db_a7ljF -> final-model-test",
171+
"qualifiedName": "default/ai/dataset/358edf08d04408846e44e83ec67f06845c5e4075b575ba3ae754194c2335aad7",
172+
"aiDatasetType": "VALIDATION",
173+
"inputs": [
174+
{
175+
"typeName": "Database",
176+
"guid": "f9503834-afdf-4cb1-a332-c5306355d134"
177+
}
178+
],
179+
"outputs": [
180+
{
181+
"typeName": "AIModel",
182+
"guid": "d1bb9565-bc37-4bdf-8133-7aad425c60bb"
183+
}
184+
],
185+
"__state": "ACTIVE"
186+
}
187+
},
188+
{
189+
"typeName": "Process",
190+
"attributes": {
191+
"name": "final-model-test -> /CITIES/_TEST_3",
192+
"qualifiedName": "default/ai/dataset/93fa4fdd5e3fd2db40f232f5e98c17073325377037a27536902a3b6148fc5d84",
193+
"aiDatasetType": "OUTPUT",
194+
"inputs": [
195+
{
196+
"typeName": "AIModel",
197+
"guid": "d1bb9565-bc37-4bdf-8133-7aad425c60bb"
198+
}
199+
],
200+
"outputs": [
201+
{
202+
"typeName": "Table",
203+
"guid": "57b97883-06c5-4565-9f0f-2cc8bcfb773e"
204+
}
205+
],
206+
"__state": "ACTIVE"
207+
}
208+
},
209+
{
210+
"typeName": "Process",
211+
"attributes": {
212+
"name": "final-model-test -> /cities_test",
213+
"qualifiedName": "default/ai/dataset/ca638688b0fef19faff963859c177bd5f720d906270eaf0750c13497fcbd933b",
214+
"aiDatasetType": "OUTPUT",
215+
"inputs": [
216+
{
217+
"typeName": "AIModel",
218+
"guid": "d1bb9565-bc37-4bdf-8133-7aad425c60bb"
219+
}
220+
],
221+
"outputs": [
222+
{
223+
"typeName": "View",
224+
"guid": "be0d1b1e-2b6d-40d7-85f0-56a0f785dce8"
225+
}
226+
],
227+
"__state": "ACTIVE"
228+
}
229+
}
230+
]
231+
}

0 commit comments

Comments
 (0)