195 lines
7.7 KiB
Python
195 lines
7.7 KiB
Python
"""Opt-in model timing remains independent of busy tools and slow log sinks."""
|
|
|
|
import asyncio
|
|
import logging
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from types import SimpleNamespace
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from langchain_core.language_models.chat_models import BaseChatModel
|
|
from langchain_core.messages import AIMessageChunk
|
|
from langchain_core.outputs import ChatGenerationChunk
|
|
|
|
from deerflow.config import runtime_performance_config
|
|
from deerflow.models import timing_callback
|
|
from deerflow.models.timing_callback import ModelTimingCallback, _TimingLogDispatcher
|
|
|
|
|
|
@pytest.fixture
|
|
def timing_logs(monkeypatch, caplog):
|
|
monkeypatch.setattr(runtime_performance_config, "get_runtime_performance_config", lambda: SimpleNamespace(model_timing_nonblocking_enabled=True))
|
|
dispatcher = _TimingLogDispatcher()
|
|
monkeypatch.setattr(timing_callback, "_timing_logs", dispatcher)
|
|
caplog.set_level(logging.INFO, logger=timing_callback.logger.name)
|
|
yield dispatcher
|
|
dispatcher.close(timeout=2)
|
|
|
|
|
|
class _NativeAsyncModel(BaseChatModel):
|
|
@property
|
|
def _llm_type(self):
|
|
return "timing-callback-test"
|
|
|
|
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
|
|
raise AssertionError("The model must use native async streaming")
|
|
|
|
async def _astream(self, messages, stop=None, run_manager=None, **kwargs):
|
|
yield ChatGenerationChunk(message=AIMessageChunk(content="first"))
|
|
yield ChatGenerationChunk(message=AIMessageChunk(content="second"))
|
|
|
|
|
|
class _SlowHandler(logging.Handler):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.entered = threading.Event()
|
|
self.release_event = threading.Event()
|
|
self.worker_threads = []
|
|
|
|
def emit(self, record):
|
|
self.worker_threads.append(threading.current_thread().name)
|
|
self.entered.set()
|
|
self.release_event.wait(timeout=5)
|
|
|
|
|
|
async def _wait_until_set(event):
|
|
while not event.is_set():
|
|
await asyncio.sleep(0.001)
|
|
|
|
|
|
def test_default_off_keeps_synchronous_logging_and_executor_dispatch(monkeypatch, caplog):
|
|
config = runtime_performance_config.RuntimePerformanceConfig()
|
|
monkeypatch.setattr(runtime_performance_config, "get_runtime_performance_config", lambda: config)
|
|
caplog.set_level(logging.INFO, logger=timing_callback.logger.name)
|
|
callback = ModelTimingCallback(model_name="test", base_url=None)
|
|
assert callback.run_inline is False
|
|
|
|
def reject_queue(*args):
|
|
raise AssertionError("Default-off callbacks must not use the background log queue")
|
|
|
|
monkeypatch.setattr(timing_callback._timing_logs, "emit", reject_queue)
|
|
run_id = uuid4()
|
|
callback.on_chat_model_start({}, [[]], run_id=run_id)
|
|
assert "LLM_CONNECT_START" in caplog.text
|
|
callback.on_llm_new_token("token", run_id=run_id)
|
|
assert "LLM_CONNECT_OK" in caplog.text
|
|
callback.on_llm_end(None, run_id=run_id)
|
|
assert "LLM_REQUEST_DONE" in caplog.text
|
|
|
|
|
|
def test_stream_progresses_with_saturated_tool_pool_and_blocked_log_handler(timing_logs, caplog):
|
|
"""Exercise LangChain dispatch, not just direct callback method calls."""
|
|
sink = _SlowHandler()
|
|
timing_callback.logger.addHandler(sink)
|
|
tool_started = threading.Event()
|
|
release_tool = threading.Event()
|
|
|
|
def long_tool():
|
|
tool_started.set()
|
|
release_tool.wait(timeout=5)
|
|
|
|
async def scenario():
|
|
loop = asyncio.get_running_loop()
|
|
# A fresh asyncio.run loop owns this executor and cleans it up.
|
|
loop.set_default_executor(ThreadPoolExecutor(max_workers=1, thread_name_prefix="busy-tool"))
|
|
blocked_tool = loop.run_in_executor(None, long_tool)
|
|
try:
|
|
await asyncio.wait_for(_wait_until_set(tool_started), timeout=1)
|
|
callback = ModelTimingCallback(model_name="test", base_url="https://models.invalid/v1")
|
|
model = _NativeAsyncModel(callbacks=[callback])
|
|
|
|
async def consume():
|
|
return [chunk.content async for chunk in model.astream("private prompt")]
|
|
|
|
chunks = await asyncio.wait_for(consume(), timeout=1)
|
|
await asyncio.wait_for(_wait_until_set(sink.entered), timeout=1)
|
|
assert "".join(chunks) == "firstsecond"
|
|
assert not blocked_tool.done()
|
|
assert not sink.release_event.is_set()
|
|
assert callback._calls == {}
|
|
finally:
|
|
release_tool.set()
|
|
sink.release_event.set()
|
|
await blocked_tool
|
|
|
|
try:
|
|
asyncio.run(scenario())
|
|
timing_logs.close(timeout=2)
|
|
finally:
|
|
release_tool.set()
|
|
sink.release_event.set()
|
|
timing_callback.logger.removeHandler(sink)
|
|
assert sink.worker_threads
|
|
assert set(sink.worker_threads) == {"deerflow-model-timing-log"}
|
|
assert "LLM_CONNECT_START" in caplog.text
|
|
assert "LLM_CONNECT_OK" in caplog.text
|
|
assert "LLM_REQUEST_DONE" in caplog.text
|
|
assert "private prompt" not in caplog.text
|
|
|
|
|
|
def test_sync_callback_preserves_sanitized_lifecycle_and_only_logs_first_token(timing_logs, caplog):
|
|
callback = ModelTimingCallback(model_name="test", base_url="https://user:secret@models.invalid:443/private?key=secret")
|
|
run_id = uuid4()
|
|
callback.on_chat_model_start({}, [["private input"]], run_id=run_id)
|
|
callback.on_llm_new_token("", run_id=run_id)
|
|
for _ in range(100):
|
|
callback.on_llm_new_token("private output", run_id=run_id)
|
|
callback.on_llm_end(None, run_id=run_id)
|
|
timing_logs.close(timeout=2)
|
|
|
|
records = [record for record in caplog.records if record.name == timing_callback.logger.name]
|
|
assert len(records) == 3
|
|
assert "base_url=https://models.invalid:443" in caplog.text
|
|
assert "first_chunk_ms=" in caplog.text
|
|
assert f"run_id={run_id}" in caplog.text
|
|
for secret in ("private", "secret", "user:"):
|
|
assert secret not in caplog.text
|
|
assert callback._calls == {}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("error", "expected", "level"),
|
|
[
|
|
(RuntimeError("private upstream payload"), "LLM_REQUEST_FAILED", logging.ERROR),
|
|
(asyncio.CancelledError("private cancellation"), "LLM_REQUEST_CANCELLED", logging.INFO),
|
|
(GeneratorExit("private closure"), "LLM_REQUEST_CANCELLED", logging.INFO),
|
|
],
|
|
)
|
|
def test_failure_and_cancellation_are_content_free(timing_logs, caplog, error, expected, level):
|
|
callback = ModelTimingCallback(model_name="test", base_url=None)
|
|
run_id = uuid4()
|
|
callback.on_chat_model_start({}, [[]], run_id=run_id)
|
|
callback.on_llm_error(error, run_id=run_id)
|
|
timing_logs.close(timeout=2)
|
|
|
|
terminal = [record for record in caplog.records if expected in record.getMessage()]
|
|
assert len(terminal) == 1
|
|
assert terminal[0].levelno == level
|
|
assert "private" not in caplog.text
|
|
assert callback._calls == {}
|
|
|
|
|
|
def test_full_log_queue_drops_without_blocking_and_reports_recovery(monkeypatch, caplog):
|
|
dispatcher = _TimingLogDispatcher(capacity=1)
|
|
monkeypatch.setattr(timing_callback, "_timing_logs", dispatcher)
|
|
caplog.set_level(logging.INFO, logger=timing_callback.logger.name)
|
|
sink = _SlowHandler()
|
|
timing_callback.logger.addHandler(sink)
|
|
try:
|
|
dispatcher.emit(logging.INFO, "first")
|
|
assert sink.entered.wait(timeout=1)
|
|
dispatcher.emit(logging.INFO, "queued")
|
|
dispatcher.emit(logging.INFO, "dropped")
|
|
assert dispatcher._queue.qsize() == 1
|
|
assert dispatcher._dropped == 1
|
|
sink.release_event.set()
|
|
dispatcher.close(timeout=2)
|
|
finally:
|
|
sink.release_event.set()
|
|
dispatcher.close(timeout=2)
|
|
timing_callback.logger.removeHandler(sink)
|
|
assert "LLM_TIMING_LOG_DROPPED count=1 reason=queue_full" in caplog.text
|
|
assert "queued" in caplog.text
|
|
assert not any(record.getMessage() == "dropped" for record in caplog.records)
|