diff --git a/llama_stack/distribution/datatypes.py b/llama_stack/distribution/datatypes.py
index 5212e6da1..0044de09e 100644
--- a/llama_stack/distribution/datatypes.py
+++ b/llama_stack/distribution/datatypes.py
@@ -34,8 +34,8 @@ RoutableObject = Union[
 
 RoutableObjectWithProvider = Union[
     ModelDefWithProvider,
-    ShieldDef,
-    MemoryBankDef,
+    ShieldDefWithProvider,
+    MemoryBankDefWithProvider,
 ]
 
 RoutedProtocol = Union[