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 a6848aff6 [python][plan] Split chat model action into per-model
attempts (#1357)
a6848aff6 is described below
commit a6848aff617434a2ec02322c27416c96b91090fa
Author: Sangkyoon Nam <[email protected]>
AuthorDate: Fri Nov 6 15:23:09 2026 +0900
[python][plan] Split chat model action into per-model attempts (#1357)
Generated-by: Codex CLI 0.166.0 (gpt-6-astra)
---
python/flink_agents/api/chat_models/chat_model.py | 49 +++++++------
.../flink_agents/plan/actions/chat_model_action.py | 82 ++++++++++++++++------
.../tests/actions/test_chat_model_action_retry.py | 42 +++++++++++
3 files changed, 131 insertions(+), 42 deletions(-)
diff --git a/python/flink_agents/api/chat_models/chat_model.py
b/python/flink_agents/api/chat_models/chat_model.py
index 1ee890839..d2560f159 100644
--- a/python/flink_agents/api/chat_models/chat_model.py
+++ b/python/flink_agents/api/chat_models/chat_model.py
@@ -588,6 +588,34 @@ class BaseChatModelSetup(Resource):
)
self.tools = callables
+ def prepare_request_messages(
+ self,
+ messages: Sequence[ChatMessage],
+ prompt_args: Mapping[str, Any] | None = None,
+ ) -> Sequence[ChatMessage]:
+ """Render the setup prompt and skill discovery for a chat request."""
+ # Apply prompt template
+ if self.prompt is not None:
+ str_prompt_args: Dict[str, str] = (
+ {k: str(v) for k, v in prompt_args.items()} if prompt_args
else {}
+ )
+ prompt_messages =
self._get_prompt().format_messages(**str_prompt_args)
+
+ # append meaningful messages
+ # any block counts, so image-only messages survive
+ for msg in messages:
+ if len(msg.blocks) > 0 or msg.role == MessageRole.ASSISTANT:
+ prompt_messages.append(msg)
+ messages = prompt_messages
+
+ if self.skill_discovery_prompt:
+ # Right after the first system message, or at the head when there
+ # is none.
+ index = find_first_system_message(messages) + 1
+ injected = [ChatMessage.system(self.skill_discovery_prompt)]
+ messages = list(messages[:index]) + injected +
list(messages[index:])
+ return messages
+
def chat(
self,
messages: Sequence[ChatMessage],
@@ -638,26 +666,7 @@ class BaseChatModelSetup(Resource):
raise ValueError(msg)
if tools is None:
- # Apply prompt template
- if self.prompt is not None:
- str_prompt_args: Dict[str, str] = (
- {k: str(v) for k, v in prompt_args.items()} if prompt_args
else {}
- )
- prompt_messages =
self._get_prompt().format_messages(**str_prompt_args)
-
- # append meaningful messages
- # any block counts, so image-only messages survive
- for msg in messages:
- if len(msg.blocks) > 0 or msg.role ==
MessageRole.ASSISTANT:
- prompt_messages.append(msg)
- messages = prompt_messages
-
- if self.skill_discovery_prompt:
- # Right after the first system message, or at the head when
there
- # is none.
- index = find_first_system_message(messages) + 1
- injected = [ChatMessage.system(self.skill_discovery_prompt)]
- messages = list(messages[:index]) + injected +
list(messages[index:])
+ messages = self.prepare_request_messages(messages, prompt_args)
# Call chat model connection to execute chat
merged_kwargs = self.model_kwargs.copy()
diff --git a/python/flink_agents/plan/actions/chat_model_action.py
b/python/flink_agents/plan/actions/chat_model_action.py
index 596f4efae..1496e6703 100644
--- a/python/flink_agents/plan/actions/chat_model_action.py
+++ b/python/flink_agents/plan/actions/chat_model_action.py
@@ -20,6 +20,7 @@ import logging
import re
import time
from concurrent.futures import CancelledError
+from dataclasses import dataclass
from functools import wraps
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Sequence, cast
from uuid import UUID
@@ -76,9 +77,6 @@ _FINALIZE_DIRECTIVE = (
" Preserve its meaning and do not add or infer any new information.\n\n"
"Response:\n"
)
-# The loop call's own record already pins the answer this call converts, so a
fixed
-# identity is enough and stays stable whatever form the arguments take on
recovery.
-_FINALIZE_DURABLE_ID = "chat:final"
_logger = logging.getLogger(__name__)
@@ -557,25 +555,34 @@ def _accept_model_response(
return response
-async def chat(
+@dataclass
+class _ChatAttemptResult:
+ response: ChatResult
+ chat_model: "BaseChatModelSetup"
+ structured_output: Any
+
+
+class _ChatAttemptFailed(Exception):
+ """A candidate exhausted its retry budget or could not be prepared."""
+
+ def __init__(self, error: Exception) -> None:
+ super().__init__(str(error))
+ self.error = error
+
+
+async def _chat_with_retries(
initial_request_id: UUID,
model: str,
messages: List[ChatMessage],
prompt_args: Dict | None,
output_schema: OutputSchema | None,
ctx: RunnerContext,
-) -> None:
- """Chat with llm.
-
- If there is no tool call generated, we return the chat response event
directly,
- otherwise, we generate tool request event according to the tool calls in
chat model
- response, and save the request and response messages in tool call context.
+ durable_id: str = _CHAT_DURABLE_ID,
+) -> _ChatAttemptResult:
+ """Attempt one model without sending a response or request failure event.
- When ``output_schema`` is set and the chat model reports that it travels
through
- the provider's native structured output, the final answer is followed by
one more
- call without tools that carries the schema, and that call's response is
parsed.
- The extra call belongs to the same attempt, so its failures, and a failure
to
- parse its response, consume the same retry budget.
+ Provider and preparation failures raise ``_ChatAttemptFailed``.
Cancellation
+ and durable storage failures propagate unchanged to the runtime.
"""
try:
chat_model = cast(
@@ -584,8 +591,7 @@ async def chat(
except (CancelledError, InterruptedError):
raise
except Exception as error:
- _send_failure(initial_request_id, error, ctx)
- return
+ raise _ChatAttemptFailed(error) from error
# Evaluated once: the answer depends only on the chat model and the
schema, and
# a raise (a NATIVE strategy the connection cannot honor) cannot change on
retry.
@@ -600,8 +606,7 @@ async def chat(
except (CancelledError, InterruptedError):
raise
except Exception as error:
- _send_failure(initial_request_id, error, ctx)
- return
+ raise _ChatAttemptFailed(error) from error
request_metric_group = ctx.action_metric_group
@@ -669,7 +674,7 @@ async def chat(
messages,
chat_async=chat_async,
prompt_args=prompt_args,
- durable_id=_CHAT_DURABLE_ID,
+ durable_id=durable_id,
)
try:
response = _accept_model_response(
@@ -705,7 +710,7 @@ async def chat(
response.text,
output_schema,
chat_async=chat_async,
- durable_id=_FINALIZE_DURABLE_ID,
+ durable_id=f"{durable_id}:final",
)
try:
structured = _accept_model_response(
@@ -741,8 +746,41 @@ async def chat(
)
if final_error is not None:
- _send_failure(initial_request_id, final_error, ctx)
+ raise _ChatAttemptFailed(final_error) from final_error
+
+ return _ChatAttemptResult(response, chat_model, structured_output)
+
+
+async def chat(
+ initial_request_id: UUID,
+ model: str,
+ messages: List[ChatMessage],
+ prompt_args: Dict | None,
+ output_schema: OutputSchema | None,
+ ctx: RunnerContext,
+) -> None:
+ """Chat with llm.
+
+ If there is no tool call generated, we return the chat response event
directly,
+ otherwise, we generate tool request event according to the tool calls in
chat model
+ response, and save the request and response messages in tool call context.
+
+ When ``output_schema`` is set and the chat model reports that it travels
through
+ the provider's native structured output, the final answer is followed by
one more
+ call without tools that carries the schema, and that call's response is
parsed.
+ The extra call belongs to the same attempt, so its failures, and a failure
to
+ parse its response, consume the same retry budget.
+ """
+ try:
+ attempt = await _chat_with_retries(
+ initial_request_id, model, messages, prompt_args, output_schema,
ctx
+ )
+ except _ChatAttemptFailed as failure:
+ _send_failure(initial_request_id, failure.error, ctx)
return
+ response = attempt.response
+ chat_model = attempt.chat_model
+ structured_output = attempt.structured_output
if (
len(response.tool_calls) > 0
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 8169f9a2d..0ca3b9023 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
@@ -48,6 +48,8 @@ from flink_agents.api.trace import (
LLMExecutionMetadataKeys,
)
from flink_agents.plan.actions.chat_model_action import (
+ _chat_with_retries,
+ _ChatAttemptFailed,
chat,
process_chat_request_or_tool_response,
)
@@ -1373,3 +1375,43 @@ class TestNativeStructuredOutputFinalization:
assert ctx.durable_execute_async.call_count == 2
ctx.durable_execute.assert_not_called()
assert sent_events[0].structured_output.result == 42
+
+
[email protected]("chat_async", [False, True])
+def
test_candidate_attempt_exhausts_retries_without_emitting_failure(chat_async):
+ model = MagicMock()
+ model.chat.side_effect = RuntimeError("candidate unavailable")
+ ctx, events, metrics, memory = _create_mock_runner_context(
+ model, max_retries=2, retry_wait_interval_sec=0, chat_async=chat_async
+ )
+ request_id = uuid4()
+ with pytest.raises(_ChatAttemptFailed) as caught:
+ asyncio.run(_chat_with_retries(
+ request_id, "candidate", [ChatMessage.user("hello")], None, None,
+ ctx, durable_id="chat:router:candidate",
+ ))
+ assert str(caught.value.error) == "RuntimeError: candidate unavailable"
+ assert caught.value.__cause__ is caught.value.error
+ assert events == []
+ assert model.chat.call_count == 3
+ durable = ctx.durable_execute_async if chat_async else ctx.durable_execute
+ assert [c.kwargs["durable_id"] for c in durable.call_args_list] == [
+ "chat:router:candidate"
+ ] * 3
+ assert
memory.get("_RETRY_STATS_CONTEXT")[str(request_id)]["total_retry_count"] == 2
+ assert metrics.get_sub_group("model_resource",
"candidate").get_counter("retryCount").get_count() == 2
+
+
+def test_candidate_native_finalization_uses_attempt_identity():
+ model = _native_chat_model()
+ ctx, events, _, _ = _native_context(model)
+ result = asyncio.run(_chat_with_retries(
+ uuid4(), "candidate", [ChatMessage.user("hello")], None,
+ OutputSchema(output_schema=_StructuredResult), ctx,
+ durable_id="chat:router:candidate",
+ ))
+ assert result.structured_output is not None
+ assert events == []
+ assert [c.kwargs["durable_id"] for c in
ctx.durable_execute.call_args_list] == [
+ "chat:router:candidate", "chat:router:candidate:final"
+ ]
