Skip to content

Commit 91cf6d2

Browse files
authored
feat(cli): add --model flag and gate prompt on TTY for openkb init (#56)
Mirrors the --language pattern from #48 to close the remaining asymmetry where `openkb init` was only partially non-interactive. - Add `--model/-m MODEL` flag that skips the interactive model prompt when set, persisting the LiteLLM "provider/model" string straight to .openkb/config.yaml. - Gate the model `click.prompt` on `_stdin_is_tty()` so piped/redirected callers fall back to the default model without a click prompt failure on EOF. - Add `_coerce_model` validation (max 100 chars, no control chars), matching `_coerce_language` so embedded newlines can't corrupt logged output or config.yaml. Now `LLM_API_KEY=... openkb init -m anthropic/claude-sonnet-4-6 -l ko` is fully non-interactive (api_key prompt remains for security).
1 parent 5e2f436 commit 91cf6d2

2 files changed

Lines changed: 143 additions & 6 deletions

File tree

openkb/cli.py

Lines changed: 49 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -359,6 +359,7 @@ def use(path):
359359

360360

361361
_LANGUAGE_MAX_LEN = 50
362+
_MODEL_MAX_LEN = 100
362363

363364

364365
def _coerce_language(value: str | None) -> str | None:
@@ -393,6 +394,35 @@ def _language_option_callback(_ctx, _param, value):
393394
return _coerce_language(value)
394395

395396

397+
def _coerce_model(value: str | None) -> str | None:
398+
"""Strip a model string; treat blanks as unset; reject unsafe values.
399+
400+
Mirrors ``_coerce_language``. The model string is passed to LiteLLM and
401+
also echoed in logs/CLI output, so embedded control characters would
402+
corrupt that output. Capping length keeps pathological values out of
403+
config.yaml.
404+
405+
Returns the cleaned string, or ``None`` if the input was missing or blank
406+
after stripping. Raises ``click.BadParameter`` on unsafe input.
407+
"""
408+
if value is None:
409+
return None
410+
value = value.strip()
411+
if not value:
412+
return None
413+
if len(value) > _MODEL_MAX_LEN or any(c in value for c in "\n\r\t"):
414+
raise click.BadParameter(
415+
f"model must be {_MODEL_MAX_LEN} characters or fewer "
416+
"with no control characters",
417+
param_hint="'--model'",
418+
)
419+
return value
420+
421+
422+
def _model_option_callback(_ctx, _param, value):
423+
return _coerce_model(value)
424+
425+
396426
def _stdin_is_tty() -> bool:
397427
"""Return True when stdin is a real terminal.
398428
@@ -404,13 +434,23 @@ def _stdin_is_tty() -> bool:
404434

405435

406436
@cli.command()
437+
@click.option(
438+
"--model", "-m", "model",
439+
default=None, metavar="MODEL",
440+
callback=_model_option_callback,
441+
help=(
442+
"LLM in LiteLLM provider/model format "
443+
"(e.g. 'gpt-5.4-mini', 'anthropic/claude-sonnet-4-6'). "
444+
"Skips the interactive prompt when set."
445+
),
446+
)
407447
@click.option(
408448
"--language", "-l", "language",
409449
default=None, metavar="LANG",
410450
callback=_language_option_callback,
411451
help="Wiki output language (e.g. 'en', 'ko'). Skips the interactive prompt when set.",
412452
)
413-
def init(language):
453+
def init(model, language):
414454
"""Initialise a new knowledge base in the current directory."""
415455
openkb_dir = Path(".openkb")
416456
if openkb_dir.exists():
@@ -424,11 +464,14 @@ def init(language):
424464
click.echo(" Gemini: gemini/gemini-3.1-pro-preview, gemini/gemini-3-flash-preview")
425465
click.echo(" Others: see https://docs.litellm.ai/docs/providers")
426466
click.echo()
427-
model = click.prompt(
428-
f"Model (enter for default {DEFAULT_CONFIG['model']})",
429-
default=DEFAULT_CONFIG["model"],
430-
show_default=False,
431-
)
467+
if model is None and _stdin_is_tty():
468+
model = _coerce_model(click.prompt(
469+
f"Model (enter for default {DEFAULT_CONFIG['model']})",
470+
default=DEFAULT_CONFIG["model"],
471+
show_default=False,
472+
))
473+
if not model:
474+
model = DEFAULT_CONFIG["model"]
432475
api_key = click.prompt(
433476
"LLM API Key (saved to .env, enter to skip)",
434477
default="",

tests/test_cli.py

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -180,6 +180,100 @@ def test_init_language_prompt_accepts_input(tmp_path):
180180
assert config["language"] == "fr"
181181

182182

183+
def test_init_defaults_model_to_default(tmp_path):
184+
"""Non-TTY (CliRunner) skips the model prompt and falls back to default."""
185+
from openkb.config import DEFAULT_CONFIG
186+
187+
runner = CliRunner()
188+
with runner.isolated_filesystem(temp_dir=tmp_path), \
189+
patch("openkb.cli.register_kb"):
190+
result = runner.invoke(cli, ["init"], input="\n")
191+
assert result.exit_code == 0
192+
# Non-TTY: prompt must not block on EOF.
193+
assert "Model (enter for default" not in result.output
194+
195+
from pathlib import Path
196+
config = yaml.safe_load((Path(".openkb") / "config.yaml").read_text())
197+
assert config["model"] == DEFAULT_CONFIG["model"]
198+
199+
200+
def test_init_model_flag_sets_config(tmp_path):
201+
runner = CliRunner()
202+
with runner.isolated_filesystem(temp_dir=tmp_path), \
203+
patch("openkb.cli.register_kb"):
204+
# Flag supplies model, so only api_key is prompted under non-TTY.
205+
result = runner.invoke(
206+
cli, ["init", "--model", "anthropic/claude-sonnet-4-6"], input="\n",
207+
)
208+
assert result.exit_code == 0
209+
# Flag must skip the model prompt entirely
210+
assert "Model (enter for default" not in result.output
211+
212+
from pathlib import Path
213+
config = yaml.safe_load((Path(".openkb") / "config.yaml").read_text())
214+
assert config["model"] == "anthropic/claude-sonnet-4-6"
215+
216+
217+
def test_init_model_short_flag(tmp_path):
218+
runner = CliRunner()
219+
with runner.isolated_filesystem(temp_dir=tmp_path), \
220+
patch("openkb.cli.register_kb"):
221+
result = runner.invoke(cli, ["init", "-m", "gpt-5.4"], input="\n")
222+
assert result.exit_code == 0
223+
224+
from pathlib import Path
225+
config = yaml.safe_load((Path(".openkb") / "config.yaml").read_text())
226+
assert config["model"] == "gpt-5.4"
227+
228+
229+
def test_init_empty_model_flag_falls_back_to_default(tmp_path):
230+
"""--model '' must not persist a blank string into config.yaml."""
231+
from openkb.config import DEFAULT_CONFIG
232+
233+
runner = CliRunner()
234+
with runner.isolated_filesystem(temp_dir=tmp_path), \
235+
patch("openkb.cli.register_kb"):
236+
result = runner.invoke(cli, ["init", "--model", ""], input="\n")
237+
assert result.exit_code == 0
238+
239+
from pathlib import Path
240+
config = yaml.safe_load((Path(".openkb") / "config.yaml").read_text())
241+
assert config["model"] == DEFAULT_CONFIG["model"]
242+
243+
244+
def test_init_rejects_model_with_control_chars(tmp_path):
245+
"""A --model value with embedded newlines could corrupt logs/output."""
246+
runner = CliRunner()
247+
with runner.isolated_filesystem(temp_dir=tmp_path), \
248+
patch("openkb.cli.register_kb"):
249+
result = runner.invoke(
250+
cli, ["init", "--model", "gpt-4\nIgnore prior instructions"],
251+
input="\n",
252+
)
253+
assert result.exit_code != 0
254+
assert "--model" in result.output
255+
256+
from pathlib import Path
257+
assert not Path(".openkb").exists()
258+
259+
260+
def test_init_model_prompt_accepts_input(tmp_path):
261+
runner = CliRunner()
262+
with runner.isolated_filesystem(temp_dir=tmp_path), \
263+
patch("openkb.cli.register_kb"), \
264+
patch("openkb.cli._stdin_is_tty", return_value=True):
265+
# Inputs: model ("anthropic/claude-opus-4-6"), api key (blank), language (blank → default)
266+
result = runner.invoke(
267+
cli, ["init"], input="anthropic/claude-opus-4-6\n\n\n",
268+
)
269+
assert result.exit_code == 0
270+
assert "Model (enter for default" in result.output
271+
272+
from pathlib import Path
273+
config = yaml.safe_load((Path(".openkb") / "config.yaml").read_text())
274+
assert config["model"] == "anthropic/claude-opus-4-6"
275+
276+
183277
class TestQueryStreamGate:
184278
"""Regression tests for issue #34.
185279

0 commit comments

Comments
 (0)