Make Safety test work, other cleanup

llama_stack/apis/memory/client.py

@@ -85,7 +85,6 @@ class MemoryClient(Memory):
 
 
 async def run_main(host: str, port: int, stream: bool):
-    client = MemoryClient(f"http://{host}:{port}")
     banks_client = MemoryBanksClient(f"http://{host}:{port}")
 
     bank = VectorMemoryBankDef(
@@ -95,7 +94,7 @@ async def run_main(host: str, port: int, stream: bool):
         chunk_size_in_tokens=512,
         overlap_size_in_tokens=64,
     )
-    await client.register_memory_bank(bank)
+    await banks_client.register_memory_bank(bank)
 
     retrieved_bank = await banks_client.get_memory_bank(bank.identifier)
     assert retrieved_bank is not None
@@ -130,6 +129,8 @@ async def run_main(host: str, port: int, stream: bool):
         for i, path in enumerate(files)
     ]
 
+    client = MemoryClient(f"http://{host}:{port}")
+
     # insert some documents
     await client.insert_documents(
         bank_id=bank.identifier,

llama_stack/apis/memory_banks/client.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 import asyncio
+import json
 
 from typing import Any, Dict, List, Optional
 
@@ -15,7 +16,9 @@ from termcolor import cprint
 from .memory_banks import *  # noqa: F403
 
 
-def deserialize_memory_bank_def(j: Optional[Dict[str, Any]]) -> MemoryBankDef:
+def deserialize_memory_bank_def(
+    j: Optional[Dict[str, Any]]
+) -> MemoryBankDefWithProvider:
     if j is None:
         return None
 
@@ -44,7 +47,7 @@ class MemoryBanksClient(MemoryBanks):
     async def shutdown(self) -> None:
         pass
 
-    async def list_memory_banks(self) -> List[MemoryBankDef]:
+    async def list_memory_banks(self) -> List[MemoryBankDefWithProvider]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/memory_banks/list",
@@ -53,10 +56,23 @@ class MemoryBanksClient(MemoryBanks):
             response.raise_for_status()
             return [deserialize_memory_bank_def(x) for x in response.json()]
 
+    async def register_memory_bank(
+        self, memory_bank: MemoryBankDefWithProvider
+    ) -> None:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self.base_url}/memory_banks/register",
+                json={
+                    "memory_bank": json.loads(memory_bank.json()),
+                },
+                headers={"Content-Type": "application/json"},
+            )
+            response.raise_for_status()
+
     async def get_memory_bank(
         self,
         identifier: str,
-    ) -> Optional[MemoryBankDef]:
+    ) -> Optional[MemoryBankDefWithProvider]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/memory_banks/get",

llama_stack/apis/models/client.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 import asyncio
+import json
 
 from typing import List, Optional
 
@@ -25,21 +26,32 @@ class ModelsClient(Models):
     async def shutdown(self) -> None:
         pass
 
-    async def list_models(self) -> List[ModelServingSpec]:
+    async def list_models(self) -> List[ModelDefWithProvider]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/models/list",
                 headers={"Content-Type": "application/json"},
             )
             response.raise_for_status()
-            return [ModelServingSpec(**x) for x in response.json()]
+            return [ModelDefWithProvider(**x) for x in response.json()]
 
-    async def get_model(self, core_model_id: str) -> Optional[ModelServingSpec]:
+    async def register_model(self, model: ModelDefWithProvider) -> None:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self.base_url}/models/register",
+                json={
+                    "model": json.loads(model.json()),
+                },
+                headers={"Content-Type": "application/json"},
+            )
+            response.raise_for_status()
+
+    async def get_model(self, identifier: str) -> Optional[ModelDefWithProvider]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/models/get",
                 params={
-                    "core_model_id": core_model_id,
+                    "identifier": identifier,
                 },
                 headers={"Content-Type": "application/json"},
             )
@@ -47,7 +59,7 @@ class ModelsClient(Models):
             j = response.json()
             if j is None:
                 return None
-            return ModelServingSpec(**j)
+            return ModelDefWithProvider(**j)
 
 
 async def run_main(host: str, port: int, stream: bool):

llama_stack/apis/shields/client.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 import asyncio
+import json
 
 from typing import List, Optional
 
@@ -25,16 +26,27 @@ class ShieldsClient(Shields):
     async def shutdown(self) -> None:
         pass
 
-    async def list_shields(self) -> List[ShieldSpec]:
+    async def list_shields(self) -> List[ShieldDefWithProvider]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/shields/list",
                 headers={"Content-Type": "application/json"},
             )
             response.raise_for_status()
-            return [ShieldSpec(**x) for x in response.json()]
+            return [ShieldDefWithProvider(**x) for x in response.json()]
 
-    async def get_shield(self, shield_type: str) -> Optional[ShieldSpec]:
+    async def register_shield(self, shield: ShieldDefWithProvider) -> None:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self.base_url}/shields/register",
+                json={
+                    "shield": json.loads(shield.json()),
+                },
+                headers={"Content-Type": "application/json"},
+            )
+            response.raise_for_status()
+
+    async def get_shield(self, shield_type: str) -> Optional[ShieldDefWithProvider]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/shields/get",
@@ -49,7 +61,7 @@ class ShieldsClient(Shields):
             if j is None:
                 return None
 
-            return ShieldSpec(**j)
+            return ShieldDefWithProvider(**j)
 
 
 async def run_main(host: str, port: int, stream: bool):