setup.py

"""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()