124 lines
3.9 KiB
Python
124 lines
3.9 KiB
Python
"""Concise model connection and latency diagnostics for production logs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
from uuid import UUID
|
|
|
|
from langchain_core.callbacks import BaseCallbackHandler
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _CallTiming:
|
|
started: float
|
|
first_chunk: float | None = None
|
|
|
|
|
|
class ModelTimingCallback(BaseCallbackHandler):
|
|
"""Log model-call start, first response chunk, completion, and failure."""
|
|
|
|
def __init__(self, *, model_name: str, base_url: str | None) -> None:
|
|
self._model_name = model_name
|
|
self._base_url = base_url or "-"
|
|
self._calls: dict[str, _CallTiming] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
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)
|
|
logger.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)
|
|
logger.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()
|
|
with self._lock:
|
|
timing = self._calls.pop(rid, None)
|
|
|
|
if timing is None:
|
|
logger.error(
|
|
"LLM_REQUEST_FAILED model=%s base_url=%s run_id=%s error_type=%s error=%s",
|
|
self._model_name,
|
|
self._base_url,
|
|
rid,
|
|
type(error).__name__ if error else "UnknownLifecycle",
|
|
error or "missing start event",
|
|
)
|
|
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:
|
|
logger.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
|
|
|
|
logger.error(
|
|
"LLM_REQUEST_FAILED model=%s base_url=%s run_id=%s status=error first_chunk_ms=%s total_ms=%d error_type=%s error=%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__,
|
|
error,
|
|
)
|