From bfe066f9f6b422a5bb5836bce554f892a784f0ba Mon Sep 17 00:00:00 2001 From: Guangya Liu Date: Thu, 22 Jan 2026 17:13:19 -0500 Subject: [PATCH] Update vector_dbs to vector_stores --- .../lib/cli/llama_stack_client.py | 4 +- .../lib/cli/vector_dbs/vector_dbs.py | 116 --------------- .../{vector_dbs => vector_stores}/__init__.py | 4 +- .../lib/cli/vector_stores/vector_stores.py | 139 ++++++++++++++++++ 4 files changed, 143 insertions(+), 120 deletions(-) delete mode 100644 src/llama_stack_client/lib/cli/vector_dbs/vector_dbs.py rename src/llama_stack_client/lib/cli/{vector_dbs => vector_stores}/__init__.py (74%) create mode 100644 src/llama_stack_client/lib/cli/vector_stores/vector_stores.py diff --git a/src/llama_stack_client/lib/cli/llama_stack_client.py b/src/llama_stack_client/lib/cli/llama_stack_client.py index 54c46aaa..7a96e5e2 100644 --- a/src/llama_stack_client/lib/cli/llama_stack_client.py +++ b/src/llama_stack_client/lib/cli/llama_stack_client.py @@ -25,7 +25,7 @@ from .scoring_functions import scoring_functions from .shields import shields from .toolgroups import toolgroups -from .vector_dbs import vector_dbs +from .vector_stores import vector_stores @click.group() @@ -81,7 +81,7 @@ def llama_stack_client(ctx, endpoint: str, api_key: str, config: str | None): # Register all subcommands llama_stack_client.add_command(models, "models") -llama_stack_client.add_command(vector_dbs, "vector_dbs") +llama_stack_client.add_command(vector_stores, "vector_stores") llama_stack_client.add_command(shields, "shields") llama_stack_client.add_command(eval_tasks, "eval_tasks") llama_stack_client.add_command(providers, "providers") diff --git a/src/llama_stack_client/lib/cli/vector_dbs/vector_dbs.py b/src/llama_stack_client/lib/cli/vector_dbs/vector_dbs.py deleted file mode 100644 index cb196942..00000000 --- a/src/llama_stack_client/lib/cli/vector_dbs/vector_dbs.py +++ /dev/null @@ -1,116 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the terms described in the LICENSE file in -# the root directory of this source tree. - -from typing import Optional - -import click -import yaml -from rich.console import Console -from rich.table import Table - -from ..common.utils import handle_client_errors - - -@click.group() -@click.help_option("-h", "--help") -def vector_dbs(): - """Manage vector databases.""" - - -@click.command("list") -@click.help_option("-h", "--help") -@click.pass_context -@handle_client_errors("list vector dbs") -def list(ctx): - """Show available vector dbs on distribution endpoint""" - - client = ctx.obj["client"] - console = Console() - vector_dbs_list_response = client.vector_dbs.list() - - if vector_dbs_list_response: - table = Table() - # Add our specific columns - table.add_column("identifier") - table.add_column("provider_id") - table.add_column("provider_resource_id") - table.add_column("vector_db_type") - table.add_column("params") - - for item in vector_dbs_list_response: - # 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", "")) - vector_db_type = str(item_dict.pop("vector_db_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, vector_db_type, params) - - console.print(table) - - -@vector_dbs.command() -@click.help_option("-h", "--help") -@click.argument("vector-db-id") -@click.option("--provider-id", help="Provider ID for the vector db", default=None) -@click.option("--provider-vector-db-id", help="Provider's vector db ID", default=None) -@click.option( - "--embedding-model", - type=str, - help="Embedding model (for vector type)", - default="all-MiniLM-L6-v2", -) -@click.option( - "--embedding-dimension", - type=int, - help="Embedding dimension (for vector type)", - default=384, -) -@click.pass_context -@handle_client_errors("register vector db") -def register( - ctx, - vector_db_id: str, - provider_id: Optional[str], - provider_vector_db_id: Optional[str], - embedding_model: Optional[str], - embedding_dimension: Optional[int], -): - """Create a new vector db""" - client = ctx.obj["client"] - - response = client.vector_dbs.register( - vector_db_id=vector_db_id, - provider_id=provider_id, - provider_vector_db_id=provider_vector_db_id, - embedding_model=embedding_model, - embedding_dimension=embedding_dimension, - ) - if response: - click.echo(yaml.dump(response.dict())) - - -@vector_dbs.command() -@click.help_option("-h", "--help") -@click.argument("vector-db-id") -@click.pass_context -@handle_client_errors("delete vector db") -def unregister(ctx, vector_db_id: str): - """Delete a vector db""" - client = ctx.obj["client"] - client.vector_dbs.unregister(vector_db_id=vector_db_id) - click.echo(f"Vector db '{vector_db_id}' deleted successfully") - - -# Register subcommands -vector_dbs.add_command(list) -vector_dbs.add_command(register) -vector_dbs.add_command(unregister) diff --git a/src/llama_stack_client/lib/cli/vector_dbs/__init__.py b/src/llama_stack_client/lib/cli/vector_stores/__init__.py similarity index 74% rename from src/llama_stack_client/lib/cli/vector_dbs/__init__.py rename to src/llama_stack_client/lib/cli/vector_stores/__init__.py index 62e1cd65..bb43c8d1 100644 --- a/src/llama_stack_client/lib/cli/vector_dbs/__init__.py +++ b/src/llama_stack_client/lib/cli/vector_stores/__init__.py @@ -4,6 +4,6 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -from .vector_dbs import vector_dbs +from .vector_stores import vector_stores -__all__ = ["vector_dbs"] +__all__ = ["vector_stores"] diff --git a/src/llama_stack_client/lib/cli/vector_stores/vector_stores.py b/src/llama_stack_client/lib/cli/vector_stores/vector_stores.py new file mode 100644 index 00000000..128f3f2e --- /dev/null +++ b/src/llama_stack_client/lib/cli/vector_stores/vector_stores.py @@ -0,0 +1,139 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + +from typing import Optional + +import click +import yaml +from rich.console import Console +from rich.table import Table + +from ..common.utils import handle_client_errors + + +@click.group() +@click.help_option("-h", "--help") +def vector_stores(): + """Manage vector stores.""" + + +@click.command("list") +@click.help_option("-h", "--help") +@click.pass_context +@handle_client_errors("list vector stores") +def list(ctx): + """Show available vector stores on distribution endpoint""" + + client = ctx.obj["client"] + console = Console() + vector_stores_list_response = client.vector_stores.list() + + if vector_stores_list_response: + table = Table() + # Add our specific columns + table.add_column("id") + table.add_column("name") + table.add_column("provider_id") + table.add_column("provider_vector_store_id") + table.add_column("params") + + for item in vector_stores_list_response: + metadata = dict(item.metadata or {}) + provider_id = str(metadata.pop("provider_id", "")) + provider_vector_store_id = str( + metadata.pop("provider_vector_store_id", metadata.pop("provider_vector_db_id", "")) + ) + params = { + "status": item.status, + "usage_bytes": item.usage_bytes, + "embedding_model": metadata.pop("embedding_model", None), + "embedding_dimension": metadata.pop("embedding_dimension", None), + "metadata": metadata or None, + } + params = yaml.dump(params, default_flow_style=False) + + table.add_row( + str(item.id), + str(item.name or ""), + provider_id, + provider_vector_store_id, + params, + ) + + console.print(table) + + +@vector_stores.command() +@click.help_option("-h", "--help") +@click.argument("vector-db-id") +@click.option("--provider-id", help="Provider ID for the vector db", default=None) +@click.option("--provider-vector-db-id", help="Provider's vector db ID", default=None) +@click.option( + "--embedding-model", + type=str, + help="Embedding model (for vector type)", + default="all-MiniLM-L6-v2", +) +@click.option( + "--embedding-dimension", + type=int, + help="Embedding dimension (for vector type)", + default=384, +) +@click.pass_context +@handle_client_errors("register vector store") +def register( + ctx, + vector_db_id: str, + provider_id: Optional[str], + provider_vector_db_id: Optional[str], + embedding_model: Optional[str], + embedding_dimension: Optional[int], +): + """Create a new vector store""" + client = ctx.obj["client"] + + metadata = {"vector_db_id": vector_db_id} + if provider_id: + metadata["provider_id"] = provider_id + if provider_vector_db_id: + metadata["provider_vector_store_id"] = provider_vector_db_id + if embedding_model: + metadata["embedding_model"] = embedding_model + if embedding_dimension is not None: + metadata["embedding_dimension"] = embedding_dimension + + response = client.vector_stores.create(name=vector_db_id, metadata=metadata) + if response: + click.echo(yaml.dump(response.to_dict())) + + +@vector_stores.command() +@click.help_option("-h", "--help") +@click.argument("vector-db-id") +@click.pass_context +@handle_client_errors("delete vector store") +def unregister(ctx, vector_db_id: str): + """Delete a vector store""" + client = ctx.obj["client"] + vector_store_id = vector_db_id + for store in client.vector_stores.list(): + if store.id == vector_db_id or store.name == vector_db_id: + vector_store_id = store.id + break + metadata = store.metadata or {} + if metadata.get("vector_db_id") == vector_db_id: + vector_store_id = store.id + break + + client.vector_stores.delete(vector_store_id=vector_store_id) + click.echo(f"Vector store '{vector_db_id}' deleted successfully") + + +# Register subcommands +vector_stores.add_command(list) +vector_stores.add_command(register) +vector_stores.add_command(unregister)