"""Interactive setup flow for provider, model, and API key configuration.""" from getpass import getpass from rich.console import Console from rich.prompt import Prompt from rich.table import Table from .config import PROVIDERS, ProviderConfig from .state import find_env_api_key, save_state def select_provider(console: Console) -> ProviderConfig: """Display provider choices and return selected provider configuration.""" table = Table(title="Choose Provider") table.add_column("Key", style="cyan", justify="right") table.add_column("Provider", style="green") table.add_row("1", "Gemini") table.add_row("2", "Groq") table.add_row("3", "Nvidia") table.add_row("4", "Ollama") console.print(table) choice = Prompt.ask("Enter choice", choices=["1", "2", "3", "4"], console=console) config = PROVIDERS.get(choice) if not config: raise RuntimeError("Invalid provider choice. Pick a number from 1 to 4.") return config def resolve_api_key(config: ProviderConfig) -> str: """Resolve API key from prompt input, env vars, or sensible local defaults.""" env_key = find_env_api_key(config.env_key_names) if config.name == "ollama": # Ollama typically runs locally and does not require a secret key. entered = getpass("API key (optional for Ollama; press Enter to skip): ").strip() return entered or env_key or "ollama" entered = getpass("API key (hidden input): ").strip() if entered: return entered if env_key: return env_key env_hint = " or ".join(config.env_key_names) raise RuntimeError(f"Missing API key. Provide one in prompt or set {env_hint}.") def run_setup(console: Console) -> dict[str, str]: """Run end-to-end setup and persist selected provider credentials/state.""" config = select_provider(console) console.print(f"\nSelected: [bold green]{config.name}[/bold green]") model = Prompt.ask("Model", default=config.default_model, console=console).strip() or config.default_model api_key = resolve_api_key(config) save_state(config.name, api_key, model) console.print(f"Saved setup for provider '{config.name}'.") return {"provider": config.name, "api_key": api_key, "model": model}