Skip to content

Commit 9b37700

Browse files
VipulMascarenhasVipulMascarenhas
authored
Websocket changes for 1.0.3 (#893)
2 parents 3d5b790 + ac49a3e commit 9b37700

File tree

6 files changed

+239
-10
lines changed

6 files changed

+239
-10
lines changed

‎ads/aqua/extension/common_ws_msg_handler.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
import json
2+
from typing import List, Union
3+
4+
from ads.aqua import ODSC_MODEL_COMPARTMENT_OCID, fetch_service_compartment
5+
from ads.aqua.common.decorator import handle_exceptions
6+
from ads.aqua.common.errors import AquaResourceAccessError
7+
from ads.aqua.common.utils import known_realm
8+
from ads.aqua.extension.aqua_ws_msg_handler import AquaWSMsgHandler
9+
from ads.aqua.extension.models.ws_models import RequestResponseType, AdsVersionResponse, AdsVersionRequest, \
10+
CompatibilityCheckResponse
11+
from importlib import metadata
12+
13+
14+
class AquaCommonWsMsgHandler(AquaWSMsgHandler):
15+
16+
@staticmethod
17+
def get_message_types() -> List[RequestResponseType]:
18+
return [RequestResponseType.AdsVersion, RequestResponseType.CompatibilityCheck]
19+
20+
def __init__(self, message: Union[str, bytes]):
21+
super().__init__(message)
22+
23+
@handle_exceptions
24+
def process(self) -> AdsVersionResponse | CompatibilityCheckResponse:
25+
request = json.loads(self.message)
26+
if request.get('kind') == 'AdsVersion':
27+
version = metadata.version("oracle_ads")
28+
response = AdsVersionResponse(
29+
message_id=request.get("message_id"),
30+
kind=RequestResponseType.AdsVersion,
31+
data=version)
32+
return response
33+
if request.get('kind') == 'CompatibilityCheck':
34+
if ODSC_MODEL_COMPARTMENT_OCID or fetch_service_compartment():
35+
return CompatibilityCheckResponse(message_id=request.get("message_id"),
36+
kind=RequestResponseType.CompatibilityCheck,
37+
data={'status': 'ok'})
38+
elif known_realm():
39+
return CompatibilityCheckResponse(message_id=request.get("message_id"),
40+
kind=RequestResponseType.CompatibilityCheck,
41+
data={'status': 'compatible'})
42+
else:
43+
raise AquaResourceAccessError(
44+
f"The AI Quick actions extension is not compatible in the given region."
45+
)

‎ads/aqua/extension/deployment_ws_msg_handler.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
import json
2+
from typing import List, Union
3+
4+
from ads.aqua.common.decorator import handle_exceptions
5+
from ads.aqua.extension.aqua_ws_msg_handler import AquaWSMsgHandler
6+
from ads.aqua.extension.models.ws_models import RequestResponseType, ListDeploymentResponse, ListDeploymentRequest, \
7+
ModelDeploymentDetailsResponse
8+
from ads.aqua.modeldeployment import AquaDeploymentApp
9+
from ads.config import COMPARTMENT_OCID
10+
11+
12+
class AquaDeploymentWSMsgHandler(AquaWSMsgHandler):
13+
14+
def __init__(self, message: Union[str, bytes]):
15+
super().__init__(message)
16+
17+
@staticmethod
18+
def get_message_types() -> List[RequestResponseType]:
19+
return [RequestResponseType.ListDeployments, RequestResponseType.DeploymentDetails]
20+
21+
@handle_exceptions
22+
def process(self) -> ListDeploymentResponse | ModelDeploymentDetailsResponse:
23+
request = json.loads(self.message)
24+
if request.get("kind") == "ListDeployments":
25+
deployment_list = AquaDeploymentApp().list(
26+
compartment_id=request.get("compartment_id") or COMPARTMENT_OCID,
27+
project_id=request.get("project_id"),
28+
)
29+
response = ListDeploymentResponse(
30+
message_id=request.get("message_id"),
31+
kind=RequestResponseType.ListDeployments,
32+
data=deployment_list,
33+
)
34+
return response
35+
elif request.get("kind") == "DeploymentDetails":
36+
deployment_details = AquaDeploymentApp().get(request.get("model_deployment_id"))
37+
response = ModelDeploymentDetailsResponse(message_id=request.get("message_id"),
38+
kind=RequestResponseType.DeploymentDetails,
39+
data=deployment_details)
40+
return response

‎ads/aqua/extension/evaluation_ws_msg_handler.py‎

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,14 @@
33
# Copyright (c) 2024 Oracle and/or its affiliates.
44
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
55

6+
import json
67
from typing import List, Union
78

89
from ads.aqua.common.decorator import handle_exceptions
910
from ads.aqua.evaluation import AquaEvaluationApp
1011
from ads.aqua.extension.aqua_ws_msg_handler import AquaWSMsgHandler
1112
from ads.aqua.extension.models.ws_models import (
12-
ListEvaluationsRequest,
13+
EvaluationDetailsResponse,
1314
ListEvaluationsResponse,
1415
RequestResponseType,
1516
)
@@ -19,21 +20,42 @@
1920
class AquaEvaluationWSMsgHandler(AquaWSMsgHandler):
2021
@staticmethod
2122
def get_message_types() -> List[RequestResponseType]:
22-
return [RequestResponseType.ListEvaluations]
23+
return [
24+
RequestResponseType.ListEvaluations,
25+
RequestResponseType.EvaluationDetails,
26+
]
2327

2428
def __init__(self, message: Union[str, bytes]):
2529
super().__init__(message)
2630

2731
@handle_exceptions
28-
def process(self) -> ListEvaluationsResponse:
29-
list_eval_request = ListEvaluationsRequest.from_json(self.message)
32+
def process(self) -> ListEvaluationsResponse | EvaluationDetailsResponse:
33+
request = json.loads(self.message)
34+
if request["kind"] == "ListEvaluations":
35+
return self.list_evaluations(request)
36+
if request["kind"] == "EvaluationDetails":
37+
return self.evaluation_details(request)
3038

39+
@staticmethod
40+
def list_evaluations(request) -> ListEvaluationsResponse:
3141
eval_list = AquaEvaluationApp().list(
32-
list_eval_request.compartment_id or COMPARTMENT_OCID,
42+
request.get("compartment_id") or COMPARTMENT_OCID
3343
)
3444
response = ListEvaluationsResponse(
35-
message_id=list_eval_request.message_id,
45+
message_id=request["message_id"],
3646
kind=RequestResponseType.ListEvaluations,
3747
data=eval_list,
3848
)
3949
return response
50+
51+
@staticmethod
52+
def evaluation_details(request) -> EvaluationDetailsResponse:
53+
evaluation_details = AquaEvaluationApp().get(
54+
eval_id=request.get("evaluation_id")
55+
)
56+
response = EvaluationDetailsResponse(
57+
message_id=request.get("message_id"),
58+
kind=RequestResponseType.EvaluationDetails,
59+
data=evaluation_details,
60+
)
61+
return response

‎ads/aqua/extension/models/ws_models.py‎

Lines changed: 78 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,23 +7,30 @@
77
from dataclasses import dataclass
88
from typing import List, Optional
99

10-
from ads.aqua.evaluation.entities import AquaEvaluationSummary
11-
from ads.aqua.model.entities import AquaModelSummary
10+
from ads.aqua.evaluation.entities import AquaEvaluationSummary, AquaEvaluationDetail
11+
from ads.aqua.model.entities import AquaModelSummary, AquaModel
12+
from ads.aqua.modeldeployment.entities import AquaDeployment, AquaDeploymentDetail
1213
from ads.common.extended_enum import ExtendedEnumMeta
1314
from ads.common.serializer import DataClassSerializable
1415

1516

1617
class RequestResponseType(str, metaclass=ExtendedEnumMeta):
1718
ListEvaluations = "ListEvaluations"
19+
EvaluationDetails = "EvaluationDetails"
20+
ListDeployments = "ListDeployments"
21+
DeploymentDetails = "DeploymentDetails"
1822
ListModels = "ListModels"
23+
ModelDetails = "ModelDetails"
24+
AdsVersion = "AdsVersion"
25+
CompatibilityCheck = "CompatibilityCheck"
1926
Error = "Error"
2027

2128

2229
@dataclass
2330
class BaseResponse(DataClassSerializable):
2431
message_id: str
2532
kind: RequestResponseType
26-
data: object
33+
data: Optional[object]
2734

2835

2936
@dataclass
@@ -40,21 +47,89 @@ class ListEvaluationsRequest(BaseRequest):
4047
kind = RequestResponseType.ListEvaluations
4148

4249

50+
@dataclass
51+
class EvaluationDetailsRequest(BaseRequest):
52+
kind = RequestResponseType.EvaluationDetails
53+
evaluation_id: str
54+
55+
4356
@dataclass
4457
class ListModelsRequest(BaseRequest):
4558
compartment_id: Optional[str] = None
59+
project_id: Optional[str] = None
60+
model_type: Optional[str] = None
61+
kind = RequestResponseType.ListDeployments
62+
63+
64+
@dataclass
65+
class ModelDetailsRequest(BaseRequest):
66+
kind = RequestResponseType.ModelDetails
67+
model_id: str
68+
69+
70+
@dataclass
71+
class ListDeploymentRequest(BaseRequest):
72+
compartment_id: str
73+
project_id: Optional[str] = None
74+
kind = RequestResponseType.ListDeployments
75+
76+
77+
@dataclass
78+
class DeploymentDetailsRequest(BaseRequest):
79+
model_deployment_id: str
80+
kind = RequestResponseType.DeploymentDetails
4681

4782

4883
@dataclass
4984
class ListEvaluationsResponse(BaseResponse):
5085
data: List[AquaEvaluationSummary]
5186

5287

88+
@dataclass
89+
class EvaluationDetailsResponse(BaseResponse):
90+
data: AquaEvaluationDetail
91+
92+
93+
@dataclass
94+
class ListDeploymentResponse(BaseResponse):
95+
data: List[AquaDeployment]
96+
97+
98+
@dataclass
99+
class ModelDeploymentDetailsResponse(BaseResponse):
100+
data: AquaDeploymentDetail
101+
102+
53103
@dataclass
54104
class ListModelsResponse(BaseResponse):
55105
data: List[AquaModelSummary]
56106

57107

108+
@dataclass
109+
class ModelDetailsResponse(BaseResponse):
110+
data: AquaModel
111+
112+
113+
@dataclass
114+
class AdsVersionRequest(BaseRequest):
115+
kind: RequestResponseType.AdsVersion
116+
117+
118+
@dataclass
119+
class AdsVersionResponse(BaseResponse):
120+
data: str
121+
122+
123+
@dataclass
124+
class CompatibilityCheckRequest(BaseRequest):
125+
kind: RequestResponseType.CompatibilityCheck
126+
127+
128+
@dataclass
129+
class CompatibilityCheckResponse(BaseResponse):
130+
data: object
131+
132+
58133
@dataclass
59134
class AquaWsError(DataClassSerializable):
60135
status: str

‎ads/aqua/extension/models_ws_msg_handler.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
import json
2+
from typing import List, Union
3+
4+
from ads.aqua.common.decorator import handle_exceptions
5+
from ads.aqua.extension.aqua_ws_msg_handler import AquaWSMsgHandler
6+
from ads.aqua.extension.models.ws_models import RequestResponseType, ListModelsResponse, ListModelsRequest, \
7+
ModelDetailsResponse
8+
from ads.aqua.model import AquaModelApp
9+
10+
11+
class AquaModelWSMsgHandler(AquaWSMsgHandler):
12+
13+
def __init__(self, message: Union[str, bytes]):
14+
super().__init__(message)
15+
16+
@staticmethod
17+
def get_message_types() -> List[RequestResponseType]:
18+
return [RequestResponseType.ListModels,RequestResponseType.ModelDetails]
19+
20+
@handle_exceptions
21+
def process(self) -> ListModelsResponse | ModelDetailsResponse:
22+
request = json.loads(self.message)
23+
if request.get('kind') == 'ListModels':
24+
models_list = AquaModelApp().list(
25+
compartment_id=request.get("compartment_id"),
26+
project_id=request.get("project_id"),
27+
model_type=request.get("model_type")
28+
)
29+
response = ListModelsResponse(
30+
message_id=request.get("message_id"),
31+
kind=RequestResponseType.ListModels,
32+
data=models_list,
33+
)
34+
return response
35+
elif request.get('kind') == 'ModelDetails':
36+
model_id=request.get("model_id")
37+
response=AquaModelApp().get(model_id)
38+
return ModelDetailsResponse(message_id=request.get("message_id"),
39+
kind=RequestResponseType.ModelDetails,
40+
data=response)
41+

‎ads/aqua/extension/ui_websocket_handler.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414

1515
from ads.aqua import logger
1616
from ads.aqua.extension.aqua_ws_msg_handler import AquaWSMsgHandler
17+
from ads.aqua.extension.common_ws_msg_handler import AquaCommonWsMsgHandler
18+
from ads.aqua.extension.deployment_ws_msg_handler import AquaDeploymentWSMsgHandler
1719
from ads.aqua.extension.evaluation_ws_msg_handler import AquaEvaluationWSMsgHandler
1820
from ads.aqua.extension.models.ws_models import (
1921
AquaWsError,
@@ -22,6 +24,7 @@
2224
ErrorResponse,
2325
RequestResponseType,
2426
)
27+
from ads.aqua.extension.models_ws_msg_handler import AquaModelWSMsgHandler
2528

2629
MAX_WORKERS = 20
2730

@@ -43,7 +46,10 @@ def get_aqua_internal_error_response(message_id: str) -> ErrorResponse:
4346
class AquaUIWebSocketHandler(WebSocketHandler):
4447
"""Handler for Aqua Websocket."""
4548

46-
_handlers_: List[Type[AquaWSMsgHandler]] = [AquaEvaluationWSMsgHandler]
49+
_handlers_: List[Type[AquaWSMsgHandler]] = [AquaEvaluationWSMsgHandler,
50+
AquaDeploymentWSMsgHandler,
51+
AquaModelWSMsgHandler,
52+
AquaCommonWsMsgHandler]
4753

4854
thread_pool: ThreadPoolExecutor
4955

0 commit comments

Comments
 (0)