Ashfaqbs commented on issue #1065:
URL: https://github.com/apache/flink-agents/issues/1065#issuecomment-5535112444
## Proposed design
Traced where the enforcement hook naturally fits with the current code:
- Chat calls already extract prompt/completion tokens per call in
`ChatModelAction.recordChatTokenMetrics()`, feeding
`BaseChatModelSetup.recordTokenMetrics()` — but that's Flink-metric counters
for dashboards only, not state that can gate the next call (and metrics aren't
reliably queryable/restorable for decision logic across checkpoints).
- Every chat call already funnels through
`ChatModelInvoker.chatWithRetries()`, on the durable/keyed execution path
(`ctx.durableExecute`). That's the natural enforcement point.
Sketch:
1. **New config options** in `AgentExecutionOptions`, matching the existing
`ConfigOption<>` pattern (`MAX_RETRIES`, `CHAT_ASYNC`, etc.):
- `chat.budget.max-tokens` (or cost) — ceiling per window, disabled by
default so it's zero-impact for existing users
- `chat.budget.window-sec` — sliding/fixed window size
- `chat.budget.on-exceeded` — `FAIL` (raise, letting the existing
`ErrorHandlingStrategy` handle it the same way any other attempt failure
already does) or `WARN` (emit an observability event only)
2. **Keyed durable state**, not Flink metrics, for the running counter —
each Flink Agents key already gets its own Agent instance on the
durable/checkpointed execution path, so a `ValueState`-backed total (tokens
spent + window start) checkpoints/restores correctly, where a metric counter
would not.
3. **Check-then-call** at the top of `chatWithRetries`, before
`durableExecute`: reset if the window has expired; if spend would exceed
budget, either throw (folds into the existing `ChatAttemptFailed` →
`ErrorHandlingStrategy` path, reusing machinery instead of adding a new failure
mode) or emit a warning and proceed.
4. Increment the counter alongside the existing `recordChatTokenMetrics`
call, right after a successful attempt.
This is meant to be a thin policy layer on infrastructure that already
exists — no new call path, reuses the existing retry/error-handling machinery,
opt-in by default.
Before I put together a PR: this touches checkpointed-state semantics
(window reset behavior across job restarts, what "per agent" scopes to under
keyed parallelism/rescaling). Wanted to get a read on that from a maintainer
before writing code, rather than build it cold and risk a rewrite. Does the
sketch above look right, or would you scope the state/window semantics
differently?
--
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]
