Skip to content

Commit a2f2e25

Browse files
committed
Adding API key support
Summary: Test Plan: Config client: ``` llama-stack-client configure --endpoint=https://api.hostname.com --api-key='_API_KEY_' Done! You can now use the Llama Stack Client CLI with endpoint https://api.hostname.com ``` Try listing the models ``` llama-stack-client models list ┏━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━━┓ ┃ identifier ┃ provider_id ┃ provider_resource_id ┃ metadata ┃ model_type ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━━┩ │ llama3.1-8b-instruct │ meta-llama │ llama3.1-8b-instruct │ {} │ llm │ │ llama3.3-70b-instruct │ meta-llama │ llama3.3-70b-instruct │ {} │ llm │ │ llama3.4-17b-instruct │ meta-llama │ llama3.4-17b-instruct │ {} │ llm │ └───────────────────────┴─────────────┴───────────────────────┴──────────┴────────────┘ ```
1 parent 6b6be35 commit a2f2e25

File tree

2 files changed

+38
-27
lines changed

2 files changed

+38
-27
lines changed

src/llama_stack_client/lib/cli/configure.py

Lines changed: 26 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -24,42 +24,42 @@ def get_config():
2424

2525

2626
@click.command()
27-
@click.option("--host", type=str, help="Llama Stack distribution host")
28-
@click.option("--port", type=str, help="Llama Stack distribution port number")
29-
@click.option("--endpoint", type=str, help="Llama Stack distribution endpoint")
30-
def configure(host: str | None, port: str | None, endpoint: str | None):
27+
@click.option("--endpoint", type=str, help="Llama Stack distribution endpoint", default="")
28+
@click.option("--api-key", type=str, help="Llama Stack distribution API key", default="")
29+
def configure(endpoint: str | None, api_key: str | None):
3130
"""Configure Llama Stack Client CLI"""
3231
os.makedirs(LLAMA_STACK_CLIENT_CONFIG_DIR, exist_ok=True)
3332
config_path = get_config_file_path()
3433

35-
if endpoint:
34+
if endpoint != "":
3635
final_endpoint = endpoint
3736
else:
38-
if host and port:
39-
final_endpoint = f"http://{host}:{port}"
40-
else:
41-
host = prompt(
42-
"> Enter the host name of the Llama Stack distribution server: ",
43-
validator=Validator.from_callable(
44-
lambda x: len(x) > 0,
45-
error_message="Host cannot be empty, please enter a valid host",
46-
),
47-
)
48-
port = prompt(
49-
"> Enter the port number of the Llama Stack distribution server: ",
50-
validator=Validator.from_callable(
51-
lambda x: x.isdigit(),
52-
error_message="Please enter a valid port number",
53-
),
54-
)
55-
final_endpoint = f"http://{host}:{port}"
37+
final_endpoint = prompt(
38+
"> Enter the endpoint of the Llama Stack distribution server: ",
39+
validator=Validator.from_callable(
40+
lambda x: len(x) > 0,
41+
error_message="Endpoint cannot be empty, please enter a valid endpoint",
42+
),
43+
)
44+
45+
if api_key != "":
46+
final_api_key = api_key
47+
else:
48+
final_api_key = prompt(
49+
"> Enter the API key (leave empty if no key is needed): ",
50+
)
51+
52+
# Prepare config dict before writing it
53+
config_dict = {
54+
"endpoint": final_endpoint,
55+
}
56+
if final_api_key != "":
57+
config_dict["api_key"] = final_api_key
5658

5759
with open(config_path, "w") as f:
5860
f.write(
5961
yaml.dump(
60-
{
61-
"endpoint": final_endpoint,
62-
},
62+
config_dict,
6363
sort_keys=True,
6464
)
6565
)

src/llama_stack_client/lib/cli/llama_stack_client.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,12 @@
3535
@click.option(
3636
"--endpoint", type=str, help="Llama Stack distribution endpoint", default=""
3737
)
38+
@click.option(
39+
"--api-key", type=str, help="Llama Stack distribution API key", default=""
40+
)
3841
@click.option("--config", type=str, help="Path to config file", default=None)
3942
@click.pass_context
40-
def cli(ctx, endpoint: str, config: str | None):
43+
def cli(ctx, endpoint: str, api_key: str, config: str | None):
4144
"""Welcome to the LlamaStackClient CLI"""
4245
ctx.ensure_object(dict)
4346

@@ -55,20 +58,28 @@ def cli(ctx, endpoint: str, config: str | None):
5558
with open(config, "r") as f:
5659
config_dict = yaml.safe_load(f)
5760
endpoint = config_dict.get("endpoint", endpoint)
61+
api_key = config_dict.get("api_key", "")
5862
except Exception as e:
5963
click.echo(f"Error loading config from {config}: {str(e)}", err=True)
6064
click.echo("Falling back to HTTP client with endpoint", err=True)
6165

6266
if endpoint == "":
6367
endpoint = "http://localhost:8321"
6468

69+
default_headers = {}
70+
if api_key != "":
71+
default_headers = {
72+
"Authorization": f"Bearer {api_key}",
73+
}
74+
6575
client = LlamaStackClient(
6676
base_url=endpoint,
6777
provider_data={
6878
"fireworks_api_key": os.environ.get("FIREWORKS_API_KEY", ""),
6979
"together_api_key": os.environ.get("TOGETHER_API_KEY", ""),
7080
"openai_api_key": os.environ.get("OPENAI_API_KEY", ""),
7181
},
82+
default_headers=default_headers,
7283
)
7384
ctx.obj = {"client": client}
7485

0 commit comments

Comments
 (0)