File size: 4,568 Bytes
45f97d1 | 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 | """
Chain-of-Thought (CoT) controller for Qwen3.5-35B-A3B.
Manages thinking behavior at three control levels:
1. Hard switch (enable_thinking parameter)
2. Thinking budget (token-level budget enforcement)
3. Thinking visibility & streaming
"""
from __future__ import annotations
import structlog
from inference.config import (
DEFAULT_THINKING_BUDGET,
THINKING_BUDGET_PRESETS,
SAMPLING_THINKING_ON,
SAMPLING_THINKING_OFF,
)
logger = structlog.get_logger(__name__)
class CoTController:
"""
Manages Chain-of-Thought behavior at three control levels.
Level 1 - Hard Switch: enable_thinking parameter via chat_template_kwargs
Level 2 - Thinking Budget: token-level budget enforcement via logits processor
Level 3 - Thinking Visibility: controls what the banker sees of thinking
"""
def __init__(self):
self.default_mode = "on"
self.default_budget = DEFAULT_THINKING_BUDGET
self.default_visibility = "collapsed"
def build_request_params(self, cot_config: dict) -> dict:
"""
Build SGLang-compatible request parameters from CoT configuration.
Args:
cot_config: Dict with keys: mode, budget, visibility
Returns:
Dict to merge into the SGLang API call.
"""
mode = cot_config.get("mode", self.default_mode)
budget = cot_config.get("budget", self.default_budget)
visibility = cot_config.get("visibility", self.default_visibility)
params: dict = {}
# Level 1: Hard switch
if mode == "off":
params["chat_template_kwargs"] = {"enable_thinking": False}
params.update(SAMPLING_THINKING_OFF)
logger.debug("cot_mode_off")
else:
params["chat_template_kwargs"] = {"enable_thinking": True}
params.update(SAMPLING_THINKING_ON)
logger.debug("cot_mode_on")
# Level 2: Budget (resolved to token count)
resolved_budget = self._resolve_budget(budget)
params["_thinking_budget"] = resolved_budget
logger.debug("cot_budget_set", budget=budget, resolved=resolved_budget)
# Level 3: Visibility (handled by response streaming layer)
params["_thinking_visibility"] = visibility
return params
def _resolve_budget(self, budget: str | int) -> int:
"""Resolve a budget preset name or raw int to a token count."""
if isinstance(budget, int):
return max(budget, -1)
return THINKING_BUDGET_PRESETS.get(budget, THINKING_BUDGET_PRESETS["standard"])
@staticmethod
def get_presets() -> list[dict]:
"""Return all available CoT presets for the UI."""
descriptions = {
"none": "No thinking — equivalent to disabling CoT",
"minimal": "128 tokens — quick sanity check only",
"short": "512 tokens — brief reasoning",
"standard": "2048 tokens — balanced reasoning (default)",
"extended": "8192 tokens — complex multi-step analysis",
"deep": "32768 tokens — maximum reasoning depth",
"unlimited": "No limit — think until done",
}
return [
{
"name": name,
"budget": tokens,
"description": descriptions.get(name, ""),
}
for name, tokens in THINKING_BUDGET_PRESETS.items()
]
@staticmethod
def get_workflow_presets() -> list[dict]:
"""Pre-configured presets for common banking workflows."""
return [
{
"name": "Quick Lookup",
"description": "Fast response, no thinking",
"mode": "off",
"budget": "none",
"visibility": "hidden",
},
{
"name": "Standard Analysis",
"description": "Balanced reasoning, collapsed view",
"mode": "on",
"budget": "standard",
"visibility": "collapsed",
},
{
"name": "Deep Review",
"description": "Extended reasoning, streaming view",
"mode": "on",
"budget": "deep",
"visibility": "streaming",
},
{
"name": "Debug Mode",
"description": "Unlimited thinking, full observability",
"mode": "on",
"budget": "unlimited",
"visibility": "full",
},
]
|