"""Interactive setup wizard — generates config.yaml and .env.""" import getpass import os import sys from pathlib import Path import yaml _ENV_KEY_MAP = { "openai": "OPENAI_API_KEY", "anthropic": "ANTHROPIC_API_KEY", "gemini": "GEMINI_API_KEY", "hugging-face": "HF_API_KEY", } _PROVIDER_CHOICES = {"1": "openai", "2": "anthropic", "3": "gemini", "4": "meta-llama"} def generate_config( bot_name: str, domain: str, provider: str, model: str, web_search: bool, ) -> str: """Return a YAML string with all config sections.""" config = { "chatbot": { "name": bot_name, "domain": domain, }, "llm": { "provider": provider, "model": model, "temperature": 0.0, "max_tokens": 8192, }, "api_keys": { "openai": "", "anthropic": "", "gemini": "", "meta-llama": "", }, "embeddings": { "provider": "local", "openai_model": "text-embedding-3-small", "emb_model": "sentence-transformers/all-mpnet-base-v2", }, "retrieval": { "chunk_size": 1000, "chunk_overlap": 100, "top_k": 20, "max_distance": 0.55, "max_context_chars": 12000, }, "web_search": { "enabled": web_search, "backend": "semantic_scholar", "max_results": 5, }, "query_understanding": { "enabled": True, "max_history": 6, "max_clarifications": 1, }, "verification": { "enabled": True, "max_iterations": 3, "strict_mode": True, }, "sql": { "enabled": True, "max_rows": 200, }, "paths": { "knowledge_base": "knowledge_base", "vector_db": "chroma_db", "sql_db": "sql_db", }, } return yaml.dump(config, default_flow_style=False, sort_keys=False) def generate_env(provider: str, api_key: str, existing_env_path: str = None) -> str: """Return .env file content, merging with existing keys if present.""" env_var = _ENV_KEY_MAP.get(provider, f"{provider.upper()}_API_KEY") # Preserve existing keys from a prior .env file existing: dict[str, str] = {} if existing_env_path and os.path.exists(existing_env_path): with open(existing_env_path, "r") as f: for line in f: line = line.strip() if line and not line.startswith("#") and "=" in line: k, v = line.split("=", 1) v = v.strip() # Unwrap one matched pair of quotes if (v.startswith('"') and v.endswith('"')) or \ (v.startswith("'") and v.endswith("'")): v = v[1:-1] # Unescape previously escaped characters v = v.replace('\\"', '"').replace('\\\\', '\\') existing[k.strip()] = v # Update with the new key existing[env_var] = api_key lines = ["# Auto-generated by setup wizard"] for k, v in sorted(existing.items()): v_escaped = v.replace('\\', '\\\\').replace('"', '\\"') lines.append(f'{k}="{v_escaped}"') lines.append("") return "\n".join(lines) def run_wizard(): """Five-step interactive setup flow.""" project_root = Path(__file__).resolve().parent print("=" * 50) print(" RAG Research Chatbot — Setup Wizard") print("=" * 50) print() # Step 1: Bot name bot_name = input("Step 1/5 — Bot name [Research Assistant]: ").strip() if not bot_name: bot_name = "Research Assistant" # Step 2: Domain domain = input("Step 2/5 — Domain / topic description: ").strip() if not domain: domain = "general research" # Step 3: Provider print("\nStep 3/5 — LLM provider:") print(" 1) OpenAI") print(" 2) Anthropic") print(" 3) Google Gemini") print(" 4) Meta Llama") provider_choice = input("Choose [1]: ").strip() or "1" provider = _PROVIDER_CHOICES.get(provider_choice, "openai") # Step 3b: API key api_key = getpass.getpass(f"Enter your {provider} API key: ") if not api_key: print(f" Warning: No API key entered for {provider}.") print(f" Set it later by editing .env or re-running: python setup.py") # Step 4: Model selection — fetch available models print(f"\nFetching available {provider} models...") try: from src.llm import list_models models = list_models(provider, api_key) except Exception: print(f" Warning: Could not validate API key for {provider}. Using default model list.") fallback = { "openai": ["gpt-4.1", "gpt-4.1-mini", "gpt-4.1-nano"], "anthropic": ["claude-sonnet-4-6", "claude-haiku-4-5", "claude-opus-4-6"], "gemini": ["gemini-2.5-flash", "gemini-2.5-pro", "gemini-2.0-flash"], "meta": ["meta-llama/Llama-3.3-70B-Instruct"], } models = fallback.get(provider, ["default-model"]) print("\nStep 4/5 — Choose a model:") for i, m in enumerate(models, 1): print(f" {i}) {m}") model_choice = input(f"Choose [1]: ").strip() or "1" try: idx = int(model_choice) - 1 if 0 <= idx < len(models): model = models[idx] else: model = models[0] except (ValueError, IndexError): model = models[0] # Step 5: Web search print("\nStep 5/5 — Enable web search (Semantic Scholar)?") print(" 1) Yes") print(" 2) No") ws_choice = input("Choose [1]: ").strip() or "1" web_search = ws_choice == "1" # Write config.yaml config_path = project_root / "config.yaml" if config_path.exists(): overwrite = input(f"\n{config_path} already exists. Overwrite? [y/N]: ").strip().lower() if overwrite != 'y': print(" Keeping existing config.yaml.") else: config_str = generate_config(bot_name, domain, provider, model, web_search) config_path.write_text(config_str) print(f"\n Wrote {config_path}") else: config_str = generate_config(bot_name, domain, provider, model, web_search) config_path.write_text(config_str) print(f"\n Wrote {config_path}") # Write .env (merges with existing keys if present) env_path = project_root / ".env" env_str = generate_env(provider, api_key, existing_env_path=str(env_path)) env_path.write_text(env_str) print(f" Wrote {env_path}") # Create knowledge_base/ directory kb_dir = project_root / "knowledge_base" kb_dir.mkdir(exist_ok=True) print(f" Created {kb_dir}/") # Next steps print("\n" + "=" * 50) print(" Setup complete! Next steps:") print("=" * 50) print(f" 1. Add documents to {kb_dir}/") print(" 2. Run: python ingest.py") print(" 3. Run: python app_cli.py (or: streamlit run app_web.py)") print() if __name__ == "__main__": run_wizard()