memory bank unregister and fixes

Author: Dinesh Yeduguru

src/llama_stack_client/lib/cli/memory_banks/memory_banks.py

@@ -26,30 +26,58 @@ def list(ctx):
     client = ctx.obj["client"]
     console = Console()
     memory_banks_list_response = client.memory_banks.list()
-    headers = []
-    if memory_banks_list_response and len(memory_banks_list_response) > 0:
-        headers = sorted(memory_banks_list_response[0].__dict__.keys())
 
     if memory_banks_list_response:
         table = Table()
-        for header in headers:
-            table.add_column(header)
+        # Add our specific columns
+        table.add_column("identifier")
+        table.add_column("provider_id")
+        table.add_column("provider_resource_id")
+        table.add_column("memory_bank_type")
+        table.add_column("params")
 
         for item in memory_banks_list_response:
-            table.add_row(*[str(getattr(item, header)) for header in headers])
+            # Create a dict of all attributes
+            item_dict = item.__dict__
+
+            # Extract our main columns
+            identifier = str(item_dict.pop("identifier", ""))
+            provider_id = str(item_dict.pop("provider_id", ""))
+            provider_resource_id = str(item_dict.pop("provider_resource_id", ""))
+            memory_bank_type = str(item_dict.pop("memory_bank_type", ""))
+            # Convert remaining attributes to YAML string for params column
+            params = yaml.dump(item_dict, default_flow_style=False)
+
+            table.add_row(identifier, provider_id, provider_resource_id, memory_bank_type, params)
+
         console.print(table)
 
 
 @memory_banks.command()
-@click.option("--memory-bank-id", required=True, help="Id of the memory bank")
+@click.argument("memory-bank-id")
 @click.option("--type", type=click.Choice(["vector", "keyvalue", "keyword", "graph"]), required=True)
 @click.option("--provider-id", help="Provider ID for the memory bank", default=None)
 @click.option("--provider-memory-bank-id", help="Provider's memory bank ID", default=None)
-@click.option("--chunk-size", type=int, help="Chunk size in tokens (for vector type)", default=512)
-@click.option("--embedding-model", type=str, help="Embedding model (for vector type)", default="all-MiniLM-L6-v2")
-@click.option("--overlap-size", type=int, help="Overlap size in tokens (for vector type)", default=64)
+@click.option(
+    "--chunk-size",
+    type=int,
+    help="Chunk size in tokens (for vector type)",
+    default=512,
+)
+@click.option(
+    "--embedding-model",
+    type=str,
+    help="Embedding model (for vector type)",
+    default="all-MiniLM-L6-v2",
+)
+@click.option(
+    "--overlap-size",
+    type=int,
+    help="Overlap size in tokens (for vector type)",
+    default=64,
+)
 @click.pass_context
-def create(
+def register(
     ctx,
     memory_bank_id: str,
     type: str,
@@ -65,18 +93,24 @@ def create(
     config = None
     if type == "vector":
         config = {
-            "type": "vector",
+            "memory_bank_type": "vector",
             "chunk_size_in_tokens": chunk_size,
             "embedding_model": embedding_model,
         }
         if overlap_size:
             config["overlap_size_in_tokens"] = overlap_size
     elif type == "keyvalue":
-        config = {"type": "keyvalue"}
+        config = {"memory_bank_type": "keyvalue"}
     elif type == "keyword":
-        config = {"type": "keyword"}
+        config = {"memory_bank_type": "keyword"}
     elif type == "graph":
-        config = {"type": "graph"}
+        config = {"memory_bank_type": "graph"}
+
+    from rich import print as rprint
+    from rich.pretty import pprint
+
+    rprint("\n[bold blue]Memory Bank Configuration:[/bold blue]")
+    pprint(config, expand_all=True)
 
     response = client.memory_banks.register(
         memory_bank_id=memory_bank_id,
@@ -88,6 +122,17 @@ def create(
         click.echo(yaml.dump(response.dict()))
 
 
+@memory_banks.command()
+@click.argument("memory-bank-id")
+@click.pass_context
+def unregister(ctx, memory_bank_id: str):
+    """Delete a memory bank"""
+    client = ctx.obj["client"]
+    client.memory_banks.unregister(memory_bank_id=memory_bank_id)
+    click.echo(f"Memory bank '{memory_bank_id}' deleted successfully")
+
+
 # Register subcommands
 memory_banks.add_command(list)
-memory_banks.add_command(create)
+memory_banks.add_command(register)
+memory_banks.add_command(unregister)

src/llama_stack_client/lib/cli/models/models.py

@@ -88,37 +88,14 @@ def register_model(
         click.echo(f"Failed to register model: {str(e)}")
 
 
-@click.command(name="update", help="Update an existing model at distribution endpoint")
-@click.argument("model_id")
-@click.option("--provider-id", help="Provider ID for the model", default=None)
-@click.option("--provider-model-id", help="Provider's model ID", default=None)
-@click.option("--metadata", help="JSON metadata for the model", default=None)
-@click.pass_context
-def update_model(
-    ctx, model_id: str, provider_id: Optional[str], provider_model_id: Optional[str], metadata: Optional[str]
-):
-    """Update an existing model at distribution endpoint"""
-    client = ctx.obj["client"]
-
-    try:
-        response = client.models.update(
-            model_id=model_id, provider_id=provider_id, provider_model_id=provider_model_id, metadata=metadata
-        )
-        if response:
-            click.echo(f"Successfully updated model {model_id}")
-    except Exception as e:
-        click.echo(f"Failed to update model: {str(e)}")
-
-
-@click.command(name="delete", help="Delete a model from distribution endpoint")
+@click.command(name="unregister", help="Unregister a model from distribution endpoint")
 @click.argument("model_id")
 @click.pass_context
-def delete_model(ctx, model_id: str):
-    """Delete a model from distribution endpoint"""
+def unregister_model(ctx, model_id: str):
     client = ctx.obj["client"]
 
     try:
-        response = client.models.delete(model_id=model_id)
+        response = client.models.unregister(model_id=model_id)
         if response:
             click.echo(f"Successfully deleted model {model_id}")
     except Exception as e:
@@ -129,5 +106,4 @@ def delete_model(ctx, model_id: str):
 models.add_command(list_models)
 models.add_command(get_model)
 models.add_command(register_model)
-models.add_command(update_model)
-models.add_command(delete_model)
+models.add_command(unregister_model)