The agent class

This page documents the Agent class, which runs the main loop of the agent. To learn about the configuration objects used to specify the behavior of an agent, see the agent configuration reference page.

sweagent.agent.agents.Agent

Agent(*, templates: TemplateConfig, tools: ToolHandler, history_processors: list[HistoryProcessor], model: AbstractModel, max_requeries: int = 3, name: str = 'main', _catch_errors: bool = True, _always_require_zero_exit_code: bool = False)

The agent handles the behaviour of the model and how it interacts with the environment.

To run the agent, either call self.run or self.setup and then self.step in a loop.

Source code in sweagent/agent/agents.py

def __init__(
    self,
    *,
    templates: TemplateConfig,
    tools: ToolHandler,
    history_processors: list[HistoryProcessor],
    model: AbstractModel,
    max_requeries: int = 3,
    name: str = "main",
    _catch_errors: bool = True,
    _always_require_zero_exit_code: bool = False,
):
    """The agent handles the behaviour of the model and how it interacts with the environment.

    To run the agent, either call `self.run` or `self.setup` and then `self.step` in a loop.
    """
    self._catch_errors = _catch_errors
    self._always_require_zero_exit_code = _always_require_zero_exit_code
    self.name = name
    self.model = model
    self.templates = templates
    self.tools = tools
    self.history_processors = history_processors
    self.max_requeries = max_requeries
    self.logger = get_logger("swea-agent", emoji="🤠")

    # Set in run method
    self._env: SWEEnv | None = None
    self._problem_statement: ProblemStatement | ProblemStatementConfig | None = None
    self.traj_path: Path | None = None

    #: Number of attempts to solve the issue when using a review loop
    self._i_attempt: int = 0

    #: The following three attributes collect the information about how the agent
    #: solved the problem.
    self._history_by_attempt: dict[int, list] = defaultdict(list)
    self._trajectory_by_attempt: dict[int, Trajectory] = defaultdict(list)
    self._info_by_attempt: dict[int, AgentInfo] = defaultdict(dict)  # type: ignore

    #: Variables to be referenced in the templates that are forwarded from one
    #: solution attempt to the next
    self._forwarded_vars: dict[str, Any] = {}

    self._chook = CombinedAgentHook()

    self._replay_config: BaseModel | None = None
    """This can be set to a RunSingleConfig from the Run instance whenever possible.
    It can be used to replay the agent's trajectory in an environment.
    """

history property writable

history: History

History that is passed on to the model. Use _append_history to modify.

history_processors instance-attribute

history_processors = history_processors

info property writable

info: AgentInfo

Information about the agent's run

logger instance-attribute

logger = get_logger('swea-agent', emoji='🤠')

max_requeries instance-attribute

max_requeries = max_requeries

messages property

messages: list[dict[str, str]]

Return the history of the agent since the last reset, processed through all history processors.

model instance-attribute

model = model

name instance-attribute

name = name

replay_config property writable

replay_config: BaseModel | None

templates instance-attribute

templates = templates

tools instance-attribute

tools = tools

traj_path instance-attribute

traj_path: Path | None = None

trajectory property writable

trajectory: Trajectory

Trajectory of the agent for the current instance. In contrast to history, this is mostly for the informational value of how the agent interacted with the environment and is also what is being used when replaying the trajectory

add_demonstrations_to_history

add_demonstrations_to_history() -> None

Add demonstrations to history

Source code in sweagent/agent/agents.py

def add_demonstrations_to_history(self) -> None:
    """Add demonstrations to history"""
    for demonstration_path in self.templates.demonstrations:
        self._add_demonstration_to_history(demonstration_path)

add_hook

add_hook(hook: AbstractAgentHook) -> None

Add hook to agent

Source code in sweagent/agent/agents.py

def add_hook(self, hook: AbstractAgentHook) -> None:
    """Add hook to agent"""
    hook.on_init(agent=self)
    self._chook.add_hook(hook)

add_instance_template_to_history

add_instance_template_to_history(state: dict[str, str]) -> None

Add observation to history, as well as the instance template or demonstrations if we're at the start of a new attempt.

Source code in sweagent/agent/agents.py

