Spaces:
Sleeping
Sleeping
| """Gradio chat client for a Hugging Face Inference Endpoint.""" | |
| import itertools | |
| import json | |
| import logging | |
| import os | |
| import sys | |
| import time | |
| from secrets import compare_digest | |
| from typing import NamedTuple | |
| import gradio as gr | |
| import requests | |
| # Root stays at WARNING so LOG_LEVEL doesn't also enable httpx/urllib3. | |
| logging.basicConfig( | |
| level=logging.WARNING, | |
| format="%(asctime)s %(levelname)-7s %(name)s %(message)s", | |
| stream=sys.stdout, | |
| ) | |
| logger = logging.getLogger("gemma-endpoint") | |
| logger.setLevel(os.environ.get("LOG_LEVEL", "INFO").upper()) | |
| _turn_ids = itertools.count(1) | |
| class EndpointWaking(Exception): | |
| """Raised before any tokens, so retrying the request is safe.""" | |
| REQUEST_TIMEOUT_SECONDS = 300 | |
| MAX_TOKENS = 4096 # reasoning models spend most of a smaller budget thinking | |
| # A cold endpoint rejects outright instead of holding the request. | |
| WAKE_STATUSES = (502, 503, 504) | |
| AUTH_STATUSES = (401, 403) | |
| WAKE_TIMEOUT_SECONDS = int(os.environ.get("WAKE_TIMEOUT_SECONDS", "300")) | |
| WAKE_RETRY_SECONDS = 5 | |
| WAKE_TICK_SECONDS = 1 # shorter than the retry so Stop stays responsive | |
| # `status: pending` renders a spinner. It needs text: reasoning_tags drops | |
| # empty messages, spinner included. | |
| WAITING_METADATA = {"title": "Waiting for the endpoint…", "status": "pending"} | |
| WAITING_TEXT = "If the endpoint was asleep, the first reply can take a few minutes." | |
| INPUT_PLACEHOLDER = "Message the model" | |
| # llama.cpp sends thinking as reasoning_content; these tags make it collapsible. | |
| THINK_OPEN = "<think>" | |
| THINK_CLOSE = "</think>" | |
| # The default 512-byte reads deliver the whole reply at once. | |
| STREAM_CHUNK_SIZE = 1 | |
| GEMMA_MODEL = "google/gemma-4-12B-it-qat-q4_0-gguf" | |
| QWEN_MODEL = "unsloth/Qwen3.5-9B-GGUF" | |
| DEEPSEEK_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B" | |
| MODEL_LABELS = { | |
| GEMMA_MODEL: "Gemma 4 12B", | |
| QWEN_MODEL: "Qwen 3.5 9B", | |
| DEEPSEEK_MODEL: "DeepSeek R1 14B", | |
| } | |
| def _required_env(name): | |
| value = os.environ.get(name, "").strip() | |
| if not value: | |
| raise RuntimeError( | |
| f"{name} is not set. Set it as a Space secret (or export it " | |
| f"locally) before starting the app." | |
| ) | |
| return value | |
| class Endpoint(NamedTuple): | |
| url: str | |
| token: str | |
| token_env: str # secret name, for logs and errors; never the token itself | |
| def _endpoint(url_env, token_env): | |
| """Per-org token, falling back to HF_TOKEN.""" | |
| token = os.environ.get(token_env, "").strip() | |
| return Endpoint( | |
| url=_required_env(url_env).rstrip("/"), | |
| token=token or HF_TOKEN, | |
| token_env=token_env if token else "HF_TOKEN", | |
| ) | |
| HF_TOKEN = _required_env("HF_TOKEN") | |
| ENDPOINT_URL = _required_env("ENDPOINT_URL").rstrip("/") | |
| LOGIN_PASSWORD = os.environ.get("LOGIN_PASSWORD", "") | |
| MODELS = { | |
| GEMMA_MODEL: Endpoint(ENDPOINT_URL, HF_TOKEN, "HF_TOKEN"), | |
| QWEN_MODEL: _endpoint("QWEN_ENDPOINT_URL", "QWEN_HF_TOKEN"), | |
| DEEPSEEK_MODEL: _endpoint("DEEPSEEK_ENDPOINT_URL", "DEEPSEEK_HF_TOKEN"), | |
| } | |
| AVAILABLE_MODELS = list(MODELS) | |
| DEFAULT_MODEL = GEMMA_MODEL | |
| def _log_config(): | |
| for model, endpoint in MODELS.items(): | |
| logger.info( | |
| "configured model=%s url=%s token=%s", | |
| model, | |
| endpoint.url, | |
| endpoint.token_env, | |
| ) | |
| logger.info( | |
| "configured auth=%s timeout=%ss max_tokens=%s", | |
| "on" if LOGIN_PASSWORD else "OFF", | |
| REQUEST_TIMEOUT_SECONDS, | |
| MAX_TOKENS, | |
| ) | |
| _log_config() | |
| def endpoint_for(model_name): | |
| return MODELS.get(model_name) or Endpoint(ENDPOINT_URL, HF_TOKEN, "HF_TOKEN") | |
| def message_text(content): | |
| """Gradio content is a string or a list of typed parts; keep the text.""" | |
| if isinstance(content, str): | |
| return content | |
| if isinstance(content, (list, tuple)): | |
| return "".join( | |
| part.get("text") or "" | |
| for part in content | |
| if isinstance(part, dict) and part.get("type") == "text" | |
| ) | |
| return "" | |
| def normalize_history(history): | |
| """Plain text turns, without spinners or other metadata messages.""" | |
| turns = [] | |
| for item in history or []: | |
| if not isinstance(item, dict) or item.get("role") not in ("user", "assistant"): | |
| continue | |
| if (item.get("metadata") or {}).get("title"): | |
| continue | |
| text = message_text(item.get("content")) | |
| if text: | |
| turns.append({"role": item["role"], "content": text}) | |
| return turns | |
| def build_api_messages(history, message): | |
| return normalize_history(history) + [{"role": "user", "content": message}] | |
| def stream_completion(messages, model_name, turn="-"): | |
| """Yield SSE chunks. reasoning_content is wrapped in think tags.""" | |
| endpoint = endpoint_for(model_name) | |
| prompt_chars = sum(len(m["content"]) for m in messages) | |
| logger.info( | |
| "turn=%s POST %s/v1/chat/completions model=%s token=%s messages=%d " | |
| "prompt_chars=%d", | |
| turn, | |
| endpoint.url, | |
| model_name, | |
| endpoint.token_env, | |
| len(messages), | |
| prompt_chars, | |
| ) | |
| logger.debug("turn=%s prompt=%s", turn, messages) | |
| started = time.monotonic() | |
| with requests.post( | |
| f"{endpoint.url}/v1/chat/completions", | |
| headers={ | |
| "Authorization": f"Bearer {endpoint.token}", | |
| "Content-Type": "application/json", | |
| }, | |
| json={ | |
| "model": model_name, | |
| "messages": messages, | |
| "max_tokens": MAX_TOKENS, | |
| "stream": True, | |
| }, | |
| timeout=REQUEST_TIMEOUT_SECONDS, | |
| stream=True, | |
| ) as response: | |
| logger.info( | |
| "turn=%s response status=%s after %.2fs", | |
| turn, | |
| response.status_code, | |
| time.monotonic() - started, | |
| ) | |
| if response.status_code in WAKE_STATUSES: | |
| raise EndpointWaking(response.status_code) | |
| if response.status_code in AUTH_STATUSES: | |
| raise gr.Error( | |
| f"{MODEL_LABELS.get(model_name, model_name)} rejected the token " | |
| f"(HTTP {response.status_code}). Check that {endpoint.token_env} " | |
| f"has access to the organisation that owns this endpoint." | |
| ) | |
| response.raise_for_status() | |
| # requests guesses ISO-8859-1 when the charset is missing. | |
| response.encoding = "utf-8" | |
| chunks = 0 | |
| thinking = False | |
| for line in response.iter_lines( | |
| chunk_size=STREAM_CHUNK_SIZE, decode_unicode=True | |
| ): | |
| if not line or not line.startswith("data:"): | |
| continue | |
| payload = line[len("data:") :].strip() | |
| if payload == "[DONE]": | |
| logger.info( | |
| "turn=%s stream done chunks=%d in %.2fs", | |
| turn, | |
| chunks, | |
| time.monotonic() - started, | |
| ) | |
| break | |
| choices = json.loads(payload).get("choices") or [{}] | |
| delta = choices[0].get("delta") or {} | |
| reasoning = delta.get("reasoning_content") | |
| content = delta.get("content") | |
| piece = "" | |
| if reasoning: | |
| if not thinking: | |
| piece += THINK_OPEN | |
| thinking = True | |
| piece += reasoning | |
| if content: | |
| if thinking: | |
| piece += THINK_CLOSE | |
| thinking = False | |
| piece += content | |
| if piece: | |
| chunks += 1 | |
| if chunks == 1: | |
| logger.info( | |
| "turn=%s first token after %.2fs", | |
| turn, | |
| time.monotonic() - started, | |
| ) | |
| yield piece | |
| else: | |
| logger.warning( | |
| "turn=%s stream ended without [DONE] after %d chunks", turn, chunks | |
| ) | |
| if thinking: | |
| yield THINK_CLOSE | |
| def waiting_turn(title=None): | |
| metadata = dict(WAITING_METADATA) | |
| if title: | |
| metadata["title"] = title | |
| return {"role": "assistant", "content": WAITING_TEXT, "metadata": metadata} | |
| def stream_turn(messages, model_name, turn="-"): | |
| """Yield ("waking", seconds) while the endpoint is cold, then ("token", text).""" | |
| started = time.monotonic() | |
| deadline = started + WAKE_TIMEOUT_SECONDS | |
| attempt = 0 | |
| while True: | |
| attempt += 1 | |
| try: | |
| for chunk in stream_completion(messages, model_name, turn=turn): | |
| yield "token", chunk | |
| return | |
| except EndpointWaking as waking: | |
| waited = time.monotonic() - started | |
| if time.monotonic() >= deadline: | |
| raise gr.Error( | |
| f"The endpoint did not come up within " | |
| f"{WAKE_TIMEOUT_SECONDS}s (last status {waking}). Check that " | |
| f"it is Running and not scaled to zero." | |
| ) from None | |
| logger.info( | |
| "turn=%s endpoint waking status=%s attempt=%d waited=%.0fs", | |
| turn, | |
| waking, | |
| attempt, | |
| waited, | |
| ) | |
| for _ in range(WAKE_RETRY_SECONDS // WAKE_TICK_SECONDS): | |
| if time.monotonic() >= deadline: | |
| break | |
| yield "waking", int(time.monotonic() - started) | |
| time.sleep(WAKE_TICK_SECONDS) | |
| def respond(message, history, model_name): | |
| history = list(history or []) | |
| message = (message or "").strip() | |
| if not message: | |
| logger.debug("ignoring an empty submission") | |
| yield history, *idle_controls() | |
| return | |
| turn = next(_turn_ids) | |
| started = time.monotonic() | |
| api_messages = build_api_messages(history, message) | |
| logger.info( | |
| "turn=%s start model=%s history_turns=%d message_chars=%d", | |
| turn, | |
| model_name or DEFAULT_MODEL, | |
| len(api_messages) - 1, | |
| len(message), | |
| ) | |
| history.append({"role": "user", "content": message}) | |
| history.append(waiting_turn()) | |
| yield list(history), *busy_controls() | |
| reply = "" | |
| try: | |
| for kind, value in stream_turn( | |
| api_messages, model_name or DEFAULT_MODEL, turn=turn | |
| ): | |
| if kind == "waking": | |
| history[-1] = waiting_turn(f"Waking up the endpoint… {value}s") | |
| else: | |
| reply += value | |
| history[-1] = {"role": "assistant", "content": reply} | |
| yield list(history), *busy_controls() | |
| if not reply: | |
| raise gr.Error("The endpoint returned an empty reply.") | |
| except Exception as error: | |
| if isinstance(error, gr.Error): | |
| detail = str(error) | |
| elif isinstance(error, requests.exceptions.Timeout): | |
| detail = "The endpoint timed out. Check that it is in the Running state." | |
| elif isinstance(error, requests.exceptions.RequestException): | |
| detail = f"Request to the endpoint failed: {_error_detail(error)}" | |
| elif isinstance(error, (KeyError, IndexError, ValueError)): | |
| detail = f"Unexpected response from the endpoint: {error}" | |
| else: | |
| logger.exception("turn=%s unexpected failure", turn) | |
| raise | |
| logger.error( | |
| "turn=%s failed after %.2fs with %d chars streamed: %s", | |
| turn, | |
| time.monotonic() - started, | |
| len(reply), | |
| detail, | |
| ) | |
| if reply: | |
| yield list(history), *idle_controls() | |
| else: | |
| del history[-2:] | |
| yield list(history), *idle_controls(message) | |
| raise gr.Error(detail, print_exception=False) from None | |
| logger.info( | |
| "turn=%s complete reply_chars=%d in %.2fs", | |
| turn, | |
| len(reply), | |
| time.monotonic() - started, | |
| ) | |
| logger.debug("turn=%s reply=%s", turn, reply) | |
| yield list(history), *idle_controls() | |
| def busy_controls(value=""): | |
| return ( | |
| gr.update(value=value, interactive=False, placeholder="Replying…"), | |
| gr.update(interactive=False, visible=False), | |
| gr.update(interactive=True, visible=True), | |
| ) | |
| def idle_controls(value=""): | |
| return ( | |
| gr.update(value=value, interactive=True, placeholder=INPUT_PLACEHOLDER), | |
| gr.update(interactive=True, visible=True), | |
| gr.update(interactive=False, visible=False), | |
| ) | |
| def cancel_stream(history): | |
| logger.info("stopped by the user") | |
| return normalize_history(history), *idle_controls() | |
| def new_chat(): | |
| logger.info("new chat") | |
| return [], *idle_controls() | |
| def model_caption(model_name): | |
| return f"`{model_name or DEFAULT_MODEL}`" | |
| def example_text(event: gr.SelectData): | |
| return event.value["text"] | |
| def _error_detail(error): | |
| if error.response is None: | |
| return str(error) | |
| try: | |
| return error.response.text or str(error) | |
| except requests.exceptions.RequestException: | |
| return str(error) | |
| def verify_login(username, password): | |
| allowed = compare_digest(password or "", LOGIN_PASSWORD) | |
| logger.log( | |
| logging.INFO if allowed else logging.WARNING, | |
| "login %s for username=%r", | |
| "succeeded" if allowed else "REJECTED", | |
| username, | |
| ) | |
| return allowed | |
| EXAMPLES = [ | |
| {"display_text": "Reply in Japanese", "text": "Introduce yourself in Japanese."}, | |
| { | |
| "display_text": "Reason step by step", | |
| "text": "Is 1,001 a prime number? Work it out step by step.", | |
| }, | |
| { | |
| "display_text": "Write some code", | |
| "text": "Write a Python function that reverses a singly linked list.", | |
| }, | |
| { | |
| "display_text": "Keep it short", | |
| "text": "Explain what a hash map is in two sentences.", | |
| }, | |
| ] | |
| EMPTY_STATE = ( | |
| "**Pick a model and send a message**\n\n" | |
| "A sleeping endpoint can take a few minutes to answer the first message." | |
| ) | |
| THEME = gr.themes.Base( | |
| primary_hue=gr.themes.colors.teal, | |
| neutral_hue=gr.themes.colors.slate, | |
| font=[gr.themes.GoogleFont("IBM Plex Sans"), "system-ui", "sans-serif"], | |
| font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace"], | |
| radius_size=gr.themes.sizes.radius_md, | |
| ).set( | |
| body_background_fill="#f2f4f6", | |
| body_background_fill_dark="#0e1216", | |
| block_border_width="0px", | |
| block_shadow="none", | |
| button_primary_background_fill="#0b6e79", | |
| button_primary_background_fill_hover="#095a63", | |
| button_primary_background_fill_dark="#14909d", | |
| button_primary_background_fill_hover_dark="#18a5b3", | |
| button_primary_text_color="#ffffff", | |
| ) | |
| CSS = """ | |
| /* Gradio caps .main at 640px, which squeezes the chat. */ | |
| .gradio-container, .gradio-container > .main { | |
| max-width: 880px !important; width: 100%; margin: 0 auto; } | |
| #header { align-items: center; gap: 12px; } | |
| #header h1 { font-size: 1.35rem; font-weight: 600; margin: 0; } | |
| #header p { margin: 2px 0 0; color: var(--body-text-color-subdued); } | |
| #header button { background: var(--block-background-fill); | |
| border: 1px solid var(--border-color-primary); border-radius: 8px; } | |
| #header button:hover { border-color: var(--button-primary-background-fill); } | |
| #model-picker .wrap { gap: 4px; padding: 4px; border-radius: 10px; | |
| background: var(--block-background-fill); | |
| border: 1px solid var(--border-color-primary); width: fit-content; } | |
| #model-picker label { border: 0 !important; box-shadow: none !important; | |
| background: transparent !important; padding: 6px 14px !important; | |
| border-radius: 7px !important; cursor: pointer; } | |
| #model-picker label.selected { background: var(--button-primary-background-fill) !important; | |
| color: var(--button-primary-text-color) !important; } | |
| #model-picker input[type="radio"] { position: absolute; opacity: 0; | |
| width: 1px; height: 1px; pointer-events: none; } | |
| #model-picker label:has(input:focus-visible) { | |
| outline: 2px solid var(--button-primary-background-fill); outline-offset: 2px; } | |
| #model-id { margin-top: -6px; } | |
| #model-id p, #model-id code { font-size: 0.8rem; color: var(--body-text-color-subdued); } | |
| #model-id code { background: none; border: 0; padding: 0; } | |
| #chat { | |
| background: var(--background-fill-primary) !important; | |
| border: 1px solid var(--border-color-primary) !important; | |
| border-radius: 12px !important; } | |
| #chat .bubble-wrap { background: transparent; } | |
| #chat .placeholder-content { justify-content: center; gap: 20px; padding: 24px; } | |
| #chat .placeholder { flex: 0 0 auto !important; height: auto !important; | |
| text-align: center; } | |
| #chat .placeholder p { margin: 4px 0; color: var(--body-text-color-subdued); } | |
| #chat .placeholder strong { color: var(--body-text-color); font-size: 1.05rem; } | |
| #chat .examples { width: 100%; max-width: 560px; margin: 0 auto; gap: 10px; } | |
| #chat .example { background: var(--background-fill-secondary) !important; | |
| border: 1px solid var(--border-color-primary) !important; | |
| border-radius: 10px !important; text-align: left; } | |
| #chat .example:hover { border-color: var(--button-primary-background-fill) !important; } | |
| #composer { align-items: stretch; gap: 8px; } | |
| #composer textarea { font-size: 0.975rem; padding: 12px 14px !important; | |
| background: var(--block-background-fill) !important; | |
| border: 1px solid var(--border-color-primary) !important; | |
| border-radius: 10px !important; } | |
| #composer textarea:focus { border-color: var(--button-primary-background-fill) !important; | |
| box-shadow: 0 0 0 3px color-mix(in srgb, | |
| var(--button-primary-background-fill) 20%, transparent) !important; } | |
| #composer button { min-width: 88px; border-radius: 10px; } | |
| """ | |
| def create_app(): | |
| with gr.Blocks(title="Endpoint chat") as demo: | |
| with gr.Row(elem_id="header", equal_height=True): | |
| gr.Markdown( | |
| "# Endpoint chat\nChat with your Hugging Face Inference Endpoints.", | |
| ) | |
| new_chat_button = gr.Button("New chat", size="sm", scale=0, min_width=110) | |
| model_selector = gr.Radio( | |
| choices=[(MODEL_LABELS[model], model) for model in AVAILABLE_MODELS], | |
| value=DEFAULT_MODEL, | |
| label="Model", | |
| show_label=False, | |
| container=False, | |
| elem_id="model-picker", | |
| ) | |
| model_id = gr.Markdown(model_caption(DEFAULT_MODEL), elem_id="model-id") | |
| chatbot = gr.Chatbot( | |
| show_label=False, | |
| height="calc(100dvh - 330px)", | |
| min_height=320, | |
| placeholder=EMPTY_STATE, | |
| examples=EXAMPLES, | |
| reasoning_tags=[(THINK_OPEN, THINK_CLOSE)], | |
| buttons=["copy", "copy_all"], | |
| feedback_options=None, | |
| elem_id="chat", | |
| ) | |
| with gr.Row(elem_id="composer"): | |
| msg = gr.Textbox( | |
| label="Message", | |
| show_label=False, | |
| placeholder=INPUT_PLACEHOLDER, | |
| lines=1, | |
| max_lines=8, | |
| autofocus=True, | |
| container=False, | |
| scale=1, | |
| ) | |
| submit_button = gr.Button("Send", variant="primary", scale=0) | |
| stop_button = gr.Button( | |
| "Stop", variant="stop", scale=0, interactive=False, visible=False | |
| ) | |
| controls = [msg, submit_button, stop_button] | |
| chat_outputs = [chatbot, *controls] | |
| respond_inputs = [msg, chatbot, model_selector] | |
| send_events = [ | |
| trigger(fn=respond, inputs=respond_inputs, outputs=chat_outputs) | |
| for trigger in (submit_button.click, msg.submit) | |
| ] | |
| send_events.append( | |
| chatbot.example_select(fn=example_text, outputs=msg).then( | |
| fn=respond, inputs=respond_inputs, outputs=chat_outputs | |
| ) | |
| ) | |
| # Cancelling skips the generator's cleanup, so Stop does it itself. | |
| stop_button.click( | |
| fn=cancel_stream, | |
| inputs=[chatbot], | |
| outputs=chat_outputs, | |
| cancels=send_events, | |
| ) | |
| new_chat_button.click(fn=new_chat, outputs=chat_outputs, cancels=send_events) | |
| model_selector.change(fn=model_caption, inputs=model_selector, outputs=model_id) | |
| return demo | |
| # Spaces' reloader looks up `demo` in __main__. | |
| demo = create_app() | |
| if __name__ == "__main__": | |
| if not LOGIN_PASSWORD: | |
| print( | |
| "WARNING: LOGIN_PASSWORD is not set - the app is starting without " | |
| "authentication and anyone who can reach it can spend endpoint quota." | |
| ) | |
| demo.launch(auth=verify_login if LOGIN_PASSWORD else None, theme=THEME, css=CSS) | |