From b8ef97990953a19c29c8d40d0bfd8751b3640e9b Mon Sep 17 00:00:00 2001
From: dltn <6599399+dltn@users.noreply.github.com>
Date: Thu, 5 Dec 2024 17:18:22 -0800
Subject: [PATCH] add register_memory_bank call to direct client test

---
 src/llama_stack_client/lib/direct/test.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/src/llama_stack_client/lib/direct/test.py b/src/llama_stack_client/lib/direct/test.py
index 82d37c8..da25be7 100644
--- a/src/llama_stack_client/lib/direct/test.py
+++ b/src/llama_stack_client/lib/direct/test.py
@@ -19,6 +19,18 @@ async def main(config_path: str):
     print("\nChat completion response:")
     print(response)
 
+    response = await client.memory_banks.register(
+        memory_bank_id="test_memory_bank",
+        params={
+            "embedding_model": "all-MiniLM-L6-v2",
+            "chunk_size_in_tokens": 512,
+            "overlap_size_in_tokens": 64,
+        },
+        provider_id="test_memory_bank",
+    )
+    print("\nRegister memory bank response:")
+    print(response)
+
 
 if __name__ == "__main__":
     import asyncio