This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 7d2f52bca [python] Key built-in chat model calls by an explicit
durable id (#1240)
7d2f52bca is described below
commit 7d2f52bcade2178c926a423d10d6a5c7721f7cd1
Author: Sangkyoon Nam <[email protected]>
AuthorDate: Tue Oct 13 16:36:43 2026 +0900
[python] Key built-in chat model calls by an explicit durable id (#1240)
Generated-by: Claude Code 2.1.296 (Claude Opus 5.5)
---
.../flink_agents/plan/actions/chat_model_action.py | 12 ++++-
.../tests/actions/test_chat_model_action_retry.py | 58 +++++++++++++++++++++-
2 files changed, 68 insertions(+), 2 deletions(-)
diff --git a/python/flink_agents/plan/actions/chat_model_action.py
b/python/flink_agents/plan/actions/chat_model_action.py
index 08ea6e3ec..6f4dd713a 100644
--- a/python/flink_agents/plan/actions/chat_model_action.py
+++ b/python/flink_agents/plan/actions/chat_model_action.py
@@ -59,6 +59,11 @@ if TYPE_CHECKING:
_TOOL_CALL_CONTEXT = "_TOOL_CALL_CONTEXT"
_TOOL_REQUEST_EVENT_CONTEXT = "_TOOL_REQUEST_EVENT_CONTEXT"
_RETRY_STATS_CONTEXT = "_RETRY_STATS_CONTEXT"
+# The identity of each chat model attempt, the same as the Java chat action's
call id
+# for a plain chat model. Without it the identity is a digest of the pickled
messages,
+# which can change between Python processes (#1190); attempts are told apart
by their
+# position in the action's durable call sequence.
+_CHAT_DURABLE_ID = "chat"
_PROMPT_ARGS = "prompt_args"
_FINISH_REASON = "finish_reason"
_TRUNCATED_FINISH_REASON = "length"
@@ -614,7 +619,12 @@ async def chat(
# Persistence/recovery failures must escape the request-failure
boundary,
# so durable calls stay outside the try blocks below.
response, failure = await _durable_invoke(
- ctx, invoke, messages, chat_async=chat_async,
prompt_args=prompt_args
+ ctx,
+ invoke,
+ messages,
+ chat_async=chat_async,
+ prompt_args=prompt_args,
+ durable_id=_CHAT_DURABLE_ID,
)
try:
response = _accept_model_response(
diff --git
a/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py
b/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py
index 44d3d42ad..a2ec36280 100644
--- a/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py
+++ b/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py
@@ -19,7 +19,7 @@
import asyncio
import time
-from typing import Any, Sequence
+from typing import Any, Callable, Sequence
from unittest.mock import MagicMock, call
from uuid import uuid4
@@ -181,6 +181,7 @@ def _create_mock_runner_context(
ctx.action_metric_group = metric_group
ctx.send_event = MagicMock(side_effect=lambda e: sent_events.append(e))
ctx.get_resource = MagicMock(return_value=chat_model)
+ # Like the runtime, the reserved durable_id keyword is not forwarded to fn.
ctx.durable_execute = MagicMock(
side_effect=lambda fn, *args, durable_id=None, **kwargs: fn(*args,
**kwargs)
)
@@ -321,6 +322,61 @@ class TestChatModelActionRetry:
)
ctx.report_execution_failed.assert_not_called()
+ @pytest.mark.parametrize("chat_async", [False, True])
+ def test_every_attempt_uses_the_chat_durable_id(self, chat_async: bool) ->
None:
+ """Each attempt is keyed by "chat", not by a digest of the messages
(#1190)."""
+ chat_model = MagicMock()
+ chat_model.chat = MagicMock(
+ side_effect=[
+ RuntimeError("transient"),
+ _response_fixture(MessageRole.ASSISTANT, "hello"),
+ ]
+ )
+ ctx, sent_events, _, _ = _create_mock_runner_context(
+ chat_model, max_retries=1, retry_wait_interval_sec=0
+ )
+ if chat_async:
+ async_option = AgentExecutionOptions.CHAT_ASYNC
+ original_get = ctx.config.get.side_effect
+ ctx.config.get = MagicMock(
+ side_effect=lambda option: True
+ if option is async_option
+ else original_get(option)
+ )
+
+ async def durable_execute_async(
+ fn: Callable[..., Any],
+ *args: Any,
+ durable_id: str | None = None,
+ **kwargs: Any,
+ ) -> Any:
+ return ctx.durable_execute(fn, *args, durable_id=durable_id,
**kwargs)
+
+ ctx.durable_execute_async =
MagicMock(side_effect=durable_execute_async)
+
+ asyncio.run(
+ chat(
+ uuid4(),
+ chat_model.connection,
+ [ChatMessage.of(MessageRole.USER, "hi")],
+ {},
+ None,
+ ctx,
+ )
+ )
+
+ assert len(sent_events) == 1
+ event = sent_events[0]
+ assert isinstance(event, ChatResponseEvent)
+ assert event.is_success
+ assert event.response.text == "hello"
+ assert chat_model.chat.call_count == 2
+ durable_call = ctx.durable_execute_async if chat_async else
ctx.durable_execute
+ assert [call.kwargs["durable_id"] for call in
durable_call.call_args_list] == [
+ "chat",
+ "chat",
+ ]
+
def test_chat_retries_with_exponential_backoff(self) -> None:
"""Fail once then succeed: 1s interval, 1 retry -> wait 1s (1 *
2^0)."""
call_count = 0