File size: 26,076 Bytes
4abc9ca
 
 
 
8f0acda
7124ced
3240ef1
fead5fe
 
4abc9ca
3240ef1
 
 
 
 
4abc9ca
3240ef1
8f0acda
 
 
3240ef1
 
 
8f0acda
3240ef1
 
 
 
 
 
 
 
 
968f928
 
 
 
 
 
 
 
 
 
 
 
3240ef1
 
8f0acda
7878e0b
3240ef1
 
7878e0b
 
 
 
 
 
3240ef1
7878e0b
3240ef1
 
 
 
 
8f0acda
3240ef1
 
 
932067e
 
3240ef1
 
 
 
e0a218a
 
 
 
 
 
 
 
010f4ed
 
e0a218a
010f4ed
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3240ef1
8f0acda
 
 
 
 
 
 
 
 
 
 
 
 
4abc9ca
3240ef1
8f0acda
 
 
 
 
 
 
3240ef1
8f0acda
3240ef1
 
 
 
 
8f0acda
 
3240ef1
 
4abc9ca
3240ef1
4abc9ca
7124ced
 
3240ef1
7124ced
 
 
 
 
3240ef1
4abc9ca
3240ef1
 
 
 
 
 
8f0acda
3240ef1
4abc9ca
8f0acda
3240ef1
4abc9ca
3240ef1
 
4abc9ca
3240ef1
 
4abc9ca
3240ef1
4abc9ca
3240ef1
 
7124ced
 
 
3240ef1
 
 
 
8f0acda
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fead5fe
 
 
 
 
 
 
 
 
5f3ab1c
fead5fe
5f3ab1c
fead5fe
 
5f3ab1c
fead5fe
 
 
 
 
 
 
 
 
 
7124ced
f11081b
fead5fe
 
 
 
 
94d58b7
 
 
fead5fe
 
 
 
 
 
b1f4f76
fead5fe
3da89a5
fead5fe
 
 
b1f4f76
fead5fe
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3da89a5
 
 
 
 
 
 
 
 
 
 
94d58b7
 
 
 
b1f4f76
 
 
 
 
3da89a5
b1f4f76
 
 
 
 
 
 
 
 
 
 
94d58b7
3da89a5
fead5fe
 
7124ced
fead5fe
3240ef1
7124ced
 
 
 
 
 
 
 
5f3ab1c
 
 
 
 
 
 
fead5fe
3240ef1
 
 
 
fead5fe
8f0acda
3240ef1
fead5fe
8f0acda
3240ef1
 
 
7124ced
 
 
3240ef1
7124ced
8f0acda
3240ef1
 
 
fead5fe
8f0acda
 
3240ef1
7124ced
 
 
 
fead5fe
7124ced
 
3240ef1
 
 
8f0acda
3240ef1
 
 
 
 
 
 
 
 
8f0acda
fead5fe
3240ef1
 
 
fead5fe
8f0acda
fead5fe
3240ef1
 
 
 
 
5f3ab1c
3240ef1
fead5fe
3240ef1
 
 
 
7124ced
3240ef1
 
 
 
 
8f0acda
3240ef1
 
 
 
 
 
 
 
fead5fe
3240ef1
 
 
 
 
 
 
 
8f0acda
3240ef1
 
 
 
fead5fe
3240ef1
 
 
 
 
 
 
 
 
 
 
5f3ab1c
3240ef1
 
8f0acda
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fead5fe
 
 
 
 
5f3ab1c
 
 
 
 
 
 
 
 
 
 
 
fead5fe
5f3ab1c
 
 
 
 
 
 
 
 
fead5fe
 
 
 
803c552
 
 
 
 
5f3ab1c
 
 
fead5fe
 
94d58b7
 
fead5fe
 
5f3ab1c
 
 
fead5fe
 
 
 
 
 
 
5f3ab1c
 
fead5fe
 
 
 
 
 
 
 
8f0acda
fead5fe
 
 
 
 
 
5f3ab1c
fead5fe
5f3ab1c
fead5fe
 
 
 
 
 
 
7124ced
fead5fe
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
import os
import fnmatch
import urllib.request
import json
import time
import threading
import asyncio
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor
from huggingface_hub import hf_hub_download, HfApi
import uvicorn
from fastapi import FastAPI, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel

# ── 1. Curated Model Catalogue ─────────────────────────────────
# Each entry is also a role's fallback: the model this Space's role loads
# on boot and falls back to if a Hub refresh (see discover_best_model_for_role)
# doesn't find anything usable. One role per deployed Space β€” see SPACE_ROLE.
CURATED_MODELS = {
    "deepseek": {
        "id":          "deepseek",
        "role":        "reasoning",
        "label":       "DeepSeek-R1-Distill Qwen 7B",
        "emoji":       "🧠",
        "category":    "Reasoning / Coding",
        "repo":        "bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF",
        "file":        "*Q4_K_M.gguf",
        "ram_gb":      4.8,
        "n_ctx":       8192,
        "description": "Distilled DeepSeek reasoning model β€” punches far above its size. Best for logic, math, coding.",
    },
    "wizardvicuna": {
        "id":          "wizardvicuna",
        "role":        "uncensored-alt",
        "label":       "Wizard-Vicuna 7B Uncensored",
        "emoji":       "πŸ§™",
        "category":    "Uncensored",
        "repo":        "TheBloke/Wizard-Vicuna-7B-Uncensored-GGUF",
        "file":        "*Q4_K_M.gguf",
        "ram_gb":      4.8,
        "n_ctx":       4096,
        "description": "Classic uncensored Llama-2 7B β€” reliable, well-tested, same size class as the other curated models so it stays within the free CPU tier's time budget.",
    },
    "dolphin": {
        "id":          "dolphin",
        "role":        "uncensored",
        "label":       "Dolphin 2.9.3 Mistral 7B",
        "emoji":       "🐬",
        "category":    "Uncensored",
        # Was Dolphin 2.9.4 Llama 3.1 8B β€” measurably too slow on the free
        # CPU tier (regularly blew past the server's own 90s self-heal
        # timeout even for a 10-token reply, unlike the two 7B models here
        # which consistently finish in under 35s). This 7B Mistral-based
        # Dolphin variant matches the size class of the working models.
        "repo":        "bartowski/dolphin-2.9.3-mistral-7B-32k-GGUF",
        "file":        "*Q4_K_M.gguf",
        "ram_gb":      4.8,
        "n_ctx":       8192,
        "description": "Eric Hartford's Dolphin series β€” gold standard uncensored model. Follows all instructions without refusal.",
    },
    "qwen": {
        "id":          "qwen",
        "role":        "general",
        "label":       "Qwen 2.5 7B Instruct",
        "emoji":       "⚑",
        "category":    "Best General",
        "repo":        "bartowski/Qwen2.5-7B-Instruct-GGUF",
        "file":        "*Q4_K_M.gguf",
        "ram_gb":      4.8,
        "n_ctx":       8192,
        "description": "Alibaba's flagship 7B model β€” extremely capable, multilingual, fast. Best all-rounder for daily use.",
    },
    # NOT the default (role deliberately not "reasoning") β€” live-tested after
    # deploy and it timed out at 90s on a trivial 2-token-answer prompt.
    # "Thinking"-tuned models ruminate at length before answering regardless
    # of base model size, which is exactly the failure mode that got Dolphin
    # downsized before. Smaller parameter count didn't save it. Left in the
    # catalogue as a manual, opt-in choice only β€” "deepseek" stays the
    # default reasoning model since it's the one actually proven to finish
    # in time on this infra.
    "qwen3thinking": {
        "id":          "qwen3thinking",
        "role":        "reasoning-alt",
        "label":       "Qwen3 4B Thinking (2507)",
        "emoji":       "πŸ€”",
        "category":    "Reasoning / Coding",
        "repo":        "unsloth/Qwen3-4B-Thinking-2507-GGUF",
        "file":        "*Q4_K_M.gguf",
        "ram_gb":      2.8,
        "n_ctx":       8192,
        "description": "Purpose-built reasoning/thinking model, newer generation than DeepSeek-R1-Distill-Qwen-7B and roughly half the size β€” faster on CPU with comparable reasoning quality.",
    },
    # Same pattern as above for the uncensored role: newer official Dolphin
    # release (3.0) on a much smaller 3B Llama-3.2 base, replacing the 7B
    # Mistral-based Dolphin 2.9.3 as the uncensored Space's default. "dolphin"
    # and "wizardvicuna" above remain available as manual fallbacks.
    "dolphin3": {
        "id":          "dolphin3",
        "role":        "uncensored",
        "label":       "Dolphin 3.0 Llama 3.2 3B",
        "emoji":       "🐬",
        "category":    "Uncensored",
        "repo":        "bartowski/Dolphin3.0-Llama3.2-3B-GGUF",
        "file":        "*Q4_K_M.gguf",
        "ram_gb":      2.2,
        "n_ctx":       8192,
        "description": "Official Dolphin 3.0 release from the cognitivecomputations/dphn team β€” newer generation, less than half the size of Dolphin 2.9.3, follows instructions without refusal.",
    },
}
ROLE_TO_MODEL_ID = {spec["role"]: mid for mid, spec in CURATED_MODELS.items()}

# Hub search terms used by discover_best_model_for_role() when you hit
# /jolly/refresh-model. "prefer" keywords just re-rank same-popularity
# candidates; they don't gate eligibility.
ROLE_SEARCH = {
    "reasoning":  {"query": "reasoning instruct gguf",   "prefer": ["r1", "reasoning", "qwq", "thinking", "o1"]},
    "uncensored": {"query": "uncensored abliterated gguf", "prefer": ["uncensored", "abliterated", "dolphin"]},
    "general":    {"query": "instruct chat gguf",         "prefer": []},
}
# Skip repos naming a parameter count above what the free CPU tier (16GB RAM)
# can load at Q4 β€” keeps auto-discovery from picking something that OOMs.
_SIZE_BLOCKLIST = ("70b", "72b", "65b", "34b", "32b", "30b", "27b", "24b", "22b", "20b", "14b", "13b", "405b")

# ── 2. Config (env vars override defaults) ─────────────────────
# SPACE_ROLE is the one thing you set per deployed Space β€” it picks which
# curated model this instance boots into automatically, no MODEL_ID needed.
SPACE_ROLE    = os.environ.get("SPACE_ROLE", "general").strip().lower()
if SPACE_ROLE not in ROLE_TO_MODEL_ID:
    print(f"[Config] Unknown SPACE_ROLE '{SPACE_ROLE}', defaulting to 'general'")
    SPACE_ROLE = "general"

CONFIG_URL    = os.environ.get("CONFIG_URL", "")
MODEL_ID      = os.environ.get("MODEL_ID", ROLE_TO_MODEL_ID[SPACE_ROLE])  # role's model unless overridden
MODEL_REPO    = os.environ.get("MODEL_REPO", "")         # override repo if set
MODEL_FILE    = os.environ.get("MODEL_FILE", "")         # override file if set
N_CTX         = int(os.environ.get("N_CTX",  "8192"))
N_THREADS     = int(os.environ.get("N_THREADS", "4"))
API_KEY       = os.environ.get("API_KEY", "")
MODEL_SOURCE  = "fallback"   # "fallback" (curated catalogue) or "hub-discovered" (via /jolly/refresh-model)
LAST_REFRESHED = None        # epoch seconds of last successful /jolly/refresh-model

# Pull optional cluster config
if CONFIG_URL:
    print(f"[Config] Fetching cluster config from: {CONFIG_URL}")
    try:
        with urllib.request.urlopen(urllib.request.Request(CONFIG_URL), timeout=10) as r:
            cfg = json.loads(r.read().decode())
            MODEL_ID   = cfg.get("MODEL_ID",   cfg.get("model_id",   MODEL_ID))
            MODEL_REPO = cfg.get("MODEL_REPO", cfg.get("model_repo", MODEL_REPO))
            MODEL_FILE = cfg.get("MODEL_FILE", cfg.get("model_file", MODEL_FILE))
            if "N_CTX"     in cfg or "n_ctx"     in cfg: N_CTX     = int(cfg.get("N_CTX",     cfg.get("n_ctx",     N_CTX)))
            if "N_THREADS" in cfg or "n_threads" in cfg: N_THREADS = int(cfg.get("N_THREADS", cfg.get("n_threads", N_THREADS)))
            API_KEY = cfg.get("API_KEY", cfg.get("api_key", API_KEY))
            print("[Config] Cluster config synced.")
    except Exception as e:
        print(f"[Config] Fetch failed ({e}), using local env vars.")

# Resolve active model spec from catalogue or raw env overrides
def resolve_model_spec(model_id, repo_override="", file_override=""):
    if repo_override:
        return repo_override, file_override or "*.gguf", N_CTX
    spec = CURATED_MODELS.get(model_id, CURATED_MODELS[ROLE_TO_MODEL_ID[SPACE_ROLE]])
    return spec["repo"], spec["file"], spec.get("n_ctx", N_CTX)

ACTIVE_MODEL_ID = MODEL_ID if MODEL_ID in CURATED_MODELS else ROLE_TO_MODEL_ID[SPACE_ROLE]
active_repo, active_file, active_ctx = resolve_model_spec(ACTIVE_MODEL_ID, MODEL_REPO, MODEL_FILE)

# ── 3. Download helper ─────────────────────────────────────────
def download_model(repo, file_pattern):
    try:
        if "*" not in file_pattern:
            return hf_hub_download(repo_id=repo, filename=file_pattern)
    except Exception as e:
        print(f"[Download] Direct download failed: {e}, searching…")
    api = HfApi()
    files = list(api.list_repo_files(repo_id=repo))
    matches = [f for f in files if fnmatch.fnmatch(f.lower(), file_pattern.lower())]
    if not matches:
        matches = [f for f in files if f.endswith(".gguf")]
    if not matches:
        raise FileNotFoundError(f"No GGUF found in {repo}")
    print(f"[Download] Selected: {matches[0]}")
    return hf_hub_download(repo_id=repo, filename=matches[0])

# ── 3b. Hub auto-discovery (manual trigger via /jolly/refresh-model) ────
def discover_best_model_for_role(role):
    """Search the Hub for the most-downloaded GGUF model matching this role.
    Returns (repo, file, label) or None if nothing usable is found β€” callers
    must leave the currently-loaded model running on None rather than clobber it."""
    spec = ROLE_SEARCH.get(role, ROLE_SEARCH["general"])
    api = HfApi()
    try:
        candidates = list(api.list_models(search=spec["query"], sort="downloads", direction=-1, limit=25))
    except Exception as e:
        print(f"[Discover] Hub search failed: {e}")
        return None

    def score(m):
        name = m.id.lower()
        return sum(1 for kw in spec["prefer"] if kw in name)
    candidates.sort(key=score, reverse=True)

    for m in candidates:
        name = m.id.lower()
        if any(b in name for b in _SIZE_BLOCKLIST):
            continue
        try:
            files = list(api.list_repo_files(m.id))
        except Exception:
            continue
        match = next((f for f in files if fnmatch.fnmatch(f.lower(), "*q4_k_m.gguf")), None)
        if not match:
            match = next((f for f in files if f.lower().endswith(".gguf") and "q4" in f.lower()), None)
        if not match:
            continue
        label = m.id.split("/")[-1]
        print(f"[Discover] Role '{role}' -> {m.id} / {match}")
        return m.id, match, label

    print(f"[Discover] No suitable model found for role '{role}'")
    return None

# ── 4. Inference runs in a dedicated child process ──────────────
# llama.cpp's generation call does not hand control back to the asyncio
# event loop while it runs β€” even in a background thread, it holds the GIL
# long enough that the ENTIRE app (every route, including /jolly/health)
# freezes for the whole duration of generation. A real OS process has its
# own GIL, so the main server stays responsive no matter how slow a
# generation is. The tradeoff: no token-by-token streaming animation
# (results come back as one completed chunk), which is a fair price for
# "the app never appears to hang."
PUBLIC_PORT = 7860
_mp_ctx = mp.get_context("spawn")

# Only meaningful inside the worker process.
_worker_llama = None

def _init_worker(model_path, ctx, n_threads):
    global _worker_llama
    from llama_cpp import Llama
    print(f"[Worker] Loading {model_path}")
    _worker_llama = Llama(model_path=model_path, n_ctx=ctx, n_threads=n_threads, verbose=False)
    print(f"[Worker] Ready: {model_path}")

def _worker_generate(messages, max_tokens, temperature):
    return _worker_llama.create_chat_completion(
        messages=messages, max_tokens=max_tokens, temperature=temperature, stream=False
    )

def _worker_noop():
    # Submitted lambdas aren't picklable (required for ProcessPoolExecutor
    # with the spawn start method) β€” a plain module-level function is.
    return True

class WorkerBusy(Exception):
    pass

class _Worker:
    """One persistent worker process holding the loaded model. Recreated on model switch."""
    def __init__(self):
        self.executor = None
        self.ready = threading.Event()
        self.model_path = ""
        self.ctx = N_CTX
        self.lock = threading.Lock()
        self.queue_lock = asyncio.Lock()

    def start(self, model_path, ctx):
        self.ready.clear()
        self.ctx = ctx
        old = self.executor
        self.executor = ProcessPoolExecutor(
            max_workers=1, mp_context=_mp_ctx,
            initializer=_init_worker, initargs=(model_path, ctx, N_THREADS),
        )
        self.model_path = model_path
        if old:
            old.shutdown(wait=False, cancel_futures=True)

        # ProcessPoolExecutor's initializer runs in the child as soon as the
        # pool spins up; submit a no-op through it so we only flip `ready`
        # once that child process has actually finished loading the model.
        def _mark_ready():
            try:
                self.executor.submit(_worker_noop).result()
                self.ready.set()
                print(f"[Worker] Model ready: {model_path}")
            except Exception as e:
                print(f"[Worker] Failed to become ready: {e}")

        threading.Thread(target=_mark_ready, daemon=True).start()

    async def generate(self, messages, max_tokens, temperature):
        # Queue behind whatever's already generating instead of rejecting
        # outright β€” a single worker process only ever handles one request
        # at a time anyway, so a second request arriving mid-generation was
        # getting an instant 429 for no good reason (previously: fail-fast
        # via a busy flag). Waiting in line is strictly better for clients
        # that don't retry cleverly. Bounded so a truly stuck queue still
        # surfaces an error instead of hanging forever.
        try:
            await asyncio.wait_for(self.queue_lock.acquire(), timeout=100)
        except asyncio.TimeoutError:
            raise WorkerBusy("Server has a long queue right now β€” please wait and try again.")
        try:
            loop = asyncio.get_running_loop()
            with self.lock:
                executor = self.executor
            try:
                # A long system prompt (diary + user context can be several
                # thousand tokens) makes prompt processing itself slow, on top
                # of generation β€” a client giving up and retrying doesn't stop
                # the actual work on the worker process, which otherwise holds
                # the queue locked indefinitely for whoever asks next. Bound it,
                # and if it's blown through, the worker process is unrecoverable
                # for this request β€” kill and replace it rather than leave a
                # zombie generation running forever.
                return await asyncio.wait_for(
                    loop.run_in_executor(executor, _worker_generate, messages, max_tokens, temperature),
                    timeout=90,
                )
            except asyncio.TimeoutError:
                print("[Worker] Generation exceeded 90s β€” killing and restarting the worker process")
                self.start(self.model_path, self.ctx)
                raise RuntimeError("Generation timed out (prompt may be too long for this model's context window) β€” the worker was restarted, try again with a shorter message.")
        finally:
            self.queue_lock.release()

worker = _Worker()

# ── 5. FastAPI app ──────────────────────────────────────────────
proxy = FastAPI(title="JollyDay LLM Space β€” Multi-Model")

proxy.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

def _check_api_key(request: Request):
    if not API_KEY:
        return
    got = request.headers.get("authorization", "")
    if got != f"Bearer {API_KEY}":
        raise HTTPException(status_code=401, detail="Invalid API key")

# ── 6. /jolly/* control endpoints ─────────────────────────────

@proxy.get("/jolly/health")
async def jolly_health():
    return JSONResponse({
        "status":      "ready" if worker.ready.is_set() else "loading",
        "role":        SPACE_ROLE,
        "model_id":    ACTIVE_MODEL_ID,
        "model_file":  os.path.basename(worker.model_path),
        "model_source": MODEL_SOURCE,
        "n_ctx":       N_CTX,
        "n_threads":   N_THREADS,
    })

@proxy.get("/jolly/info")
async def jolly_info():
    spec = CURATED_MODELS.get(ACTIVE_MODEL_ID, {})
    return JSONResponse({
        "role":        SPACE_ROLE,
        "model_id":    ACTIVE_MODEL_ID,
        "model_repo":  active_repo,
        "model_file":  active_file,
        "model_path":  os.path.basename(worker.model_path),
        "model_source": MODEL_SOURCE,
        "last_refreshed": LAST_REFRESHED,
        "label":       spec.get("label", ""),
        "n_ctx":       N_CTX,
        "n_threads":   N_THREADS,
        "api_key_set": bool(API_KEY),
        "config_url":  CONFIG_URL,
        "server_ready": worker.ready.is_set(),
    })

@proxy.get("/jolly/models")
async def jolly_models():
    return JSONResponse({
        "role":      SPACE_ROLE,
        "active_id": ACTIVE_MODEL_ID,
        "models":    list(CURATED_MODELS.values()),
    })

@proxy.get("/jolly/active")
async def jolly_active():
    spec = CURATED_MODELS.get(ACTIVE_MODEL_ID, {})
    return JSONResponse({
        "id":          ACTIVE_MODEL_ID,
        "role":        SPACE_ROLE,
        "label":       spec.get("label", os.path.basename(worker.model_path)),
        "emoji":       spec.get("emoji", "πŸ€–"),
        "category":    spec.get("category", ""),
        "description": spec.get("description", ""),
        "model_file":  os.path.basename(worker.model_path),
        "model_source": MODEL_SOURCE,
        "server_ready": worker.ready.is_set(),
    })

@proxy.get("/jolly/wake")
async def jolly_wake():
    """Lightweight wake endpoint β€” just touching this URL wakes a sleeping HF Space."""
    spec = CURATED_MODELS.get(ACTIVE_MODEL_ID, {})
    return JSONResponse({
        "status":  "ready" if worker.ready.is_set() else "loading",
        "model_id": ACTIVE_MODEL_ID,
        "label":   spec.get("label", ACTIVE_MODEL_ID),
        "emoji":   spec.get("emoji", "πŸ€–"),
    })

class SwitchRequest(BaseModel):
    model_id: str

@proxy.post("/jolly/switch")
async def jolly_switch(req: SwitchRequest):
    global ACTIVE_MODEL_ID, active_repo, active_file, active_ctx, MODEL_SOURCE

    model_id = req.model_id.lower().strip()
    if model_id not in CURATED_MODELS:
        raise HTTPException(
            status_code=400,
            detail=f"Unknown model_id '{model_id}'. Valid options: {list(CURATED_MODELS.keys())}"
        )

    if model_id == ACTIVE_MODEL_ID and worker.ready.is_set():
        spec = CURATED_MODELS[model_id]
        return JSONResponse({"status": "already_active", "model_id": model_id, "label": spec["label"]})

    spec = CURATED_MODELS[model_id]
    ACTIVE_MODEL_ID = model_id
    active_repo     = spec["repo"]
    active_file     = spec["file"]
    active_ctx      = spec.get("n_ctx", N_CTX)
    MODEL_SOURCE    = "fallback"

    def _switch():
        try:
            path = download_model(active_repo, active_file)
            worker.start(path, active_ctx)
        except Exception as e:
            print(f"[Switch] ERROR: {e}")

    threading.Thread(target=_switch, daemon=True).start()

    return JSONResponse({
        "status":      "switching",
        "model_id":    model_id,
        "label":       spec["label"],
        "emoji":       spec["emoji"],
        "description": spec["description"],
        "note":        "Model is downloading/loading. Poll /jolly/health until status is 'ready'.",
    })

@proxy.post("/jolly/refresh-model")
async def jolly_refresh_model():
    """Manual 'update models' button: re-query the Hub for the current best
    model for this Space's role (SPACE_ROLE) and hot-swap to it if found.
    Leaves the running model untouched if the search comes up empty."""
    global ACTIVE_MODEL_ID, active_repo, active_file, active_ctx, MODEL_SOURCE, LAST_REFRESHED

    found = discover_best_model_for_role(SPACE_ROLE)
    if not found:
        return JSONResponse({
            "status": "no_change",
            "role":   SPACE_ROLE,
            "detail": "No suitable Hub model found for this role β€” keeping the current model.",
            "active_model_id": ACTIVE_MODEL_ID,
        })

    repo, file, label = found
    ACTIVE_MODEL_ID = f"hub:{label}"
    active_repo     = repo
    active_file     = file
    active_ctx      = N_CTX
    MODEL_SOURCE    = "hub-discovered"
    LAST_REFRESHED  = time.time()

    def _switch():
        try:
            path = download_model(active_repo, active_file)
            worker.start(path, active_ctx)
        except Exception as e:
            print(f"[Refresh] ERROR: {e}")

    threading.Thread(target=_switch, daemon=True).start()

    return JSONResponse({
        "status":     "switching",
        "role":       SPACE_ROLE,
        "model_repo": repo,
        "model_file": file,
        "label":      label,
        "note":       "Discovered from Hub search. Poll /jolly/health until status is 'ready'.",
    })

# ── 7. OpenAI-style routes ──────────────────────────────────────
# No true token streaming (see note above) β€” a streaming request still gets
# a single SSE chunk with the full text, then [DONE]. Obsidian's client
# just accumulates delta.content, so this renders identically to a real
# stream, just without the progressive animation.

@proxy.get("/v1/models")
async def list_models(request: Request):
    _check_api_key(request)
    return JSONResponse({
        "object": "list",
        "data": [{"id": ACTIVE_MODEL_ID, "object": "model", "owned_by": "me"}],
    })

@proxy.post("/v1/chat/completions")
async def chat_completions(request: Request):
    _check_api_key(request)
    if not worker.ready.is_set():
        return JSONResponse(
            {"error": "Model is loading, please retry in a moment.", "status": "loading"},
            status_code=503,
        )
    try:
        body = await request.json()
    except Exception:
        return JSONResponse({"error": "Invalid JSON body"}, status_code=400)

    messages = body.get("messages", [])
    # Hard-capped regardless of what the client asks for. Measured throughput
    # on this free CPU tier is ~5 tokens/sec, and responses no longer stream
    # token-by-token (see note above) β€” the client waits the full duration
    # with zero visible feedback. 180 tokens (~35s) was still landing "All
    # providers failed" in practice β€” almost certainly some client-side
    # network timeout (independent of our own AbortController) kicking in
    # around there. 80 tokens (~15s) leaves real margin.
    max_tokens  = min(int(body.get("max_tokens", 80) or 80), 80)
    temperature = body.get("temperature", 0.7)
    stream      = bool(body.get("stream", False))

    try:
        result = await worker.generate(messages, max_tokens, temperature)
    except WorkerBusy as e:
        return JSONResponse({"error": str(e), "status": "busy"}, status_code=429)
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)

    if stream:
        async def _gen():
            content = result.get("choices", [{}])[0].get("message", {}).get("content", "")
            chunk = {
                "id": result.get("id", "chatcmpl"), "object": "chat.completion.chunk",
                "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": None}],
            }
            yield f"data: {json.dumps(chunk)}\n\n"
            yield "data: [DONE]\n\n"
        return StreamingResponse(_gen(), media_type="text/event-stream")

    return JSONResponse(result)

# ── 8. Boot + launch ─────────────────────────────────────────────
# Guarded so the `spawn` multiprocessing start method (re-imports this file
# fresh in each worker process) doesn't recursively re-run the boot/serve
# logic β€” only route/function definitions above this line execute in the
# worker; nothing here does.
if __name__ == "__main__":
    print(f"[Boot] Space role : {SPACE_ROLE}")
    print(f"[Boot] Model ID   : {ACTIVE_MODEL_ID}")
    print(f"[Boot] Model repo : {active_repo}")
    print(f"[Boot] Model file : {active_file}")
    print(f"[Boot] Context    : {active_ctx}  Threads: {N_THREADS}")
    print(f"[Boot] API key set: {'yes' if API_KEY else 'no'}")

    try:
        boot_path = download_model(active_repo, active_file)
    except Exception as e:
        print(f"[WARN] {e} β€” falling back to Qwen 2.5 1.5B")
        boot_path = hf_hub_download(
            repo_id="Qwen/Qwen2.5-1.5B-Instruct-GGUF",
            filename="qwen2.5-1.5b-instruct-q4_k_m.gguf"
        )

    worker.start(boot_path, active_ctx)

    print(f"[Proxy] Starting on :{PUBLIC_PORT}…")
    uvicorn.run(proxy, host="0.0.0.0", port=PUBLIC_PORT, log_level="info")