Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .env.example
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,7 @@
# OpenAI: LLM_API_KEY=sk-...
# Anthropic: LLM_API_KEY=sk-ant-...
# Gemini: LLM_API_KEY=AIza...
# DeepSeek: LLM_API_KEY=sk-...
# Cohere: LLM_API_KEY=...
# Mistral: LLM_API_KEY=...
LLM_API_KEY=your-key-here
51 changes: 49 additions & 2 deletions openkb/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,27 @@ def filter(self, record: logging.LogRecord) -> bool:
load_dotenv() # load from cwd (covers running inside the KB dir)


_KNOWN_PROVIDER_KEYS = (
"OPENAI_API_KEY", "ANTHROPIC_API_KEY", "GEMINI_API_KEY",
"DEEPSEEK_API_KEY", "COHERE_API_KEY", "MISTRAL_API_KEY",
"TOGETHER_API_KEY", "GROQ_API_KEY",
)


def _extract_provider(model: str) -> str | None:
"""Extract the LiteLLM provider name from a model string.

``model`` uses ``provider/model`` LiteLLM format.
OpenAI models can omit the prefix; default to ``"openai"``.
"""
model = model.strip()
if not model:
return None
if "/" in model:
return model.split("/")[0].lower()
return "openai"


def _setup_llm_key(kb_dir: Path | None = None) -> None:
"""Set LiteLLM API key from LLM_API_KEY env var if present.

Expand All @@ -62,6 +83,8 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None:

Also propagates to provider-specific env vars (OPENAI_API_KEY, etc.)
so that the Agents SDK litellm provider can pick them up.
Provider is auto-detected from the KB config when available; otherwise
a common provider set is used as a fallback.
"""
if kb_dir is not None:
env_file = kb_dir / ".env"
Expand All @@ -74,9 +97,23 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None:
load_dotenv(global_env, override=False)

api_key = os.environ.get("LLM_API_KEY", "")

# Try to resolve the active provider from the KB config
provider: str | None = None
if kb_dir is not None:
config_path = kb_dir / ".openkb" / "config.yaml"
if config_path.exists():
config = load_config(config_path)
model = config.get("model", "")
provider = _extract_provider(str(model))

if not api_key:
# Check if any provider key is already set
has_key = any(os.environ.get(k) for k in ("OPENAI_API_KEY", "ANTHROPIC_API_KEY", "GEMINI_API_KEY"))
check_keys = (
(f"{provider.upper()}_API_KEY",) if provider
else _KNOWN_PROVIDER_KEYS
)
has_key = any(os.environ.get(k) for k in check_keys)
if not has_key:
click.echo(
"Warning: No LLM API key found. Set one of:\n"
Expand All @@ -86,7 +123,16 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None:
)
else:
litellm.api_key = api_key
for env_var in ("OPENAI_API_KEY", "ANTHROPIC_API_KEY", "GEMINI_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
for env_var in _KNOWN_PROVIDER_KEYS:
if not os.environ.get(env_var):
os.environ[env_var] = api_key

Expand Down Expand Up @@ -352,6 +398,7 @@ def init(language):
click.echo(" OpenAI: gpt-5.4-mini, gpt-5.4")
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-0417")
click.echo(" Others: see https://docs.litellm.ai/docs/providers")
click.echo()
model = click.prompt(
Expand Down