def add_instance_template_to_history(self, state: dict[str, str]) -> None:
    """Add observation to history, as well as the instance template or demonstrations if we're
    at the start of a new attempt.
    """
    templates: list[str] = []
    # Determine observation template based on what prior observation was
    assert self.history[-1]["role"] == "system" or self.history[-1].get("is_demo", False)
    # Show instance template if prev. obs. was initial system message
    templates = [self.templates.instance_template]
    if self.templates.strategy_template is not None:
        templates.append(self.templates.strategy_template)

    self._add_templated_messages_to_history(templates, **state)

add_step_to_history

add_step_to_history(step: StepOutput) -> None

Adds a step (command that was run and output) to the model history

Source code in sweagent/agent/agents.py

def add_step_to_history(self, step: StepOutput) -> None:
    """Adds a step (command that was run and output) to the model history"""
    self._append_history(
        {
            "role": "assistant",
            "content": step.output,
            "thought": step.thought,
            "action": step.action,
            "agent": self.name,
            "tool_calls": step.tool_calls,
            "message_type": "action",
        },
    )

    if step.observation is None or step.observation.strip() == "":
        # Show no output template if observation content was empty
        templates = [self.templates.next_step_no_output_template]
    else:
        # Show standard output template if there is observation content
        templates = [self.templates.next_step_template]
    self._add_templated_messages_to_history(
        templates,
        observation=step.observation,
        tool_call_ids=step.tool_call_ids,
        **step.state,
    )

add_step_to_trajectory

add_step_to_trajectory(step: StepOutput) -> None

Source code in sweagent/agent/agents.py

def add_step_to_trajectory(self, step: StepOutput) -> None:
    trajectory_step = TrajectoryStep(
        {
            "action": step.action,
            "observation": step.observation,
            "response": step.output,
            "thought": step.thought,
            "execution_time": step.execution_time,
            "state": step.state,
            "messages": self.messages,
        },
    )
    self.trajectory.append(trajectory_step)

add_system_message_to_history

add_system_message_to_history() -> None

Add system message to history

Source code in sweagent/agent/agents.py

def add_system_message_to_history(self) -> None:
    """Add system message to history"""
    assert self._problem_statement is not None
    system_msg = Template(self.templates.system_template).render(**self._get_format_dict())
    self.logger.info(f"SYSTEM ({self.name})\n{system_msg}")
    self._append_history(
        {"role": "system", "content": system_msg, "agent": self.name, "message_type": "system_prompt"}
    )

attempt_autosubmission_after_error

attempt_autosubmission_after_error(step: StepOutput) -> StepOutput

For most exceptions, we attempt to still extract the patch and submit that. This means we send the submit command to the runtime and parse the output.

Source code in sweagent/agent/agents.py

def attempt_autosubmission_after_error(self, step: StepOutput) -> StepOutput:
    """For most exceptions, we attempt to still extract the patch and submit that.
    This means we send the `submit` command to the runtime and parse the output.
    """
    step = step.model_copy()
    step.done = True
    assert self._env is not None
    try:
        self._env.interrupt_session()
    except Exception as e:
        self.logger.debug("Failed to interrupt session before autosubmit: %s. Ignoring.", e)
    try:
        observation = self._env.communicate(input=self.tools.config.submit_command)
    except Exception as e:
        self.logger.debug("Failed to submit after error, got %s", e)
        return step
    step = self.handle_submission(step, observation=observation)
    if step.submission:
        self.logger.info("Exiting with autosubmission")
        step.observation = "Exited (autosubmitted)"
    return step

forward

forward(history: list[dict[str, str]]) -> StepOutput

Forward the model without handling errors.

All exceptions raised will contain the StepOutput object with some of the attributes set.

Parameters:

Name	Type	Description	Default
history	list[dict[str, str]]	history to query the model with	required

Returns:

Name	Type	Description
step_output	StepOutput	step output

Source code in sweagent/agent/agents.py

