Skip to content

Commit a0853f8

Browse files
committed
chore: add search keywords to header
1 parent feccd23 commit a0853f8

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

src/sagemaker/jumpstart/types.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,7 @@ def __init__(
130130
class JumpStartModelHeader(JumpStartDataHolderType):
131131
"""Data class JumpStart model header."""
132132

133-
__slots__ = ["model_id", "version", "min_version", "spec_key"]
133+
__slots__ = ["model_id", "version", "min_version", "spec_key", "search_keywords"]
134134

135135
def __init__(self, header: Dict[str, str]):
136136
"""Initializes a JumpStartModelHeader object from its json representation.
@@ -142,7 +142,11 @@ def __init__(self, header: Dict[str, str]):
142142

143143
def to_json(self) -> Dict[str, str]:
144144
"""Returns json representation of JumpStartModelHeader object."""
145-
json_obj = {att: getattr(self, att) for att in self.__slots__ if hasattr(self, att)}
145+
json_obj = {
146+
att: getattr(self, att)
147+
for att in self.__slots__
148+
if getattr(self, att, None) is not None
149+
}
146150
return json_obj
147151

148152
def from_json(self, json_obj: Dict[str, str]) -> None:
@@ -155,6 +159,7 @@ def from_json(self, json_obj: Dict[str, str]) -> None:
155159
self.version: str = json_obj["version"]
156160
self.min_version: str = json_obj["min_version"]
157161
self.spec_key: str = json_obj["spec_key"]
162+
self.search_keywords: Optional[List[str]] = json_obj.get("search_keywords")
158163

159164

160165
class JumpStartECRSpecs(JumpStartDataHolderType):

0 commit comments

Comments
 (0)