def save_screenshot(memory_step: ActionStep, agent: CodeAgent) -> None:
    """ 
    Takes screenshots and writes to observations.
"""
  png_bytes = driver.get_screenshot_as_png()
        memory_step.observations_images = [image.copy()]  # persist images to memory_step
    url_info = f"Current url: {driver.current_url}"
    memory_step.observations = (
        url_info if memory_step.observations is None else memory_step.observations + "\n" + url_info
    )
    return

agent = CodeAgent(
    tools=[go_back, close_popups, search_item_ctrl_f], # pass navigation tools
    model=model,
    additional_authorized_imports=["helium"],
    step_callbacks=[save_screenshot], # pass callback
)