1105 lines
44 KiB
Python
1105 lines
44 KiB
Python
import base64
|
|
import importlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import random
|
|
import re
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from collections.abc import Generator
|
|
from typing import Any
|
|
|
|
import httpx
|
|
import valkey
|
|
from dotenv import load_dotenv
|
|
from flask import Flask, Request, Response, request, stream_with_context
|
|
|
|
if os.getenv("GUNICORN_WORKER_ID"):
|
|
from gevent import monkey
|
|
|
|
monkey.patch_all()
|
|
|
|
lxml_html: Any = importlib.import_module("lxml.html")
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s [%(levelname)s] %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
load_dotenv()
|
|
app = Flask(__name__)
|
|
|
|
API_URL = os.getenv("OPENAI_API_URL", "https://api.openai.com/v1/chat/completions")
|
|
API_KEY = os.getenv("OPENAI_API_KEY", "")
|
|
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]
|
|
|
|
worker_running = False
|
|
worker_thread: threading.Thread | None = None
|
|
|
|
logger.info(f"Using model: {MODEL}")
|
|
logger.info(f"API URL: {API_URL}")
|
|
logger.info(f"API key present: {bool(API_KEY)}")
|
|
|
|
|
|
LOADING_PAGE = """<!DOCTYPE html>
|
|
<html lang="en">
|
|
<head>
|
|
<meta charset="UTF-8">
|
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
|
<title>Generating your page...</title>
|
|
<style>
|
|
* { margin: 0; padding: 0; box-sizing: border-box; }
|
|
body {
|
|
min-height: 100vh;
|
|
display: flex;
|
|
flex-direction: column;
|
|
align-items: center;
|
|
justify-content: center;
|
|
background: #0a0a0f;
|
|
color: #e0e0e0;
|
|
font-family: system-ui, -apple-system, sans-serif;
|
|
overflow: hidden;
|
|
}
|
|
.spinner {
|
|
width: 80px;
|
|
height: 80px;
|
|
border: 4px solid rgba(100, 100, 255, 0.15);
|
|
border-top-color: #6366f1;
|
|
border-radius: 50%;
|
|
animation: spin 0.8s linear infinite;
|
|
}
|
|
@keyframes spin { to { transform: rotate(360deg); } }
|
|
.title {
|
|
margin-top: 24px;
|
|
font-size: 1.5rem;
|
|
font-weight: 600;
|
|
letter-spacing: 0.02em;
|
|
}
|
|
.subtitle {
|
|
margin-top: 8px;
|
|
font-size: 0.9rem;
|
|
color: #888;
|
|
}
|
|
.stats {
|
|
margin-top: 20px;
|
|
font-size: 0.85rem;
|
|
color: #6366f1;
|
|
font-variant-numeric: tabular-nums;
|
|
}
|
|
.retry-btn {
|
|
margin-top: 20px;
|
|
padding: 10px 24px;
|
|
background: rgba(99, 102, 241, 0.2);
|
|
color: #e0e0e0;
|
|
border: 1px solid rgba(99, 102, 241, 0.5);
|
|
border-radius: 8px;
|
|
font-size: 0.95rem;
|
|
cursor: pointer;
|
|
font-family: system-ui, -apple-system, sans-serif;
|
|
display: none;
|
|
transition: all 0.15s ease;
|
|
}
|
|
.retry-btn:hover {
|
|
background: rgba(99, 102, 241, 0.4);
|
|
border-color: #6366f1;
|
|
color: #fff;
|
|
}
|
|
.particles {
|
|
position: fixed;
|
|
inset: 0;
|
|
pointer-events: none;
|
|
overflow: hidden;
|
|
}
|
|
.particle {
|
|
position: absolute;
|
|
width: 3px;
|
|
height: 3px;
|
|
background: #6366f1;
|
|
border-radius: 50%;
|
|
opacity: 0;
|
|
animation: float 3s ease-in-out infinite;
|
|
}
|
|
@keyframes float {
|
|
0% { opacity: 0; transform: translateY(100vh) scale(0); }
|
|
50% { opacity: 0.6; }
|
|
100% { opacity: 0; transform: translateY(-20vh) scale(1); }
|
|
}
|
|
</style>
|
|
</head>
|
|
<body>
|
|
<div class="particles" id="particles"></div>
|
|
<div class="spinner"></div>
|
|
<div class="title">Crafting your page...</div>
|
|
<div class="subtitle">The AI is designing something unique for you</div>
|
|
<div class="stats" id="stats">Characters: 0</div>
|
|
<button class="retry-btn" id="retryBtn">Retry</button>
|
|
<script>
|
|
const stats = document.getElementById('stats');
|
|
const titleEl = document.querySelector('.title');
|
|
const retryBtn = document.getElementById('retryBtn');
|
|
const particles = document.getElementById('particles');
|
|
for (let i = 0; i < 30; i++) {
|
|
const p = document.createElement('div');
|
|
p.className = 'particle';
|
|
p.style.left = Math.random() * 100 + '%';
|
|
p.style.animationDelay = Math.random() * 3 + 's';
|
|
p.style.animationDuration = (2 + Math.random() * 2) + 's';
|
|
particles.appendChild(p);
|
|
}
|
|
retryBtn.onclick = () => { window.location.href = '/'; };
|
|
const start = Date.now();
|
|
const es = new EventSource('/stream' + window.location.search + window.location.hash);
|
|
let chars = 0;
|
|
let htmlChunks = [];
|
|
let htmlSize = 0;
|
|
let mode = 'idle';
|
|
let requestId = null;
|
|
let heartbeatInterval = null;
|
|
let waitingConfirmStart = null;
|
|
let statusStart = null;
|
|
const interval = setInterval(() => {
|
|
if (mode === 'generating') {
|
|
stats.textContent = 'Characters: ' + chars.toLocaleString() + ' | ' + ((Date.now() - start) / 1000).toFixed(1) + 's';
|
|
}
|
|
if (mode === 'waiting_confirmation' && waitingConfirmStart && Date.now() - waitingConfirmStart > 10000) {
|
|
es.close();
|
|
clearInterval(interval);
|
|
clearInterval(heartbeatInterval);
|
|
clearTimeout(timeout);
|
|
titleEl.textContent = 'Connection lost';
|
|
document.querySelector('.subtitle').textContent = 'Could not connect to worker. Please refresh.';
|
|
retryBtn.style.display = 'inline-block';
|
|
}
|
|
if (mode === 'status' && statusStart && Date.now() - statusStart > 60000) {
|
|
es.close();
|
|
clearInterval(interval);
|
|
clearInterval(heartbeatInterval);
|
|
clearTimeout(timeout);
|
|
titleEl.textContent = 'Connection lost';
|
|
document.querySelector('.subtitle').textContent = 'Could not connect to worker. Please refresh.';
|
|
retryBtn.style.display = 'inline-block';
|
|
}
|
|
}, 100);
|
|
const timeout = setTimeout(() => {
|
|
es.close();
|
|
clearInterval(interval);
|
|
clearInterval(heartbeatInterval);
|
|
titleEl.textContent = 'Request timed out';
|
|
document.querySelector('.subtitle').textContent = 'Generation took too long. Please try again.';
|
|
retryBtn.style.display = 'inline-block';
|
|
}, 120000);
|
|
window.__wotf_html = '';
|
|
es.onmessage = (e) => {
|
|
if (e.data === '[DONE]') {
|
|
clearInterval(heartbeatInterval);
|
|
heartbeatInterval = null;
|
|
titleEl.textContent = 'Validating HTML...';
|
|
} else if (e.data === '[RENDER]') {
|
|
clearInterval(heartbeatInterval);
|
|
heartbeatInterval = null;
|
|
clearTimeout(timeout);
|
|
clearInterval(interval);
|
|
es.close();
|
|
titleEl.textContent = 'Rendering...';
|
|
const fullHtml = htmlChunks.join('');
|
|
window.__wotf_html = fullHtml;
|
|
document.open();
|
|
document.write(fullHtml);
|
|
document.close();
|
|
injectMenu();
|
|
} else if (e.data === '[CLEAR]') {
|
|
htmlChunks = [];
|
|
htmlSize = 0;
|
|
chars = 0;
|
|
} else if (e.data.startsWith('[ID:')) {
|
|
requestId = e.data.slice(4, -1);
|
|
} else if (e.data.startsWith('[STATUS:')) {
|
|
const msg = e.data.slice(8);
|
|
const cleanMsg = msg.endsWith(']') ? msg.slice(0, -1) : msg;
|
|
if (cleanMsg === 'waiting_confirmation') {
|
|
mode = 'waiting_confirmation';
|
|
waitingConfirmStart = Date.now();
|
|
if (requestId && !window.__confirmed) {
|
|
window.__confirmed = true;
|
|
fetch('/confirm/' + requestId, { method: 'POST' }).catch(() => {
|
|
es.close();
|
|
clearInterval(interval);
|
|
clearInterval(heartbeatInterval);
|
|
clearTimeout(timeout);
|
|
titleEl.textContent = 'Connection to worker failed';
|
|
document.querySelector('.subtitle').textContent = 'Could not connect to generation worker. Please refresh.';
|
|
retryBtn.style.display = 'inline-block';
|
|
});
|
|
}
|
|
} else if (mode !== 'generating') {
|
|
stats.textContent = cleanMsg;
|
|
mode = 'status';
|
|
waitingConfirmStart = null;
|
|
statusStart = Date.now();
|
|
}
|
|
} else if (e.data === '[ERROR]') {
|
|
clearInterval(heartbeatInterval);
|
|
heartbeatInterval = null;
|
|
clearTimeout(timeout);
|
|
clearInterval(interval);
|
|
es.close();
|
|
titleEl.textContent = 'Error generating page';
|
|
document.querySelector('.subtitle').textContent = 'Please refresh and try again';
|
|
retryBtn.style.display = 'inline-block';
|
|
} else {
|
|
mode = 'generating';
|
|
statusStart = null;
|
|
if (!heartbeatInterval && requestId) {
|
|
heartbeatInterval = setInterval(() => {
|
|
fetch('/heartbeat/' + requestId, { method: 'POST' });
|
|
}, 3000);
|
|
}
|
|
const bytes = Uint8Array.from(atob(e.data), c => c.charCodeAt(0));
|
|
const decoded = new TextDecoder().decode(bytes);
|
|
htmlChunks.push(decoded);
|
|
htmlSize += decoded.length;
|
|
chars = htmlSize;
|
|
}
|
|
};
|
|
es.onerror = () => {
|
|
clearInterval(heartbeatInterval);
|
|
heartbeatInterval = null;
|
|
clearTimeout(timeout);
|
|
clearInterval(interval);
|
|
es.close();
|
|
titleEl.textContent = 'Connection error';
|
|
document.querySelector('.subtitle').textContent = 'Please refresh and try again';
|
|
retryBtn.style.display = 'inline-block';
|
|
};
|
|
|
|
function injectMenu() {
|
|
const style = document.createElement('style');
|
|
style.textContent = `
|
|
#wotf-menu-trigger {
|
|
position: fixed;
|
|
top: 0;
|
|
left: 0;
|
|
right: 0;
|
|
height: 12px;
|
|
z-index: 2147483646;
|
|
cursor: default;
|
|
}
|
|
#wotf-menu {
|
|
position: fixed;
|
|
top: -52px;
|
|
left: 0;
|
|
right: 0;
|
|
height: 52px;
|
|
background: rgba(10, 10, 15, 0.95);
|
|
backdrop-filter: blur(12px);
|
|
display: flex;
|
|
align-items: center;
|
|
justify-content: center;
|
|
gap: 16px;
|
|
padding: 0 24px;
|
|
z-index: 2147483647;
|
|
transition: top 0.25s ease;
|
|
border-bottom: 1px solid rgba(99, 102, 241, 0.3);
|
|
}
|
|
#wotf-menu.visible {
|
|
top: 0;
|
|
}
|
|
#wotf-menu button {
|
|
background: rgba(99, 102, 241, 0.15);
|
|
color: #e0e0e0;
|
|
border: 1px solid rgba(99, 102, 241, 0.4);
|
|
padding: 6px 16px;
|
|
border-radius: 6px;
|
|
font-size: 0.85rem;
|
|
cursor: pointer;
|
|
transition: all 0.15s ease;
|
|
font-family: system-ui, -apple-system, sans-serif;
|
|
}
|
|
#wotf-menu button:hover {
|
|
background: rgba(99, 102, 241, 0.35);
|
|
border-color: #6366f1;
|
|
color: #fff;
|
|
}
|
|
#wotf-menu .wotf-label {
|
|
font-size: 0.75rem;
|
|
color: #666;
|
|
letter-spacing: 0.05em;
|
|
text-transform: uppercase;
|
|
}
|
|
`;
|
|
document.head.appendChild(style);
|
|
|
|
const trigger = document.createElement('div');
|
|
trigger.id = 'wotf-menu-trigger';
|
|
|
|
const menu = document.createElement('div');
|
|
menu.id = 'wotf-menu';
|
|
|
|
const label = document.createElement('span');
|
|
label.className = 'wotf-label';
|
|
label.textContent = 'Web On The Fly';
|
|
|
|
const exportBtn = document.createElement('button');
|
|
exportBtn.textContent = 'Export Code';
|
|
exportBtn.onclick = () => {
|
|
const html = window.__wotf_html || '<!-- No HTML available -->';
|
|
const blob = new Blob([html], { type: 'text/html' });
|
|
const url = URL.createObjectURL(blob);
|
|
const a = document.createElement('a');
|
|
a.href = url;
|
|
a.download = 'generated-page.html';
|
|
a.click();
|
|
URL.revokeObjectURL(url);
|
|
};
|
|
|
|
const newBtn = document.createElement('button');
|
|
newBtn.textContent = 'New Page';
|
|
newBtn.onclick = () => {
|
|
window.location.href = '/';
|
|
};
|
|
|
|
menu.appendChild(label);
|
|
menu.appendChild(exportBtn);
|
|
menu.appendChild(newBtn);
|
|
document.body.appendChild(trigger);
|
|
document.body.appendChild(menu);
|
|
|
|
let hideTimer = null;
|
|
trigger.addEventListener('mouseenter', () => {
|
|
clearTimeout(hideTimer);
|
|
menu.classList.add('visible');
|
|
});
|
|
menu.addEventListener('mouseenter', () => {
|
|
clearTimeout(hideTimer);
|
|
menu.classList.add('visible');
|
|
});
|
|
trigger.addEventListener('mouseleave', () => {
|
|
hideTimer = setTimeout(() => menu.classList.remove('visible'), 150);
|
|
});
|
|
menu.addEventListener('mouseleave', (e) => {
|
|
if (e.clientY > 52) {
|
|
menu.classList.remove('visible');
|
|
}
|
|
});
|
|
}
|
|
</script>
|
|
</body>
|
|
</html>"""
|
|
|
|
|
|
def make_httpx_client() -> httpx.Client:
|
|
return httpx.Client(timeout=300)
|
|
|
|
|
|
def call_llm(prompt: str, system: str = "You are a helpful assistant.") -> str:
|
|
headers = {"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"}
|
|
payload: dict[str, Any] = {
|
|
"model": MODEL,
|
|
"messages": [
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
"temperature": 0.7,
|
|
"stream": False,
|
|
}
|
|
try:
|
|
with make_httpx_client() as client:
|
|
resp = client.post(API_URL, headers=headers, json=payload)
|
|
resp.raise_for_status()
|
|
data: dict[str, Any] = resp.json()
|
|
return data["choices"][0]["message"]["content"]
|
|
except BaseException as e: # noqa: BLE001
|
|
logger.error(f"call_llm failed: {e}")
|
|
return ""
|
|
|
|
|
|
def validate_html(html: str) -> list[str]:
|
|
errors: list[str] = []
|
|
|
|
if not html.strip().lower().startswith("<!doctype"):
|
|
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 <html> tag, found {len(html_tags)}")
|
|
if len(body_tags) != 1:
|
|
errors.append(f"Expected 1 <body> tag, found {len(body_tags)}")
|
|
if len(head_tags) != 1:
|
|
errors.append(f"Expected 1 <head> tag, found {len(head_tags)}")
|
|
except BaseException as e: # noqa: BLE001
|
|
errors.append(f"HTML parse error: {e}")
|
|
|
|
style_blocks = re.findall(r"<style[^>]*>(.*?)</style>", html, re.DOTALL | re.IGNORECASE)
|
|
for i, block in enumerate(style_blocks):
|
|
if block.count("{") != block.count("}"):
|
|
errors.append(f"CSS block {i + 1}: mismatched braces")
|
|
|
|
script_blocks = re.findall(r"<script[^>]*>(.*?)</script>", html, re.DOTALL | re.IGNORECASE)
|
|
for i, block in enumerate(script_blocks):
|
|
if block.count("{") != block.count("}"):
|
|
errors.append(f"JS block {i + 1}: mismatched braces")
|
|
if block.count("(") != block.count(")"):
|
|
errors.append(f"JS block {i + 1}: mismatched parentheses")
|
|
if block.count("[") != block.count("]"):
|
|
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<code>\n{js_text}\n</code>"
|
|
)
|
|
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}")
|
|
|
|
return errors
|
|
|
|
|
|
def fix_html(html: str, errors: list[str]) -> str:
|
|
error_text = "\n".join(f"- {e}" for e in errors)
|
|
prompt = f"""Fix this HTML document. Issues found:
|
|
{error_text}
|
|
|
|
Rules:
|
|
- Return ONLY the complete corrected HTML document (<!DOCTYPE html> ... </html>)
|
|
- Do NOT wrap in markdown code blocks
|
|
- Do NOT add explanations
|
|
- Preserve the design and content, only fix errors
|
|
|
|
HTML to fix:
|
|
{html}"""
|
|
|
|
fixed = call_llm(prompt, system="You are an expert HTML/CSS/JS developer. You fix broken code.")
|
|
|
|
fixed = re.sub(r"^```html?\s*", "", fixed.strip(), flags=re.MULTILINE)
|
|
fixed = re.sub(r"\s*```$", "", fixed.strip(), flags=re.MULTILINE)
|
|
|
|
if fixed and "<html" in fixed.lower():
|
|
return fixed
|
|
logger.warning("fix_html: LLM returned invalid output, returning original")
|
|
return html
|
|
|
|
|
|
def summarize_errors(errors: list[str]) -> str:
|
|
seen: list[str] = []
|
|
for e in errors:
|
|
short = e[:60]
|
|
if short not in seen:
|
|
seen.append(short)
|
|
if len(seen) >= 3:
|
|
break
|
|
return ", ".join(seen)
|
|
|
|
|
|
def get_user_context(req: Request) -> dict[str, Any]:
|
|
user_agent = req.headers.get("User-Agent", "")
|
|
lang = req.headers.get("Accept-Language", "")
|
|
ip = req.remote_addr or ""
|
|
path = req.path
|
|
referer = req.headers.get("Referer", "")
|
|
cookies = dict(req.cookies)
|
|
|
|
return {
|
|
"user_agent": user_agent,
|
|
"language": lang,
|
|
"ip": ip,
|
|
"path": path,
|
|
"referer": referer,
|
|
"cookies": cookies,
|
|
}
|
|
|
|
|
|
THEMES = [
|
|
("fake_store", "A quirky online store selling unusual products (e.g., bottled laughter, cloud-shaped pillows)"),
|
|
("fake_store", "A vintage record store homepage with vinyl listings and album reviews"),
|
|
("fake_store", "A plant shop selling exotic and rare houseplants"),
|
|
("fake_store", "A retro gaming store selling refurbished consoles and cartridges"),
|
|
("fake_store", "A boutique selling handmade candles with weird scents"),
|
|
("fake_store", "A pet supply store for unusual pets (lizards, tarantulas, etc.)"),
|
|
("fake_homepage", "A personal portfolio site for a fictional photographer"),
|
|
("fake_homepage", "A musician's homepage with tour dates and album art"),
|
|
("fake_homepage", "A food blogger's site with recipe cards and photos"),
|
|
("fake_homepage", "A travel vlogger's site showcasing weird destinations"),
|
|
("fake_homepage", "A developer's portfolio with fun project descriptions"),
|
|
("fake_homepage", "An artist's gallery site with abstract paintings"),
|
|
("fake_social", "A social media profile page for a fictional influencer"),
|
|
("fake_social", "A Reddit-style thread about a bizarre conspiracy theory"),
|
|
("fake_social", "A Twitter/X feed from a fictional account"),
|
|
("fake_social", "A dating app profile page with funny bios"),
|
|
("fake_social", "A forum discussion about the best fictional pizza toppings"),
|
|
("fake_wiki", "A Wikipedia article about a fictional historical event"),
|
|
("fake_wiki", "A Wikipedia article about a fictional animal species"),
|
|
("fake_wiki", "A Wikipedia article about a fictional technology"),
|
|
("fake_wiki", "A wiki page for a fictional video game"),
|
|
("fake_wiki", "A wiki page for a fictional band or music group"),
|
|
("fake_forum", "A 2000s-era forum page with threads about tech support"),
|
|
("fake_forum", "A gaming forum with strategy discussions"),
|
|
("fake_forum", "A cooking forum with recipe debates and tips"),
|
|
("fake_forum", "A conspiracy forum with wild theories"),
|
|
("fake_blog", "A tech blog reviewing fictional gadgets"),
|
|
("fake_blog", "A lifestyle blog about minimalist living"),
|
|
("fake_blog", "A horror story blog with creepy tales"),
|
|
("fake_blog", "A science blog explaining weird phenomena"),
|
|
("fake_news", "A local news article about a quirky town event"),
|
|
("fake_news", "A sports article about a fictional team"),
|
|
("fake_news", "A celebrity gossip article"),
|
|
("fake_app", "A landing page for a fictional mobile app"),
|
|
("fake_app", "A SaaS dashboard for a fake productivity tool"),
|
|
("fake_app", "A streaming service page with fictional shows"),
|
|
("fake_other", "A fake email inbox from the year 2000"),
|
|
("fake_other", "A fake search engine results page for a weird query"),
|
|
("fake_other", "A fake 404 error page with an interactive game"),
|
|
("fake_other", "A fake loading screen for a fictional video game"),
|
|
("fake_other", "A fake terminal emulator with ASCII art and commands"),
|
|
("web_game", "A browser-based snake game with score tracking and increasing difficulty"),
|
|
("web_game", "A memory card matching game with colorful icons and timer"),
|
|
("web_game", "A clicker/idle game where you build a fictional empire"),
|
|
("web_game", "A whack-a-mole style arcade game with combos and high scores"),
|
|
("web_game", "A platformer game with a cute character and multiple levels"),
|
|
("web_game", "A typing speed test game with word races and accuracy stats"),
|
|
("web_game", "A 2048-style number sliding puzzle game"),
|
|
("web_game", "A reaction time test game with random color flashes"),
|
|
("web_game", "A virtual pet game where you feed, play, and care for a creature"),
|
|
("web_game", "A maze runner game with procedurally generated levels"),
|
|
]
|
|
|
|
|
|
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
|
|
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]
|
|
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_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"
|
|
while worker_running:
|
|
request_id: str | None = None
|
|
try:
|
|
lock_holder = r.get("wotf:processing") # type: ignore[union-attr]
|
|
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
|
|
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]
|
|
else:
|
|
time.sleep(0.5)
|
|
continue
|
|
|
|
acquired = r.set("wotf:processing", worker_id, ex=120, nx=True) # type: ignore[union-attr]
|
|
if not acquired:
|
|
time.sleep(0.5)
|
|
continue
|
|
|
|
r.set(hb_key, worker_id, ex=15) # type: ignore[union-attr]
|
|
|
|
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
|
|
if request_id is None:
|
|
r.delete("wotf:processing") # type: ignore[union-attr]
|
|
r.delete(hb_key) # type: ignore[union-attr]
|
|
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]
|
|
|
|
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]
|
|
|
|
r.delete("wotf:processing") # type: ignore[union-attr]
|
|
r.delete(hb_key) # type: ignore[union-attr]
|
|
|
|
confirmed = False
|
|
for _ in range(150):
|
|
time.sleep(0.1)
|
|
if r.get(f"wotf:req:{request_id}:confirmed"): # type: ignore[union-attr]
|
|
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]
|
|
time.sleep(0.5)
|
|
continue
|
|
|
|
acquired = r.set("wotf:processing", worker_id, ex=120, nx=True) # type: ignore[union-attr]
|
|
if not acquired:
|
|
time.sleep(0.5)
|
|
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]
|
|
|
|
ctx_raw = r.hgetall(f"wotf:req:{request_id}:ctx") # type: ignore[union-attr]
|
|
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:
|
|
ctx[k] = v
|
|
|
|
start = time.time()
|
|
hb_start = time.time()
|
|
hb_check_start = time.time()
|
|
html_buffer: list[str] = []
|
|
aborted = False
|
|
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]
|
|
hb_start = time.time()
|
|
|
|
if time.time() - hb_check_start >= 2:
|
|
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]
|
|
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]
|
|
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
|
|
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]
|
|
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]
|
|
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]
|
|
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]
|
|
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]
|
|
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]
|
|
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]
|
|
encoded = base64.b64encode(chunk.encode("utf-8")).decode("ascii")
|
|
r.rpush(f"wotf:stream:{request_id}:chunks", encoded) # type: ignore[union-attr]
|
|
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]
|
|
elapsed = time.time() - start
|
|
r.set(f"wotf:req:{request_id}:status", "done") # type: ignore[union-attr]
|
|
|
|
n_raw = r.get("wotf:stats:gen_count") # type: ignore[union-attr]
|
|
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_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]
|
|
|
|
r.delete("wotf:processing") # type: ignore[union-attr]
|
|
r.delete(hb_key) # type: ignore[union-attr]
|
|
logger.info(f"[worker] Generated {request_id} in {elapsed:.1f}s (avg={new_avg:.1f}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]
|
|
time.sleep(1)
|
|
|
|
|
|
def build_prompt(context: dict[str, Any]) -> str:
|
|
site_type, theme = random.choice(THEMES)
|
|
mood = random.choice(["playful", "mysterious", "cozy", "epic", "dreamy", "energetic", "melancholic", "whimsical"])
|
|
|
|
return f"""You are a web designer creating a realistic-looking fake website.
|
|
|
|
Here is information about the visitor:
|
|
- Preferred language: {context["language"]}
|
|
|
|
Site type: {site_type}
|
|
Theme: {theme}
|
|
Mood: {mood}
|
|
|
|
Generate a COMPLETE HTML document (<!DOCTYPE html> ... </html>) that looks like a REAL website of this type:
|
|
|
|
- Make it look authentic and believable as if it actually existed on the internet
|
|
- Include realistic details: fake names, fake content, fake comments, fake usernames, fake timestamps
|
|
- Use appropriate layout and conventions for the site type (e.g., a wiki should look like Wikipedia, a forum should have threads, etc.)
|
|
- Use inline CSS for styling - make it polished and visually appealing
|
|
- Include subtle animations or interactive elements where appropriate
|
|
- Make the content fun, interesting, and creative - not boring lorem ipsum
|
|
- Use the visitor's language for text when possible
|
|
- Responsive design that works on mobile
|
|
|
|
Examples:
|
|
- fake_store: product listings, prices, cart button, reviews
|
|
- fake_social: profile pic, followers count, posts, likes, comments
|
|
- fake_wiki: infobox, table of contents, citations, edit links
|
|
- fake_forum: thread list, usernames, post counts, timestamps
|
|
- fake_blog: article with author, date, comments section
|
|
- fake_homepage: about section, portfolio/work, contact info
|
|
|
|
Return ONLY the HTML, no markdown, no explanations."""
|
|
|
|
|
|
def stream_llm(context: dict[str, Any]) -> Generator[str]:
|
|
prompt = build_prompt(context)
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {API_KEY}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
payload: dict[str, Any] = {
|
|
"model": MODEL,
|
|
"messages": [
|
|
{"role": "system", "content": "You are a creative web designer."},
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
"temperature": 0.9,
|
|
"stream": True,
|
|
}
|
|
|
|
req_id = uuid.uuid4().hex[:8]
|
|
start = time.time()
|
|
logger.info(f"[{req_id}] Streaming LLM call for IP={context['ip']}, lang={context['language']}")
|
|
|
|
try:
|
|
with make_httpx_client() as client, client.stream("POST", API_URL, headers=headers, json=payload) as resp:
|
|
elapsed = time.time() - start
|
|
resp.raise_for_status()
|
|
buffer = ""
|
|
for chunk in resp.iter_text():
|
|
buffer += chunk
|
|
while "\n" in buffer:
|
|
line, buffer = buffer.split("\n", 1)
|
|
line = line.strip()
|
|
if not line or line == "data: [DONE]":
|
|
continue
|
|
if line.startswith("data: "):
|
|
data_str = line[6:]
|
|
try:
|
|
data: dict[str, Any] = json.loads(data_str)
|
|
delta = data.get("choices", [{}])[0].get("delta", {}).get("content", "")
|
|
if delta:
|
|
yield delta
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
elapsed = time.time() - start
|
|
logger.info(f"[{req_id}] Stream complete in {elapsed:.1f}s")
|
|
except httpx.HTTPStatusError as e:
|
|
elapsed = time.time() - start
|
|
logger.error(f"[{req_id}] HTTP error in {elapsed:.1f}s: {e.response.status_code} {e.response.text[:200]}")
|
|
yield f"\n\n<!-- Error: {e.response.status_code} -->"
|
|
except Exception as e: # noqa: BLE001
|
|
elapsed = time.time() - start
|
|
logger.error(f"[{req_id}] Stream failed in {elapsed:.1f}s: {e}")
|
|
yield f"\n\n<!-- Error: {e} -->"
|
|
|
|
|
|
def sse_stream(context: dict[str, Any]) -> Generator[str]:
|
|
buffer: list[str] = []
|
|
|
|
def status(msg: str) -> str:
|
|
return f"data: [STATUS:{msg}]\n\n"
|
|
|
|
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"
|
|
|
|
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"
|
|
|
|
|
|
@app.route("/", defaults={"path": ""})
|
|
@app.route("/<path:path>")
|
|
def index(path: str) -> Response:
|
|
logger.info(f"Request from {request.remote_addr}: {request.path} (UA: {request.headers.get('User-Agent', '')[:80]})")
|
|
return Response(LOADING_PAGE, mimetype="text/html")
|
|
|
|
|
|
@app.route("/stream")
|
|
def stream() -> Response:
|
|
context = get_user_context(request)
|
|
logger.info(f"Stream request from {request.remote_addr}")
|
|
|
|
request_id, position = enqueue_request(context)
|
|
logger.info(f"[{request_id}] Enqueued at position {position}")
|
|
|
|
def generate() -> Generator[str]:
|
|
yield "data: [CLEAR]\n\n"
|
|
|
|
last_idx = 0
|
|
while True:
|
|
status_raw = r.get(f"wotf:req:{request_id}:status") # type: ignore[union-attr]
|
|
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_msg = error_raw if isinstance(error_raw, str) else "Unknown error"
|
|
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: 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"
|
|
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: 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"
|
|
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: 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"
|
|
break
|
|
|
|
info = get_queue_info(request_id)
|
|
pos = info["position"]
|
|
eta = info["eta"]
|
|
if pos == 0:
|
|
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"
|
|
time.sleep(1)
|
|
|
|
while True:
|
|
chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1) # type: ignore[union-attr]
|
|
chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else [] # type: ignore[arg-type]
|
|
if chunks:
|
|
for chunk in chunks:
|
|
if chunk == "[RENDER]":
|
|
yield f"data: {chunk}\n\n"
|
|
return
|
|
yield f"data: {chunk}\n\n"
|
|
last_idx += len(chunks)
|
|
time.sleep(0.05)
|
|
else:
|
|
status_raw = r.get(f"wotf:req:{request_id}:status") # type: ignore[union-attr]
|
|
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: 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_msg = error_raw if isinstance(error_raw, str) else "Unknown error"
|
|
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: [ERROR]\n\n"
|
|
return
|
|
time.sleep(0.2)
|
|
|
|
return Response(
|
|
stream_with_context(generate()),
|
|
mimetype="text/event-stream",
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"X-Accel-Buffering": "no",
|
|
"Connection": "keep-alive",
|
|
},
|
|
)
|
|
|
|
|
|
@app.route("/confirm/<request_id>", methods=["POST"])
|
|
def confirm(request_id: str) -> Response:
|
|
ctx = r.exists(f"wotf:req:{request_id}:ctx") # type: ignore[union-attr]
|
|
if not ctx:
|
|
return Response("not found", status=404)
|
|
r.set(f"wotf:req:{request_id}:confirmed", "1") # type: ignore[union-attr]
|
|
return Response("ok", status=200)
|
|
|
|
|
|
@app.route("/heartbeat/<request_id>", methods=["POST"])
|
|
def heartbeat(request_id: str) -> Response:
|
|
status = r.get(f"wotf:req:{request_id}:status") # type: ignore[union-attr]
|
|
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]
|
|
return Response("ok", status=200)
|
|
|
|
|
|
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]
|
|
if claimed:
|
|
worker_running = True
|
|
if os.getenv("GUNICORN_WORKER_ID"):
|
|
from gevent import spawn
|
|
|
|
worker_thread = spawn(worker_loop) # type: ignore[assignment]
|
|
logger.info("[startup] Worker greenlet started in this process")
|
|
else:
|
|
worker_thread = threading.Thread(target=worker_loop, daemon=True)
|
|
worker_thread.start()
|
|
logger.info("[startup] Worker thread started in this process")
|
|
else:
|
|
logger.info("[startup] Worker thread already running in another process")
|
|
|
|
|
|
def gunicorn_post_fork(server: Any, worker: Any) -> None:
|
|
start_worker_if_needed()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
worker_running = True
|
|
worker_thread = threading.Thread(target=worker_loop, daemon=True)
|
|
worker_thread.start()
|
|
app.run(host="0.0.0.0", port=5000, debug=True)
|