diff --git a/.env.example b/.env.example index bda78e3d..60cbcf32 100644 --- a/.env.example +++ b/.env.example @@ -4,3 +4,11 @@ # Gemini: LLM_API_KEY=AIza... # DeepSeek: LLM_API_KEY=sk-... LLM_API_KEY=your-key-here + +# Optional: API base URL override for self-hosted / proxied providers. +# `openkb init` writes the right one for the chosen model automatically. +# You can also set these by hand — LiteLLM reads them per provider: +# OPENAI_API_BASE=http://localhost:11434/v1 # Ollama, vLLM, LM Studio, ... +# ANTHROPIC_API_BASE=https://your-proxy.example.com +# OPENROUTER_API_BASE=https://openrouter.ai/api/v1 +# Omit for the official endpoint of each provider. diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index af3ac1bd..8c45057e 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -1942,7 +1942,7 @@ async def compile_short_doc( source_path: Path, kb_dir: Path, model: str, - max_concurrency: int = DEFAULT_COMPILE_CONCURRENCY, + max_concurrency: int | None = None, ) -> None: """Compile a short document using a multi-step LLM pipeline with caching. @@ -1956,6 +1956,10 @@ async def compile_short_doc( language: str = config.get("language", "en") entity_types = resolve_entity_types(config) + # Resolve concurrency: explicit param > config > hard-coded default. + if max_concurrency is None: + max_concurrency = config.get("compile_concurrency", DEFAULT_COMPILE_CONCURRENCY) + wiki_dir = kb_dir / "wiki" schema_md = get_agents_md(wiki_dir) content = source_path.read_text(encoding="utf-8") @@ -2005,7 +2009,7 @@ async def compile_long_doc( kb_dir: Path, model: str, doc_description: str = "", - max_concurrency: int = DEFAULT_COMPILE_CONCURRENCY, + max_concurrency: int | None = None, ) -> None: """Compile a long (PageIndex) document's concepts and index. @@ -2019,6 +2023,10 @@ async def compile_long_doc( language: str = config.get("language", "en") entity_types = resolve_entity_types(config) + # Resolve concurrency: explicit param > config > hard-coded default. + if max_concurrency is None: + max_concurrency = config.get("compile_concurrency", DEFAULT_COMPILE_CONCURRENCY) + wiki_dir = kb_dir / "wiki" schema_md = get_agents_md(wiki_dir) summary_content = summary_path.read_text(encoding="utf-8") diff --git a/openkb/cli.py b/openkb/cli.py index a34a0b95..924b7998 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -61,6 +61,7 @@ def filter(self, record: logging.LogRecord) -> bool: "OPENAI_API_KEY", "ANTHROPIC_API_KEY", "GEMINI_API_KEY", "DEEPSEEK_API_KEY", "MISTRAL_API_KEY", "MOONSHOT_API_KEY", "ZHIPUAI_API_KEY", "DASHSCOPE_API_KEY", + "MINIMAX_API_KEY", ) # Providers that authenticate via OAuth device flow (subscription login @@ -68,6 +69,113 @@ def filter(self, record: logging.LogRecord) -> bool: # missing-key warning would be a false alarm for them. _OAUTH_PROVIDERS = {"chatgpt", "github_copilot"} +# Public providers with well-known official endpoints — for these the user +# never needs to set a base URL, so ``openkb init`` skips the prompt. +# Anything outside this set (``ollama/``, ``vllm/``, ``openrouter/``, +# ``custom/``, ...) is presumed self-hosted / proxied and triggers the +# base-URL prompt. +# +# ``minimax`` is a special case: it ships two regional endpoints (global / +# China). It lives here so the generic base-URL prompt doesn't fire, but +# ``openkb init`` runs a dedicated region picker for it instead. +_KNOWN_PUBLIC_PROVIDERS: frozenset[str] = frozenset({ + "openai", "anthropic", "gemini", "google", "deepseek", "mistral", + "moonshot", "zhipuai", "dashscope", "minimax", +}) + +# MiniMax regional endpoints. Both expose an OpenAI-compatible /v1 API +# under the same ``minimax/`` LiteLLM provider prefix; the base URL is +# the only differentiator. Defaults to global when no choice is made. +_MINIMAX_GLOBAL_URL = "https://api.minimax.io/v1" +_MINIMAX_CHINA_URL = "https://api.minimaxi.com/v1" + +# Maps each LiteLLM provider prefix to the env var LiteLLM reads for +# its API key. ``openkb init`` uses this to write the *right* variable +# into .env — not the generic ``LLM_API_KEY`` — so the file matches the +# provider the user just picked. A value of ``None`` means the provider +# runs locally and doesn't need a key (ollama, vllm by default). +# +# Sources (LiteLLM docs, verified): +# OPENAI_API_KEY docs.litellm.ai/docs/set_keys +# ANTHROPIC_API_KEY docs.litellm.ai/docs/providers/anthropic +# GEMINI_API_KEY docs.litellm.ai/docs/providers/gemini +# DEEPSEEK_API_KEY docs.litellm.ai/docs/providers/deepseek +# MISTRAL_API_KEY docs.litellm.ai/docs/providers/mistral +# MOONSHOT_API_KEY docs.litellm.ai/docs/providers/moonshot +# DASHSCOPE_API_KEY docs.litellm.ai/docs/providers/dashscope +# OPENROUTER_API_KEY docs.litellm.ai/docs/providers/openrouter +# MINIMAX_API_KEY user-specified +# ZHIPUAI_API_KEY the LiteLLM provider was renamed to "Z.AI" but +# still uses the ``zhipuai/`` prefix, so the env +# var is unchanged. +_PROVIDER_KEY_ENV: dict[str, str | None] = { + "openai": "OPENAI_API_KEY", + "anthropic": "ANTHROPIC_API_KEY", + "gemini": "GEMINI_API_KEY", + "google": "GOOGLE_API_KEY", + "deepseek": "DEEPSEEK_API_KEY", + "mistral": "MISTRAL_API_KEY", + "moonshot": "MOONSHOT_API_KEY", + "zhipuai": "ZHIPUAI_API_KEY", + "dashscope": "DASHSCOPE_API_KEY", + "minimax": "MINIMAX_API_KEY", + "openrouter": "OPENROUTER_API_KEY", + "ollama": None, # local server, no key + "vllm": "HOSTED_VLLM_API_KEY", # optional per LiteLLM docs +} + + +def _key_env_for_provider(provider: str | None) -> str | None: + """Return the LiteLLM ``*_API_KEY`` env var for ``provider``. + + Returns ``None`` for providers that don't use a key (ollama, vllm) + AND for unknown providers — the latter is intentionally treated as + "no canonical name known", so the caller can decide whether to fall + back to the generic ``LLM_API_KEY``. + """ + if provider is None: + return None + if provider in _PROVIDER_KEY_ENV: + return _PROVIDER_KEY_ENV[provider] + return None + +# LiteLLM reads these per-provider env vars to override the base URL. +# Used by ``openkb init`` to map a user-supplied base URL into the right +# ``*_API_BASE`` key in the KB's .env. LiteLLM also accepts ``api_base=`` +# per-call and ``litellm.api_base`` globally, but the env-var route is +# provider-agnostic and survives model switches without code changes. +_PROVIDER_TO_BASE_ENV: dict[str, str] = { + "openai": "OPENAI_API_BASE", + "anthropic": "ANTHROPIC_API_BASE", + "gemini": "GEMINI_API_BASE", + "google": "GOOGLE_API_BASE", + "deepseek": "DEEPSEEK_API_BASE", + "mistral": "MISTRAL_API_BASE", + "moonshot": "MOONSHOT_API_BASE", + "zhipuai": "ZHIPUAI_API_BASE", + "dashscope": "DASHSCOPE_API_BASE", + # Common proxies / aggregators. Users with truly custom providers can + # always edit .env by hand. + "openrouter": "OPENROUTER_API_BASE", + "ollama": "OLLAMA_API_BASE", + "vllm": "OPENAI_API_BASE", # vLLM exposes an OpenAI-compatible endpoint + "minimax": "MINIMAX_API_BASE", +} + + +def _base_url_env_for_provider(provider: str | None) -> str | None: + """Return the LiteLLM ``*_API_BASE`` env var name for ``provider``. + + Falls back to ``OPENAI_API_BASE`` for unknown prefixes — most local / + proxy servers (vLLM, LM Studio, xinference, etc.) expose an + OpenAI-compatible endpoint, so this covers the common case. + """ + if not provider: + return None + if provider in _PROVIDER_TO_BASE_ENV: + return _PROVIDER_TO_BASE_ENV[provider] + return "OPENAI_API_BASE" + def _extract_provider(model: str) -> str | None: """Extract the LiteLLM provider name from a model string. @@ -106,10 +214,11 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None: if global_env.exists(): load_dotenv(global_env, override=False) - api_key = os.environ.get("LLM_API_KEY", "") - - # Try to resolve the active provider, extra headers, and request timeout - # from the KB config + # Resolve the active provider first — its key env var (e.g. + # ``OPENAI_API_KEY``) takes priority over the generic ``LLM_API_KEY`` + # so users who set the provider-specific name directly (the format + # ``openkb init`` now writes to .env) get picked up without having to + # also export a generic catch-all. provider: str | None = None extra_headers: dict[str, str] = {} timeout: float | None = None @@ -124,37 +233,57 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None: set_extra_headers(extra_headers) set_timeout(timeout) + provider_key_env = _key_env_for_provider(provider) + api_key = "" + if provider_key_env: + api_key = os.environ.get(provider_key_env, "").strip() + if not api_key: + api_key = os.environ.get("LLM_API_KEY", "").strip() + if not api_key: - # Check if any provider key is already set. OAuth-based providers - # (ChatGPT subscription, GitHub Copilot) don't use API keys at all, - # so the warning is skipped for them. - check_keys = ( - (f"{provider.upper()}_API_KEY",) if provider - else _KNOWN_PROVIDER_KEYS + # No key found under either the provider-specific name or the + # generic ``LLM_API_KEY``. OAuth-based providers (ChatGPT + # subscription, GitHub Copilot) don't use API keys at all, so the + # warning is skipped for them. Keyless self-hosted providers + # (ollama, vllm) likewise don't need one. + keyless = provider is not None and ( + provider in _OAUTH_PROVIDERS + or _PROVIDER_KEY_ENV.get(provider) is None ) - has_key = any(os.environ.get(k) for k in check_keys) - if not has_key and provider not in _OAUTH_PROVIDERS: + if not keyless: + key_hint = provider_key_env or "LLM_API_KEY" click.echo( "Warning: No LLM API key found. Set one of:\n" - f" 1. {kb_dir / '.env' if kb_dir else '/.env'} — LLM_API_KEY=sk-...\n" - f" 2. {GLOBAL_CONFIG_DIR / '.env'} — LLM_API_KEY=sk-...\n" - " 3. Export LLM_API_KEY in your shell profile" + f" 1. {kb_dir / '.env' if kb_dir else '/.env'} — {key_hint}=...\n" + f" 2. {GLOBAL_CONFIG_DIR / '.env'} — {key_hint}=...\n" + f" 3. Export {key_hint} in your shell profile" ) else: litellm.api_key = api_key - # Dynamically set the provider-specific env var when possible - if provider: - provider_env = f"{provider.upper()}_API_KEY" - if not os.environ.get(provider_env): - os.environ[provider_env] = api_key - - # Fallback: also set common provider keys so multi-provider - # configs (e.g. PageIndex Cloud) still work + # Propagate the key to every provider env var we know about. + # This keeps multi-provider setups (e.g. PageIndex Cloud, agent + # calls that use a different provider than compile) working + # without requiring the user to duplicate the key in .env. + if provider_key_env and not os.environ.get(provider_key_env): + os.environ[provider_key_env] = api_key for env_var in _KNOWN_PROVIDER_KEYS: if not os.environ.get(env_var): os.environ[env_var] = api_key + # Base URL: pick up the provider-specific *_API_BASE env var (written + # by `openkb init` for self-hosted / proxied providers) and apply it + # to litellm.api_base so LiteLLM uses it on every request. LiteLLM + # already reads e.g. OPENAI_API_BASE natively for some providers, but + # setting litellm.api_base makes the override reliable across + # providers and call paths. + base_env = _base_url_env_for_provider(provider) + base_url = "" + if base_env: + base_url = os.environ.get(base_env, "").strip() + if base_url: + litellm.api_base = base_url + # Supported document extensions for the `add` command SUPPORTED_EXTENSIONS = { ".pdf", ".md", ".markdown", ".docx", ".pptx", ".xlsx", ".xls", @@ -524,6 +653,149 @@ def _model_option_callback(_ctx, _param, value): return _coerce_model(value) +_BASE_URL_MAX_LEN = 2048 + + +def _coerce_base_url(value: str | None) -> str | None: + """Strip a base URL; treat blanks as unset; reject unsafe values. + + Mirrors ``_coerce_model``. The URL is written verbatim to ``.env`` and + may be echoed in CLI output, so embedded control characters would + corrupt that file / output and are rejected. Capping length keeps + pathological values out of the .env file. + """ + if value is None: + return None + value = value.strip() + if not value: + return None + if len(value) > _BASE_URL_MAX_LEN or any(c in value for c in "\n\r\t"): + raise click.BadParameter( + f"base URL must be {_BASE_URL_MAX_LEN} characters or fewer " + "with no control characters", + param_hint="'--base-url'", + ) + return value + + +def _base_url_option_callback(_ctx, _param, value): + return _coerce_base_url(value) + + +def _build_env_content( + env_writes: dict[str, str], provider: str | None, +) -> str: + """Build the KB-local .env content. + + Always emits the file — even when the user skipped both the API-key + prompt and the base-URL prompt — so that the user has a discoverable + place to drop credentials later. + + Variable naming follows the actual LiteLLM provider (e.g. + ``OPENAI_API_KEY`` for OpenAI, ``MINIMAX_API_KEY`` for MiniMax, + ``ANTHROPIC_API_BASE`` for Anthropic), not a generic catch-all — so + the file reads naturally to someone familiar with the provider. For + unknown / local providers, fall back to ``LLM_API_KEY``. + + ``env_writes`` maps env-var name → value for fields that should be + active (uncommented). Any field not present in ``env_writes`` is + rendered as a commented placeholder. + """ + key_env = _key_env_for_provider(provider) + # Unknown provider (custom/self-hosted) OR no provider context at + # all → fall back to LLM_API_KEY so the value still propagates + # through _setup_llm_key. Keyless providers are detected below by + # checking ``_PROVIDER_KEY_ENV`` directly. + if key_env is None: + key_env = "LLM_API_KEY" + # Truly keyless provider (ollama, vllm) → skip the key section + # entirely; emitting a placeholder would mislead the user. + skip_key_section = ( + provider is not None + and provider in _PROVIDER_KEY_ENV + and _PROVIDER_KEY_ENV[provider] is None + ) + + lines: list[str] = [ + "# OpenKB environment configuration", + "# Generated by `openkb init` — edit as needed. See .env.example", + "# for the full list of supported variables.", + "", + ] + if not skip_key_section: + key_label = key_env or "LLM_API_KEY" + lines += [ + f"# {key_label} — used by LiteLLM to authenticate to " + f"{provider or 'your provider'}.", + "# Uncomment and paste your key below.", + ] + if key_env and key_env in env_writes: + lines.append(f"{key_env}={env_writes[key_env]}") + elif "LLM_API_KEY" in env_writes and key_env != "LLM_API_KEY": + # Caller stored the key under the generic name but the + # active env var should be the provider-specific one — emit + # the generic line commented so the user can move it. + lines.append(f"# {key_env}=sk-...") + lines.append(f"LLM_API_KEY={env_writes['LLM_API_KEY']}") + elif "LLM_API_KEY" in env_writes: + lines.append(f"LLM_API_KEY={env_writes['LLM_API_KEY']}") + else: + lines.append(f"# {key_env}=sk-...") + lines.append("") + + base_env_var = _base_url_env_for_provider(provider) + if base_env_var is not None: + if base_env_var in env_writes: + lines.append( + f"{base_env_var}={env_writes[base_env_var]}" + ) + else: + label = ( + "API base URL" + if provider is None or provider not in _KNOWN_PUBLIC_PROVIDERS + else f"{provider} endpoint" + ) + lines += [ + f"# Optional: {label} override.", + f"# Uncomment to point at a self-hosted / proxied server.", + f"# {base_env_var}=https://your-endpoint/v1", + ] + return "\n".join(lines) + "\n" + + +def _prompt_minimax_region() -> str: + """Prompt the user to pick a MiniMax regional endpoint. + + MiniMax ships two endpoints under the same ``minimax/`` LiteLLM + provider prefix — global and China — and the only way to disambiguate + them is the base URL. Returns one of ``_MINIMAX_GLOBAL_URL`` or + ``_MINIMAX_CHINA_URL``. Accepts ``1``/``global`` (default) or + ``2``/``china``; anything else re-prompts so a typo never silently + routes the user to the wrong region. + + The output uses blank lines + a clear heading so the picker can't be + mistaken for a continuation of the model / API-key prompts around it + (which would silently route the user to whichever default happens to + win — usually the wrong region). + """ + click.echo() + click.echo("── MiniMax region ─────────────────────────────────") + click.echo("MiniMax has two regional endpoints under the same") + click.echo("`minimax/` LiteLLM prefix — pick one:") + click.echo(f" [1] Global ({_MINIMAX_GLOBAL_URL}) [default]") + click.echo(f" [2] China ({_MINIMAX_CHINA_URL})") + click.echo("──────────────────────────────────────────────────") + while True: + choice = click.prompt( + "Region (1=Global, 2=China)", default="1", show_default=False, + ).strip().lower() + if choice in ("", "1", "global"): + return _MINIMAX_GLOBAL_URL + if choice in ("2", "china"): + return _MINIMAX_CHINA_URL + click.echo(f"Unknown choice {choice!r}; please enter 1 or 2.") + + def _stdin_is_tty() -> bool: """Return True when stdin is a real terminal. @@ -551,7 +823,19 @@ def _stdin_is_tty() -> bool: callback=_language_option_callback, help="Wiki output language (e.g. 'en', 'ko'). Skips the interactive prompt when set.", ) -def init(model, language): +@click.option( + "--base-url", "-u", "base_url", + default=None, metavar="URL", + callback=_base_url_option_callback, + help=( + "LLM API base URL (for self-hosted / proxied providers, e.g. " + "'http://localhost:11434' for Ollama). When the chosen model is a " + "public provider (OpenAI, Anthropic, Gemini, DeepSeek, ...) the " + "interactive prompt is skipped automatically. Stored in .env as " + "the provider-specific *_API_BASE variable." + ), +) +def init(model, language, base_url): """Initialise a new knowledge base in the current directory.""" openkb_dir = Path(".openkb") if openkb_dir.exists(): @@ -564,6 +848,7 @@ def init(model, language): click.echo(" Anthropic: anthropic/claude-sonnet-4-6, anthropic/claude-opus-4-6") click.echo(" Gemini: gemini/gemini-3.1-pro-preview, gemini/gemini-3-flash-preview") click.echo(" DeepSeek: deepseek/deepseek-v4-flash, deepseek/deepseek-v4-pro") + click.echo(" MiniMax: minimax/MiniMax-M2.7, minimax/MiniMax-M3") click.echo(" Others: see https://docs.litellm.ai/docs/providers") click.echo() if model is None and _stdin_is_tty(): @@ -574,6 +859,27 @@ def init(model, language): )) if not model: model = DEFAULT_CONFIG["model"] + # Only ask for a base URL when the chosen model isn't a known public + # provider (i.e. it's self-hosted, proxied, or otherwise needs a + # non-default endpoint). The --base-url flag overrides this gate. + provider = _extract_provider(model) + # MiniMax is a known public provider but ships two regional endpoints + # under the same prefix — the only differentiator is the base URL, so + # we run a dedicated region picker instead of the generic prompt. + # Non-TTY (scripted init) falls back to global silently; users who + # need China there can pass --base-url explicitly. + if base_url is None and provider == "minimax" and _stdin_is_tty(): + base_url = _coerce_base_url(_prompt_minimax_region()) + elif base_url is None and provider == "minimax": + base_url = _MINIMAX_GLOBAL_URL + if base_url is None and _stdin_is_tty() and provider not in _KNOWN_PUBLIC_PROVIDERS: + base_url = _coerce_base_url(click.prompt( + "API base URL (for self-hosted / proxied providers, enter to skip)", + default="", + show_default=False, + )) + if not base_url: + base_url = None api_key = click.prompt( "LLM API Key (saved to .env, enter to skip)", default="", @@ -606,24 +912,57 @@ def init(model, language): "model": model, "language": language, "pageindex_threshold": DEFAULT_CONFIG["pageindex_threshold"], + "eval_concurrency": DEFAULT_CONFIG["eval_concurrency"], + "compile_concurrency": DEFAULT_CONFIG["compile_concurrency"], } save_config(openkb_dir / "config.yaml", config) atomic_write_json(openkb_dir / "hashes.json", {}) - # Write API key to KB-local .env (0600) if the user provided one + # Write the KB-local .env (0600). The file is always created — even + # when the user skipped both the API-key prompt and the base-URL + # prompt — so there's a discoverable place to drop credentials later. + # Missing fields are written as commented placeholders naming the + # right env-var for the chosen provider. + env_writes: dict[str, str] = {} if api_key: - env_path = Path(".env") - if env_path.exists(): - click.echo(".env already exists, skipping write. Add LLM_API_KEY manually if needed.") + # Write under the provider-specific env var (e.g. OPENAI_API_KEY) + # so the file matches what LiteLLM actually reads; fall back to + # the generic name for unknown / keyless providers. + key_env = _key_env_for_provider(provider) or "LLM_API_KEY" + env_writes[key_env] = api_key + if base_url: + base_env_var = _base_url_env_for_provider(provider) + if base_env_var: + env_writes[base_env_var] = base_url + env_path = Path(".env") + if env_path.exists(): + if env_writes: + click.echo( + ".env already exists, skipping write. Add the missing " + "entries manually if needed: " + ", ".join(env_writes), + ) + else: + env_path.write_text( + _build_env_content(env_writes, provider), + encoding="utf-8", + ) + os.chmod(env_path, 0o600) + if env_writes: + click.echo("Saved to .env: " + ", ".join(env_writes)) else: - env_path.write_text(f"LLM_API_KEY={api_key}\n", encoding="utf-8") - os.chmod(env_path, 0o600) - click.echo("Saved LLM API key to .env.") + click.echo( + "Created .env with commented placeholders — " + "fill in your API key before running compile." + ) # Register this KB in the global config register_kb(Path.cwd()) click.echo("Knowledge base initialized.") + click.echo( + f" • Review .env and {openkb_dir / 'config.yaml'} if anything " + "needs adjusting." + ) @cli.command() diff --git a/openkb/config.py b/openkb/config.py index d5489688..04c4b728 100644 --- a/openkb/config.py +++ b/openkb/config.py @@ -17,6 +17,8 @@ "model": "gpt-5.4-mini", "language": "en", "pageindex_threshold": 20, + "eval_concurrency": 8, + "compile_concurrency": 5, } # Default entity-type vocabulary. Overridable per-KB via the optional diff --git a/openkb/skill/evaluator.py b/openkb/skill/evaluator.py index 223966cb..906641ca 100644 --- a/openkb/skill/evaluator.py +++ b/openkb/skill/evaluator.py @@ -416,12 +416,18 @@ async def run_eval( content = _skill_content_block(skill_dir) result = EvalResult(prompts=eval_set) + # Concurrency: fall back to config file, then hard-coded default. + from openkb.config import load_config + openkb_dir = skill_dir.parent.parent.parent / ".openkb" # /.openkb + config = load_config(openkb_dir / "config.yaml") + eval_concurrency = config.get("eval_concurrency", EVAL_CONCURRENCY) + # Run grading concurrently. Each prompt is independent — graders read # the same `desc`/`content` strings and produce results that are then # appended to `result` in eval_set order below, so concurrent # execution is correctness-preserving. A semaphore caps simultaneous # LLM calls to avoid hitting provider rate limits. - sem = asyncio.Semaphore(EVAL_CONCURRENCY) + sem = asyncio.Semaphore(eval_concurrency) async def _trigger(p: EvalPrompt) -> Literal["trigger", "no-trigger"]: async with sem: diff --git a/tests/test_cli.py b/tests/test_cli.py index 3f727138..0f8343e6 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -276,6 +276,540 @@ def test_init_model_prompt_accepts_input(tmp_path): assert config["model"] == "anthropic/claude-opus-4-6" +# --------------------------------------------------------------------------- +# Base URL prompt + .env wiring +# --------------------------------------------------------------------------- + + +def test_init_public_provider_skips_base_url_prompt(tmp_path): + """OpenAI / Anthropic / etc. use official endpoints — no prompt.""" + from openkb.cli import _KNOWN_PUBLIC_PROVIDERS # noqa: F401 (sanity) + + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + # Inputs: model (gpt-5.4), api key (blank), language (blank) + result = runner.invoke(cli, ["init"], input="gpt-5.4\n\n\n") + assert result.exit_code == 0, result.output + # The base-URL prompt must NOT appear for a public provider. + assert "API base URL" not in result.output + + +def test_init_custom_provider_prompts_for_base_url(tmp_path): + """A non-public provider (e.g. custom/...) must trigger the prompt.""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + # Inputs: model (custom/my-model), base url, api key (blank), language (blank) + result = runner.invoke( + cli, ["init"], + input="custom/my-model\nhttp://localhost:8080/v1\n\n\n", + ) + assert result.exit_code == 0, result.output + assert "API base URL" in result.output + + from pathlib import Path + env_content = Path(".env").read_text() + # custom/ is unknown → falls back to OPENAI_API_BASE (most + # proxies are OAI-compatible) and the generic LLM_API_KEY. + assert "OPENAI_API_BASE=http://localhost:8080/v1" in env_content + # User skipped the key — it must appear only as a COMMENTED + # placeholder, never as an active assignment. + assert "# LLM_API_KEY=" in env_content + for line in env_content.splitlines(): + stripped = line.lstrip() + assert not stripped.startswith("LLM_API_KEY="), ( + f"LLM_API_KEY must not be active when user skipped: {line!r}" + ) + + +def test_init_ollama_provider_no_key_section(tmp_path): + """Ollama runs locally and doesn't take an API key — .env must not + mislead the user with a placeholder. + """ + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + result = runner.invoke( + cli, ["init"], + input="ollama/llama3\nhttp://localhost:11434\n\n\n", + ) + assert result.exit_code == 0, result.output + + from pathlib import Path + env_content = Path(".env").read_text() + assert "OLLAMA_API_BASE=http://localhost:11434" in env_content + # No key section at all — ollama doesn't need one. + assert "_API_KEY=" not in env_content + + +def test_init_base_url_flag_writes_env(tmp_path): + """--base-url on the CLI sets the URL without prompting.""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + result = runner.invoke( + cli, [ + "init", + "--model", "openai/gpt-5.4-mini", + "--base-url", "https://proxy.example.com/v1", + ], + input="\n\n", # api key, language + ) + assert result.exit_code == 0, result.output + # Public provider but --base-url forced it: prompt should NOT fire. + assert "API base URL" not in result.output + + from pathlib import Path + env_content = Path(".env").read_text() + # openai/ → OPENAI_API_BASE. + assert "OPENAI_API_BASE=https://proxy.example.com/v1" in env_content + + +def test_init_base_url_and_key_written_together(tmp_path): + """Both LLM_API_KEY and the base URL land in .env when provided.""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + result = runner.invoke( + cli, [ + "init", + "--model", "vllm/custom-llama", + "--base-url", "http://gpu-host:8000/v1", + ], + input="sk-test-key\n\n", # api key, language + ) + assert result.exit_code == 0, result.output + + from pathlib import Path + env_content = Path(".env").read_text() + # vllm → HOSTED_VLLM_API_KEY (LiteLLM), OPENAI_API_BASE for URL. + assert "HOSTED_VLLM_API_KEY=sk-test-key" in env_content + assert "OPENAI_API_BASE=http://gpu-host:8000/v1" in env_content + + # chmod 600 was applied. + import stat + mode = Path(".env").stat().st_mode + assert stat.S_IMODE(mode) == 0o600 + + +def test_init_base_url_blank_prompt_still_writes_env_with_comments(tmp_path): + """When the user provides nothing, .env is still created with + commented placeholders so the file exists as a discoverable target + for the user to drop their credentials into later. + """ + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + # Use anthropic so both the key and base URL are present as + # commented placeholders. (ollama is keyless — its .env has no + # key section at all, which we test separately.) + result = runner.invoke( + cli, ["init", "--model", "anthropic/claude-sonnet-4-6"], + input="\n\n\n", # blank key, blank language + ) + assert result.exit_code == 0, result.output + + from pathlib import Path + env_path = Path(".env") + assert env_path.exists(), "init must always create .env" + content = env_path.read_text() + + # No active assignments should leak in for fields the user skipped. + for line in content.splitlines(): + stripped = line.lstrip() + assert not stripped.startswith("ANTHROPIC_API_KEY="), ( + f"ANTHROPIC_API_KEY must not be active when user skipped: {line!r}" + ) + assert not stripped.startswith("ANTHROPIC_API_BASE="), ( + f"ANTHROPIC_API_BASE must not be active when user skipped: {line!r}" + ) + # Both placeholders appear as comments so the user knows what to set. + assert "# ANTHROPIC_API_KEY=" in content + assert "# ANTHROPIC_API_BASE=" in content + + # chmod 600 still applied even when content is mostly comments. + import stat + assert stat.S_IMODE(env_path.stat().st_mode) == 0o600 + + +def test_init_existing_env_preserved(tmp_path): + """If .env already exists, init must not clobber it; user is told.""" + from pathlib import Path + + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path) as cwd: + # Pre-existing .env. + Path(".env").write_text("EXISTING=keep-me\n", encoding="utf-8") + with patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + result = runner.invoke( + cli, ["init", "--model", "ollama/llama3"], + input="http://localhost:11434\n\n\n", + ) + assert result.exit_code == 0, result.output + + # Original content preserved verbatim; new key was NOT appended. + assert Path(".env").read_text() == "EXISTING=keep-me\n" + assert "skipping write" in result.output + + +def test_init_rejects_base_url_with_control_chars(tmp_path): + """A --base-url value with embedded newlines is unsafe (would corrupt .env).""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"): + result = runner.invoke( + cli, [ + "init", + "--base-url", "http://x\nLLM_API_KEY=stolen", + ], + input="\n\n", + ) + assert result.exit_code != 0 + assert "--base-url" in result.output + + from pathlib import Path + # Init must abort before writing any KB state. + assert not Path(".openkb").exists() + + +def test_init_rejects_overly_long_base_url(tmp_path): + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"): + result = runner.invoke( + cli, ["init", "--base-url", "x" * 3000], + input="\n\n", + ) + assert result.exit_code != 0 + assert "--base-url" in result.output + + +def test_init_emits_post_init_reminder(tmp_path): + """After init succeeds, the user is pointed at .env and config.yaml.""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"): + result = runner.invoke(cli, ["init"], input="\n\n") + assert result.exit_code == 0, result.output + assert "Review .env" in result.output + assert "config.yaml" in result.output + + +# --------------------------------------------------------------------------- +# MiniMax region picker (global / China) +# --------------------------------------------------------------------------- + + +def test_init_minimax_global_region_writes_env(tmp_path): + """Interactive choice 1 ⇒ MINIMAX_API_BASE points to the global endpoint.""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + # Inputs: model (flag), region (1 = global), api key, language + result = runner.invoke( + cli, ["init", "--model", "minimax/MiniMax-M2.7"], + input="1\n\n\n", + ) + assert result.exit_code == 0, result.output + # The picker must be visually distinct (heading + bracketed + # options) so it can't be mistaken for a continuation of the + # surrounding model / API-key prompts. + assert "MiniMax region" in result.output + assert "[1] Global" in result.output + assert "[2] China" in result.output + + from pathlib import Path + env_content = Path(".env").read_text() + assert "MINIMAX_API_BASE=https://api.minimax.io/v1" in env_content + + +def test_init_minimax_china_region_writes_env(tmp_path): + """Interactive choice 2 ⇒ MINIMAX_API_BASE points to the China endpoint.""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + result = runner.invoke( + cli, ["init", "--model", "minimax/MiniMax-M3"], + input="2\n\n\n", + ) + assert result.exit_code == 0, result.output + + from pathlib import Path + env_content = Path(".env").read_text() + assert "MINIMAX_API_BASE=https://api.minimaxi.com/v1" in env_content + + +def test_init_minimax_picker_fires_for_typed_model(tmp_path): + """The picker must fire when the user TYPES the model interactively, + not only when --model is passed. Regression: a previous version + silently skipped the picker unless --model was explicit, which made + MiniMax users end up with no MINIMAX_API_BASE in .env. + """ + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + # Inputs: model (typed), region, api key, language + result = runner.invoke( + cli, ["init"], + input="minimax/MiniMax-M2.7\n1\nsk-test\nen\n", + ) + assert result.exit_code == 0, result.output + assert "MiniMax region" in result.output + + from pathlib import Path + env_content = Path(".env").read_text() + assert "MINIMAX_API_BASE=https://api.minimax.io/v1" in env_content + assert "MINIMAX_API_KEY=sk-test" in env_content + + +def test_init_minimax_default_to_global_under_non_tty(tmp_path): + """Scripted (non-TTY) init falls back to the global endpoint silently.""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"): + # CliRunner is non-TTY by default; region picker must NOT fire. + result = runner.invoke( + cli, ["init", "--model", "minimax/MiniMax-M2.7"], + input="\n\n", + ) + assert result.exit_code == 0, result.output + assert "MiniMax region" not in result.output + + from pathlib import Path + env_content = Path(".env").read_text() + assert "MINIMAX_API_BASE=https://api.minimax.io/v1" in env_content + + +def test_init_minimax_base_url_flag_overrides_region_picker(tmp_path): + """An explicit --base-url bypasses the region picker entirely.""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + result = runner.invoke( + cli, [ + "init", + "--model", "minimax/MiniMax-M2.7", + "--base-url", "https://my-proxy.example.com/v1", + ], + input="\n\n", + ) + assert result.exit_code == 0, result.output + assert "MiniMax region" not in result.output # picker skipped + + from pathlib import Path + env_content = Path(".env").read_text() + assert "MINIMAX_API_BASE=https://my-proxy.example.com/v1" in env_content + # Neither built-in endpoint should have leaked into the file. + assert "api.minimax.io" not in env_content + assert "api.minimaxi.com" not in env_content + + +def test_init_minimax_invalid_choice_reprompts(tmp_path): + """An unrecognised region entry re-prompts instead of silently defaulting.""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + result = runner.invoke( + cli, ["init", "--model", "minimax/MiniMax-M2.7"], + input="99\n2\n\n\n", # bad, then China, then api key, then lang + ) + assert result.exit_code == 0, result.output + assert "Unknown choice '99'" in result.output + + from pathlib import Path + env_content = Path(".env").read_text() + # The second prompt answer (2 = China) wins. + assert "MINIMAX_API_BASE=https://api.minimaxi.com/v1" in env_content + + +def test_init_minimax_key_and_url_written_together(tmp_path): + """Both LLM_API_KEY and MINIMAX_API_BASE land in .env when provided.""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + result = runner.invoke( + cli, ["init", "--model", "minimax/MiniMax-M2.7"], + input="1\nsk-minimax-key\n\n", + ) + assert result.exit_code == 0, result.output + + from pathlib import Path + env_content = Path(".env").read_text() + assert "MINIMAX_API_KEY=sk-minimax-key" in env_content + assert "MINIMAX_API_BASE=https://api.minimax.io/v1" in env_content + + +def test_known_provider_keys_includes_minimax(): + """``_KNOWN_PROVIDER_KEYS`` must list MINIMAX_API_KEY so that + ``_setup_llm_key`` propagates a generic LLM_API_KEY to it — otherwise + the Agents-SDK litellm provider wouldn't see the credential for + ``minimax/``-prefixed models. + """ + from openkb.cli import _KNOWN_PROVIDER_KEYS + assert "MINIMAX_API_KEY" in _KNOWN_PROVIDER_KEYS + + +def test_provider_to_base_env_includes_minimax(): + """``_PROVIDER_TO_BASE_ENV`` must map ``minimax`` to its env var.""" + from openkb.cli import _PROVIDER_TO_BASE_ENV + assert _PROVIDER_TO_BASE_ENV["minimax"] == "MINIMAX_API_BASE" + + +def test_init_minimax_no_key_writes_env_with_placeholder(tmp_path): + """MiniMax region picked, but no key: .env still created with both + the active URL line and a commented LLM_API_KEY placeholder. + """ + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path), \ + patch("openkb.cli.register_kb"), \ + patch("openkb.cli._stdin_is_tty", return_value=True): + result = runner.invoke( + cli, ["init", "--model", "minimax/MiniMax-M2.7"], + input="2\n\n\n", # region=china, blank key, blank language + ) + assert result.exit_code == 0, result.output + + from pathlib import Path + content = Path(".env").read_text() + assert "MINIMAX_API_BASE=https://api.minimaxi.com/v1" in content + # Key placeholder present as a comment under the provider-specific + # name (MINIMAX_API_KEY), never as an active assignment. + assert "# MINIMAX_API_KEY=" in content + for line in content.splitlines(): + assert not line.lstrip().startswith("MINIMAX_API_KEY=") + assert not line.lstrip().startswith("LLM_API_KEY=") + + +def test_build_env_content_no_provider_no_base_url(): + """When no provider context is known, the env builder still emits + a valid LLM_API_KEY placeholder and no spurious *_API_BASE section. + """ + from openkb.cli import _build_env_content + content = _build_env_content({}, provider=None) + # provider=None → generic LLM_API_KEY placeholder. + assert "# LLM_API_KEY=" in content + # No provider ⇒ no base URL section at all (no misleading hint). + assert "_API_BASE=" not in content + + +def test_build_env_content_active_key_and_placeholder_url(): + from openkb.cli import _build_env_content + content = _build_env_content( + {"ANTHROPIC_API_KEY": "sk-test"}, provider="anthropic", + ) + # Active key written under the provider-specific name. + assert "ANTHROPIC_API_KEY=sk-test" in content + # No generic LLM_API_KEY leaks in for a known provider. + assert "LLM_API_KEY=sk-test" not in content + # Base URL placeholder for anthropic is present but commented. + assert "# ANTHROPIC_API_BASE=" in content + # No active (uncommented) assignment leaks the placeholder URL. + for line in content.splitlines(): + assert not line.lstrip().startswith("ANTHROPIC_API_BASE=") + + +@pytest.mark.parametrize("provider,key_env,key_value", [ + ("openai", "OPENAI_API_KEY", "sk-openai"), + ("anthropic", "ANTHROPIC_API_KEY", "sk-ant"), + ("gemini", "GEMINI_API_KEY", "AIza-test"), + ("deepseek", "DEEPSEEK_API_KEY", "sk-ds"), + ("mistral", "MISTRAL_API_KEY", "mistral-key"), + ("moonshot", "MOONSHOT_API_KEY", "ms-key"), + ("dashscope", "DASHSCOPE_API_KEY", "ds-key"), + ("openrouter", "OPENROUTER_API_KEY", "or-key"), + ("minimax", "MINIMAX_API_KEY", "minimax-key"), + ("zhipuai", "ZHIPUAI_API_KEY", "zhipu-key"), +]) +def test_build_env_content_per_provider_key_naming(provider, key_env, key_value): + """Regression: each LiteLLM provider has its own *_API_KEY env var, + and ``openkb init`` must write the right one — not the generic + ``LLM_API_KEY`` — so the file reads naturally to anyone familiar + with that provider. + """ + from openkb.cli import _build_env_content + content = _build_env_content({key_env: key_value}, provider) + assert f"{key_env}={key_value}" in content + # The active line must be uncommented. + active_lines = [ + line for line in content.splitlines() + if line.startswith(f"{key_env}=") + ] + assert active_lines == [f"{key_env}={key_value}"] + + +def test_key_env_for_provider_known_and_unknown(): + from openkb.cli import _key_env_for_provider + assert _key_env_for_provider("openai") == "OPENAI_API_KEY" + assert _key_env_for_provider("minimax") == "MINIMAX_API_KEY" + # ollama has no key (None, not the generic fallback). + assert _key_env_for_provider("ollama") is None + # Unknown provider also returns None (caller decides fallback). + assert _key_env_for_provider("custom-thing") is None + assert _key_env_for_provider(None) is None + + +def test_setup_llm_key_reads_provider_specific_env_var(tmp_path): + """``_setup_llm_key`` must pick up the provider-specific env var + (the format ``openkb init`` now writes) without requiring a + generic ``LLM_API_KEY`` fallback. + """ + from pathlib import Path + from openkb import cli as cli_mod + + monkeypatch = pytest.MonkeyPatch() + # Simulate what openkb init now writes: only the provider-specific + # env var is set; LLM_API_KEY is empty. + monkeypatch.setenv("OPENAI_API_KEY", "sk-direct-openai") + monkeypatch.delenv("LLM_API_KEY", raising=False) + try: + kb_dir = tmp_path / "kb" + kb_dir.mkdir() + (kb_dir / ".openkb").mkdir() + (kb_dir / ".openkb/config.yaml").write_text( + "model: openai/gpt-5.4-mini\n", encoding="utf-8", + ) + cli_mod._setup_llm_key(kb_dir) + assert cli_mod.litellm.api_key == "sk-direct-openai" + finally: + monkeypatch.undo() + + +def test_setup_llm_key_applies_minimax_base_url(tmp_path): + """``_setup_llm_key`` reads MINIMAX_API_BASE and sets litellm.api_base.""" + from pathlib import Path + + from openkb import cli as cli_mod + + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setenv("MINIMAX_API_BASE", "https://api.minimaxi.com/v1") + monkeypatch.setenv("LLM_API_KEY", "sk-test") + try: + kb_dir = tmp_path / "kb" + kb_dir.mkdir() + (kb_dir / ".openkb").mkdir() + (kb_dir / ".openkb/config.yaml").write_text( + "model: minimax/MiniMax-M2.7\n", encoding="utf-8", + ) + cli_mod._setup_llm_key(kb_dir) + assert cli_mod.litellm.api_base == "https://api.minimaxi.com/v1" + finally: + monkeypatch.undo() + + class TestQueryStreamGate: """Regression tests for issue #34.