71 lines
3.0 KiB
Python
71 lines
3.0 KiB
Python
import asyncio
|
|
import logging
|
|
from types import SimpleNamespace
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
|
|
from deerflow.config import runtime_performance_config
|
|
from deerflow.models.timing_callback import ModelTimingCallback
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def synchronous_timing_logs(monkeypatch):
|
|
# These assertions inspect logs immediately. Keep the legacy path isolated
|
|
# from a developer's opt-in config; the queued path has its own regression
|
|
# suite in test_model_timing_nonblocking.py.
|
|
monkeypatch.setattr(
|
|
runtime_performance_config,
|
|
"get_runtime_performance_config",
|
|
lambda: SimpleNamespace(model_timing_nonblocking_enabled=False),
|
|
)
|
|
|
|
|
|
def test_model_timing_callback_logs_connection_and_completion(caplog):
|
|
callback = ModelTimingCallback(model_name="qwen", base_url="http://model/v1")
|
|
run_id = uuid4()
|
|
|
|
with caplog.at_level(logging.INFO, logger="deerflow.models.timing_callback"):
|
|
callback.on_chat_model_start({}, [[object()]], run_id=run_id)
|
|
callback.on_llm_new_token("", run_id=run_id, chunk=object())
|
|
callback.on_llm_end(object(), run_id=run_id)
|
|
|
|
messages = [record.getMessage() for record in caplog.records]
|
|
assert any("LLM_CONNECT_START" in message for message in messages)
|
|
assert any("LLM_CONNECT_OK" in message and "first_chunk_ms=" in message for message in messages)
|
|
assert any("LLM_REQUEST_DONE" in message and "total_ms=" in message for message in messages)
|
|
|
|
|
|
def test_model_timing_callback_logs_failure(caplog):
|
|
callback = ModelTimingCallback(
|
|
model_name="qwen",
|
|
base_url="https://user:API_KEY_SHOULD_NOT_LOG@model.example:8443/v1?token=secret",
|
|
)
|
|
run_id = uuid4()
|
|
|
|
with caplog.at_level(logging.INFO, logger="deerflow.models.timing_callback"):
|
|
callback.on_chat_model_start({}, [[object()]], run_id=run_id)
|
|
callback.on_llm_error(TimeoutError("PROMPT_OR_TOOL_BODY_SHOULD_NOT_LOG"), run_id=run_id)
|
|
|
|
message = caplog.records[-1].getMessage()
|
|
assert "LLM_REQUEST_FAILED" in message
|
|
assert "TimeoutError" in message
|
|
assert "base_url=https://model.example:8443" in message
|
|
assert "API_KEY_SHOULD_NOT_LOG" not in message
|
|
assert "PROMPT_OR_TOOL_BODY_SHOULD_NOT_LOG" not in message
|
|
|
|
|
|
def test_model_timing_callback_logs_cancellation_as_info(caplog):
|
|
callback = ModelTimingCallback(model_name="qwen", base_url="http://model/v1")
|
|
run_id = uuid4()
|
|
|
|
with caplog.at_level(logging.INFO, logger="deerflow.models.timing_callback"):
|
|
callback.on_chat_model_start({}, [[object()]], run_id=run_id)
|
|
callback.on_llm_error(asyncio.CancelledError(), run_id=run_id)
|
|
|
|
cancellation_records = [record for record in caplog.records if "LLM_REQUEST_CANCELLED" in record.getMessage()]
|
|
assert len(cancellation_records) == 1
|
|
assert cancellation_records[0].levelno == logging.INFO
|
|
assert "status=cancelled" in cancellation_records[0].getMessage()
|
|
assert not any("LLM_REQUEST_FAILED" in record.getMessage() for record in caplog.records)
|