| import queue |
| import threading |
| import time |
| import logging |
|
|
| from bot.telegram_api import ( |
| send_message, |
| edit_message, |
| delete_message, |
| md_to_telegram_html, |
| split_for_telegram, |
| ) |
| from model.qwen import generate_stream |
| from db import models as dbm |
|
|
| log = logging.getLogger("streaming") |
|
|
| EDIT_THROTTLE_SEC = 1.0 |
| CURSOR = " ▌" |
| PREVIEW_LIMIT = 3500 |
| QUEUE_TIMEOUT = 300 |
| FALLBACK_PLAIN_LIMIT = 3900 |
|
|
|
|
| async def stream_reply(chat_id, user_id, user_text): |
| try: |
| window = int(dbm.get_setting("conversation_window", "6")) |
| except Exception: |
| window = 6 |
| try: |
| max_new = int(dbm.get_setting("max_new_tokens", "1024")) |
| except Exception: |
| max_new = 1024 |
| try: |
| temperature = float(dbm.get_setting("temperature", "0.7")) |
| except Exception: |
| temperature = 0.7 |
| try: |
| top_p = float(dbm.get_setting("top_p", "0.9")) |
| except Exception: |
| top_p = 0.9 |
| system_prompt = dbm.get_setting("system_prompt", "You are a helpful assistant.") |
|
|
| history = dbm.get_history(user_id, window) |
| messages = [{"role": "system", "content": system_prompt}] + history |
|
|
| sent = await send_message(chat_id, CURSOR, parse_mode=None) |
| if not sent.get("ok"): |
| log.error(f"sendMessage placeholder failed: {sent}") |
| return |
| placeholder_id = sent["result"]["message_id"] |
|
|
| buffer = "" |
| last_edit = 0.0 |
| gpu_start = time.time() |
| first_token_at = None |
| error = None |
|
|
| q = queue.Queue() |
|
|
| def _producer(): |
| try: |
| for chunk in generate_stream(messages, max_new, temperature, top_p): |
| q.put(("chunk", chunk)) |
| except Exception as e: |
| q.put(("error", repr(e))) |
| finally: |
| q.put(("done", None)) |
|
|
| threading.Thread(target=_producer, daemon=True).start() |
|
|
| try: |
| while True: |
| try: |
| kind, payload = q.get(timeout=QUEUE_TIMEOUT) |
| except queue.Empty: |
| error = "timeout" |
| break |
|
|
| if kind == "chunk": |
| if first_token_at is None: |
| first_token_at = time.time() |
| buffer += payload |
| now = time.time() |
| if now - last_edit >= EDIT_THROTTLE_SEC: |
| preview = buffer |
| if len(preview) > PREVIEW_LIMIT: |
| preview = preview[-PREVIEW_LIMIT:] |
| await _safe_edit_plain(chat_id, placeholder_id, preview + CURSOR) |
| last_edit = now |
| elif kind == "error": |
| error = payload |
| break |
| elif kind == "done": |
| break |
| except Exception as e: |
| error = repr(e) |
| log.exception("stream_reply failure") |
|
|
| gpu_ms = int((time.time() - gpu_start) * 1000) |
| latency_ms = int((time.time() - (first_token_at or gpu_start)) * 1000) |
|
|
| if error: |
| await _safe_edit_plain(chat_id, placeholder_id, f"⚠️ Error: {error}") |
| dbm.log_gpu(user_id, gpu_ms, "error") |
| return |
|
|
| if not buffer.strip(): |
| await _safe_edit_plain(chat_id, placeholder_id, "⚠️ Empty response from model.") |
| dbm.log_gpu(user_id, gpu_ms, "empty") |
| return |
|
|
| await delete_message(chat_id, placeholder_id) |
|
|
| delivered = await _deliver_final(chat_id, buffer) |
| if not delivered: |
| log.error("Final delivery failed entirely") |
| await send_message( |
| chat_id, |
| "⚠️ Could not deliver the response (formatting error). Full text below:", |
| parse_mode=None, |
| ) |
| for chunk in _chunk_plain(buffer, FALLBACK_PLAIN_LIMIT): |
| await send_message(chat_id, chunk, parse_mode=None) |
|
|
| |
| dbm.add_message( |
| user_id, "assistant", buffer, |
| latency_ms=latency_ms, gpu_time_ms=gpu_ms, |
| ) |
| dbm.archive_append(user_id, "assistant", buffer, latency_ms=latency_ms) |
| dbm.increment_user_turns(user_id) |
| dbm.log_gpu(user_id, gpu_ms, "ok") |
|
|
| user = dbm.get_user(user_id) |
| if user and user.get("token"): |
| dbm.bump_token_use(user["token"]) |
|
|
|
|
| async def _deliver_final(chat_id, raw_text: str) -> bool: |
| try: |
| html_text = md_to_telegram_html(raw_text) |
| chunks = split_for_telegram(html_text) |
| except Exception: |
| log.exception("md/split failed for final delivery") |
| return False |
|
|
| sent_count = 0 |
| for i, chunk in enumerate(chunks): |
| r = await send_message(chat_id, chunk, parse_mode="HTML") |
| if r.get("ok"): |
| sent_count += 1 |
| continue |
| if i == 0: |
| log.warning(f"Final chunk 0 HTML failed: {r}") |
| return False |
| log.warning(f"Final chunk {i} HTML failed, retrying plain: {r}") |
| r2 = await send_message(chat_id, _strip_html(chunk), parse_mode=None) |
| if r2.get("ok"): |
| sent_count += 1 |
| continue |
| log.error(f"Plain fallback for chunk {i} also failed: {r2}") |
| return False |
|
|
| return sent_count == len(chunks) |
|
|
|
|
| def _strip_html(s: str) -> str: |
| import re |
| return re.sub(r"<[^>]+>", "", s) |
|
|
|
|
| def _chunk_plain(text: str, max_len: int): |
| chunks = [] |
| remaining = text |
| while len(remaining) > max_len: |
| cut = remaining.rfind("\n", 0, max_len) |
| if cut == -1 or cut < max_len * 0.5: |
| cut = max_len |
| chunks.append(remaining[:cut]) |
| remaining = remaining[cut:].lstrip("\n") |
| if remaining: |
| chunks.append(remaining) |
| return chunks |
|
|
|
|
| async def _safe_edit_plain(chat_id, message_id, text): |
| try: |
| r = await edit_message( |
| chat_id, message_id, text, |
| parse_mode=None, |
| ) |
| if not r.get("ok"): |
| log.warning(f"edit (plain) failed: {r}") |
| except Exception: |
| log.exception("plain edit failed") |