def forward(self, history: list[dict[str, str]]) -> StepOutput:
    """Forward the model without handling errors.

    All exceptions raised will contain the `StepOutput` object
    with some of the attributes set.

    Args:
        history: history to query the model with

    Returns:
        step_output: step output
    """
    # we continuously add actions, output etc. to the step object
    # because some of the specific exception handling requires some of these
    # attributes (e.g., if we want to requery the model for a bash syntax error, we
    # need to have the previous model output to format the requery template)
    step = StepOutput()
    try:
        # Forward model and get actions
        self._chook.on_model_query(messages=history, agent=self.name)
        output = self.model.query(history)  # type: ignore
        step.output = output["message"]
        if isinstance(self.model, HumanThoughtModel):
            # TODO: This might be a bit hacky
            # consider changing sweagent/tools/tools.py:ToolConfig to enforce this.
            step.thought, step.action = ThoughtActionParser()(output, self.tools.config.commands)
        elif isinstance(self.model, HumanModel):
            step.thought, step.action = "", output["message"]
        else:
            step.thought, step.action = self.tools.parse_actions(output)
        if output.get("tool_calls") is not None:
            step.tool_call_ids = [call["id"] for call in output["tool_calls"]]
            step.tool_calls = output["tool_calls"]
        self.logger.info(f"💭 THOUGHT\n{step.thought}\n\n🎬 ACTION\n{step.action.strip()}")
        self._chook.on_actions_generated(step=step)
        return self.handle_action(step)
    except Exception as e:
        # Attach the step object to the exception
        e.step = step  # type: ignore
        raise

forward_with_handling

forward_with_handling(history: list[dict[str, str]]) -> StepOutput

Forward the model and handle errors, requerying the model if we can. For example, if the model outputs a bash command that has syntax errors, we will not execute it but requery the model for a corrected command.

Note: This will update the trajectory, but not the history.

Parameters:

Name	Type	Description	Default
history	list[dict[str, str]]	history to forward	required

Returns:

Name	Type	Description
step_output	StepOutput	step output

Source code in sweagent/agent/agents.py

def forward_with_handling(self, history: list[dict[str, str]]) -> StepOutput:
    """Forward the model and handle errors, requerying the model if we can.
    For example, if the model outputs a bash command that has syntax errors,
    we will not execute it but requery the model for a corrected command.

    Note: This will update the trajectory, but not the history.

    Args:
        history: history to forward

    Returns:
        step_output: step output
    """

    def handle_error_with_autosubmission(exit_status: str, message: str) -> StepOutput:
        """Attempts to autosubmit (extract patch from the environment) and stops the loop."""
        self.logger.warning(message)
        return self.attempt_autosubmission_after_error(
            StepOutput(
                thought=message,
                exit_status=exit_status,
                output=message,
                done=True,
            )
        )

    def handle_error_with_retry(exception: Exception, template: str) -> list[dict[str, str]]:
        """Requeries the model if the error is a format/blocklist/bash syntax error."""
        self.logger.warning("Requerying model (%s)", type(exception).__name__)
        step: StepOutput = getattr(exception, "step", StepOutput())
        self.add_step_to_trajectory(step)
        exception_message = getattr(exception, "message", "")
        if not exception_message:
            try:
                exception_message = exception.args[0]
            except (IndexError, AttributeError):
                pass
        return self.get_model_requery_history(
            error_template=template,
            **step.to_template_format_dict(),
            **getattr(exception, "extra_info", {}),
            exception_message=exception_message,
        )

    n_format_fails = 0
    while n_format_fails < self.max_requeries:
        try:
            return self.forward(history)

        # Errors that are raised

        except KeyboardInterrupt:
            raise

        # Errors that cause requery

        except FormatError as e:
            n_format_fails += 1
            history = handle_error_with_retry(exception=e, template=self.tools.config.format_error_template)
        except _BlockedActionError as e:
            n_format_fails += 1
            history = handle_error_with_retry(
                exception=e, template=self.tools.config.filter.blocklist_error_template
            )
        except BashIncorrectSyntaxError as e:
            n_format_fails += 1
            history = handle_error_with_retry(
                exception=e,
                template=self.templates.shell_check_error_template,
            )
        except _RetryWithOutput as e:
            n_format_fails = 0
            history = handle_error_with_retry(
                exception=e,
                template=self.templates.next_step_template,
            )
        except _RetryWithoutOutput:
            n_format_fails = 0
            # Requery with the same template as the last step

        # Errors that cause exit

        except ContextWindowExceededError:
            return handle_error_with_autosubmission(
                "exit_context",
                "Exit due to context window",
            )
        except CostLimitExceededError:
            return handle_error_with_autosubmission(
                "exit_cost",
                "Exit due to cost limit",
            )
        except RetryError as e:
            self.logger.exception(f"Exiting due to retry error: {e}", exc_info=True)
            return handle_error_with_autosubmission(
                "exit_api",
                f"Exit due to retry error: {e}",
            )
        except SwerexException as e:
            self.logger.exception(f"Exiting due to environment error: {e}", exc_info=True)
            return handle_error_with_autosubmission(
                "exit_environment_error",
                f"Exit due to environment error: {e}",
            )
        except RuntimeError as e:
            self.logger.exception(f"Exiting due to runtime error: {e}", exc_info=True)
            return handle_error_with_autosubmission(
                "exit_error",
                f"Exit due to runtime error: {e}",
            )
        except Exception as e:
            self.logger.exception(f"Exiting due to unknown error: {e}", exc_info=True)
            return handle_error_with_autosubmission(
                "exit_error",
                f"Exit due to unknown error: {e}",
            )
    self.logger.exception(
        "Exit due to repeated format/blocklist/bash syntax errors",
        exc_info=True,
    )
    return handle_error_with_autosubmission(
        "exit_format",
        "Exit due to repeated format/blocklist/bash syntax errors",
    )

from_config classmethod

from_config(config: AgentConfig) -> Self

Source code in sweagent/agent/agents.py

@classmethod
def from_config(cls, config: AgentConfig) -> Self:
    model = get_model(config.model, config.tools)
    return cls(
        templates=config.templates,
        tools=ToolHandler(config.tools),
        history_processors=config.history_processors,
        model=model,
        max_requeries=config.max_requeries,
    )

get_model_requery_history

get_model_requery_history(error_template: str, *, output: str, **kwargs: str | int | float | bool | None) -> list[dict[str, str]]

Ask the model to correct after a hitting one of the following errors:

1. Malformatted output (could not parse action)
2. Blocked action (command is on the blocklist)
3. Bash command syntax error

At the time this function is called, the proposed action and observation are not part of the history yet.

This function adds temporary history based on the error template and queries the model. If the model is able to correct itself, the records of the mistakes will not be part of the history (but they are saved in the trajectory).

Parameters:

Name	Type	Description	Default
error_template	str	error template	required
output	str	model output	required
**kwargs	str | int | float | bool | None	keyword arguments to be passed to the error template	{}

Returns:

Type	Description
list[dict[str, str]]	model output after requery

Source code in sweagent/agent/agents.py

def get_model_requery_history(
    self, error_template: str, *, output: str, **kwargs: str | int | float | bool | None
) -> list[dict[str, str]]:
    """Ask the model to correct after a hitting one of the following errors:

    1. Malformatted output (could not parse action)
    2. Blocked action (command is on the blocklist)
    3. Bash command syntax error

    At the time this function is called, the proposed action and observation are not part of the history
    yet.

    This function adds temporary history based on the error template and queries the model.
    If the model is able to correct itself, the records of the mistakes will not be part of the history
    (but they are saved in the trajectory).

    Args:
        error_template: error template
        output: model output
        **kwargs: keyword arguments to be passed to the error template

    Returns:
        model output after requery
    """
    format_dict = {**kwargs, **self._get_format_dict()}
    error_template = Template(error_template).render(**format_dict)

    self.logger.warning(f"{error_template}")

    return self.messages + [
        {"role": "assistant", "content": output, "agent": self.name},
        {"role": "user", "content": error_template, "agent": self.name},
    ]

handle_action

handle_action(step: StepOutput) -> StepOutput

Runs an action proposed by the agent in the environment and returns the corresponding output.

Parameters:

Name	Type	Description	Default
action		command to run in bash shell	required
output		output from model (only used for error handling)	required

Returns:

Name	Type	Description
action_execution_output	StepOutput	action execution output

Source code in sweagent/agent/agents.py

def handle_action(self, step: StepOutput) -> StepOutput:
    """Runs an action proposed by the agent in the environment and returns the corresponding output.

    Args:
        action: command to run in bash shell
        output: output from model (only used for error handling)

    Returns:
        action_execution_output: action execution output
    """
    if self.tools.should_block_action(step.action):
        raise _BlockedActionError()

    if step.action.strip() == "exit":
        self.logger.info("Exiting agent")
        step.done = True
        step.observation = "Exited"
        step.exit_status = "exit_command"
        assert self._env is not None
        step.state = self.tools.get_state(env=self._env)  # for history
        return step

    assert self._env is not None
    self._chook.on_action_started(step=step)
    execution_t0 = time.perf_counter()
    run_action: str = self.tools.guard_multiline_input(step.action).strip()
    try:
        step.observation = self._env.communicate(
            input=run_action,
            timeout=self.tools.config.execution_timeout,
            set_last_action=True,
            check="raise" if self._always_require_zero_exit_code else "ignore",
        )
    except CommandTimeoutError:
        try:
            self._env.interrupt_session()
        except Exception as f:
            self.logger.exception("Failed to interrupt session after command timeout: %s", f, exc_info=True)
            raise
        step.observation = Template(self.templates.command_cancelled_timeout_template).render(
            **self._get_format_dict(),
            timeout=self.tools.config.execution_timeout,
            command=run_action,
        )

    step.execution_time = time.perf_counter() - execution_t0
    self._chook.on_action_executed(step=step)
    step.state = self.tools.get_state(env=self._env)

    if RETRY_WITH_OUTPUT_TOKEN in step.observation:
        step.observation = step.observation.replace(RETRY_WITH_OUTPUT_TOKEN, "")
        raise _RetryWithOutput()
    elif RETRY_WITHOUT_OUTPUT_TOKEN in step.observation:
        step.observation = step.observation.replace(RETRY_WITHOUT_OUTPUT_TOKEN, "")
        raise _RetryWithoutOutput()

    return self.handle_submission(step)

handle_submission

handle_submission(step: StepOutput, *, observation='') -> StepOutput

Check if there was a submission in the observation and handle it.

Parameters:

Name	Type	Description	Default
step	StepOutput		required
observation		If specified, will use this rather than stepobservation	''

Source code in sweagent/agent/agents.py

def handle_submission(self, step: StepOutput, *, observation="") -> StepOutput:
    """Check if there was a submission in the observation and handle it.

    Args:
        step:
        observation: If specified, will use this rather than stepobservation
    """
    step = step.model_copy()
    assert self.tools is not None
    submission = self.tools.parse_submission_cmd_output(observation or step.observation)
    if submission is not None:
        if submission.strip() != "":
            step.submission = submission
        else:
            step.submission = None
        step.observation = submission
        if not step.exit_status:
            step.exit_status = "submitted"
        else:
            step.exit_status = f"submitted ({step.exit_status})"
        step.done = True
        self.logger.info(f"Found submission: {submission}")
    return step

run

run(env: SWEEnv, problem_statement: ProblemStatement | ProblemStatementConfig, output_dir: Path = Path('.')) -> AgentRunResult

Run the agent on a problem instance. This method contains the main loop that repeatedly calls self._step until the problem is solved.

Parameters:

Name	Type	Description	Default
setup_args		Arguments to pass to the agent's setup method.	required
env	SWEEnv	The environment to run the agent on.	required
traj_dir		Directory to save the trajectory to	required

Source code in sweagent/agent/agents.py

def run(
    self,
    env: SWEEnv,
    problem_statement: ProblemStatement | ProblemStatementConfig,
    output_dir: Path = Path("."),
) -> AgentRunResult:
    """Run the agent on a problem instance. This method contains the
    main loop that repeatedly calls `self._step` until the problem is solved.

    Args:
        setup_args: Arguments to pass to the agent's setup method.
        env: The environment to run the agent on.
        traj_dir: Directory to save the trajectory to
    """
    output_dir.mkdir(parents=True, exist_ok=True)
    self.setup(env=env, problem_statement=problem_statement, output_dir=output_dir)

    # Run action/observation loop
    self._chook.on_run_start()
    step_output = StepOutput()
    while not step_output.done:
        step_output = self.step()
        self.save_trajectory()
    self._chook.on_run_done(trajectory=self.trajectory, info=self.info)

    self.logger.info("Trajectory saved to %s", self.traj_path)

    return AgentRunResult(info=self.info, trajectory=self.trajectory)

save_trajectory

save_trajectory() -> None

Save the trajectory to disk. This includes the history, the environment state, and the model stats.

Source code in sweagent/agent/agents.py

def save_trajectory(
    self,
) -> None:
    """Save the trajectory to disk.
    This includes the history, the environment state, and the model stats.
    """

    def get_attempt_data(attempt_idx: int) -> dict[str, Any]:
        """Get data saved for every attempt"""
        assert self._env is not None
        # The deepcopy here is important because else the
        # data["info"]["model_stats"] update will create havoc!
        attempt_data = copy.deepcopy(
            {
                "trajectory": self._trajectory_by_attempt[attempt_idx],
                "history": self._history_by_attempt[attempt_idx],
                "info": self._info_by_attempt[attempt_idx],
            }
        )
        attempt_data["replay_config"] = (
            self.replay_config.model_dump_json() if self.replay_config is not None else None
        )
        attempt_data["environment"] = self._env.name
        return attempt_data

    data = {
        **get_attempt_data(0),
    }

    assert self.traj_path is not None
    self.traj_path.write_text(json.dumps(data, indent=2))

setup

setup(env: SWEEnv, problem_statement: ProblemStatement | ProblemStatementConfig, output_dir: Path = Path('.')) -> None

Setup the agent for a new instance. This includes formatting the system message and adding demonstrations to the history.

This method is called by self.run.

Parameters:

Name	Type	Description	Default
instance_args		Arguments for the instance	required

Source code in sweagent/agent/agents.py

def setup(
    self,
    env: SWEEnv,
    problem_statement: ProblemStatement | ProblemStatementConfig,
    output_dir: Path = Path("."),
) -> None:
    """Setup the agent for a new instance. This includes
    formatting the system message and adding demonstrations to the history.

    This method is called by `self.run`.

    Args:
        instance_args: Arguments for the instance
    """
    self._problem_statement = problem_statement
    self._env = env

    # Save/reset some attributes
    self.info = AgentInfo()
    self.info["swe_agent_hash"] = get_agent_commit_hash()
    self.info["swe_agent_version"] = __version__
    self.info["swe_rex_version"] = get_rex_version()
    self.info["swe_rex_hash"] = get_rex_commit_hash()
    self.traj_path = output_dir / (self._problem_statement.id + ".traj")
    self.logger.info("Trajectory will be saved to %s", self.traj_path)

    self._i_attempt = 0
    self._history_by_attempt = defaultdict(list)
    self._trajectory_by_attempt = defaultdict(list)
    self._info_by_attempt = defaultdict(dict)  # type: ignore
    self._forwarded_vars = {}
    # if self._rloop is not None:
    #     self._forwarded_vars = self._rloop.get_forwarded_vars()
    self._chook.on_tools_installation_started()
    self.tools.install(self._env)
    self.setup_attempt()

    self._chook.on_setup_done()

setup_attempt

setup_attempt() -> None

Setup the agent for a new attempt. This includes resetting the model stats.

Source code in sweagent/agent/agents.py

def setup_attempt(self) -> None:
    """Setup the agent for a new attempt. This includes resetting the model stats."""
    assert self._env is not None
    if self._i_attempt > 0:
        self._env.reset()
    self.model.reset_stats()
    self.add_system_message_to_history()
    self.add_demonstrations_to_history()
    self.add_instance_template_to_history(state=self.tools.get_state(self._env))

step

step() -> StepOutput

Run a step of the agent. This is a wrapper around self.forward_with_handling with additional bookkeeping:

1. Update message history with performed action and observation
2. Update trajectory with the final executed result
3. Update the info dictionary

Returns:

Name	Type	Description
step_output	StepOutput	step output (same as the output of self.forward_with_handling)

Source code in sweagent/agent/agents.py

def step(self) -> StepOutput:
    """Run a step of the agent. This is a wrapper around `self.forward_with_handling`
    with additional bookkeeping:

    1. Update message history with performed action and observation
    2. Update trajectory with the final executed result
    3. Update the info dictionary

    Returns:
        step_output: step output (same as the output of `self.forward_with_handling`)
    """

    assert self._env is not None
    self._chook.on_step_start()

    n_step = len(self.trajectory) + 1
    self.logger.info("=" * 25 + f" STEP {n_step} " + "=" * 25)
    step_output = self.forward_with_handling(self.messages)
    self.add_step_to_history(step_output)

    self.info["submission"] = step_output.submission
    self.info["exit_status"] = step_output.exit_status  # type: ignore
    self.info.update(self._get_edited_files_with_context(patch=step_output.submission or ""))  # type: ignore
    model_stats: InstanceStats = self.model.stats
    self.info["model_stats"] = model_stats.model_dump()

    self.add_step_to_trajectory(step_output)

    self._chook.on_step_done(step=step_output, info=self.info)
    return step_output