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, )