Push registration methods onto the backing providers

llama_stack/apis/agents/agents.py

@@ -261,7 +261,7 @@ class Session(BaseModel):
     turns: List[Turn]
     started_at: datetime
 
-    memory_bank: Optional[MemoryBank] = None
+    memory_bank: Optional[MemoryBankDef] = None
 
 
 class AgentConfigCommon(BaseModel):

llama_stack/apis/inference/inference.py

@@ -14,6 +14,7 @@ from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
+from llama_stack.apis.models import *  # noqa: F403
 
 
 class LogProbConfig(BaseModel):
@@ -203,3 +204,12 @@ class Inference(Protocol):
         model: str,
         contents: List[InterleavedTextMedia],
     ) -> EmbeddingsResponse: ...
+
+    @webmethod(route="/inference/register_model")
+    async def register_model(self, model: ModelDef) -> None: ...
+
+    @webmethod(route="/inference/list_models")
+    async def list_models(self) -> List[ModelDef]: ...
+
+    @webmethod(route="/inference/get_model")
+    async def get_model(self, identifier: str) -> Optional[ModelDef]: ...

llama_stack/apis/memory/memory.py

@@ -15,6 +15,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
+from llama_stack.apis.memory_banks import *  # noqa: F403
 
 
 @json_schema_type
@@ -76,3 +77,12 @@ class Memory(Protocol):
         bank_id: str,
         document_ids: List[str],
     ) -> None: ...
+
+    @webmethod(route="/memory/register_memory_bank")
+    async def register_memory_bank(self, memory_bank: MemoryBankDef) -> None: ...
+
+    @webmethod(route="/memory/list_memory_banks")
+    async def list_memory_banks(self) -> List[MemoryBankDef]: ...
+
+    @webmethod(route="/memory/get_memory_bank")
+    async def get_memory_bank(self, identifier: str) -> Optional[MemoryBankDef]: ...

llama_stack/apis/safety/safety.py

@@ -11,6 +11,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
+from llama_stack.apis.shields import *  # noqa: F403
 
 
 @json_schema_type
@@ -42,3 +43,12 @@ class Safety(Protocol):
     async def run_shield(
         self, shield_type: str, messages: List[Message], params: Dict[str, Any] = None
     ) -> RunShieldResponse: ...
+
+    @webmethod(route="/safety/register_shield")
+    async def register_shield(self, shield: ShieldDef) -> None: ...
+
+    @webmethod(route="/safety/list_shields")
+    async def list_shields(self) -> List[ShieldDef]: ...
+
+    @webmethod(route="/safety/get_shield")
+    async def get_shield(self, identifier: str) -> Optional[ShieldDef]: ...