Skip to content

Commit 862e900

Browse files
gyliu513gyliu513
authored
Update vector_dbs to vector_stores (#316)
Fixed #315 ```console (llama-stack) (base) gualiu@gualiu-mac llama-stack-client-python % llama-stack-client vector_stores list INFO:httpx:HTTP Request: GET http://localhost:8321/v1/vector_stores "HTTP/1.1 200 OK" ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ id ┃ name ┃ provider_id ┃ provider_vector_store_id ┃ params ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ vs_8472b5ef-3855-44eb-838e-9e1fd53aa… │ test_vector_store_4041f46f-0edb-4692-… │ faiss │ vs_8472b5ef-3855-44eb-838e-9e1fd53aa… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ │ vs_b3829f97-8aa9-47b3-8606-086f69a48… │ test_vector_store_1a0fea8d-3d16-4150-… │ faiss │ vs_b3829f97-8aa9-47b3-8606-086f69a48… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ │ vs_7e16d792-431d-4c03-b6ba-e810dea29… │ live_bank │ faiss │ vs_7e16d792-431d-4c03-b6ba-e810dea29… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ │ vs_e496a34e-0d31-420e-8518-8c397ec48… │ vector_store │ faiss │ vs_e496a34e-0d31-420e-8518-8c397ec48… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ │ vs_1dde3765-dd40-47b7-bfab-70de7d334… │ vector_store │ faiss │ vs_1dde3765-dd40-47b7-bfab-70de7d334… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ │ vs_40910714-7cd5-4426-945b-2253b6c0e… │ live_bank │ faiss │ vs_40910714-7cd5-4426-945b-2253b6c0e… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ │ vs_97c7c129-18ec-4c30-a25e-e0fa93901… │ vector_store │ faiss │ vs_97c7c129-18ec-4c30-a25e-e0fa93901… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ │ vs_2a8740f2-20ae-4ea9-b45b-14fc23efc… │ live_bank │ faiss │ vs_2a8740f2-20ae-4ea9-b45b-14fc23efc… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ │ vs_ace492f1-4cd1-4ab9-9295-6ab47f01a… │ interio_bank │ faiss │ vs_ace492f1-4cd1-4ab9-9295-6ab47f01a… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ │ vs_239df24b-7fb9-4c03-bb15-d2f7da710… │ test_vector_store_22ae7bc0-a5a9-4426-… │ faiss │ vs_239df24b-7fb9-4c03-bb15-d2f7da710… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ │ vs_99fe883e-7c79-441c-9f72-b2283dc42… │ test_vector_store_d603de47-60f9-4a62-… │ faiss │ vs_99fe883e-7c79-441c-9f72-b2283dc42… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ │ vs_563486f2-f565-4bc1-af53-83b9a21cf… │ test_vector_store_1ce59bf1-ab85-470f-… │ faiss │ vs_563486f2-f565-4bc1-af53-83b9a21cf… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ │ vs_c4e8c818-a5d6-4131-b287-4cca75a35… │ test_vector_store_ba923690-cb84-4ccc-… │ faiss │ vs_c4e8c818-a5d6-4131-b287-4cca75a35… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ │ vs_0404cb9c-cad6-4a1c-92eb-43e098449… │ test_vector_store_54d918db-1713-48dc-… │ faiss │ vs_0404cb9c-cad6-4a1c-92eb-43e098449… │ embedding_dimension: '768' │ │ │ │ │ │ embedding_model: │ │ │ │ │ │ sentence-transformers/nomic-ai/nomic-… │ │ │ │ │ │ metadata: null │ │ │ │ │ │ status: completed │ │ │ │ │ │ usage_bytes: 0 │ │ │ │ │ │ │ └───────────────────────────────────────┴────────────────────────────────────────┴─────────────┴───────────────────────────────────────┴────────────────────────────────────────┘ ```
1 parent acd5e64 commit 862e900

File tree

4 files changed

+143
-120
lines changed

4 files changed

+143
-120
lines changed

‎src/llama_stack_client/lib/cli/llama_stack_client.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
from .scoring_functions import scoring_functions
2626
from .shields import shields
2727
from .toolgroups import toolgroups
28-
from .vector_dbs import vector_dbs
28+
from .vector_stores import vector_stores
2929

3030

3131
@click.group()
@@ -81,7 +81,7 @@ def llama_stack_client(ctx, endpoint: str, api_key: str, config: str | None):
8181

8282
# Register all subcommands
8383
llama_stack_client.add_command(models, "models")
84-
llama_stack_client.add_command(vector_dbs, "vector_dbs")
84+
llama_stack_client.add_command(vector_stores, "vector_stores")
8585
llama_stack_client.add_command(shields, "shields")
8686
llama_stack_client.add_command(eval_tasks, "eval_tasks")
8787
llama_stack_client.add_command(providers, "providers")

‎src/llama_stack_client/lib/cli/vector_dbs/vector_dbs.py‎

Lines changed: 0 additions & 116 deletions
This file was deleted.

src/llama_stack_client/lib/cli/vector_dbs/__init__.py renamed to src/llama_stack_client/lib/cli/vector_stores/__init__.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,6 @@
44
# This source code is licensed under the terms described in the LICENSE file in
55
# the root directory of this source tree.
66

7-
from .vector_dbs import vector_dbs
7+
from .vector_stores import vector_stores
88

9-
__all__ = ["vector_dbs"]
9+
__all__ = ["vector_stores"]

‎src/llama_stack_client/lib/cli/vector_stores/vector_stores.py‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the terms described in the LICENSE file in
5+
# the root directory of this source tree.
6+
7+
from typing import Optional
8+
9+
import click
10+
import yaml
11+
from rich.console import Console
12+
from rich.table import Table
13+
14+
from ..common.utils import handle_client_errors
15+
16+
17+
@click.group()
18+
@click.help_option("-h", "--help")
19+
def vector_stores():
20+
"""Manage vector stores."""
21+
22+
23+
@click.command("list")
24+
@click.help_option("-h", "--help")
25+
@click.pass_context
26+
@handle_client_errors("list vector stores")
27+
def list(ctx):
28+
"""Show available vector stores on distribution endpoint"""
29+
30+
client = ctx.obj["client"]
31+
console = Console()
32+
vector_stores_list_response = client.vector_stores.list()
33+
34+
if vector_stores_list_response:
35+
table = Table()
36+
# Add our specific columns
37+
table.add_column("id")
38+
table.add_column("name")
39+
table.add_column("provider_id")
40+
table.add_column("provider_vector_store_id")
41+
table.add_column("params")
42+
43+
for item in vector_stores_list_response:
44+
metadata = dict(item.metadata or {})
45+
provider_id = str(metadata.pop("provider_id", ""))
46+
provider_vector_store_id = str(
47+
metadata.pop("provider_vector_store_id", metadata.pop("provider_vector_db_id", ""))
48+
)
49+
params = {
50+
"status": item.status,
51+
"usage_bytes": item.usage_bytes,
52+
"embedding_model": metadata.pop("embedding_model", None),
53+
"embedding_dimension": metadata.pop("embedding_dimension", None),
54+
"metadata": metadata or None,
55+
}
56+
params = yaml.dump(params, default_flow_style=False)
57+
58+
table.add_row(
59+
str(item.id),
60+
str(item.name or ""),
61+
provider_id,
62+
provider_vector_store_id,
63+
params,
64+
)
65+
66+
console.print(table)
67+
68+
69+
@vector_stores.command()
70+
@click.help_option("-h", "--help")
71+
@click.argument("vector-db-id")
72+
@click.option("--provider-id", help="Provider ID for the vector db", default=None)
73+
@click.option("--provider-vector-db-id", help="Provider's vector db ID", default=None)
74+
@click.option(
75+
"--embedding-model",
76+
type=str,
77+
help="Embedding model (for vector type)",
78+
default="all-MiniLM-L6-v2",
79+
)
80+
@click.option(
81+
"--embedding-dimension",
82+
type=int,
83+
help="Embedding dimension (for vector type)",
84+
default=384,
85+
)
86+
@click.pass_context
87+
@handle_client_errors("register vector store")
88+
def register(
89+
ctx,
90+
vector_db_id: str,
91+
provider_id: Optional[str],
92+
provider_vector_db_id: Optional[str],
93+
embedding_model: Optional[str],
94+
embedding_dimension: Optional[int],
95+
):
96+
"""Create a new vector store"""
97+
client = ctx.obj["client"]
98+
99+
metadata = {"vector_db_id": vector_db_id}
100+
if provider_id:
101+
metadata["provider_id"] = provider_id
102+
if provider_vector_db_id:
103+
metadata["provider_vector_store_id"] = provider_vector_db_id
104+
if embedding_model:
105+
metadata["embedding_model"] = embedding_model
106+
if embedding_dimension is not None:
107+
metadata["embedding_dimension"] = embedding_dimension
108+
109+
response = client.vector_stores.create(name=vector_db_id, metadata=metadata)
110+
if response:
111+
click.echo(yaml.dump(response.to_dict()))
112+
113+
114+
@vector_stores.command()
115+
@click.help_option("-h", "--help")
116+
@click.argument("vector-db-id")
117+
@click.pass_context
118+
@handle_client_errors("delete vector store")
119+
def unregister(ctx, vector_db_id: str):
120+
"""Delete a vector store"""
121+
client = ctx.obj["client"]
122+
vector_store_id = vector_db_id
123+
for store in client.vector_stores.list():
124+
if store.id == vector_db_id or store.name == vector_db_id:
125+
vector_store_id = store.id
126+
break
127+
metadata = store.metadata or {}
128+
if metadata.get("vector_db_id") == vector_db_id:
129+
vector_store_id = store.id
130+
break
131+
132+
client.vector_stores.delete(vector_store_id=vector_store_id)
133+
click.echo(f"Vector store '{vector_db_id}' deleted successfully")
134+
135+
136+
# Register subcommands
137+
vector_stores.add_command(list)
138+
vector_stores.add_command(register)
139+
vector_stores.add_command(unregister)

0 commit comments

Comments
 (0)