feat: add stop_run State key and TokenBudgetHook - #12411
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub. 1 Skipped Deployment
|
github-actions
Bot
added
topic:tests
type:documentation
Coverage report
Click to see where and how coverage changed
This report was generated by python-coverage-comment-action |
||||||||||||||||||||||||||||||||||||||||||||||||||||||
anakin87
changed the title
stop_run State key and TokenBudgetHook
| assert agent.chat_generator.run.call_count == 1 | ||
|
|
||
|
|
||
| class TestStopRun: |
There was a problem hiding this comment.
to exactly understand how the Agent behaves with this key (in relation to hooks and exit reasons), I recommend taking a quick look at the following tests
| if total_tokens >= self.max_total_tokens: | ||
| state.set("stop_run", "token_budget_exceeded") |
There was a problem hiding this comment.
I also wonder if optionally setting a final message in state.data["messages"] would be useful. E.g. the final message could be "Agent stopped because token budget was exceeded".
This could be an optional init param of this hook. WDYT?
There was a problem hiding this comment.
I don't have strong opinions about this point.
Do you think that this would be helpful? Who would be the final user of this?
There was a problem hiding this comment.
I'm imagining the case if someone is hosting an Agent in Hayhooks and that whatever the last message is in State is what the end user will see. Right now it would be a tool call result which could be confusing.
I think it can be solved here in the Hook or we would need to add a check for that in the Hayhooks app to provide an actionable message to the user
There was a problem hiding this comment.
OK, I see. What about a final_message: ChatMessage | bool = False
False: do nothingTrue: default final messageChatMessage: provided message
Just an idea, I can improve the design (if you agree on the direction).
There was a problem hiding this comment.
Looks good! And I agree with the off by default. I also realized that users anyways may prefer to create an after_run hook that manages what the final message should be based on reading the exit reason
There was a problem hiding this comment.
Yes, agree on after_run hook.
| total_tokens = _first_numeric(usage, ("total_tokens",)) | ||
| if not total_tokens: | ||
| total_tokens = _first_numeric(usage, _INPUT_TOKEN_KEYS) + _first_numeric(usage, _OUTPUT_TOKEN_KEYS) | ||
| if total_tokens >= self.max_total_tokens: |
There was a problem hiding this comment.
I think adding a logger warning message here could be helpful so it appears in logs if the Agent is running over the token budget.
2 participants
Related Issues
Proposed Changes:
stop_runStatekeyAgent._run_stepbefore the LLM callexit_reasonexit_reasonis the original oneTokenBudgetHookto test the abstraction aboveHow did you test it?
CI, new tests
Checklist
fix:,feat:,build:,chore:,ci:,docs:,style:,refactor:,perf:,test:and added!in case the PR includes breaking changes.