context model2

Author: yzhu0

src/sagemaker/lineage/context.py

@@ -283,6 +283,29 @@ def models(self) -> List[association.Association]:
         ]
         return model_list
 
+    def models_v2(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.BOTH
+    ) -> List[Artifact]:
+        """Get artifacts representing models from the context lineage.
+
+        Returns:
+            list of Artifacts: Artifacts representing a model.
+
+        """
+        query_filter = LineageFilter(
+            entities=[LineageEntityEnum.ARTIFACT], sources=[LineageSourceEnum.MODEL]
+        )
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.context_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        model_artifacts = []
+        for vertex in query_result.vertices:
+            model_artifacts.append(vertex.to_lineage_object())
+        return model_artifacts
+
     def dataset_artifacts(
         self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
     ) -> List[Artifact]: