Skip to content

feat: add lifecycle hooks for agent execution flow#1959

Draft
wingding12 wants to merge 1 commit intohuggingface:mainfrom
wingding12:feat/lifecycle-hooks-1883
Draft

feat: add lifecycle hooks for agent execution flow#1959
wingding12 wants to merge 1 commit intohuggingface:mainfrom
wingding12:feat/lifecycle-hooks-1883

Conversation

@wingding12
Copy link

Summary

Fixes #1883 - Add lifecycle hooks for CodeAgent execution flow

Adds lifecycle callbacks that fire at specific points during agent execution, enabling finer-grained control than step_callbacks (which only fire after steps complete). This addresses the use case of needing to execute callbacks at different points in the execution flow, such as sending LLM output to a TTS service before code execution.

New Lifecycle Events

Event When it fires Use cases
MODEL_OUTPUT After LLM generates output, before parsing TTS queuing, early logging, output validation
PRE_EXECUTION Before code/tool execution, after parsing Human-in-the-loop approval, custom code validation, security checks

Usage Example

from smolagents import CodeAgent, LifecycleEvent, ModelOutputEvent, PreExecutionEvent

def on_model_output(event: ModelOutputEvent, agent):
    # Send to TTS service before code execution starts
    tts_service.queue(event.model_output)

def on_pre_execution(event: PreExecutionEvent, agent):
    # Validate code before execution
    if "dangerous_function" in (event.code_action or ""):
        raise ValueError("Dangerous code blocked!")

agent = CodeAgent(
    tools=[...],
    model=model,
    lifecycle_callbacks={
        LifecycleEvent.MODEL_OUTPUT: on_model_output,
        LifecycleEvent.PRE_EXECUTION: on_pre_execution,
    }
)

Files Changed

File Changes
src/smolagents/memory.py Add LifecycleEvent enum, ModelOutputEvent/PreExecutionEvent dataclasses, LifecycleCallbackRegistry class
src/smolagents/agents.py Add lifecycle_callbacks parameter to MultiStepAgent, trigger events in CodeAgent._step_stream and ToolCallingAgent._step_stream
tests/test_agents.py Add 6 comprehensive tests for lifecycle hooks

Test plan

  • Test setup of lifecycle callbacks registry
  • Test invalid lifecycle event type raises error
  • Test MODEL_OUTPUT callback in CodeAgent
  • Test PRE_EXECUTION callback in CodeAgent (verifies order: model → callback → execution)
  • Test lifecycle callbacks in ToolCallingAgent
  • Test single-argument callback signature (backwards compatible)
  • All existing tests still pass

Adds lifecycle callbacks that fire at specific points during agent
execution, enabling finer-grained control than step_callbacks (which
only fire after steps complete).

New lifecycle events:
- MODEL_OUTPUT: Triggered after LLM generates output, before parsing.
  Use case: Send to TTS service before code execution starts.
- PRE_EXECUTION: Triggered before code/tool execution, after parsing.
  Use case: Human-in-the-loop approval, custom validation.

Usage:
```python
def on_model_output(event: ModelOutputEvent, agent):
    tts_service.queue(event.model_output)

def on_pre_execution(event: PreExecutionEvent, agent):
    if "dangerous" in (event.code_action or ""):
        raise ValueError("Blocked!")

agent = CodeAgent(
    tools=[...],
    model=model,
    lifecycle_callbacks={
        LifecycleEvent.MODEL_OUTPUT: on_model_output,
        LifecycleEvent.PRE_EXECUTION: on_pre_execution,
    }
)
```

Changes:
- memory.py: Add LifecycleEvent enum, ModelOutputEvent/PreExecutionEvent
  dataclasses, and LifecycleCallbackRegistry class
- agents.py: Add lifecycle_callbacks parameter to MultiStepAgent,
  trigger events in CodeAgent._step_stream and ToolCallingAgent._step_stream
- test_agents.py: Add comprehensive tests for lifecycle hooks

Fixes huggingface#1883
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

ENH: Add lifecycle hooks for CodeAgent execution flow

1 participant