diff --git a/README.md b/README.md
index c7764ac..b3b4114 100644
--- a/README.md
+++ b/README.md
@@ -35,6 +35,7 @@ Loaded from `.env` via `python-dotenv`:
| `OPENAI_API_KEY` | Yes | — | API key for the LLM provider |
| `OPENAI_API_URL` | No | `https://api.openai.com/v1/chat/completions` | OpenAI-compatible endpoint |
| `MODEL` | No | `gpt-4o-mini` | Model name to use |
+| `VALKEY_URL` | No | `redis://localhost:6379/0` | Valkey/Redis URL |
## Site types
diff --git a/app.py b/app.py
index 0f8371c..3ba13de 100644
--- a/app.py
+++ b/app.py
@@ -1,10 +1,11 @@
import base64
-import importlib
import json
import logging
import os
import random
import re
+import signal
+import sys
import threading
import time
import uuid
@@ -21,7 +22,7 @@ if os.getenv("GUNICORN_WORKER_ID"):
monkey.patch_all()
-lxml_html: Any = importlib.import_module("lxml.html")
+from lxml import html as lxml_html # type: ignore[import-untyped]
logging.basicConfig(
level=logging.INFO,
@@ -40,6 +41,68 @@ MODEL = os.getenv("MODEL", "gpt-4o-mini")
VALKEY_URL = os.getenv("VALKEY_URL", "redis://localhost:6379/0")
r = valkey.Valkey.from_url(VALKEY_URL, decode_responses=True) # type: ignore[assignment]
+
+class ValkeyClient:
+ def __init__(self, client: valkey.Valkey) -> None:
+ self._c = client
+
+ def get(self, key: str) -> str | None:
+ v = self._c.get(key) # type: ignore[union-attr]
+ return str(v) if isinstance(v, bytes) else v # type: ignore[return-value]
+
+ def set(self, key: str, value: str, ex: int | None = None, nx: bool = False, xx: bool = False) -> bool:
+ return bool(self._c.set(key, value, ex=ex, nx=nx, xx=xx)) # type: ignore[union-attr]
+
+ def delete(self, *keys: str) -> int:
+ return int(self._c.delete(*keys)) # type: ignore[union-attr]
+
+ def exists(self, key: str) -> bool:
+ return bool(self._c.exists(key)) # type: ignore[union-attr]
+
+ def rpush(self, key: str, *values: str) -> int:
+ return int(self._c.rpush(key, *values)) # type: ignore[union-attr]
+
+ def lpop(self, key: str) -> str | None:
+ v = self._c.lpop(key) # type: ignore[union-attr]
+ return str(v) if isinstance(v, bytes) else v # type: ignore[return-value]
+
+ def lrange(self, key: str, start: int, end: int) -> list[str]:
+ v = self._c.lrange(key, start, end) # type: ignore[union-attr]
+ return [str(x) for x in v] # type: ignore[arg-type]
+
+ def llen(self, key: str) -> int:
+ return int(self._c.llen(key)) # type: ignore[union-attr]
+
+ def ttl(self, key: str) -> int:
+ return int(self._c.ttl(key)) # type: ignore[union-attr]
+
+ def hset(self, key: str, mapping: dict[str, str]) -> int:
+ return int(self._c.hset(key, mapping=mapping)) # type: ignore[union-attr]
+
+ def hgetall(self, key: str) -> dict[str, str]:
+ v = self._c.hgetall(key) # type: ignore[union-attr]
+ return {str(k): str(vv) for k, vv in v.items()} # type: ignore[arg-type]
+
+ def incr(self, key: str) -> int:
+ return int(self._c.incr(key)) # type: ignore[union-attr]
+
+ def expire(self, key: str, seconds: int) -> bool:
+ return bool(self._c.expire(key, seconds)) # type: ignore[union-attr]
+
+
+rv = ValkeyClient(r)
+
+CHUNK_SIZE = 1000
+CONFIRMATION_TIMEOUT_ITERATIONS = 150
+CONFIRMATION_POLL_INTERVAL = 0.1
+HEARTBEAT_TIMEOUT_S = 10
+HEARTBEAT_CHECK_INTERVAL_S = 2
+WORKER_LOCK_TTL_S = 300
+WORKER_LOCK_HEARTBEAT_INTERVAL_S = 60
+PROCESSING_LOCK_TTL_S = 120
+HB_KEY_TTL_S = 15
+REQUEST_CTX_TTL_S = 30
+
worker_running = False
worker_thread: threading.Thread | None = None
@@ -427,16 +490,16 @@ def validate_html(html: str) -> list[str]:
errors.append("Missing DOCTYPE")
try:
- doc: Any = lxml_html.fromstring(html)
- html_tags: list[Any] = doc.xpath("//html")
- body_tags: list[Any] = doc.xpath("//body")
- head_tags: list[Any] = doc.xpath("//head")
- if len(html_tags) != 1:
- errors.append(f"Expected 1 tag, found {len(html_tags)}")
- if len(body_tags) != 1:
- errors.append(f"Expected 1 tag, found {len(body_tags)}")
- if len(head_tags) != 1:
- errors.append(f"Expected 1 tag, found {len(head_tags)}")
+ doc: Any = lxml_html.fromstring(html) # type: ignore[union-attr]
+ html_tags: list[Any] = doc.xpath("//html") # type: ignore[assignment,union-attr]
+ body_tags: list[Any] = doc.xpath("//body") # type: ignore[assignment,union-attr]
+ head_tags: list[Any] = doc.xpath("//head") # type: ignore[assignment,union-attr]
+ if len(html_tags) != 1: # type: ignore[arg-type]
+ errors.append(f"Expected 1 tag, found {len(html_tags)}") # type: ignore[arg-type]
+ if len(body_tags) != 1: # type: ignore[arg-type]
+ errors.append(f"Expected 1 tag, found {len(body_tags)}") # type: ignore[arg-type]
+ if len(head_tags) != 1: # type: ignore[arg-type]
+ errors.append(f"Expected 1 tag, found {len(head_tags)}") # type: ignore[arg-type]
except BaseException as e: # noqa: BLE001
errors.append(f"HTML parse error: {e}")
@@ -455,15 +518,18 @@ def validate_html(html: str) -> list[str]:
errors.append(f"JS block {i + 1}: mismatched brackets")
if script_blocks:
- js_text = "\n".join(script_blocks)
- review = call_llm(
- f"Review this JavaScript for runtime errors (undefined variables, wrong selectors, syntax issues). List each issue concisely. If no issues, say 'No issues found.'\n\n\n{js_text}\n"
- )
- if review and "no issue" not in review.lower():
- for line in review.strip().split("\n"):
- line = line.strip()
- if line and not line.startswith("<"):
- errors.append(f"JS review: {line}")
+ try:
+ js_text = "\n".join(script_blocks)
+ review = call_llm(
+ f"Review this JavaScript for runtime errors (undefined variables, wrong selectors, syntax issues). List each issue concisely. If no issues, say 'No issues found.'\n\n\n{js_text}\n"
+ )
+ if review and "no issue" not in review.lower():
+ for line in review.strip().split("\n"):
+ line = line.strip()
+ if line and not line.startswith("<"):
+ errors.append(f"JS review: {line}")
+ except BaseException as e: # noqa: BLE001
+ logger.warning(f"JS review failed: {e}")
return errors
@@ -580,117 +646,111 @@ THEMES = [
def enqueue_request(context: dict[str, Any]) -> tuple[str, int]:
request_id = uuid.uuid4().hex[:12]
ctx_key = f"wotf:req:{request_id}:ctx"
- r.hset(ctx_key, mapping={k: json.dumps(v) for k, v in context.items()}) # type: ignore[union-attr]
- r.set(f"wotf:req:{request_id}:status", "pending") # type: ignore[union-attr]
- r.rpush("wotf:queue", request_id) # type: ignore[union-attr]
- position_raw = r.llen("wotf:queue") # type: ignore[union-attr]
- position: int = int(position_raw) if isinstance(position_raw, (int, float)) else 0
+ rv.hset(ctx_key, mapping={k: json.dumps(v) for k, v in context.items()})
+ rv.set(f"wotf:req:{request_id}:status", "pending")
+ rv.rpush("wotf:queue", request_id)
+ position = rv.llen("wotf:queue")
return (request_id, position)
def get_queue_info(request_id: str) -> dict[str, Any]:
- queue_items_raw = r.lrange("wotf:queue", 0, -1) # type: ignore[union-attr]
- queue_items: list[str] = list(queue_items_raw) if isinstance(queue_items_raw, list) else [] # type: ignore[arg-type]
+ queue_items = rv.lrange("wotf:queue", 0, -1)
if request_id in queue_items:
position = queue_items.index(request_id) + 1
else:
position = 0
- avg_raw = r.get("wotf:stats:avg_time") # type: ignore[union-attr]
+ avg_raw = rv.get("wotf:stats:avg_time")
avg_str: str = str(avg_raw) if avg_raw is not None else "15"
avg = float(avg_str)
eta = position * avg
return {"position": position, "avg_time": avg, "eta": eta}
-def get_generated_html(request_id: str) -> str:
- result = r.get(f"wotf:req:{request_id}:html") # type: ignore[union-attr]
- return result if isinstance(result, str) else ""
-
-
def worker_loop() -> None:
worker_id = f"worker-{os.getpid()}-{threading.current_thread().name}"
hb_key = "wotf:processing:hb"
+ lock_hb_start = time.time()
while worker_running:
request_id: str | None = None
try:
- lock_holder = r.get("wotf:processing") # type: ignore[union-attr]
+ lock_holder = rv.get("wotf:processing")
if lock_holder is not None:
- hb_ttl_raw = r.ttl(hb_key) # type: ignore[union-attr]
- hb_ttl: int = int(hb_ttl_raw) if isinstance(hb_ttl_raw, (int, float)) else -2
+ hb_ttl = rv.ttl(hb_key)
if hb_ttl == -2:
logger.info(f"[worker] Reclaiming stale lock from {lock_holder}")
- r.delete("wotf:processing") # type: ignore[union-attr]
- r.delete(hb_key) # type: ignore[union-attr]
+ rv.delete("wotf:processing")
+ rv.delete(hb_key)
else:
time.sleep(0.5)
continue
- acquired = r.set("wotf:processing", worker_id, ex=120, nx=True) # type: ignore[union-attr]
+ acquired = rv.set("wotf:processing", worker_id, ex=PROCESSING_LOCK_TTL_S, nx=True)
if not acquired:
time.sleep(0.5)
continue
- r.set(hb_key, worker_id, ex=15) # type: ignore[union-attr]
+ rv.set(hb_key, worker_id, ex=HB_KEY_TTL_S)
- request_id_raw = r.lpop("wotf:queue") # type: ignore[union-attr]
- if request_id_raw is None:
- r.delete("wotf:processing") # type: ignore[union-attr]
- r.delete(hb_key) # type: ignore[union-attr]
- time.sleep(0.5)
- continue
-
- request_id = request_id_raw if isinstance(request_id_raw, str) else None
+ request_id = rv.lpop("wotf:queue")
if request_id is None:
- r.delete("wotf:processing") # type: ignore[union-attr]
- r.delete(hb_key) # type: ignore[union-attr]
+ rv.delete("wotf:processing")
+ rv.delete(hb_key)
time.sleep(0.5)
continue
- r.set(f"wotf:req:{request_id}:status", "waiting_confirmation") # type: ignore[union-attr]
- r.delete(f"wotf:stream:{request_id}:chunks") # type: ignore[union-attr]
+ rv.set(f"wotf:req:{request_id}:status", "waiting_confirmation")
+ rv.delete(f"wotf:stream:{request_id}:chunks")
- r.rpush(f"wotf:stream:{request_id}:chunks", f"[ID:{request_id}]") # type: ignore[union-attr]
- r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:waiting_confirmation]") # type: ignore[union-attr]
+ rv.rpush(f"wotf:stream:{request_id}:chunks", f"[ID:{request_id}]")
+ rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:waiting_confirmation]")
- r.delete("wotf:processing") # type: ignore[union-attr]
- r.delete(hb_key) # type: ignore[union-attr]
+ rv.delete("wotf:processing")
+ rv.delete(hb_key)
confirmed = False
- for _ in range(150):
- time.sleep(0.1)
- if r.get(f"wotf:req:{request_id}:confirmed"): # type: ignore[union-attr]
+ for _ in range(CONFIRMATION_TIMEOUT_ITERATIONS):
+ time.sleep(CONFIRMATION_POLL_INTERVAL)
+ if rv.get(f"wotf:req:{request_id}:confirmed"):
confirmed = True
break
if not confirmed:
logger.info(f"[worker] No confirmation for {request_id}, marking as error")
- r.set(f"wotf:req:{request_id}:status", "error") # type: ignore[union-attr]
- r.set(f"wotf:req:{request_id}:error", "Worker did not receive confirmation from client") # type: ignore[union-attr]
- r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Connection to worker failed]") # type: ignore[union-attr]
- r.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]") # type: ignore[union-attr]
- r.delete(f"wotf:req:{request_id}:confirmed") # type: ignore[union-attr]
- r.expire(f"wotf:req:{request_id}:ctx", 30) # type: ignore[union-attr]
+ rv.set(f"wotf:req:{request_id}:status", "error")
+ rv.set(f"wotf:req:{request_id}:error", "Worker did not receive confirmation from client")
+ rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Connection to worker failed]")
+ rv.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]")
+ rv.delete(f"wotf:req:{request_id}:confirmed")
+ rv.expire(f"wotf:req:{request_id}:ctx", REQUEST_CTX_TTL_S)
time.sleep(0.5)
continue
- acquired = r.set("wotf:processing", worker_id, ex=120, nx=True) # type: ignore[union-attr]
+ acquired = False
+ for _ in range(5):
+ if rv.set("wotf:processing", worker_id, ex=PROCESSING_LOCK_TTL_S, nx=True):
+ acquired = True
+ break
+ time.sleep(0.5)
if not acquired:
- time.sleep(0.5)
+ logger.error(f"[worker] Failed to acquire lock for {request_id} after retries")
+ rv.set(f"wotf:req:{request_id}:status", "pending")
+ rv.delete(f"wotf:stream:{request_id}:chunks")
+ rv.delete(f"wotf:req:{request_id}:confirmed")
+ rv.rpush("wotf:queue", request_id)
continue
- r.set(hb_key, worker_id, ex=15) # type: ignore[union-attr]
- r.set(f"wotf:req:{request_id}:status", "processing") # type: ignore[union-attr]
- r.set(f"wotf:req:{request_id}:last_heartbeat", str(time.time())) # type: ignore[union-attr]
+ rv.set(hb_key, worker_id, ex=HB_KEY_TTL_S)
+ rv.set(f"wotf:req:{request_id}:status", "processing")
+ rv.set(f"wotf:req:{request_id}:last_heartbeat", "0")
- ctx_raw = r.hgetall(f"wotf:req:{request_id}:ctx") # type: ignore[union-attr]
+ ctx_raw = rv.hgetall(f"wotf:req:{request_id}:ctx")
ctx: dict[str, Any] = {}
- if isinstance(ctx_raw, dict):
- for k_str, v_raw in ctx_raw.items(): # type: ignore[union-attr]
- k: str = str(k_str) # type: ignore[arg-type]
- v: Any = v_raw # type: ignore[assignment]
- try:
- ctx[k] = json.loads(v) if isinstance(v, str) else v
- except json.JSONDecodeError:
+ for k_str, v_raw in ctx_raw.items():
+ k: str = k_str
+ v: Any = v_raw
+ try:
+ ctx[k] = json.loads(v) if isinstance(v, str) else v
+ except json.JSONDecodeError:
ctx[k] = v
start = time.time()
@@ -701,105 +761,107 @@ def worker_loop() -> None:
try:
for chunk in stream_llm(ctx):
if time.time() - hb_start >= 5:
- r.set(hb_key, worker_id, ex=15) # type: ignore[union-attr]
+ rv.set(hb_key, worker_id, ex=HB_KEY_TTL_S)
hb_start = time.time()
- if time.time() - hb_check_start >= 2:
+ if time.time() - hb_check_start >= HEARTBEAT_CHECK_INTERVAL_S:
hb_check_start = time.time()
- last_hb_raw = r.get(f"wotf:req:{request_id}:last_heartbeat") # type: ignore[union-attr]
- if last_hb_raw:
- last_hb = float(last_hb_raw) # type: ignore[arg-type]
- if time.time() - last_hb > 10:
- logger.info(f"[worker] Client disconnected (heartbeat stale), aborting {request_id}")
- r.set(f"wotf:req:{request_id}:status", "aborted") # type: ignore[union-attr]
- r.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]") # type: ignore[union-attr]
- r.delete("wotf:processing") # type: ignore[union-attr]
- r.delete(hb_key) # type: ignore[union-attr]
- aborted = True
- break
- else:
- logger.info(f"[worker] No heartbeat from client, aborting {request_id}")
- r.set(f"wotf:req:{request_id}:status", "aborted") # type: ignore[union-attr]
- r.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]") # type: ignore[union-attr]
- r.delete("wotf:processing") # type: ignore[union-attr]
- r.delete(hb_key) # type: ignore[union-attr]
+ last_hb_raw = rv.get(f"wotf:req:{request_id}:last_heartbeat")
+ if last_hb_raw is None:
+ logger.info(f"[worker] Heartbeat key missing, aborting {request_id}")
+ rv.set(f"wotf:req:{request_id}:status", "aborted")
+ rv.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]")
+ rv.delete("wotf:processing")
+ rv.delete(hb_key)
+ aborted = True
+ break
+ last_hb = float(last_hb_raw) # type: ignore[arg-type]
+ if last_hb > 0 and time.time() - last_hb > HEARTBEAT_TIMEOUT_S:
+ logger.info(f"[worker] Client disconnected (heartbeat stale), aborting {request_id}")
+ rv.set(f"wotf:req:{request_id}:status", "aborted")
+ rv.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]")
+ rv.delete("wotf:processing")
+ rv.delete(hb_key)
aborted = True
break
html_buffer.append(chunk)
encoded = base64.b64encode(chunk.encode("utf-8")).decode("ascii")
- r.rpush(f"wotf:stream:{request_id}:chunks", encoded) # type: ignore[union-attr]
+ rv.rpush(f"wotf:stream:{request_id}:chunks", encoded)
except BaseException as e: # noqa: BLE001
elapsed = time.time() - start
logger.error(f"[worker] stream_llm failed for {request_id} in {elapsed:.1f}s: {e}")
- retries_raw = r.incr(f"wotf:req:{request_id}:retries") # type: ignore[union-attr]
- retries: int = int(retries_raw) if isinstance(retries_raw, (int, float)) else 1
+ retries = rv.incr(f"wotf:req:{request_id}:retries")
if retries <= 2:
logger.info(f"[worker] Requeuing {request_id} (attempt {retries}/2)")
- r.set(f"wotf:req:{request_id}:status", "pending") # type: ignore[union-attr]
- r.delete(f"wotf:stream:{request_id}:chunks") # type: ignore[union-attr]
- r.rpush("wotf:queue", request_id) # type: ignore[union-attr]
+ rv.set(f"wotf:req:{request_id}:status", "pending")
+ rv.delete(f"wotf:stream:{request_id}:chunks")
+ rv.rpush("wotf:queue", request_id)
else:
logger.error(f"[worker] {request_id} failed after {retries} attempts")
- r.set(f"wotf:req:{request_id}:status", "error") # type: ignore[union-attr]
- r.set(f"wotf:req:{request_id}:error", str(e)) # type: ignore[union-attr]
- r.delete("wotf:processing") # type: ignore[union-attr]
- r.delete(hb_key) # type: ignore[union-attr]
+ rv.set(f"wotf:req:{request_id}:status", "error")
+ rv.set(f"wotf:req:{request_id}:error", str(e))
+ rv.delete("wotf:processing")
+ rv.delete(hb_key)
time.sleep(1)
continue
if aborted:
time.sleep(1)
continue
html = "".join(html_buffer)
- r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Validating HTML structure...]") # type: ignore[union-attr]
+ rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Validating HTML structure...]")
errors = validate_html(html)
if errors:
short = summarize_errors(errors)
- r.rpush(f"wotf:stream:{request_id}:chunks", f"[STATUS:Found {len(errors)} issue(s): {short}. Fixing (attempt 1/2)...]") # type: ignore[union-attr]
+ rv.rpush(f"wotf:stream:{request_id}:chunks", f"[STATUS:Found {len(errors)} issue(s): {short}. Fixing (attempt 1/2)...]")
html = fix_html(html, errors)
errors = validate_html(html)
if errors:
short = summarize_errors(errors)
- r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Still invalid. Fixing (attempt 2/2)...]") # type: ignore[union-attr]
+ rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Still invalid. Fixing (attempt 2/2)...]")
html = fix_html(html, errors)
errors = validate_html(html)
if errors:
short = summarize_errors(errors)
- r.rpush(f"wotf:stream:{request_id}:chunks", f"[STATUS:Still has {len(errors)} issue(s): {short}. Serving as-is.]") # type: ignore[union-attr]
+ rv.rpush(f"wotf:stream:{request_id}:chunks", f"[STATUS:Still has {len(errors)} issue(s): {short}. Serving as-is.]")
else:
- r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Fixed successfully.]") # type: ignore[union-attr]
- r.rpush(f"wotf:stream:{request_id}:chunks", "[CLEAR]") # type: ignore[union-attr]
- for i in range(0, len(html), 1000):
- chunk = html[i : i + 1000]
+ rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Fixed successfully.]")
+ rv.rpush(f"wotf:stream:{request_id}:chunks", "[CLEAR]")
+ for i in range(0, len(html), CHUNK_SIZE):
+ chunk = html[i : i + CHUNK_SIZE]
encoded = base64.b64encode(chunk.encode("utf-8")).decode("ascii")
- r.rpush(f"wotf:stream:{request_id}:chunks", encoded) # type: ignore[union-attr]
+ rv.rpush(f"wotf:stream:{request_id}:chunks", encoded)
else:
- r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Validation passed.]") # type: ignore[union-attr]
- r.rpush(f"wotf:stream:{request_id}:chunks", "[RENDER]") # type: ignore[union-attr]
+ rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Validation passed.]")
+ rv.rpush(f"wotf:stream:{request_id}:chunks", "[RENDER]")
elapsed = time.time() - start
- r.set(f"wotf:req:{request_id}:status", "done") # type: ignore[union-attr]
+ rv.set(f"wotf:req:{request_id}:status", "done")
- n_raw = r.get("wotf:stats:gen_count") # type: ignore[union-attr]
+ n_raw = rv.get("wotf:stats:gen_count")
n_str: str = str(n_raw) if n_raw is not None else "0"
n = int(n_str)
- avg_raw = r.get("wotf:stats:avg_time") # type: ignore[union-attr]
+ avg_raw = rv.get("wotf:stats:avg_time")
avg_str: str = str(avg_raw) if avg_raw is not None else "15"
avg = float(avg_str)
new_avg = (avg * n + elapsed) / (n + 1) if n > 0 else elapsed
- r.set("wotf:stats:avg_time", str(new_avg)) # type: ignore[union-attr]
- r.set("wotf:stats:gen_count", str(n + 1)) # type: ignore[union-attr]
+ rv.set("wotf:stats:avg_time", str(new_avg))
+ rv.set("wotf:stats:gen_count", str(n + 1))
- r.delete("wotf:processing") # type: ignore[union-attr]
- r.delete(hb_key) # type: ignore[union-attr]
+ rv.delete("wotf:processing")
+ rv.delete(hb_key)
logger.info(f"[worker] Generated {request_id} in {elapsed:.1f}s (avg={new_avg:.1f}s)")
+ if time.time() - lock_hb_start >= WORKER_LOCK_HEARTBEAT_INTERVAL_S:
+ lock_hb_start = time.time()
+ rv.set("wotf:worker_lock", "1", xx=True, ex=WORKER_LOCK_TTL_S)
+
except BaseException as e: # noqa: BLE001
logger.error(f"[worker] Error: {e}")
if request_id:
- r.set(f"wotf:req:{request_id}:status", "error") # type: ignore[union-attr]
- r.set(f"wotf:req:{request_id}:error", str(e)) # type: ignore[union-attr]
- r.delete("wotf:processing") # type: ignore[union-attr]
- r.delete(hb_key) # type: ignore[union-attr]
+ rv.set(f"wotf:req:{request_id}:status", "error")
+ rv.set(f"wotf:req:{request_id}:error", str(e))
+ rv.delete("wotf:processing")
+ rv.delete(hb_key)
time.sleep(1)
@@ -894,52 +956,27 @@ def stream_llm(context: dict[str, Any]) -> Generator[str]:
yield f"\n\n"
-def sse_stream(context: dict[str, Any]) -> Generator[str]:
- buffer: list[str] = []
+@app.route("/health")
+def health() -> Response:
+ try:
+ valkey_ok = rv.get("wotf:health_check") is not None or rv.set("wotf:health_check", "1", ex=60)
+ except BaseException: # noqa: BLE001
+ valkey_ok = False
- def status(msg: str) -> str:
- return f"data: [STATUS:{msg}]\n\n"
+ queue_depth = rv.llen("wotf:queue")
+ worker_active = rv.exists("wotf:worker_lock")
- yield status("Generating HTML...")
- for chunk in stream_llm(context):
- buffer.append(chunk)
- encoded = base64.b64encode(chunk.encode("utf-8")).decode("ascii")
- yield f"data: {encoded}\n\n"
+ status = "ok" if (valkey_ok and worker_active) else "degraded"
- html = "".join(buffer)
- yield "data: [DONE]\n\n"
-
- yield status("Validating HTML structure...")
- errors = validate_html(html)
-
- if errors:
- short = summarize_errors(errors)
- yield status(f"Found {len(errors)} issue(s): {short}. Fixing (attempt 1/2)...")
-
- html = fix_html(html, errors)
- errors = validate_html(html)
-
- if errors:
- short = summarize_errors(errors)
- yield status("Still invalid. Fixing (attempt 2/2)...")
- html = fix_html(html, errors)
- errors = validate_html(html)
-
- if errors:
- short = summarize_errors(errors)
- yield status(f"Still has {len(errors)} issue(s): {short}. Serving as-is.")
- else:
- yield status("Fixed successfully.")
-
- yield "data: [CLEAR]\n\n"
- for i in range(0, len(html), 1000):
- chunk = html[i : i + 1000]
- encoded = base64.b64encode(chunk.encode("utf-8")).decode("ascii")
- yield f"data: {encoded}\n\n"
- else:
- yield status("Validation passed.")
-
- yield "data: [RENDER]\n\n"
+ return Response(
+ json.dumps({
+ "status": status,
+ "valkey_connected": valkey_ok,
+ "worker_active": worker_active,
+ "queue_depth": queue_depth,
+ }),
+ mimetype="application/json",
+ )
@app.route("/", defaults={"path": ""})
@@ -962,41 +999,41 @@ def stream() -> Response:
last_idx = 0
while True:
- status_raw = r.get(f"wotf:req:{request_id}:status") # type: ignore[union-attr]
+ status_raw = rv.get(f"wotf:req:{request_id}:status")
status = status_raw if isinstance(status_raw, str) else "pending"
if status == "error":
- error_raw = r.get(f"wotf:req:{request_id}:error") # type: ignore[union-attr]
+ error_raw = rv.get(f"wotf:req:{request_id}:error")
error_msg = error_raw if isinstance(error_raw, str) else "Unknown error"
- yield f"data: [STATUS:Error: {error_msg}\n\n"
+ yield f"data: [STATUS:Error: {error_msg}]\n\n"
yield "data: [ERROR]\n\n"
return
if status == "waiting_confirmation":
- chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", 0, -1) # type: ignore[union-attr]
+ chunks_raw = rv.lrange(f"wotf:stream:{request_id}:chunks", 0, -1)
chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else [] # type: ignore[arg-type]
if chunks:
for c in chunks:
yield f"data: {c}\n\n"
else:
- yield "data: [STATUS:Waiting for your connection...\n\n"
+ yield "data: [STATUS:Waiting for your connection...]\n\n"
time.sleep(0.2)
continue
if status == "processing":
- chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1) # type: ignore[union-attr]
+ chunks_raw = rv.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1)
chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else [] # type: ignore[arg-type]
if chunks:
for c in chunks:
yield f"data: {c}\n\n"
last_idx += len(chunks)
else:
- yield "data: [STATUS:Generating your page...\n\n"
+ yield "data: [STATUS:Generating your page...]\n\n"
time.sleep(0.1)
continue
if status == "done":
- chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1) # type: ignore[union-attr]
+ chunks_raw = rv.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1)
chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else [] # type: ignore[arg-type]
for c in chunks:
yield f"data: {c}\n\n"
@@ -1006,13 +1043,13 @@ def stream() -> Response:
pos = info["position"]
eta = info["eta"]
if pos == 0:
- yield "data: [STATUS:Your turn is next...\n\n"
+ yield "data: [STATUS:Your turn is next...]\n\n"
else:
- yield f"data: [STATUS:Position {pos} in line | Est. wait: ~{eta:.0f}s\n\n"
+ yield f"data: [STATUS:Position {pos} in line | Est. wait: ~{eta:.0f}s]\n\n"
time.sleep(1)
while True:
- chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1) # type: ignore[union-attr]
+ chunks_raw = rv.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1)
chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else [] # type: ignore[arg-type]
if chunks:
for chunk in chunks:
@@ -1023,23 +1060,23 @@ def stream() -> Response:
last_idx += len(chunks)
time.sleep(0.05)
else:
- status_raw = r.get(f"wotf:req:{request_id}:status") # type: ignore[union-attr]
+ status_raw = rv.get(f"wotf:req:{request_id}:status")
status = status_raw if isinstance(status_raw, str) else "pending"
if status == "done":
- chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1) # type: ignore[union-attr]
+ chunks_raw = rv.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1)
chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else [] # type: ignore[arg-type]
if chunks:
for chunk in chunks:
yield f"data: {chunk}\n\n"
break
if status == "error":
- error_raw = r.get(f"wotf:req:{request_id}:error") # type: ignore[union-attr]
+ error_raw = rv.get(f"wotf:req:{request_id}:error")
error_msg = error_raw if isinstance(error_raw, str) else "Unknown error"
- yield f"data: [STATUS:Error: {error_msg}\n\n"
+ yield f"data: [STATUS:Error: {error_msg}]\n\n"
yield "data: [ERROR]\n\n"
return
if status == "aborted":
- yield "data: [STATUS:Generation aborted (connection lost)\n\n"
+ yield "data: [STATUS:Generation aborted (connection lost)]\n\n"
yield "data: [ERROR]\n\n"
return
time.sleep(0.2)
@@ -1057,19 +1094,19 @@ def stream() -> Response:
@app.route("/confirm/", methods=["POST"])
def confirm(request_id: str) -> Response:
- ctx = r.exists(f"wotf:req:{request_id}:ctx") # type: ignore[union-attr]
+ ctx = rv.exists(f"wotf:req:{request_id}:ctx")
if not ctx:
return Response("not found", status=404)
- r.set(f"wotf:req:{request_id}:confirmed", "1") # type: ignore[union-attr]
+ rv.set(f"wotf:req:{request_id}:confirmed", "1")
return Response("ok", status=200)
@app.route("/heartbeat/", methods=["POST"])
def heartbeat(request_id: str) -> Response:
- status = r.get(f"wotf:req:{request_id}:status") # type: ignore[union-attr]
+ status = rv.get(f"wotf:req:{request_id}:status")
if status is None:
return Response("not found", status=404)
- r.set(f"wotf:req:{request_id}:last_heartbeat", str(time.time())) # type: ignore[union-attr]
+ rv.set(f"wotf:req:{request_id}:last_heartbeat", str(time.time()))
return Response("ok", status=200)
@@ -1077,7 +1114,7 @@ def start_worker_if_needed() -> None:
global worker_running, worker_thread
if worker_running:
return
- claimed = r.set("wotf:worker_lock", "1", nx=True, ex=300) # type: ignore[union-attr]
+ claimed = rv.set("wotf:worker_lock", "1", nx=True, ex=WORKER_LOCK_TTL_S)
if claimed:
worker_running = True
if os.getenv("GUNICORN_WORKER_ID"):
@@ -1097,7 +1134,21 @@ def gunicorn_post_fork(server: Any, worker: Any) -> None:
start_worker_if_needed()
+def graceful_shutdown(signum: int, frame: Any) -> None:
+ global worker_running
+ sig_name = "SIGINT" if signum == signal.SIGINT else "SIGTERM"
+ logger.info(f"[shutdown] Received {sig_name}, stopping worker...")
+ worker_running = False
+ rv.delete("wotf:worker_lock")
+ rv.delete("wotf:processing")
+ rv.delete("wotf:processing:hb")
+ logger.info("[shutdown] Locks released, exiting")
+ sys.exit(0)
+
+
if __name__ == "__main__":
+ signal.signal(signal.SIGINT, graceful_shutdown)
+ signal.signal(signal.SIGTERM, graceful_shutdown)
worker_running = True
worker_thread = threading.Thread(target=worker_loop, daemon=True)
worker_thread.start()
diff --git a/pyproject.toml b/pyproject.toml
index add17a5..bb4387c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
[project]
name = "web-on-the-fly"
version = "0.1.0"
-description = "Add your description here"
+description = "Flask app that streams AI-generated fake websites via SSE"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [