serve_decide.py
24.5 KB · 467 lines · python Raw
1 #!/usr/bin/env python3
2 """vLLM OpenAI-compatible server with a built-in System 1 endpoint (POST /v1/decide) and adaptive thinking.
3
4 One engine serves both systems, for JEV-27B / JEV-27B-VL (Qwen3.8) and GEV-26B-A4B (Gemma-4):
5 System 2 the unmodified base model, through the usual OpenAI endpoints (served model name)
6 System 1 the LoRA module "jev-decision" (backbone LoRA + decision head as an lm_head LoRA), through /v1/decide
7
8 POST /v1/decide
9 kind "noul" (yes/no), "score" (0-5) or "choice" (2-256 options)
10 state string, JSON object, or a list mixing text and images: ["text", {"image": "https://... | data:..."}]
11 question string
12 options list of strings (choice only)
13 strategy choices with more options than the head's trained choice slots (16 for JEV-27B / GEV; 256 for wide heads,
14 which always read up to 256 options in one pass): "single" (one pass, labels A-P then Q-Z, AA, ...),
15 "tournament" (groups of <=16 + a final of 16), "permute" (single pass over 4 option orders, averaged);
16 default per model
17 thinking "off" (default: System 1 only), "auto" (think only when the leading option is below `threshold`),
18 "on" (always think)
19 threshold System 1 confidence below which "auto" switches thinking on (default per model)
20 reasoning controls, as for the base model's own chat API:
21 chat_template_kwargs passed to the base model's chat template (e.g. Qwen3.8: {"reasoning_effort": "low"})
22 reasoning_effort shorthand for chat_template_kwargs.reasoning_effort, if the base model has it
23 (Qwen3.8 / JEV-27B: xhigh (default), medium, low; Gemma-4 / GEV: not available)
24 think_budget maximum thinking tokens; default: no limit beyond the context window
25 return_reasoning include System 2's reasoning text
26 debug include the System 1 and System 2 distributions
27
28 Adaptive thinking: System 1 gives p1 in one pass. If thinking is switched on, the base model reasons in its thinking
29 mode over the same state, question and options; when the thinking channel closes, the answer-letter distribution p2 is
30 read in one step, and the result is p = (1 - w) * p1 + w * p2 (default w = 0.5: the fast decision and the reasoning get
31 equal weight, which keeps the calibration of System 1; p2 alone is close to one-hot and over-confident).
32
33 Response (same shape as the hosted JEV API, plus "thinking" when requested):
34 {"kind", "effective_kind", "options", "probabilities", "choice_index", "choice", "adaptation", "protocol", "model",
35 "usage", "elapsed_seconds", "num_model_requests", "thinking": {...}}
36
37 Run it like `python -m vllm.entrypoints.openai.api_server` (same flags), with --lora-modules jev-decision=<adapter dir>
38 and --trust-request-chat-template. decision_head.json is read from the adapter directory, calibration.json from its
39 parent (override with JEV_DECIDE_CALIBRATION).
40 """
41 import asyncio
42 import json
43 import math
44 import os
45 import string
46 import time
47 import zlib
48 from typing import Literal
49
50 from fastapi import APIRouter, Request
51 from fastapi.responses import JSONResponse
52 from pydantic import BaseModel
53
54 import vllm.entrypoints.launchers.api_server.entry as entry
55
56 # MoE backbones whose System 1 adapter has no routed-expert LoRA (Qwen3.5-MoE: attention, linear attention and shared
57 # expert only): JEV_DECIDE_NO_MOE_LORA=1 keeps the routed experts on the normal MoE kernels (vLLM otherwise picks its
58 # MoE-LoRA kernel for every expert layer, which has no LoRA weights here and stops at rank 128). Pair it with
59 # --lora-target-modules listing the adapter's modules.
60 if os.environ.get("JEV_DECIDE_NO_MOE_LORA") == "1":
61 from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig as _FMC
62 _fmc_post = _FMC.__post_init__
63
64 def _fmc_post_no_lora(self):
65 self.is_lora_enabled = False
66 _fmc_post(self)
67
68 _FMC.__post_init__ = _fmc_post_no_lora
69 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
70 from vllm.entrypoints.openai.completion.protocol import CompletionRequest
71
72 LORA = os.environ.get("JEV_DECIDE_LORA", "jev-decision")
73 # JEV_DECIDE_SYSTEM2=0: the served model has no usable thinking mode (e.g. a pretrained base model without instruction
74 # tuning); thinking "auto"/"on" and system2_only are refused with 400.
75 SYSTEM2 = os.environ.get("JEV_DECIDE_SYSTEM2", "1") != "0"
76 MAX_OPTIONS = 256
77 CHUNK = 128 # vLLM caps logprob_token_ids at 128 per request
78 PROTOCOL = "jev27-bare-v1"
79 PLACEHOLDER = "XQXCONTENTXQX"
80
81 # Per-model settings. threshold / mix: adaptive-thinking defaults. GEV: fitted on 1,754 questions outside the Decision
82 # Index (CommonsenseQA, OpenBookQA, AQuA-RAT, MedMCQA, LogiQA, StrategyQA). JEV-27B: same defaults, not yet validated.
83 PROFILES = {
84 "qwen": {"prefix": "", "image": "<|vision_start|><|image_pad|><|vision_end|>", "end_think": "</think>",
85 "after_think": "\n\n", "strategy": "single", "threshold": 0.8, "mix": 0.5, "thinking": "off"},
86 "gemma": {"prefix": "<bos>", "image": "<|image|>", "end_think": "<channel|>", "after_think": "",
87 "strategy": "tournament", "threshold": 0.8, "mix": 0.5, "thinking": "off"},
88 }
89 # Read the full distribution: override generation_config defaults (top_k/top_p) that would truncate processed logprobs.
90 READ = dict(max_tokens=1, temperature=1.0, top_p=1.0, top_k=0, min_p=0.0, repetition_penalty=1.0,
91 add_special_tokens=False, return_tokens_as_token_ids=True)
92 S = {}
93
94
95 class DecideRequest(BaseModel):
96 kind: Literal["noul", "score", "choice"]
97 state: str | dict | list = ""
98 question: str
99 options: list[str] | None = None
100 strategy: Literal["auto", "single", "tournament", "permute"] = "auto"
101 thinking: Literal["default", "off", "auto", "on"] = "default"
102 threshold: float | None = None
103 think_budget: int | None = None
104 reasoning_effort: str | None = None
105 chat_template_kwargs: dict | None = None
106 return_reasoning: bool = False
107 debug: bool = False
108 system2_only: bool = False # skip System 1: think and return System 2's answer distribution (for callers that already have p1)
109
110
111 # ------------------------------------------------------------------ setup
112 def setup(args):
113 from transformers import AutoConfig, AutoTokenizer
114
115 global LORA
116 merged_head = os.environ.get("JEV_DECIDE_MERGED_HEAD") # System 1 merged into the weights: no LoRA module
117 if merged_head:
118 path = os.path.dirname(os.path.abspath(merged_head))
119 head = json.load(open(merged_head))
120 names = args.served_model_name
121 LORA = (names[0] if isinstance(names, list) else names) or args.model # System 1 requests go to the base model
122 else:
123 path = next((m.path for m in (args.lora_modules or []) if m.name == LORA), None)
124 if path is None:
125 raise SystemExit(f"serve_decide: start vLLM with --lora-modules {LORA}=<adapter dir>, or set JEV_DECIDE_MERGED_HEAD")
126 head = json.load(open(os.path.join(path, "decision_head.json")))
127 calib = os.environ.get("JEV_DECIDE_CALIBRATION") or os.path.join(os.path.dirname(os.path.abspath(path)), "calibration.json")
128 temps = json.load(open(calib))["per_kind"] if os.path.exists(calib) else {"noul": 1.0, "score": 1.0, "choice": 1.0}
129 temps.setdefault("score", 1.0)
130 src = args.tokenizer or args.model
131 tok = AutoTokenizer.from_pretrained(src, trust_remote_code=args.trust_remote_code)
132 mtype = AutoConfig.from_pretrained(args.model, trust_remote_code=args.trust_remote_code).model_type
133 prof = dict(PROFILES["gemma" if "gemma" in mtype else "qwen"])
134 for k in ("strategy", "threshold", "mix", "thinking"):
135 if os.environ.get(f"JEV_DECIDE_{k.upper()}"):
136 v = os.environ[f"JEV_DECIDE_{k.upper()}"]
137 prof[k] = v if k in ("strategy", "thinking") else float(v)
138
139 def single_token_labels(context):
140 out = []
141 for lab in list(string.ascii_uppercase) + [a + b for a in string.ascii_uppercase for b in string.ascii_uppercase]:
142 t = tok.encode(lab, add_special_tokens=False)
143 if len(t) == 1 and t[0] in tok.encode(context.format(lab), add_special_tokens=False):
144 out.append((lab, t[0]))
145 if len(out) == MAX_OPTIONS:
146 break
147 return out
148
149 dec = single_token_labels("x\n{}) y") # System 1 option-line labels
150 lo, hi = head["slots"]["ranges"]["choice"]
151 assert [t for _, t in dec[: hi - lo]] == head["verbalizer_ids"][lo:hi], "first labels must be the trained A-P head"
152 base = tok.encode("Answer: (", add_special_tokens=False) # System 2 answer labels
153 ans = []
154 for lab in list(string.ascii_uppercase) + [a + b for a in string.ascii_uppercase for b in string.ascii_uppercase]:
155 ids = tok.encode(f"Answer: ({lab})", add_special_tokens=False)
156 if ids[: len(base)] == base and len(ids) == len(base) + 2:
157 ans.append((lab, ids[len(base)]))
158 if len(ans) == MAX_OPTIONS:
159 break
160 chat = tok.apply_chat_template([{"role": "user", "content": PLACEHOLDER}], tokenize=False, add_generation_prompt=True,
161 enable_thinking=True)
162 pre, post = chat.split(PLACEHOLDER)
163 raw = ("{%- for m in messages -%}{%- if m['content'] is string -%}{{ m['content'] }}{%- else -%}"
164 "{%- for c in m['content'] -%}{%- if c['type'] == 'text' -%}{{ c['text'] }}{%- else -%}" + prof["image"] +
165 "{%- endif -%}{%- endfor -%}{%- endif -%}{%- endfor -%}")
166 names = args.served_model_name
167 S.update(tok=tok, has_effort="reasoning_effort" in (tok.chat_template or ""), think_cache={},
168 max_logprobs=int(getattr(args, "max_logprobs", 20) or 20),
169 head=head, temps=temps, labels=[l for l, _ in dec], label_ids=[t for _, t in dec], ans_labels=ans, prof=prof,
170 think_pre=pre, think_post=post, raw=raw, end_think_id=tok.convert_tokens_to_ids(prof["end_think"]),
171 model=(names[0] if isinstance(names, list) else names) or args.model, model_type=mtype)
172
173
174 # ------------------------------------------------------------------ helpers
175 def _parts(state):
176 if isinstance(state, str):
177 return [{"type": "text", "text": state}], False
178 if isinstance(state, dict):
179 return [{"type": "text", "text": json.dumps(state, ensure_ascii=False)}], False
180 out, img = [], False
181 for p in state:
182 if isinstance(p, str):
183 out.append({"type": "text", "text": p})
184 elif isinstance(p, dict) and "image" in p:
185 out.append({"type": "image_url", "image_url": {"url": p["image"]}}); img = True
186 elif isinstance(p, dict) and p.get("type") == "image_url":
187 out.append(p); img = True
188 elif isinstance(p, dict) and p.get("type") == "text":
189 out.append(p)
190 else:
191 out.append({"type": "text", "text": json.dumps(p, ensure_ascii=False)})
192 return out, img
193
194
195 def _softmax(z):
196 m = max(z); e = [math.exp(x - m) for x in z]; s = sum(e)
197 return [x / s for x in e]
198
199
200 def _err(msg, code=400):
201 return JSONResponse({"error": {"message": msg, "type": "BadRequestError", "code": code}}, status_code=code)
202
203
204 class Upstream(Exception):
205 def __init__(self, resp):
206 self.resp = resp
207
208
209 def _check(out):
210 if hasattr(out, "error"):
211 raise Upstream(out)
212 return out
213
214
215 class Ctx:
216 def __init__(self, raw):
217 self.raw, self.requests, self.prompt_tokens, self.completion_tokens = raw, 0, 0, 0
218
219 def count(self, out):
220 self.requests += 1
221 if getattr(out, "usage", None):
222 self.prompt_tokens += out.usage.prompt_tokens or 0
223 self.completion_tokens += out.usage.completion_tokens or 0
224
225
226 def _readout_allowed(ids):
227 """Read-out path: restrict the next token to `ids` and return all of them as top logprobs when the server allows that
228 many (--max-logprobs); otherwise ask for `ids` explicitly with logprob_token_ids (chunks of 128). The first path also
229 works with speculative decoding; logprob_token_ids does not in this vLLM build. Both give the same distribution."""
230 return ids if len(ids) <= S["max_logprobs"] else None
231
232
233 async def _logprobs(ctx, content, has_img, ids, model, allowed=None):
234 """One-token read-out with a fallback to the other read-out path if the first one fails."""
235 try:
236 return await _logprobs_once(ctx, content, has_img, ids, model, allowed)
237 except Upstream:
238 raise
239 except Exception:
240 other = None if allowed is not None else (ids if len(ids) <= S["max_logprobs"] else None)
241 if other is None and allowed is None:
242 raise
243 return await _logprobs_once(ctx, content, has_img, ids, model, other)
244
245
246 async def _logprobs_once(ctx, content, has_img, ids, model, allowed=None):
247 """One-token read-out: {token_id: logprob} for `ids` after the content (text or chat parts)."""
248 lp = {}
249 step = CHUNK if allowed is None else len(ids)
250 for i in range(0, len(ids), step):
251 chunk = ids[i: i + step]
252 kw = dict(logprob_token_ids=chunk) if allowed is None else dict(allowed_token_ids=allowed)
253 if has_img:
254 r = ChatCompletionRequest(model=model, messages=[{"role": "user", "content": content}], chat_template=S["raw"],
255 add_generation_prompt=False, logprobs=True, top_logprobs=len(chunk) if allowed else 1, **kw, **READ)
256 out = _check(await ctx.raw.app.state.openai_serving_chat.create_chat_completion(r, ctx.raw)); ctx.count(out)
257 lp.update({int(t.token.split(":")[1]): t.logprob for t in out.choices[0].logprobs.content[0].top_logprobs})
258 else:
259 text = "".join(c["text"] for c in content)
260 r = CompletionRequest(model=model, prompt=text, logprobs=len(chunk) if allowed else 1, **kw, **READ)
261 out = _check(await ctx.raw.app.state.openai_serving_completion.create_completion(r, ctx.raw)); ctx.count(out)
262 lp.update({int(k.split(":")[1]): v for k, v in out.choices[0].logprobs.top_logprobs[0].items()})
263 if allowed is not None:
264 break
265 return lp
266
267
268 # ------------------------------------------------------------------ System 1
269 async def s1_pass(ctx, kind, parts, has_img, question, opts):
270 """One System 1 pass over `opts` (<=256). Returns probabilities aligned with opts."""
271 head, temps, prof = S["head"], S["temps"], S["prof"]
272 lo, hi = head["slots"]["ranges"][kind]
273 if kind == "choice":
274 n = len(opts); ids = S["label_ids"][:n]
275 bias = [head["bias"][lo + i] if lo + i < hi else 0.0 for i in range(n)]
276 lines = [f"{S['labels'][i]}) {o}" for i, o in enumerate(opts)]
277 else:
278 ids = head["verbalizer_ids"][lo:hi]; bias = head["bias"][lo:hi]; lines = opts
279 content = ([{"type": "text", "text": f"{prof['prefix']}[kind] {kind}\n[state] "}] + parts +
280 [{"type": "text", "text": f"\n[question] {question}\n[options]\n" + "\n".join(lines) + "\n[decision]:"}])
281 lp = await _logprobs(ctx, content, has_img, ids, LORA, allowed=_readout_allowed(ids))
282 return _softmax([(max(lp.get(t, -1e9), -1e9) + b) / temps[kind] for t, b in zip(ids, bias)])
283
284
285 def _groups(n, k=16):
286 g = math.ceil(n / k); base, extra = divmod(n, g); out, i = [], 0
287 for j in range(g):
288 size = base + (1 if j < extra else 0); out.append(list(range(i, i + size))); i += size
289 return out
290
291
292 def _native():
293 lo, hi = S["head"]["slots"]["ranges"]["choice"]
294 return hi - lo # trained choice slots: 16 (JEV-27B, GEV) or 256 (wide heads)
295
296
297 async def s1_dist(ctx, kind, parts, has_img, question, opts, strategy):
298 if kind != "choice" or len(opts) <= _native() or strategy == "single":
299 return await s1_pass(ctx, kind, parts, has_img, question, opts)
300 n = len(opts)
301 if strategy == "permute":
302 import random
303 rng = random.Random(zlib.crc32(question.encode()))
304 orders = [list(range(n))] + [rng.sample(range(n), n) for _ in range(3)]
305 res = await asyncio.gather(*(s1_pass(ctx, kind, parts, has_img, question, [opts[i] for i in o]) for o in orders))
306 p = [0.0] * n
307 for o, r in zip(orders, res):
308 for i, v in zip(o, r):
309 p[i] += v / len(orders)
310 return p
311 groups = _groups(n) # tournament: groups of <=16 in the given order (in parallel), then a final of 16
312 parts_g = await asyncio.gather(*(s1_pass(ctx, kind, parts, has_img, question, [opts[i] for i in g]) for g in groups))
313 in_group = {o: p for g, ps in zip(groups, parts_g) for o, p in zip(g, ps)}
314 chosen = [max(g, key=lambda o: (in_group[o], -o)) for g in groups]
315 rest = sorted((o for g in groups for o in g if o not in set(chosen)), key=lambda o: (-in_group[o], o))
316 fin = sorted(chosen + rest[: max(0, 16 - len(chosen))])
317 final = dict(zip(fin, await s1_pass(ctx, kind, parts, has_img, question, [opts[i] for i in fin])))
318 group_of = {o: gi for gi, g in enumerate(groups) for o in g}
319 share = [0.0] * len(groups); cap = [0.0] * len(groups)
320 for f in fin:
321 share[group_of[f]] += final[f]; cap[group_of[f]] += in_group[f]
322 among = sum(a * b for a, b in zip(share, cap))
323 p = [final[o] * among if o in final else share[group_of[o]] * in_group[o] for o in range(n)]
324 s = sum(p)
325 return [x / s for x in p]
326
327
328 # ------------------------------------------------------------------ System 2 (thinking)
329 def _think_frame(kwargs):
330 """Chat-template prefix / suffix around the user turn, thinking on, with the caller's template kwargs (cached)."""
331 key = json.dumps(kwargs or {}, sort_keys=True)
332 if key not in S["think_cache"]:
333 kw = {k: v for k, v in (kwargs or {}).items() if k != "enable_thinking"}
334 chat = S["tok"].apply_chat_template([{"role": "user", "content": PLACEHOLDER}], tokenize=False, add_generation_prompt=True,
335 enable_thinking=True, **kw)
336 S["think_cache"][key] = tuple(chat.split(PLACEHOLDER))
337 return S["think_cache"][key]
338
339
340 async def s2_dist(ctx, kind, parts, has_img, question, opts, budget, want_text, tmpl_kwargs=None):
341 """The base model thinks over the same input; the answer-letter distribution is read after the thinking channel."""
342 shown = ["Yes (true)", "No (false)"] if kind == "noul" else opts
343 labs = S["ans_labels"][: len(shown)]
344 body = "\n".join(f"({l}) {o}" for (l, _), o in zip(labs, shown))
345 tail = (f"\n\nQuestion: {question}\n\nOptions:\n{body}\n\n"
346 "Think it through carefully, then give your final answer on the last line in the form: Answer: (X)")
347 pre, post = _think_frame(tmpl_kwargs)
348 user = [{"type": "text", "text": pre}] + parts + [{"type": "text", "text": tail + post}]
349 seed = zlib.crc32((question + body).encode()) & 0x7FFFFFFF
350 t0 = time.time()
351 gen = ChatCompletionRequest(model=S["model"], messages=[{"role": "user", "content": user}], chat_template=S["raw"],
352 add_generation_prompt=False, add_special_tokens=False, max_tokens=budget, seed=seed, # None: up to the context window
353 stop_token_ids=[S["end_think_id"]], skip_special_tokens=False)
354 out = _check(await ctx.raw.app.state.openai_serving_chat.create_chat_completion(gen, ctx.raw)); ctx.count(out)
355 thought = out.choices[0].message.content or ""
356 finished = out.choices[0].finish_reason == "stop"
357 ntok = out.usage.completion_tokens if out.usage else None
358 read = user + [{"type": "text", "text": thought + S["prof"]["end_think"] + S["prof"]["after_think"] + "Answer: ("}]
359 ids = [t for _, t in labs]
360 lp = await _logprobs(ctx, read, True, ids, S["model"], allowed=_readout_allowed(ids))
361 p = _softmax([max(lp.get(t, -1e9), -1e9) for t in ids])
362 if kind == "noul":
363 p = [p[1], p[0]] # (A) yes / (B) no -> [P(false), P(true)]
364 info = {"think_tokens": ntok, "think_seconds": round(time.time() - t0, 3), "finished_within_budget": finished}
365 if want_text:
366 info["reasoning"] = thought.replace("<|channel>thought\n", "").strip()
367 return p, info
368
369
370 def _fold(p1, p2, w):
371 return [(1 - w) * x + w * y for x, y in zip(p1, p2)]
372
373
374 # ------------------------------------------------------------------ routes
375 router = APIRouter()
376
377
378 @router.get("/v1/decide/info")
379 async def info():
380 p = S["prof"]
381 return {"protocol": PROTOCOL, "model": S["model"], "model_type": S["model_type"], "max_options": len(S["labels"]),
382 "native_choice_options": len(S["head"]["slots"]["verbalizers"]) - 8, "temperatures": S["temps"],
383 "defaults": {"strategy": p["strategy"], "thinking": p["thinking"], "threshold": p["threshold"], "mix": p["mix"],
384 "think_budget": None},
385 "system2_available": SYSTEM2,
386 "reasoning_controls": {"enable_thinking": True, "reasoning_effort": ["xhigh", "medium", "low"] if S["has_effort"] else None,
387 "think_budget": "max thinking tokens (default: up to the context window)"}}
388
389
390 @router.post("/v1/decide")
391 async def decide(req: DecideRequest, raw: Request):
392 t0 = time.time()
393 if req.kind == "choice":
394 opts = req.options or []
395 if not 2 <= len(opts) <= len(S["labels"]):
396 return _err(f"choice needs 2-{len(S['labels'])} options, got {len(opts)}")
397 else:
398 opts = ["false", "true"] if req.kind == "noul" else [str(i) for i in range(6)]
399 prof = S["prof"]
400 thinking = prof["thinking"] if req.thinking == "default" else req.thinking
401 if req.kind == "score":
402 if req.thinking in ("auto", "on"):
403 return _err("thinking is supported for noul and choice")
404 thinking = "off"
405 tmpl = dict(req.chat_template_kwargs or {})
406 if req.reasoning_effort is not None:
407 if not S["has_effort"]:
408 return _err("this base model has no reasoning_effort setting; control the reasoning length with think_budget")
409 tmpl["reasoning_effort"] = req.reasoning_effort
410 strategy = prof["strategy"] if req.strategy == "auto" else req.strategy
411 tau = prof["threshold"] if req.threshold is None else req.threshold
412 parts, has_img = _parts(req.state)
413 ctx = Ctx(raw)
414 if not SYSTEM2 and (thinking in ("auto", "on") or req.system2_only):
415 return _err("thinking (System 2) is not available for this model: it is a pretrained base model without "
416 "instruction tuning; use thinking \"off\"")
417 try:
418 if req.system2_only:
419 if req.kind == "score":
420 return _err("system2_only is supported for noul and choice")
421 p2, think = await s2_dist(ctx, req.kind, parts, has_img, req.question, opts, req.think_budget, req.return_reasoning, tmpl)
422 k = max(range(len(p2)), key=p2.__getitem__)
423 return {"kind": req.kind, "options": opts, "probabilities": p2, "choice_index": k, "choice": opts[k], "system": 2,
424 "model": S["model"], "thinking": {"used": True, "budget": req.think_budget, **think},
425 "usage": {"prompt_tokens": ctx.prompt_tokens, "completion_tokens": ctx.completion_tokens},
426 "elapsed_seconds": time.time() - t0, "num_model_requests": ctx.requests}
427 p1 = await s1_dist(ctx, req.kind, parts, has_img, req.question, opts, strategy)
428 probs, think = p1, None
429 if thinking == "on" or (thinking == "auto" and max(p1) < tau):
430 p2, think = await s2_dist(ctx, req.kind, parts, has_img, req.question, opts, req.think_budget, req.return_reasoning, tmpl)
431 probs = _fold(p1, p2, prof["mix"])
432 think = {"used": True, **think}
433 if req.debug:
434 think.update(system1=p1, system2=p2)
435 elif thinking != "off":
436 think = {"used": False}
437 if req.debug:
438 think["system1"] = p1
439 except Upstream as e:
440 return JSONResponse(e.resp.model_dump(), status_code=e.resp.error.code)
441 k = max(range(len(probs)), key=probs.__getitem__)
442 native = req.kind != "choice" or len(opts) <= _native()
443 resp = {"kind": req.kind, "effective_kind": req.kind, "options": opts, "probabilities": probs, "choice_index": k,
444 "choice": opts[k], "adaptation": "native" if native else f"{strategy}", "protocol": PROTOCOL, "model": S["model"],
445 "usage": {"prompt_tokens": ctx.prompt_tokens, "completion_tokens": ctx.completion_tokens,
446 "total_tokens": ctx.prompt_tokens + ctx.completion_tokens},
447 "elapsed_seconds": time.time() - t0, "num_model_requests": ctx.requests}
448 if think is not None:
449 resp["thinking"] = {"mode": thinking, "threshold": tau, "budget": req.think_budget, **({"chat_template_kwargs": tmpl} if tmpl else {}), **think}
450 return resp
451
452
453 _build_app = entry.build_app
454
455
456 def build_app(args, *a, **kw):
457 app = _build_app(args, *a, **kw)
458 setup(args)
459 app.include_router(router)
460 return app
461
462
463 entry.build_app = build_app
464
465 if __name__ == "__main__":
466 entry.main()
467