| import asyncio |
| import json |
| import logging |
| import os |
| import re |
| import time |
| from pathlib import Path |
| from typing import Any, AsyncIterator, Iterator, Literal |
|
|
| import path_explorer |
| from fastapi import FastAPI, HTTPException |
| from fastapi.responses import FileResponse, StreamingResponse |
| from fastapi.staticfiles import StaticFiles |
| from pydantic import BaseModel, Field |
|
|
| from src.inference import ( |
| BaseInferenceBackend, |
| InferenceConfig, |
| TransformersBackend, |
| build_inference_backend, |
| ) |
| from src.pipeline.runtime_instructions import ( |
| instruction_block, |
| thinking_limit_reached_instruction, |
| ) |
|
|
|
|
| logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO")) |
| logger = logging.getLogger(__name__) |
| BASE_DIR = Path(__file__).resolve().parent |
| OUTPUTS_DIR = BASE_DIR / "outputs" |
| MODEL_JSON_DIR = BASE_DIR / "model_json" |
| DISABLE_VLLM_MODE = True |
|
|
| app = FastAPI(title="MATH Reasoning Lab") |
| app.mount("/static", StaticFiles(directory=BASE_DIR / "static"), name="static") |
|
|
|
|
| class ChatMessage(BaseModel): |
| role: Literal["system", "user", "assistant"] |
| content: str |
|
|
|
|
| class ChatRequest(BaseModel): |
| messages: list[ChatMessage] = Field(default_factory=list) |
| model: str = Field(default="peft:1K") |
| temperature: float = Field(default=0.6, ge=0.0, le=2.0) |
| max_tokens: int = Field(default=2048, ge=1, le=32768) |
| fast_mode: bool = Field(default=False) |
| thinking_effort: Literal[ |
| "none", |
| "flash", |
| "speed", |
| "low", |
| "medium", |
| "high", |
| "extreme", |
| "max", |
| ] = Field(default="medium") |
|
|
|
|
| class ModelRequest(BaseModel): |
| model: str = Field(default="peft:1K") |
|
|
|
|
| class JsonSaveRequest(BaseModel): |
| filename: str = Field(default="dpo_turns.jsonl") |
| record: dict[str, Any] = Field(default_factory=dict) |
|
|
|
|
| def discover_model_options() -> list[dict[str, Any]]: |
| options: list[dict[str, Any]] = [ |
| { |
| "id": "base:qwen2.5-3b-math", |
| "label": "Qwen2.5 3B base", |
| "kind": "base", |
| "path": None, |
| "available": True, |
| } |
| ] |
|
|
| if OUTPUTS_DIR.exists(): |
| for adapter_config in sorted(OUTPUTS_DIR.rglob("adapter_config.json")): |
| adapter_dir = adapter_config.parent |
| relative_path = adapter_dir.relative_to(OUTPUTS_DIR).as_posix() |
| checkpoint = checkpoint_name(adapter_dir) |
| family = model_family_name(adapter_dir) |
| label = adapter_label(adapter_dir) |
|
|
| options.append( |
| { |
| "id": f"peft:{relative_path}", |
| "label": label, |
| "kind": "peft", |
| "path": str(adapter_dir), |
| "relative_path": relative_path, |
| "family": family, |
| "checkpoint": checkpoint, |
| "available": True, |
| } |
| ) |
|
|
| return options |
|
|
|
|
| def checkpoint_name(adapter_dir: Path) -> str | None: |
| for part in reversed(adapter_dir.parts): |
| if part.startswith("checkpoint-"): |
| return part |
| return None |
|
|
|
|
| def checkpoint_step(option: dict[str, Any]) -> int: |
| checkpoint = option.get("checkpoint") or "" |
| match = re.search(r"checkpoint-(\d+)", checkpoint) |
| return int(match.group(1)) if match else -1 |
|
|
|
|
| def model_family_name(adapter_dir: Path) -> str: |
| relative_parts = adapter_dir.relative_to(OUTPUTS_DIR).parts |
| checkpoint = checkpoint_name(adapter_dir) |
|
|
| if checkpoint and len(relative_parts) > 1: |
| return relative_parts[-2] |
| return relative_parts[-1] |
|
|
|
|
| def adapter_label(adapter_dir: Path) -> str: |
| relative_path = adapter_dir.relative_to(OUTPUTS_DIR).as_posix() |
| checkpoint = checkpoint_name(adapter_dir) |
| family = model_family_name(adapter_dir) |
|
|
| if checkpoint: |
| checkpoint_label = checkpoint.replace("-", " ") |
| return f"{checkpoint_label} - {family}" |
| return f"{relative_path} PEFT" |
|
|
|
|
| def model_options_by_id() -> dict[str, dict[str, Any]]: |
| return {option["id"]: option for option in discover_model_options()} |
|
|
|
|
| def default_model_id() -> str: |
| options = discover_model_options() |
| peft_options = [option for option in options if option["kind"] == "peft"] |
| checkpoint_options = [option for option in peft_options if option.get("checkpoint")] |
|
|
| if checkpoint_options: |
| return max( |
| checkpoint_options, |
| key=lambda option: (checkpoint_step(option), option.get("relative_path", "")), |
| )["id"] |
|
|
| peft_by_id = {option["id"]: option for option in peft_options} |
| if "peft:1K" in peft_by_id: |
| return "peft:1K" |
| if peft_options: |
| return peft_options[0]["id"] |
| return "base:qwen2.5-3b-math" |
|
|
|
|
| MODEL_STATE: dict[str, Any] = { |
| "connected": False, |
| "loading": False, |
| "mode": "real", |
| "model": default_model_id(), |
| "label": model_options_by_id().get(default_model_id(), {}).get("label", "Qwen2.5 3B base"), |
| "error": None, |
| } |
|
|
| LOADED_MODELS: dict[str, Any] = {} |
| MODEL_LOCK = asyncio.Lock() |
| INFERENCE_BACKENDS: dict[tuple[str, str, str], BaseInferenceBackend] = {} |
| INFERENCE_LOCK = asyncio.Lock() |
| LAST_INFERENCE_HEALTH: dict[str, Any] = { |
| "ok": None, |
| "backend": None, |
| "model": None, |
| "error": None, |
| } |
| THINKING_TOKEN_BUDGETS = { |
| "none": 64, |
| "flash": 256, |
| "speed": 512, |
| "low": 1024, |
| "medium": 2048, |
| "high": 4096, |
| "extreme": 5120, |
| "max": 8192, |
| } |
|
|
| BEGIN_THOUGHT = "<|begin_of_thought|>" |
| END_THOUGHT = "<|end_of_thought|>" |
| BEGIN_SOLUTION = "<|begin_of_solution|>" |
| END_SOLUTION = "<|end_of_solution|>" |
| PROTOCOL_TAGS = (BEGIN_THOUGHT, END_THOUGHT, BEGIN_SOLUTION, END_SOLUTION) |
| THOUGHT_STOP_SEQUENCES = (END_THOUGHT, BEGIN_SOLUTION, END_SOLUTION) |
| STREAM_TEXT_BATCH_CHARS = 96 |
| STREAM_TEXT_BATCH_SECONDS = 0.04 |
| THOUGHT_SENTENCE_COMPLETION_TOKENS = 96 |
|
|
|
|
| @app.get("/") |
| async def index() -> FileResponse: |
| return FileResponse(BASE_DIR / "index.html") |
|
|
|
|
| @app.get("/api/status") |
| async def status() -> dict[str, Any]: |
| return { |
| **MODEL_STATE, |
| **runtime_info(), |
| "inference": inference_status(), |
| "models": discover_model_options(), |
| "default_model": default_model_id(), |
| } |
|
|
|
|
| @app.get("/api/models") |
| async def models() -> dict[str, Any]: |
| return {"models": discover_model_options(), "default_model": default_model_id()} |
|
|
|
|
| @app.post("/api/preload") |
| async def preload(request: ModelRequest) -> dict[str, Any]: |
| option = resolve_model_option(request.model) |
| try: |
| config = inference_config() |
| if config.backend == "vllm": |
| await check_inference_health() |
| else: |
| await get_local_model(option["id"]) |
| except Exception: |
| pass |
| return {**MODEL_STATE, **runtime_info(), "inference": inference_status()} |
|
|
|
|
| @app.get("/api/inference/health") |
| async def inference_health() -> dict[str, Any]: |
| return await check_inference_health() |
|
|
|
|
| @app.get("/api/model-json/files") |
| async def model_json_files() -> dict[str, list[str]]: |
| MODEL_JSON_DIR.mkdir(parents=True, exist_ok=True) |
| files = sorted(path.name for path in MODEL_JSON_DIR.glob("*.jsonl") if path.is_file()) |
| return {"files": files} |
|
|
|
|
| @app.post("/api/model-json/save") |
| async def save_model_json(request: JsonSaveRequest) -> dict[str, Any]: |
| MODEL_JSON_DIR.mkdir(parents=True, exist_ok=True) |
| path = model_json_path(request.filename) |
|
|
| with path.open("a", encoding="utf-8") as file: |
| file.write(json.dumps(request.record, ensure_ascii=False, separators=(",", ":")) + "\n") |
|
|
| return {"ok": True, "file": path.name} |
|
|
|
|
| def last_user_message(messages: list[ChatMessage]) -> str: |
| for message in reversed(messages): |
| if message.role == "user": |
| return message.content |
| return "Solve a sample olympiad-style algebra problem." |
|
|
|
|
| def set_model_state(**updates: Any) -> None: |
| MODEL_STATE.update(updates) |
|
|
|
|
| def model_json_path(filename: str) -> Path: |
| name = filename.strip() |
| if not name: |
| name = "dpo_turns.jsonl" |
| if not name.endswith(".jsonl"): |
| name = f"{name}.jsonl" |
| if "/" in name or "\\" in name or name in {".jsonl", ".", ".."}: |
| raise HTTPException(status_code=400, detail="Invalid JSONL filename.") |
| path = (MODEL_JSON_DIR / name).resolve() |
| if path.parent != MODEL_JSON_DIR.resolve(): |
| raise HTTPException(status_code=400, detail="Invalid JSONL filename.") |
| return path |
|
|
|
|
| def runtime_info() -> dict[str, str]: |
| try: |
| from src.device import runtime_metadata |
|
|
| return runtime_metadata() |
| except Exception as exc: |
| return { |
| "device": "unknown", |
| "device_label": "Runtime unavailable", |
| "torch_dtype": "unknown", |
| "device_error": str(exc), |
| } |
|
|
|
|
| def distillation_settings() -> dict[str, Any]: |
| try: |
| settings = path_explorer.Get("src/pipeline/distillation/deepseek/settings", raw=True) |
| return settings if isinstance(settings, dict) else {} |
| except Exception as exc: |
| logger.warning("Could not load distillation settings for inference config: %r", exc) |
| return {} |
|
|
|
|
| def local_model_settings() -> dict[str, Any]: |
| settings = distillation_settings().get("LOCAL_MODEL", {}) |
| return settings if isinstance(settings, dict) else {} |
|
|
|
|
| def qwen_system_prompt() -> str: |
| try: |
| value = path_explorer.Get("src/model/Qwen3BInstruct/settings/PROMPT") |
| return str(value or "") |
| except Exception as exc: |
| logger.warning("Could not load Qwen system prompt for vLLM backend: %r", exc) |
| return "" |
|
|
|
|
| def inference_config() -> InferenceConfig: |
| config = InferenceConfig.from_mapping(local_model_settings(), system_prompt=qwen_system_prompt()) |
| if DISABLE_VLLM_MODE and config.backend == "vllm": |
| config.backend = "transformers" |
| return config |
|
|
|
|
| def inference_status() -> dict[str, Any]: |
| config = inference_config() |
| return { |
| "backend": config.backend, |
| "vllm_model": config.vllm_model_name, |
| "vllm_base_url": config.vllm_base_url, |
| "fallback_to_transformers": config.vllm_fallback_to_transformers, |
| "health": LAST_INFERENCE_HEALTH, |
| } |
|
|
|
|
| def resolve_model_option(model_id: str) -> dict[str, Any]: |
| options = model_options_by_id() |
| if model_id == "qwen2.5-3b-math": |
| model_id = "base:qwen2.5-3b-math" |
| elif model_id in {"1K", "outputs/1K"}: |
| model_id = "peft:1K" |
| elif model_id.startswith("outputs/"): |
| model_id = f"peft:{model_id.removeprefix('outputs/')}" |
|
|
| if model_id in options: |
| return options[model_id] |
|
|
| path_matches = [ |
| option |
| for option in options.values() |
| if option.get("relative_path") == model_id |
| or option.get("path") == model_id |
| or Path(str(option.get("path"))).name == model_id |
| ] |
| if len(path_matches) == 1: |
| return path_matches[0] |
|
|
| return options[default_model_id()] |
|
|
|
|
| async def get_local_model(model_id: str) -> Any: |
| if model_id in LOADED_MODELS: |
| option = resolve_model_option(model_id) |
| set_model_state( |
| loading=False, |
| connected=True, |
| mode=option["kind"], |
| model=option["id"], |
| label=option["label"], |
| error=None, |
| ) |
| return LOADED_MODELS[model_id] |
|
|
| async with MODEL_LOCK: |
| if model_id in LOADED_MODELS: |
| option = resolve_model_option(model_id) |
| set_model_state( |
| loading=False, |
| connected=True, |
| mode=option["kind"], |
| model=option["id"], |
| label=option["label"], |
| error=None, |
| ) |
| return LOADED_MODELS[model_id] |
|
|
| option = resolve_model_option(model_id) |
|
|
| set_model_state( |
| loading=True, |
| connected=False, |
| mode=option["kind"], |
| model=option["id"], |
| label=option["label"], |
| error=None, |
| ) |
|
|
| def load() -> Any: |
| from src.model.Qwen3BInstruct import BaseModel as QwenModel |
|
|
| model = QwenModel(peft_adapter_path=option["path"]) |
| model._ensure_init() |
| return model |
|
|
| try: |
| loaded_model = await asyncio.to_thread(load) |
| LOADED_MODELS[model_id] = loaded_model |
| set_model_state(loading=False, connected=True, mode=option["kind"], error=None) |
| return loaded_model |
| except Exception as exc: |
| set_model_state(loading=False, connected=False, mode=option["kind"], error=str(exc)) |
| raise |
|
|
|
|
| def request_messages(request: ChatRequest) -> list[dict[str, str]]: |
| return [{"role": message.role, "content": message.content} for message in request.messages] |
|
|
|
|
| async def transformers_backend(model_id: str) -> TransformersBackend: |
| option = resolve_model_option(model_id) |
| model = await get_local_model(option["id"]) |
| return TransformersBackend(model=model, model_name=option["label"]) |
|
|
|
|
| async def get_inference_backend(model_id: str) -> BaseInferenceBackend: |
| config = inference_config() |
| option = resolve_model_option(model_id) |
| cache_key = ( |
| config.backend, |
| config.vllm_base_url if config.backend == "vllm" else option["id"], |
| config.vllm_model_name if config.backend == "vllm" else option["label"], |
| ) |
|
|
| if cache_key in INFERENCE_BACKENDS: |
| return INFERENCE_BACKENDS[cache_key] |
|
|
| async with INFERENCE_LOCK: |
| if cache_key in INFERENCE_BACKENDS: |
| return INFERENCE_BACKENDS[cache_key] |
|
|
| if config.backend == "vllm": |
| backend = build_inference_backend(config) |
| set_model_state( |
| loading=False, |
| connected=LAST_INFERENCE_HEALTH.get("ok") is True, |
| mode="vllm", |
| model=config.vllm_model_name, |
| label=f"vLLM {config.vllm_model_name}", |
| error=None, |
| ) |
| else: |
| backend = await transformers_backend(option["id"]) |
|
|
| INFERENCE_BACKENDS[cache_key] = backend |
| return backend |
|
|
|
|
| async def fallback_transformers_backend(model_id: str, reason: Exception) -> TransformersBackend: |
| logger.warning("vllm_fallback_to_transformers reason=%r", reason) |
| option = resolve_model_option(model_id) |
| return await transformers_backend(option["id"]) |
|
|
|
|
| async def check_inference_health() -> dict[str, Any]: |
| config = inference_config() |
| if config.backend == "transformers": |
| result = { |
| "ok": True, |
| "backend": "transformers", |
| "model": resolve_model_option(MODEL_STATE.get("model", default_model_id()))["label"], |
| } |
| LAST_INFERENCE_HEALTH.update(result, error=None) |
| return dict(LAST_INFERENCE_HEALTH) |
|
|
| backend = build_inference_backend(config) |
| try: |
| result = await asyncio.to_thread(backend.health_check) |
| LAST_INFERENCE_HEALTH.update(result, error=None) |
| set_model_state( |
| loading=False, |
| connected=True, |
| mode="vllm", |
| model=config.vllm_model_name, |
| label=f"vLLM {config.vllm_model_name}", |
| error=None, |
| ) |
| return dict(LAST_INFERENCE_HEALTH) |
| except Exception as exc: |
| result = { |
| "ok": False, |
| "backend": "vllm", |
| "model": config.vllm_model_name, |
| "base_url": config.vllm_base_url, |
| "error": str(exc), |
| } |
| LAST_INFERENCE_HEALTH.update(result) |
| set_model_state( |
| loading=False, |
| connected=False, |
| mode="vllm", |
| model=config.vllm_model_name, |
| label=f"vLLM {config.vllm_model_name}", |
| error=str(exc), |
| ) |
| if not config.vllm_fallback_to_transformers: |
| raise HTTPException(status_code=503, detail=result) |
| return dict(LAST_INFERENCE_HEALTH) |
|
|
|
|
| def mock_response(request: ChatRequest) -> str: |
| prompt = last_user_message(request.messages) |
| return f"""<|begin_of_thought|> |
| The local model is not connected, so this fallback response is being streamed to keep the UI testable. |
| |
| For the prompt "{prompt}", a simple sample computation is: |
| \\[ |
| 1 + 1 = 2. |
| \\] |
| <|end_of_thought|> |
| |
| <|begin_of_solution|> |
| Mock fallback answer: \\(\\boxed{{2}}\\). |
| <|end_of_solution|>""" |
|
|
|
|
| async def mock_model_stream(request: ChatRequest) -> AsyncIterator[str]: |
| text = mock_response(request) |
| chunk_size = 18 |
|
|
| for index in range(0, len(text), chunk_size): |
| await asyncio.sleep(0.055) |
| yield text[index : index + chunk_size] |
|
|
|
|
| async def next_stream_chunk(streamer: Iterator[str]) -> str | None: |
| sentinel = object() |
| chunk = await asyncio.to_thread(next, streamer, sentinel) |
| return None if chunk is sentinel else chunk |
|
|
|
|
| def matching_protocol_tag(buffer: str) -> str | None: |
| for tag in PROTOCOL_TAGS: |
| if buffer.startswith(tag): |
| return tag |
| return None |
|
|
|
|
| def could_be_protocol_tag(buffer: str) -> bool: |
| return buffer.startswith("<") and any(tag.startswith(buffer) for tag in PROTOCOL_TAGS) |
|
|
|
|
| async def stream_model_once( |
| request: ChatRequest, |
| max_tokens: int, |
| assistant_prefix: str | None = None, |
| state: dict[str, Any] | None = None, |
| stop_sequences: tuple[str, ...] = (), |
| ignore_eos: bool = False, |
| temperature_override: float | None = None, |
| ) -> AsyncIterator[str]: |
| if state is not None: |
| state["generated_tokens"] = 0 |
| state["stop_reason"] = None |
| state["backend"] = None |
| state["fallback_reason"] = None |
|
|
| config = inference_config() |
| backend = await get_inference_backend(request.model) |
| messages = request_messages(request) |
|
|
| async def open_stream(active_backend: BaseInferenceBackend) -> Iterator[str]: |
| effective_max_tokens = max_tokens |
| effective_temperature = request.temperature if temperature_override is None else temperature_override |
| effective_top_p = 0.95 |
| if active_backend.name == "vllm": |
| effective_max_tokens = min(max_tokens, config.vllm_max_tokens) if config.vllm_max_tokens > 0 else max_tokens |
| effective_temperature = config.vllm_temperature |
| effective_top_p = config.vllm_top_p |
|
|
| if state is not None: |
| state["active_max_tokens"] = effective_max_tokens |
|
|
| return await asyncio.to_thread( |
| active_backend.stream_generate, |
| messages, |
| effective_max_tokens, |
| effective_temperature, |
| effective_top_p, |
| list(stop_sequences), |
| assistant_prefix, |
| ignore_eos, |
| ) |
|
|
| try: |
| streamer = await open_stream(backend) |
| except Exception as exc: |
| if backend.name != "vllm" or not config.vllm_fallback_to_transformers: |
| raise |
| if state is not None: |
| state["fallback_reason"] = str(exc) |
| backend = await fallback_transformers_backend(request.model, exc) |
| streamer = await open_stream(backend) |
|
|
| if state is not None: |
| state["backend"] = backend.name |
|
|
| try: |
| while True: |
| try: |
| chunk = await next_stream_chunk(streamer) |
| except Exception as exc: |
| if backend.name != "vllm" or not config.vllm_fallback_to_transformers: |
| raise |
| if getattr(streamer, "output_text", "") or int( |
| getattr(streamer, "generated_token_count", 0) or 0, |
| ) > 0: |
| raise |
| if state is not None: |
| state["fallback_reason"] = str(exc) |
| backend = await fallback_transformers_backend(request.model, exc) |
| streamer = await open_stream(backend) |
| if state is not None: |
| state["backend"] = backend.name |
| state["generated_tokens"] = 0 |
| state["stop_reason"] = None |
| continue |
|
|
| if state is not None: |
| state["generated_tokens"] = int( |
| getattr(streamer, "generated_token_count", state.get("generated_tokens", 0)), |
| ) |
| state["stop_reason"] = getattr(streamer, "stop_reason", state.get("stop_reason")) |
| if chunk is None: |
| break |
| if chunk: |
| yield chunk |
| finally: |
| if state is not None: |
| tokens = int(getattr(streamer, "generated_token_count", state.get("generated_tokens", 0))) |
| state["generated_tokens"] = tokens |
| stop_reason = getattr(streamer, "stop_reason", state.get("stop_reason")) |
| if stop_reason == "stop": |
| stop_reason = "stop_sequence" |
| active_max_tokens = int(state.get("active_max_tokens") or max_tokens) |
| state["stop_reason"] = stop_reason or ("length" if tokens >= active_max_tokens else "eos") |
|
|
|
|
| def token_event(text: str) -> dict[str, Any]: |
| return {"type": "token", "text": text} |
|
|
|
|
| def usage_event( |
| *, |
| phase: str, |
| generated_tokens: int, |
| token_budget: int, |
| stop_reason: str | None = None, |
| done: bool = False, |
| ) -> dict[str, Any]: |
| return { |
| "type": "usage", |
| "phase": phase, |
| "thinking_tokens": generated_tokens, |
| "thinking_token_budget": token_budget, |
| "thinking_stop_reason": stop_reason, |
| "thinking_done": done, |
| } |
|
|
|
|
| def reported_thinking_tokens(generated_tokens: int, token_budget: int) -> int: |
| return min(max(0, generated_tokens), token_budget) |
|
|
|
|
| def thought_tail(thought_parts: list[str], max_chars: int = 256) -> str: |
| if not thought_parts: |
| return "" |
|
|
| tail: list[str] = [] |
| remaining = max_chars |
| for part in reversed(thought_parts): |
| if remaining <= 0: |
| break |
| segment = part[-remaining:] |
| tail.append(segment) |
| remaining -= len(segment) |
|
|
| return "".join(reversed(tail)) |
|
|
|
|
| def thought_ends_sentence(text: str) -> bool: |
| stripped = text.rstrip() |
| if not stripped: |
| return False |
| if re.search(r"""[.!?。！？]["')\]}»”’]*$""", stripped): |
| return True |
| if stripped.endswith(("\\]", "$$")): |
| return True |
| if re.search(r"""\\end\{[a-zA-Z*]+\}$""", stripped): |
| return True |
| return False |
|
|
|
|
| async def stream_thought_phase( |
| source: AsyncIterator[str], |
| thought_parts: list[str], |
| generation_state: dict[str, Any], |
| token_budget: int, |
| *, |
| allow_sentence_completion_after_budget: bool = False, |
| ) -> AsyncIterator[dict[str, Any]]: |
| buffer = "" |
| forwarding_thought = True |
|
|
| yield token_event(BEGIN_THOUGHT) |
|
|
| async for chunk in source: |
| buffer += chunk |
| output: list[str] = [] |
| stop_after_output = False |
|
|
| while buffer: |
| tag = matching_protocol_tag(buffer) |
| if tag is not None: |
| buffer = buffer[len(tag):] |
|
|
| if tag in {END_THOUGHT, BEGIN_SOLUTION, END_SOLUTION}: |
| forwarding_thought = False |
| generation_state["stop_reason"] = "stop_sequence" |
| stop_after_output = True |
| buffer = "" |
| break |
| continue |
|
|
| if could_be_protocol_tag(buffer): |
| break |
|
|
| generated_tokens = int(generation_state.get("generated_tokens") or 0) |
| budget_reached = generated_tokens >= token_budget |
| if ( |
| allow_sentence_completion_after_budget |
| and budget_reached |
| and thought_ends_sentence(thought_tail(thought_parts)) |
| ): |
| generation_state["stop_reason"] = "length" |
| stop_after_output = True |
| buffer = "" |
| break |
|
|
| char = buffer[0] |
| buffer = buffer[1:] |
|
|
| if forwarding_thought: |
| thought_parts.append(char) |
| output.append(char) |
| if ( |
| allow_sentence_completion_after_budget |
| and budget_reached |
| and thought_ends_sentence(thought_tail(thought_parts)) |
| ): |
| generation_state["stop_reason"] = "length" |
| stop_after_output = True |
| buffer = "" |
| break |
|
|
| if output: |
| yield token_event("".join(output)) |
| yield usage_event( |
| phase="thought", |
| generated_tokens=reported_thinking_tokens( |
| int(generation_state.get("generated_tokens") or 0), |
| token_budget, |
| ), |
| token_budget=token_budget, |
| ) |
|
|
| if stop_after_output: |
| return |
|
|
|
|
| async def stream_solution_phase(source: AsyncIterator[str]) -> AsyncIterator[str]: |
| buffer = "" |
| saw_end_solution = False |
|
|
| async for chunk in source: |
| buffer += chunk |
| output: list[str] = [] |
|
|
| while buffer: |
| tag = matching_protocol_tag(buffer) |
| if tag is not None: |
| buffer = buffer[len(tag):] |
|
|
| if tag == END_SOLUTION: |
| output.append(tag) |
| saw_end_solution = True |
| buffer = "" |
| break |
| continue |
|
|
| if could_be_protocol_tag(buffer): |
| break |
|
|
| output.append(buffer[0]) |
| buffer = buffer[1:] |
|
|
| if output: |
| yield "".join(output) |
|
|
| if saw_end_solution: |
| return |
|
|
| if buffer and not could_be_protocol_tag(buffer): |
| yield buffer |
|
|
| if not saw_end_solution: |
| yield f"\n{END_SOLUTION}" |
|
|
|
|
| async def real_model_events(request: ChatRequest) -> AsyncIterator[dict[str, Any]]: |
| if request.fast_mode: |
| async for event in fast_model_events(request): |
| yield event |
| return |
|
|
| thinking_limit = THINKING_TOKEN_BUDGETS[request.thinking_effort] |
| thought_parts: list[str] = [] |
| thought_state: dict[str, Any] = {} |
| thought_generation_limit = thinking_limit + THOUGHT_SENTENCE_COMPLETION_TOKENS |
|
|
| first_pass = stream_model_once( |
| request, |
| max_tokens=thought_generation_limit, |
| state=thought_state, |
| stop_sequences=THOUGHT_STOP_SEQUENCES, |
| ignore_eos=True, |
| ) |
| async for event in stream_thought_phase( |
| first_pass, |
| thought_parts, |
| thought_state, |
| thinking_limit, |
| allow_sentence_completion_after_budget=True, |
| ): |
| yield event |
|
|
| thought_stop_reason = thought_state.get("stop_reason") |
| thought_tokens = reported_thinking_tokens( |
| int(thought_state.get("generated_tokens") or 0), |
| thinking_limit, |
| ) |
|
|
| if thought_stop_reason == "length": |
| runtime_instruction = instruction_block(thinking_limit_reached_instruction()) |
| thought_parts.append(runtime_instruction) |
| yield token_event(runtime_instruction) |
|
|
| thought = "".join(thought_parts).strip() |
| assistant_prefix = f"{BEGIN_THOUGHT}\n{thought}\n{END_THOUGHT}\n\n{BEGIN_SOLUTION}\n" |
|
|
| yield usage_event( |
| phase="thought", |
| generated_tokens=thought_tokens, |
| token_budget=thinking_limit, |
| stop_reason=thought_stop_reason, |
| done=True, |
| ) |
|
|
| yield token_event(f"\n{END_THOUGHT}\n\n{BEGIN_SOLUTION}\n") |
|
|
| second_pass = stream_model_once( |
| request, |
| max_tokens=request.max_tokens, |
| assistant_prefix=assistant_prefix, |
| stop_sequences=(END_SOLUTION,), |
| ) |
| async for chunk in stream_solution_phase(second_pass): |
| yield token_event(chunk) |
|
|
|
|
| async def fast_model_events(request: ChatRequest) -> AsyncIterator[dict[str, Any]]: |
| thinking_limit = THINKING_TOKEN_BUDGETS[request.thinking_effort] |
| generation_state: dict[str, Any] = {} |
| max_tokens = thinking_limit + request.max_tokens |
|
|
| source = stream_model_once( |
| request, |
| max_tokens=max_tokens, |
| state=generation_state, |
| stop_sequences=(END_SOLUTION,), |
| temperature_override=0.0, |
| ) |
|
|
| async for chunk in source: |
| yield token_event(chunk) |
|
|
| yield usage_event( |
| phase="thought", |
| generated_tokens=reported_thinking_tokens( |
| int(generation_state.get("generated_tokens") or 0), |
| thinking_limit, |
| ), |
| token_budget=thinking_limit, |
| stop_reason=generation_state.get("stop_reason"), |
| done=True, |
| ) |
|
|
|
|
| async def mock_model_events(request: ChatRequest) -> AsyncIterator[dict[str, Any]]: |
| async for chunk in mock_model_stream(request): |
| yield token_event(chunk) |
|
|
|
|
| async def model_events(request: ChatRequest) -> AsyncIterator[dict[str, Any]]: |
| option = resolve_model_option(request.model) |
| request.model = option["id"] |
| config = inference_config() |
| if config.backend == "vllm": |
| set_model_state( |
| model=config.vllm_model_name, |
| label=f"vLLM {config.vllm_model_name}", |
| mode="vllm", |
| ) |
| else: |
| set_model_state(model=option["id"], label=option["label"], mode=option["kind"]) |
|
|
| try: |
| async for event in real_model_events(request): |
| yield event |
| except Exception as exc: |
| set_model_state(loading=False, connected=False, error=str(exc)) |
| logger.exception("local_model_events_failed backend=%s model=%s", config.backend, request.model) |
| if config.backend == "vllm" and not config.vllm_fallback_to_transformers: |
| raise |
| async for event in mock_model_events(request): |
| yield event |
|
|
|
|
| async def model_stream(request: ChatRequest) -> AsyncIterator[str]: |
| async for event in model_events(request): |
| if event.get("type") == "token": |
| yield event.get("text", "") |
|
|
|
|
| async def model_complete(request: ChatRequest) -> str: |
| chunks: list[str] = [] |
| async for chunk in model_stream(request): |
| chunks.append(chunk) |
| return "".join(chunks) |
|
|
|
|
| @app.post("/api/chat") |
| async def chat(request: ChatRequest) -> dict[str, str]: |
| return {"text": await model_complete(request)} |
|
|
|
|
| @app.post("/api/chat/stream") |
| async def chat_stream(request: ChatRequest) -> StreamingResponse: |
| async def events() -> AsyncIterator[str]: |
| text_batch: list[str] = [] |
| text_batch_len = 0 |
| last_flush = time.monotonic() |
|
|
| def flush_text_batch() -> str | None: |
| nonlocal text_batch, text_batch_len, last_flush |
| if not text_batch: |
| return None |
| payload = json.dumps({"text": "".join(text_batch)}) |
| text_batch = [] |
| text_batch_len = 0 |
| last_flush = time.monotonic() |
| return f"event: token\ndata: {payload}\n\n" |
|
|
| try: |
| async for event in model_events(request): |
| event_type = event.get("type") |
|
|
| if event_type == "token": |
| text = event.get("text", "") |
| if not text: |
| continue |
| text_batch.append(text) |
| text_batch_len += len(text) |
| if ( |
| text_batch_len >= STREAM_TEXT_BATCH_CHARS |
| or time.monotonic() - last_flush >= STREAM_TEXT_BATCH_SECONDS |
| ): |
| flushed = flush_text_batch() |
| if flushed: |
| yield flushed |
| continue |
|
|
| flushed = flush_text_batch() |
| if flushed: |
| yield flushed |
|
|
| if event_type == "usage": |
| payload = json.dumps({key: value for key, value in event.items() if key != "type"}) |
| yield f"event: usage\ndata: {payload}\n\n" |
|
|
| flushed = flush_text_batch() |
| if flushed: |
| yield flushed |
| yield "event: done\ndata: {}\n\n" |
| except asyncio.CancelledError: |
| raise |
| except Exception as exc: |
| payload = json.dumps({"message": str(exc)}) |
| yield f"event: error\ndata: {payload}\n\n" |
|
|
| return StreamingResponse( |
| events(), |
| media_type="text/event-stream", |
| headers={ |
| "Cache-Control": "no-cache", |
| "Connection": "keep-alive", |
| "X-Accel-Buffering": "no", |
| }, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| import uvicorn |
|
|
| uvicorn.run( |
| "server:app", |
| host="127.0.0.1", |
| port=8000, |
| reload=False, |
| ) |
|
|
