kchen707/wedding-bundle-builder

1"""2Lazy OpenAI/OpenRouter client factory.3 4This file uses module-level state instead of threading.local because5Gradio reassigns work across threads within a single request, which6broke the thread-local approach. For a single-user demo this is fine.7"""8 9import os10from openai import OpenAI11 12_OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"13 14_state = {"api_key": None, "llm_client": None, "embed_client": None}15 16 17class MissingAPIKeyError(RuntimeError):18    pass19 20 21def _build_clients(key: str):22    headers = {23        "HTTP-Referer": "https://huggingface.co/spaces/wedding-bundle-builder",24        "X-Title": "SD Wedding Bundle Builder",25    }26    return (27        OpenAI(base_url=_OPENROUTER_BASE_URL, api_key=key, default_headers=headers),28        OpenAI(base_url=_OPENROUTER_BASE_URL, api_key=key, default_headers=headers),29    )30 31 32def set_active_key(api_key: str) -> None:33    # Fall back to env/secret if user didn't paste anything34    if not api_key or not str(api_key).strip():35        api_key = os.environ.get("OPENROUTER_API_KEY", "")36 37    if not api_key or not str(api_key).strip():38        _state["api_key"] = None39        _state["llm_client"] = None40        _state["embed_client"] = None41        return42 43    key = str(api_key).strip()44    if _state["api_key"] == key and _state["llm_client"] is not None:45        return46 47    llm, emb = _build_clients(key)48    _state["api_key"] = key49    _state["llm_client"] = llm50    _state["embed_client"] = emb51 52 53def has_active_key() -> bool:54    return bool(_state["api_key"])55 56 57def get_llm_client() -> OpenAI:58    client = _state["llm_client"]59    if client is None:60        raise MissingAPIKeyError(61            "No OpenRouter API key set. Paste your key into the field at the top "62            "of the page before clicking the build button."63        )64    return client65 66 67def get_embed_client() -> OpenAI:68    client = _state["embed_client"]69    if client is None:70        raise MissingAPIKeyError(71            "No OpenRouter API key set. Paste your key into the field at the top "72            "of the page before clicking the build button."73        )74    return client75 76 77def validate_key(api_key: str) -> tuple[bool, str]:78    from config import CHAT_MODEL79 80    if not api_key or not str(api_key).strip():81        return False, "Key is empty."82 83    try:84        headers = {85            "HTTP-Referer": "https://huggingface.co/spaces/wedding-bundle-builder",86            "X-Title": "SD Wedding Bundle Builder",87        }88        client = OpenAI(89            base_url=_OPENROUTER_BASE_URL,90            api_key=str(api_key).strip(),91            default_headers=headers,92        )93        resp = client.chat.completions.create(94            model=CHAT_MODEL,95            messages=[{"role": "user", "content": "Say OK"}],96            max_tokens=5,97        )98        text = (resp.choices[0].message.content or "").strip()99        return True, f"Key is valid. Test response: {text or '(empty)'}"100    except Exception as e:101        return False, f"Key check failed: {e}"