32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
from __future__ import annotations
|
|
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|
|
|
from deerflow.agents.middlewares.message_timestamp_middleware import MessageTimestampMiddleware
|
|
|
|
|
|
def test_after_model_stamps_missing_message_timestamps():
|
|
middleware = MessageTimestampMiddleware()
|
|
state = {
|
|
"messages": [
|
|
HumanMessage(content="hello"),
|
|
AIMessage(content="hi"),
|
|
]
|
|
}
|
|
|
|
middleware.after_model(state, runtime=None) # type: ignore[arg-type]
|
|
|
|
timestamps = [msg.additional_kwargs.get("deerflow_created_at") for msg in state["messages"]]
|
|
assert all(isinstance(ts, str) and ts.endswith("+08:00") for ts in timestamps)
|
|
|
|
|
|
def test_after_model_keeps_existing_timestamp():
|
|
middleware = MessageTimestampMiddleware()
|
|
human = HumanMessage(content="hello")
|
|
human.additional_kwargs["deerflow_created_at"] = "2026-04-22T01:00:00.000Z"
|
|
state = {"messages": [human, AIMessage(content="hi")]}
|
|
|
|
middleware.after_model(state, runtime=None) # type: ignore[arg-type]
|
|
|
|
assert state["messages"][0].additional_kwargs["deerflow_created_at"] == "2026-04-22T01:00:00.000Z"
|