add register_memory_bank call to direct client test

meta-llama · Dec 6, 2024 · b8ef979 · b8ef979
1 parent 88d4aa2
commit b8ef979
Showing 1 changed file with 12 additions and 0 deletions.
diff --git 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