weiqingy commented on code in PR #924:
URL: https://github.com/apache/flink-agents/pull/924#discussion_r3654253847
##########
plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRetryTest.java:
##########
@@ -127,6 +128,129 @@ void chatSucceedsWithoutRetry_retryCountIsZero() throws
Exception {
verify(mockActionMetricGroup, never()).getSubGroup(anyString(),
anyString());
}
+ @Test
+ void chatReportsLlmExecution() throws Exception {
+ RunnerContext reportingCtx = reportingRunnerContext();
+ BaseChatModelSetup chatModel =
configureReportingChatContext(reportingCtx);
+ when(chatModel.chat(any(), any(), any()))
+ .thenReturn(new ChatMessage(MessageRole.ASSISTANT, "hello"));
+
+ ChatModelAction.chat(
+ UUID.randomUUID(),
+ "test-model",
+ List.of(new ChatMessage(MessageRole.USER, "hi")),
+ Map.of(),
+ null,
+ reportingCtx);
+
+ ExecutionReporter reporter = (ExecutionReporter) reportingCtx;
+ verify(reporter)
+ .reportExecutionStarted(ExecutionReporter.EntityTypes.LLM,
"test-model", Map.of());
+ verify(reporter)
+ .reportExecutionSucceeded(
+ ExecutionReporter.EntityTypes.LLM, "test-model",
Map.of());
+ }
+
+ @Test
+ void chatReportsStructuredOutputParserExecution() throws Exception {
+ RunnerContext reportingCtx = reportingRunnerContext();
+ BaseChatModelSetup chatModel =
configureReportingChatContext(reportingCtx);
+ when(chatModel.chat(any(), any(), any()))
+ .thenReturn(new ChatMessage(MessageRole.ASSISTANT,
"{\"answer\":\"42\"}"));
+
+ ChatModelAction.chat(
+ UUID.randomUUID(),
+ "test-model",
+ List.of(new ChatMessage(MessageRole.USER, "hi")),
+ Map.of(),
+ Map.class,
+ reportingCtx);
+
+ ExecutionReporter reporter = (ExecutionReporter) reportingCtx;
+ verify(reporter)
+ .reportExecutionStarted(
+ ExecutionReporter.EntityTypes.PARSER,
Agent.STRUCTURED_OUTPUT, Map.of());
+ verify(reporter)
+ .reportExecutionSucceeded(
+ ExecutionReporter.EntityTypes.PARSER,
Agent.STRUCTURED_OUTPUT, Map.of());
+ }
+
+ @Test
+ void chatRetriesStructuredOutputParseErrorWithoutFailingLlm() throws
Exception {
+ RunnerContext reportingCtx = reportingRunnerContext();
+ BaseChatModelSetup chatModel =
configureReportingChatContext(reportingCtx);
+ when(reportingCtx.getConfig())
+ .thenReturn(readableConfig(Agent.ErrorHandlingStrategy.RETRY,
1, 0));
+ when(chatModel.chat(any(), any(), any()))
+ .thenReturn(
+ new ChatMessage(MessageRole.ASSISTANT, "not-json"),
+ new ChatMessage(MessageRole.ASSISTANT,
"{\"answer\":\"42\"}"));
+
+ ChatModelAction.chat(
+ UUID.randomUUID(),
+ "test-model",
+ List.of(new ChatMessage(MessageRole.USER, "hi")),
+ Map.of(),
+ Map.class,
+ reportingCtx);
+
+ verify(chatModel, times(2)).chat(any(), any(), any());
+
+ ExecutionReporter reporter = (ExecutionReporter) reportingCtx;
+ verify(reporter, times(2))
+ .reportExecutionStarted(ExecutionReporter.EntityTypes.LLM,
"test-model", Map.of());
+ verify(reporter, times(2))
+ .reportExecutionSucceeded(
+ ExecutionReporter.EntityTypes.LLM, "test-model",
Map.of());
+ verify(reporter)
+ .reportExecutionFailed(
+ eq(ExecutionReporter.EntityTypes.PARSER),
+ eq(Agent.STRUCTURED_OUTPUT),
+ eq(Map.of()),
+ any(Exception.class),
+
eq(ExecutionReporter.ProblemCategories.MODEL_OUTPUT_PARSE_ERROR));
+ verify(reporter, never())
+ .reportExecutionFailed(
+ eq(ExecutionReporter.EntityTypes.LLM),
+ eq("test-model"),
+ eq(Map.of()),
+ any(Throwable.class),
+ any());
+ }
+
+ @Test
+ void chatReportsEachRetriedModelInvocation() throws Exception {
Review Comment:
Anchoring here since this is where the contract gets pinned:
`chatReportsEachRetriedModelInvocation` asserts `times(2)` on
`reportExecutionStarted`. `started` sits inside the retry loop
(`ChatModelAction.java:371-373`), and Python matches
(`chat_model_action.py:352-375`), so both languages emit one execution per
attempt.
On the design discussion the contract landed on one execution per framework
model invocation, with retries collapsing into it and Parser recorded
separately. Parser is implemented that way (`:457-474`); the retry half went
the other way.
Downstream, a `retryCount` of 3 on a successful request reads as 4 LLM
executions at a 25% success rate, and nothing in the record separates "attempt
2" from "a second call". `ChatResponseEvent` already carries `retryCount`, so
the two views disagree.
Which way did you land in the end? Two shapes if either helps: hoist
`started`/`succeeded` outside the loop and put attempts on the single terminal,
or keep per-attempt and add an ordinal plus a shared correlation id. Either
way, would it be worth pinning in `ExecutionReporter`'s javadoc, now that it's
permanent public API?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
