"""Concise model connection and latency diagnostics for production logs.""" from __future__ import annotations import asyncio import atexit import logging import queue import threading import time from dataclasses import dataclass from typing import Any from urllib.parse import urlsplit from uuid import UUID from langchain_core.callbacks import BaseCallbackHandler logger = logging.getLogger(__name__) class _TimingLogDispatcher: """Keep slow log handlers off both the event loop and its tool executor. Only already-sanitized, small lifecycle records enter this bounded queue. A daemon thread invokes the existing handlers; producers never wait for their I/O. Full queues drop diagnostics and report the count when the sink recovers, instead of retaining unbounded records or stalling generation. """ def __init__(self, *, capacity: int = 4096) -> None: self._queue: queue.Queue[logging.LogRecord] = queue.Queue(maxsize=capacity) self._lock = threading.Lock() self._thread: threading.Thread | None = None self._closed = False self._dropped = 0 def emit(self, level: int, message: str, *args: Any) -> None: if not logger.isEnabledFor(level): return # Preserve the event's wall-clock timestamp even if delivery is delayed. # Formatting, filters and handler I/O run only on the consumer thread. record = logger.makeRecord(logger.name, level, __file__, 0, message, args, None) with self._lock: if self._closed: return try: self._queue.put_nowait(record) except queue.Full: self._dropped += 1 return if self._thread is None: self._thread = threading.Thread(target=self._run, name="deerflow-model-timing-log", daemon=True) self._thread.start() def _run(self) -> None: while True: try: record = self._queue.get(timeout=0.1) except queue.Empty: with self._lock: if self._closed: return continue try: logger.handle(record) except Exception: # Diagnostic sinks must not fail a model call or kill the worker. pass finally: self._queue.task_done() with self._lock: dropped, self._dropped = self._dropped, 0 if dropped: try: logger.warning("LLM_TIMING_LOG_DROPPED count=%d reason=queue_full", dropped) except Exception: pass def close(self, timeout: float = 0.5) -> None: """Best-effort drain without making process exit wait on a log sink.""" with self._lock: self._closed = True thread = self._thread if thread is not None and thread is not threading.current_thread(): thread.join(timeout=timeout) _timing_logs = _TimingLogDispatcher() atexit.register(_timing_logs.close) def _safe_base_url(value: str | None) -> str: """Log only the upstream origin, never credentials, query, or URL path.""" if not value: return "-" try: parsed = urlsplit(value) except (TypeError, ValueError): return "-" if parsed.scheme not in {"http", "https"} or not parsed.hostname: return "-" host = parsed.hostname try: port = parsed.port except ValueError: return "-" return f"{parsed.scheme}://{host}:{port}" if port is not None else f"{parsed.scheme}://{host}" @dataclass(slots=True) class _CallTiming: started: float first_chunk: float | None = None class ModelTimingCallback(BaseCallbackHandler): """Optionally record timing inline and deliver logs on a dedicated queue.""" # LangChain otherwise sends each synchronous callback, including every # token's no-op after the first, through the shared tool thread pool. # When enabled, inline work is limited to timestamps, short locked state # changes and a non-blocking enqueue, never synchronous handler I/O. run_inline = False def __init__(self, *, model_name: str, base_url: str | None) -> None: from deerflow.config.runtime_performance_config import get_runtime_performance_config self.run_inline = get_runtime_performance_config().model_timing_nonblocking_enabled self._model_name = model_name self._base_url = _safe_base_url(base_url) self._calls: dict[str, _CallTiming] = {} self._lock = threading.Lock() def _log(self, level: int, message: str, *args: Any) -> None: if self.run_inline: _timing_logs.emit(level, message, *args) else: # The default-off path retains LangChain's prior callback dispatch # and synchronous logging behavior for deployment compatibility. logger.log(level, message, *args) def on_chat_model_start( self, serialized: dict[str, Any], messages: list[list[Any]], *, run_id: UUID, **kwargs: Any, ) -> None: rid = str(run_id) now = time.monotonic() with self._lock: self._calls[rid] = _CallTiming(started=now) self._log( logging.INFO, "LLM_CONNECT_START model=%s base_url=%s run_id=%s message_batches=%d message_count=%d", self._model_name, self._base_url, rid, len(messages), sum(len(batch) for batch in messages), ) def on_llm_new_token( self, token: str, *, run_id: UUID, chunk: Any | None = None, **kwargs: Any, ) -> None: """Treat the first callback as connection success, including reasoning-only chunks.""" rid = str(run_id) now = time.monotonic() with self._lock: timing = self._calls.get(rid) if timing is None or timing.first_chunk is not None: return timing.first_chunk = now first_chunk_ms = int((now - timing.started) * 1000) self._log( logging.INFO, "LLM_CONNECT_OK model=%s base_url=%s run_id=%s first_chunk_ms=%d", self._model_name, self._base_url, rid, first_chunk_ms, ) def on_llm_end(self, response: Any, *, run_id: UUID, **kwargs: Any) -> None: self._finish(run_id=run_id, error=None) def on_llm_error(self, error: BaseException, *, run_id: UUID, **kwargs: Any) -> None: self._finish(run_id=run_id, error=error) def _finish(self, *, run_id: UUID, error: BaseException | None) -> None: rid = str(run_id) now = time.monotonic() cancelled = isinstance(error, (asyncio.CancelledError, GeneratorExit)) with self._lock: timing = self._calls.pop(rid, None) if timing is None: if cancelled: self._log( logging.INFO, "LLM_REQUEST_CANCELLED model=%s base_url=%s run_id=%s status=cancelled error_type=%s", self._model_name, self._base_url, rid, type(error).__name__, ) return self._log( logging.ERROR, "LLM_REQUEST_FAILED model=%s base_url=%s run_id=%s error_type=%s", self._model_name, self._base_url, rid, type(error).__name__ if error else "UnknownLifecycle", ) return total_ms = int((now - timing.started) * 1000) first_chunk_ms = int((timing.first_chunk - timing.started) * 1000) if timing.first_chunk else None if error is None: self._log( logging.INFO, "LLM_REQUEST_DONE model=%s base_url=%s run_id=%s status=success first_chunk_ms=%s total_ms=%d", self._model_name, self._base_url, rid, first_chunk_ms if first_chunk_ms is not None else "none", total_ms, ) return if cancelled: self._log( logging.INFO, "LLM_REQUEST_CANCELLED model=%s base_url=%s run_id=%s status=cancelled first_chunk_ms=%s total_ms=%d error_type=%s", self._model_name, self._base_url, rid, first_chunk_ms if first_chunk_ms is not None else "none", total_ms, type(error).__name__, ) return self._log( logging.ERROR, "LLM_REQUEST_FAILED model=%s base_url=%s run_id=%s status=error first_chunk_ms=%s total_ms=%d error_type=%s", self._model_name, self._base_url, rid, first_chunk_ms if first_chunk_ms is not None else "none", total_ms, type(error).__name__, )