Repository: huggingface/smolagents

Issue:
[BUG] Open Deep Research Example - serpapi import error
**Describe the bug**

Using python 3.10 with serpapi 0.1.5, in the Open Deep Research example, there is an import error in [`text_web_browser.py` line 14](https://github.com/huggingface/smolagents/blob/main/examples/open_deep_research/scripts/text_web_browser.py#L14):

```from serpapi import GoogleSearch```

Produces import error below: 

```
Traceback (most recent call last):
  File "/xxx/smolagents/examples/open_deep_research/run.py", line 8, in <module>
    from scripts.text_web_browser import (
  File "/xxx/smolagents/examples/open_deep_research/scripts/text_web_browser.py", line 14, in <module>
    from serpapi import GoogleSearch
ImportError: cannot import name 'GoogleSearch' from 'serpapi' (/xxx/envs/smolagents/lib/python3.10/site-packages/serpapi/__init__.py)
```

This works:

```from serpapi.google_search import GoogleSearch```


**Code to reproduce the error**

Using python 3.10 with **serpapi                                  0.1.5**

Open the python console and enter:

```from serpapi import GoogleSearch```

Produces:

```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ImportError: cannot import name 'GoogleSearch' from 'serpapi' (xxx/envs/my_env/lib/python3.10/site-packages/serpapi/__init__.py
```

This works however:

```from serpapi.google_search import GoogleSearch```

**Error logs (if any)**

See above.

**Expected behavior**
A clear and concise description of what you expected to happen.

**Packages version:**
Run `pip freeze | grep smolagents` and paste it here.

Package                                  Version
---------------------------------------- -----------
accelerate                               1.5.2
aiofiles                                 23.2.1
aiohappyeyeballs                         2.6.1
aiohttp                                  3.11.14
aioitertools                             0.12.0
aiosignal                                1.3.2
aiosqlite                                0.21.0
alembic                                  1.15.1
annotated-types                          0.7.0
anthropic                                0.49.0
anyio                                    4.9.0
arize-phoenix                            8.17.1
arize-phoenix-client                     1.1.0
arize-phoenix-evals                      0.20.3
arize-phoenix-otel                       0.8.0
asttokens                                3.0.0
async-timeout                            5.0.1
attrs                                    25.3.0
Authlib                                  1.5.1
beautifulsoup4                           4.13.3
bio                                      1.7.1
biopython                                1.85
biothings_client                         0.4.1
cachetools                               5.5.2
certifi                                  2025.1.31
cffi                                     1.17.1
charset-normalizer                       3.4.1
chess                                    1.11.2
click                                    8.1.8
cobble                                   0.1.4
cryptography                             44.0.2
datasets                                 3.4.1
decorator                                5.2.1
defusedxml                               0.7.1
Deprecated                               1.2.18
dill                                     0.3.8
distro                                   1.9.0
dnspython                                2.7.0
docker                                   7.1.0
duckduckgo_search                        7.5.3
e2b                                      1.2.1
e2b-code-interpreter                     1.1.1
email_validator                          2.2.0
et_xmlfile                               2.0.0
exceptiongroup                           1.2.2
executing                                2.2.0
fastapi                                  0.115.11
ffmpy                                    0.5.0
filelock                                 3.18.0
frozenlist                               1.5.0
fsspec                                   2024.12.0
google-search-results                    2.4.2
googleapis-common-protos                 1.69.2
gprofiler-official                       1.0.0
gradio                                   5.22.0
gradio_client                            1.8.0
graphql-core                             3.2.6
greenlet                                 3.1.1
groovy                                   0.1.2
grpc-interceptor                         0.15.4
grpcio                                   1.71.0
h11                                      0.14.0
helium                                   5.1.1
httpcore                                 1.0.7
httpx                                    0.28.1
httpx-sse                                0.4.0
huggingface-hub                          0.29.3
idna                                     3.10
importlib_metadata                       8.6.1
iniconfig                                2.1.0
ipython                                  8.34.0
jedi                                     0.19.2
Jinja2                                   3.1.6
jiter                                    0.9.0
joblib                                   1.4.2
jsonref                                  1.1.0
jsonschema                               4.23.0
jsonschema-specifications                2024.10.1
litellm                                  1.63.12
lxml                                     5.3.1
Mako                                     1.3.9
mammoth                                  1.9.0
markdown-it-py                           3.0.0
markdownify                              1.1.0
MarkupSafe                               3.0.2
matplotlib-inline                        0.1.7
mcp                                      1.5.0
mcpadapt                                 0.0.15
mdurl                                    0.1.2
mlx                                      0.23.1
mlx-lm                                   0.22.1
mpmath                                   1.3.0
multidict                                6.2.0
multiprocess                             0.70.16
mygene                                   3.2.2
networkx                                 3.4.2
numexpr                                  2.10.2
numpy                                    2.2.4
nvidia-cublas-cu12                       12.4.5.8
nvidia-cuda-cupti-cu12                   12.4.127
nvidia-cuda-nvrtc-cu12                   12.4.127
nvidia-cuda-runtime-cu12                 12.4.127
nvidia-cudnn-cu12                        9.1.0.70
nvidia-cufft-cu12                        11.2.1.3
nvidia-curand-cu12                       10.3.5.147
nvidia-cusolver-cu12                     11.6.1.9
nvidia-cusparse-cu12                     12.3.1.170
nvidia-cusparselt-cu12                   0.6.2
nvidia-nccl-cu12                         2.21.5
nvidia-nvjitlink-cu12                    12.4.127
nvidia-nvtx-cu12                         12.4.127
openai                                   1.68.2
openinference-instrumentation            0.1.24
openinference-instrumentation-smolagents 0.1.8
openinference-semantic-conventions       0.1.15
openpyxl                                 3.1.5
opentelemetry-api                        1.31.1
opentelemetry-exporter-otlp              1.31.1
opentelemetry-exporter-otlp-proto-common 1.31.1
opentelemetry-exporter-otlp-proto-grpc   1.31.1
opentelemetry-exporter-otlp-proto-http   1.31.1
opentelemetry-instrumentation            0.52b1
opentelemetry-proto                      1.31.1
opentelemetry-sdk                        1.31.1
opentelemetry-semantic-conventions       0.52b1
orjson                                   3.10.15
outcome                                  1.3.0.post0
packaging                                24.2
pandas                                   2.2.3
parso                                    0.8.4
pathvalidate                             3.2.3
pdfminer                                 20191125
pdfminer.six                             20240706
pexpect                                  4.9.0
pillow                                   11.1.0
pip                                      22.0.2
platformdirs                             4.3.7
pluggy                                   1.5.0
pooch                                    1.8.2
primp                                    0.14.0
prompt_toolkit                           3.0.50
propcache                                0.3.0
protobuf                                 5.29.4
psutil                                   7.0.0
ptyprocess                               0.7.0
PubChemPy                                1.0.4
pure_eval                                0.2.3
puremagic                                1.28
pyarrow                                  19.0.1
pycparser                                2.22
pycryptodome                             3.22.0
pydantic                                 2.10.6
pydantic_core                            2.27.2
pydantic-settings                        2.8.1
pydub                                    0.25.1
Pygments                                 2.19.1
pypdf                                    5.4.0
PyPDF2                                   3.0.1
PySocks                                  1.7.1
pytest                                   8.3.5
pytest-datadir                           1.6.1
python-dateutil                          2.9.0.post0
python-dotenv                            1.0.1
python-multipart                         0.0.20
python-pptx                              1.0.2
pytz                                     2025.1
PyYAML                                   6.0.2
rank-bm25                                0.2.2
referencing                              0.36.2
regex                                    2024.11.6
requests                                 2.32.3
rich                                     13.9.4
rpds-py                                  0.23.1
ruff                                     0.11.2
safehttpx                                0.1.6
safetensors                              0.5.3
scikit-learn                             1.6.1
scipy                                    1.15.2
selenium                                 4.30.0
semantic-version                         2.10.0
sentencepiece                            0.2.0
**serpapi                                  0.1.5**
setuptools                               59.6.0
shellingham                              1.5.4
six                                      1.17.0
smolagents                               1.13.0.dev0
sniffio                                  1.3.1
sortedcontainers                         2.4.0
soundfile                                0.13.1
soupsieve                                2.6
SpeechRecognition                        3.14.1
SQLAlchemy                               2.0.39
sqlean.py                                3.47.0
sse-starlette                            2.2.1
stack-data                               0.6.3
starlette                                0.46.1
strawberry-graphql                       0.262.5
sympy                                    1.13.1
threadpoolctl                            3.6.0
tiktoken                                 0.9.0
tokenizers                               0.21.1
tomli                                    2.2.1
tomlkit                                  0.13.2
torch                                    2.6.0
torchvision                              0.21.0
tqdm                                     4.67.1
traitlets                                5.14.3
transformers                             4.50.0
trio                                     0.29.0
trio-websocket                           0.12.2
triton                                   3.2.0
typer                                    0.15.2
typing_extensions                        4.12.2
tzdata                                   2025.1
urllib3                                  2.3.0
uvicorn                                  0.34.0
wcwidth                                  0.2.13
websocket-client                         1.8.0
websockets                               15.0.1
wrapt                                    1.17.2
wsproto                                  1.2.0
xlrd                                     2.0.1
XlsxWriter                               3.2.2
xxhash                                   3.5.0
yarl                                     1.18.3
youtube-transcript-api                   1.0.2
zipp                                     3.21.0

**Additional context**
Add any other context about the problem here.

Hints from issue comments:
@albertvillanova: You need to install all the requirements: https://github.com/huggingface/smolagents/blob/main/examples/open_deep_research/requirements.txt

More concretely, you need `google-search-results`: https://github.com/serpapi/google-search-results-python@edlee123: Correct, I'd already installed it.  What fixed it for me was changing to `from serpapi.google_search import GoogleSearch` 



> You need to install all the requirements: https://github.com/huggingface/smolagents/blob/main/examples/open_deep_research/requirements.txt
> 
> More concretely, you need `google-search-results`: https://github.com/serpapi/google-search-results-python

@albertvillanova: Thank you for sharing your solution! However, please note that:  
- We have been using `from serpapi import GoogleSearch` without any issues.  
- This is also the approach recommended in the official SerpAPI documentation: [SerpAPI Quick Start](https://github.com/serpapi/google-search-results-python?tab=readme-ov-file#quick-start).  

For reference, the documented usage is:  

```python
from serpapi import GoogleSearch

search = GoogleSearch({
    "q": "coffee", 
    "location": "Austin, Texas",
    "api_key": "<your secret api key>"
})

result = search.get_dict()
```  

If you encountered a case where `from serpapi.google_search import GoogleSearch` was necessary, I think you should open an issue in their repository: https://github.com/serpapi/google-search-results-python@edlee123: Thank you, not a biggie.  Looks like there was one https://github.com/serpapi/google-search-results-python/issues/63

Solution patch:
diff --git a/examples/agent_from_any_llm.py b/examples/agent_from_any_llm.py
index bc42127..d5e33f0 100644
--- a/examples/agent_from_any_llm.py
+++ b/examples/agent_from_any_llm.py
@@ -52,6 +52,6 @@ agent = ToolCallingAgent(tools=[get_weather], model=model, verbosity_level=2)
 
 print("ToolCallingAgent:", agent.run("What's the weather like in Paris?"))
 
-agent = CodeAgent(tools=[get_weather], model=model, verbosity_level=2)
+agent = CodeAgent(tools=[get_weather], model=model, verbosity_level=2, stream_outputs=True)
 
 print("CodeAgent:", agent.run("What's the weather like in Paris?"))
diff --git a/examples/gradio_ui.py b/examples/gradio_ui.py
index 81c56a1..87f5326 100644
--- a/examples/gradio_ui.py
+++ b/examples/gradio_ui.py
@@ -1,25 +1,15 @@
-from io import BytesIO
-
-import requests
-from PIL import Image
-
 from smolagents import CodeAgent, GradioUI, InferenceClientModel
 
 
-def add_agent_image(memory_step, agent):
-    url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/smolagents.png"
-    response = requests.get(url)
-    memory_step.observations_images = [Image.open(BytesIO(response.content))]
-
-
 agent = CodeAgent(
     tools=[],
     model=InferenceClientModel(),
     verbosity_level=1,
     planning_interval=3,
     name="example_agent",
-    description="This is an example agent that has not tool but will always see an agent at the end of its step.",
-    step_callbacks=[add_agent_image],
+    description="This is an example agent.",
+    step_callbacks=[],
+    stream_outputs=False,
 )
 
 GradioUI(agent, file_upload_folder="./data").launch()
diff --git a/examples/multi_llm_agent.py b/examples/multi_llm_agent.py
index 186fa06..e46354e 100644
--- a/examples/multi_llm_agent.py
+++ b/examples/multi_llm_agent.py
@@ -3,10 +3,7 @@ import os
 from smolagents import CodeAgent, DuckDuckGoSearchTool, LiteLLMRouterModel
 
 
-os.environ["OPENAI_API_KEY"] = ""
-os.environ["AWS_ACCESS_KEY_ID"] = ""
-os.environ["AWS_SECRET_ACCESS_KEY"] = ""
-os.environ["AWS_REGION"] = ""
+# Make sure to setup the necessary environment variables!
 
 llm_loadbalancer_model_list = [
     {
@@ -42,6 +39,6 @@ model = LiteLLMRouterModel(
     model_list=llm_loadbalancer_model_list,
     client_kwargs={"routing_strategy": "simple-shuffle"},
 )
-agent = CodeAgent(tools=[DuckDuckGoSearchTool()], model=model)
+agent = CodeAgent(tools=[DuckDuckGoSearchTool()], model=model, stream_outputs=True)
 
 agent.run("How many seconds would it take for a leopard at full speed to run through Pont des Arts?")
diff --git a/examples/multiple_tools.py b/examples/multiple_tools.py
index a268554..1a56e35 100644
--- a/examples/multiple_tools.py
+++ b/examples/multiple_tools.py
@@ -243,11 +243,12 @@ agent = CodeAgent(
         search_wikipedia,
     ],
     model=model,
+    stream_outputs=True,
 )
 
 # Uncomment the line below to run the agent with a specific query
 
-agent.run("5000 dollars to Euros")
+agent.run("Convert 5000 dollars to Euros")
 # agent.run("What is the weather in New York?")
 # agent.run("Give me the top news headlines")
 # agent.run("Tell me a joke")
diff --git a/examples/open_deep_research/README.md b/examples/open_deep_research/README.md
index c2c7996..bfd372f 100644
--- a/examples/open_deep_research/README.md
+++ b/examples/open_deep_research/README.md
@@ -51,4 +51,14 @@ For example, to use the default `o1` model, you need to set the `OPENAI_API_KEY`
 Then you're good to go! Run the run.py script, as in:
 ```bash
 python run.py --model-id "o1" "Your question here!"
-```
\ No newline at end of file
+```
+
+## Full reproducibility of results
+
+The data used in our submissions to GAIA was augmented in this way:
+ -  For each single-page .pdf or .xls file, it was opened in a file reader (MacOS Sonoma Numbers or Preview), and a ".png" screenshot was taken and added to the folder.
+- Then for any file used in a question, the file loading system checks if there is a ".png" extension version of the file, and loads it instead of the original if it exists.
+
+This process was done manually but could be automatized.
+
+After processing, the annotated was uploaded to a [new dataset](https://huggingface.co/datasets/smolagents/GAIA-annotated). You need to request access (granted instantly).
\ No newline at end of file
diff --git a/examples/open_deep_research/analysis.ipynb b/examples/open_deep_research/analysis.ipynb
index ccb6a1d..53e7bd4 100644
--- a/examples/open_deep_research/analysis.ipynb
+++ b/examples/open_deep_research/analysis.ipynb
@@ -6,7 +6,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# !pip install plotly kaleido datasets nbformat -U -q"
+    "!pip install plotly kaleido datasets nbformat -U -q"
    ]
   },
   {
@@ -28,7 +28,7 @@
     "\n",
     "pd.set_option(\"max_colwidth\", None)\n",
     "\n",
-    "OUTPUT_DIR = \"../../output\""
+    "OUTPUT_DIR = \"output\""
    ]
   },
   {
@@ -51,7 +51,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 88,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -134,31 +134,6 @@
     "# result_df[\"tool_calls\"] = result_df[\"intermediate_steps\"].apply(sum_tool_calls)"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def get_thoughts(x):\n",
-    "    try:\n",
-    "        output = x[0][\"task\"]\n",
-    "        for y in x[1:]:\n",
-    "            try:\n",
-    "                if \"observation\" in y:\n",
-    "                    output += y[\"llm_output\"] + \"\\nObservation:\" + y[\"observation\"]\n",
-    "                else:\n",
-    "                    output += y[\"llm_output\"] + r\"\\Error:\" + str(y[\"error\"])\n",
-    "            except Exception:\n",
-    "                pass\n",
-    "        return output\n",
-    "    except Exception:\n",
-    "        return None\n",
-    "\n",
-    "\n",
-    "result_df[\"thoughts\"] = result_df[\"intermediate_steps\"].apply(lambda x: get_thoughts(x))"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -460,9 +435,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "test",
+   "display_name": "agents",
    "language": "python",
-   "name": "test"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
diff --git a/examples/open_deep_research/requirements.txt b/examples/open_deep_research/requirements.txt
index 4fe0e0e..fe6c98e 100644
--- a/examples/open_deep_research/requirements.txt
+++ b/examples/open_deep_research/requirements.txt
@@ -20,7 +20,6 @@ pypdf>=5.1.0
 python-dotenv>=1.0.1
 python_pptx>=1.0.2
 Requests>=2.32.3
-serpapi>=0.1.5
 tqdm>=4.66.4
 torch>=2.2.2
 torchvision>=0.17.2
@@ -37,4 +36,4 @@ PyPDF2
 python-pptx
 torch
 xlrd
-SpeechRecognition
\ No newline at end of file
+SpeechRecognition
diff --git a/examples/open_deep_research/run_gaia.py b/examples/open_deep_research/run_gaia.py
index 1920817..9c7bacd 100644
--- a/examples/open_deep_research/run_gaia.py
+++ b/examples/open_deep_research/run_gaia.py
@@ -1,4 +1,4 @@
-# EXAMPLE COMMAND: python examples/open_deep_research/run_gaia.py --concurrency 32 --run-name generate-traces-03-apr-noplanning --model-id gpt-4o
+# EXAMPLE COMMAND: from folder examples/open_deep_research, run: python run_gaia.py --concurrency 32 --run-name generate-traces-03-apr-noplanning --model-id gpt-4o
 import argparse
 import json
 import os
@@ -6,11 +6,12 @@ import threading
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime
 from pathlib import Path
+from typing import Any
 
 import datasets
 import pandas as pd
 from dotenv import load_dotenv
-from huggingface_hub import login
+from huggingface_hub import login, snapshot_download
 from scripts.reformulator import prepare_response
 from scripts.run_agents import (
     get_single_file_description,
@@ -49,35 +50,18 @@ def parse_args():
     parser.add_argument("--concurrency", type=int, default=8)
     parser.add_argument("--model-id", type=str, default="o1")
     parser.add_argument("--run-name", type=str, required=True)
+    parser.add_argument("--set-to-run", type=str, default="validation")
+    parser.add_argument("--use-open-models", type=bool, default=False)
+    parser.add_argument("--use-raw-dataset", action="store_true")
     return parser.parse_args()
 
 
 ### IMPORTANT: EVALUATION SWITCHES
 
-print("Make sure you deactivated Tailscale VPN, else some URLs will be blocked!")
-
-USE_OPEN_MODELS = False
-
-SET = "validation"
+print("Make sure you deactivated any VPN like Tailscale, else some URLs will be blocked!")
 
 custom_role_conversions = {"tool-call": "assistant", "tool-response": "user"}
 
-### LOAD EVALUATION DATASET
-
-eval_ds = datasets.load_dataset("gaia-benchmark/GAIA", "2023_all")[SET]
-eval_ds = eval_ds.rename_columns({"Question": "question", "Final answer": "true_answer", "Level": "task"})
-
-
-def preprocess_file_paths(row):
-    if len(row["file_name"]) > 0:
-        row["file_name"] = f"data/gaia/{SET}/" + row["file_name"]
-    return row
-
-
-eval_ds = eval_ds.map(preprocess_file_paths)
-eval_df = pd.DataFrame(eval_ds)
-print("Loaded evaluation dataset:")
-print(eval_df["task"].value_counts())
 
 user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"
 
@@ -142,17 +126,54 @@ def create_agent_team(model: Model):
     return manager_agent
 
 
+def load_gaia_dataset(use_raw_dataset: bool, set_to_run: str) -> datasets.Dataset:
+    if not os.path.exists("data/gaia"):
+        if use_raw_dataset:
+            snapshot_download(
+                repo_id="gaia-benchmark/GAIA",
+                repo_type="dataset",
+                local_dir="data/gaia",
+                ignore_patterns=[".gitattributes", "README.md"],
+            )
+        else:
+            # WARNING: this dataset is gated: make sure you visit the repo to require access.
+            snapshot_download(
+                repo_id="smolagents/GAIA-annotated",
+                repo_type="dataset",
+                local_dir="data/gaia",
+                ignore_patterns=[".gitattributes", "README.md"],
+            )
+
+    def preprocess_file_paths(row):
+        if len(row["file_name"]) > 0:
+            row["file_name"] = f"data/gaia/{set_to_run}/" + row["file_name"]
+        return row
+
+    eval_ds = datasets.load_dataset(
+        "data/gaia/GAIA.py",
+        name="2023_all",
+        split=set_to_run,
+        # data_files={"validation": "validation/metadata.jsonl", "test": "test/metadata.jsonl"},
+    )
+
+    eval_ds = eval_ds.rename_columns({"Question": "question", "Final answer": "true_answer", "Level": "task"})
+    eval_ds = eval_ds.map(preprocess_file_paths)
+    return eval_ds
+
+
 def append_answer(entry: dict, jsonl_file: str) -> None:
-    jsonl_file = Path(jsonl_file)
-    jsonl_file.parent.mkdir(parents=True, exist_ok=True)
+    jsonl_path = Path(jsonl_file)
+    jsonl_path.parent.mkdir(parents=True, exist_ok=True)
     with append_answer_lock, open(jsonl_file, "a", encoding="utf-8") as fp:
         fp.write(json.dumps(entry) + "\n")
-    assert os.path.exists(jsonl_file), "File not found!"
-    print("Answer exported to file:", jsonl_file.resolve())
+    assert jsonl_path.exists(), "File not found!"
+    print("Answer exported to file:", jsonl_path.resolve())
 
 
-def answer_single_question(example, model_id, answers_file, visual_inspection_tool):
-    model_params = {
+def answer_single_question(
+    example: dict, model_id: str, answers_file: str, visual_inspection_tool: TextInspectorTool
+) -> None:
+    model_params: dict[str, Any] = {
         "model_id": model_id,
         "custom_role_conversions": custom_role_conversions,
     }
@@ -162,15 +183,16 @@ def answer_single_question(example, model_id, answers_file, visual_inspection_to
     else:
         model_params["max_tokens"] = 4096
     model = LiteLLMModel(**model_params)
-    # model = InferenceClientModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct", provider="together", max_tokens=4096)
+    # model = InferenceClientModel(model_id="Qwen/Qwen3-32B", provider="novita", max_tokens=4096)
     document_inspection_tool = TextInspectorTool(model, 100000)
 
     agent = create_agent_team(model)
 
     augmented_question = """You have one question to answer. It is paramount that you provide a correct answer.
-Give it all you can: I know for a fact that you have access to all the relevant tools to solve it and find the correct answer (the answer does exist). Failure or 'I cannot answer' or 'None found' will not be tolerated, success will be rewarded.
-Run verification steps if that's needed, you must make sure you find the correct answer!
-Here is the task:
+Give it all you can: I know for a fact that you have access to all the relevant tools to solve it and find the correct answer (the answer does exist).
+Failure or 'I cannot answer' or 'None found' will not be tolerated, success will be rewarded.
+Run verification steps if that's needed, you must make sure you find the correct answer! Here is the task:
+
 """ + example["question"]
 
     if example["file_name"]:
@@ -180,7 +202,7 @@ Here is the task:
                 example["file_name"], example["question"], visual_inspection_tool, document_inspection_tool
             )
         else:
-            prompt_use_files = "\n\nTo solve the task above, you will have to use this attached file:"
+            prompt_use_files = "\n\nTo solve the task above, you will have to use this attached file:\n"
             prompt_use_files += get_single_file_description(
                 example["file_name"], example["question"], visual_inspection_tool, document_inspection_tool
             )
@@ -241,7 +263,7 @@ Here is the task:
     append_answer(annotated_example, answers_file)
 
 
-def get_examples_to_answer(answers_file, eval_ds) -> list[dict]:
+def get_examples_to_answer(answers_file: str, eval_ds: datasets.Dataset) -> list[dict]:
     print(f"Loading answers from {answers_file}...")
     try:
         done_questions = pd.read_json(answers_file, lines=True)["question"].tolist()
@@ -250,14 +272,18 @@ def get_examples_to_answer(answers_file, eval_ds) -> list[dict]:
         print("Error when loading records: ", e)
         print("No usable records! ▶️ Starting new.")
         done_questions = []
-    return [line for line in eval_ds.to_list() if line["question"] not in done_questions]
+    return [line for line in eval_ds.to_list() if line["question"] not in done_questions and line["file_name"]]
 
 
 def main():
     args = parse_args()
     print(f"Starting run with arguments: {args}")
 
-    answers_file = f"output/{SET}/{args.run_name}.jsonl"
+    eval_ds = load_gaia_dataset(args.use_raw_dataset, args.set_to_run)
+    print("Loaded evaluation dataset:")
+    print(pd.DataFrame(eval_ds)["task"].value_counts())
+
+    answers_file = f"output/{args.set_to_run}/{args.run_name}.jsonl"
     tasks_to_run = get_examples_to_answer(answers_file, eval_ds)
 
     with ThreadPoolExecutor(max_workers=args.concurrency) as exe:
diff --git a/examples/open_deep_research/scripts/run_agents.py b/examples/open_deep_research/scripts/run_agents.py
index 37da8a4..e2e020c 100644
--- a/examples/open_deep_research/scripts/run_agents.py
+++ b/examples/open_deep_research/scripts/run_agents.py
@@ -38,12 +38,13 @@ def get_single_file_description(file_path: str, question: str, visual_inspection
         )
         return file_description
     elif file_extension in ["pdf", "xls", "xlsx", "docx", "doc", "xml"]:
-        file_description = f" - Attached document: {file_path}"
         image_path = file_path.split(".")[0] + ".png"
         if os.path.exists(image_path):
             description = get_image_description(image_path, question, visual_inspection_tool)
+            file_path = image_path
         else:
             description = get_document_description(file_path, question, document_inspection_tool)
+        file_description = f" - Attached document: {file_path}"
         file_description += f"\n     -> File description: {description}"
         return file_description
     elif file_extension in ["mp3", "m4a", "wav"]:
diff --git a/examples/rag.py b/examples/rag.py
index 3ff572f..6f378d1 100644
--- a/examples/rag.py
+++ b/examples/rag.py
@@ -58,9 +58,10 @@ from smolagents import CodeAgent, InferenceClientModel
 retriever_tool = RetrieverTool(docs_processed)
 agent = CodeAgent(
     tools=[retriever_tool],
-    model=InferenceClientModel(model_id="meta-llama/Llama-3.3-70B-Instruct"),
+    model=InferenceClientModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
     max_steps=4,
     verbosity_level=2,
+    stream_outputs=True,
 )
 
 agent_output = agent.run("For a transformers model training, which is slower, the forward or the backward pass?")
diff --git a/src/smolagents/agents.py b/src/smolagents/agents.py
index ec32101..d052f14 100644
--- a/src/smolagents/agents.py
+++ b/src/smolagents/agents.py
@@ -23,7 +23,6 @@ import tempfile
 import textwrap
 import time
 from abc import ABC, abstractmethod
-from collections import deque
 from collections.abc import Callable, Generator
 from logging import getLogger
 from pathlib import Path
@@ -57,7 +56,7 @@ from .memory import (
     TaskStep,
     ToolCall,
 )
-from .models import ChatMessage, MessageRole, Model, parse_json_if_needed
+from .models import ChatMessage, ChatMessageStreamDelta, MessageRole, Model, parse_json_if_needed
 from .monitoring import (
     YELLOW_HEX,
     AgentLogger,
@@ -248,6 +247,7 @@ class MultiStepAgent(ABC):
         self.monitor = Monitor(self.model, self.logger)
         self.step_callbacks = step_callbacks if step_callbacks is not None else []
         self.step_callbacks.append(self.monitor.update_metrics)
+        self.stream_outputs = False
 
     def _validate_name(self, name: str | None) -> str | None:
         if name is not None and not is_valid_name(name):
@@ -346,11 +346,11 @@ You have been provided with these additional arguments, that you can access usin
 
         if stream:
             # The steps are returned as they are executed through a generator to iterate on.
-            return self._run(task=self.task, max_steps=max_steps, images=images)
+            return self._run_stream(task=self.task, max_steps=max_steps, images=images)
         # Outputs are returned only at the end. We only look at the last step.
-        return deque(self._run(task=self.task, max_steps=max_steps, images=images), maxlen=1)[0].final_answer
+        return list(self._run_stream(task=self.task, max_steps=max_steps, images=images))[-1].final_answer
 
-    def _run(
+    def _run_stream(
         self, task: str, max_steps: int, images: list["PIL.Image.Image"] | None = None
     ) -> Generator[ActionStep | PlanningStep | FinalAnswerStep]:
         final_answer = None
@@ -362,16 +362,18 @@ You have been provided with these additional arguments, that you can access usin
             if self.planning_interval is not None and (
                 self.step_number == 1 or (self.step_number - 1) % self.planning_interval == 0
             ):
-                planning_step = self._generate_planning_step(
+                for element in self._generate_planning_step(
                     task, is_first_step=(self.step_number == 1), step=self.step_number
-                )
-                self.memory.steps.append(planning_step)
-                yield planning_step
+                ):
+                    yield element
+                self.memory.steps.append(element)
             action_step = ActionStep(
                 step_number=self.step_number, start_time=step_start_time, observations_images=images
             )
             try:
-                final_answer = self._execute_step(task, action_step)
+                for el in self._execute_step(action_step):
+                    yield el
+                final_answer = el
             except AgentGenerationError as e:
                 # Agent generation errors are not caused by a Model error but an implementation error: so we should raise them and exit.
                 raise e
@@ -389,12 +391,15 @@ You have been provided with these additional arguments, that you can access usin
             yield action_step
         yield FinalAnswerStep(handle_agent_output_types(final_answer))
 
-    def _execute_step(self, task: str, memory_step: ActionStep) -> None | Any:
+    def _execute_step(self, memory_step: ActionStep) -> Generator[Any]:
         self.logger.log_rule(f"Step {self.step_number}", level=LogLevel.INFO)
-        final_answer = self.step(memory_step)
+        final_answer = None
+        for el in self._step_stream(memory_step):
+            final_answer = el
+            yield el
         if final_answer is not None and self.final_answer_checks:
             self._validate_final_answer(final_answer)
-        return final_answer
+        yield final_answer
 
     def _validate_final_answer(self, final_answer: Any):
         for check_function in self.final_answer_checks:
@@ -427,7 +432,9 @@ You have been provided with these additional arguments, that you can access usin
             )
         return final_answer
 
-    def _generate_planning_step(self, task, is_first_step: bool, step: int) -> PlanningStep:
+    def _generate_planning_step(
+        self, task, is_first_step: bool, step: int
+    ) -> Generator[ChatMessageStreamDelta, PlanningStep]:
         if is_first_step:
             input_messages = [
                 {
@@ -443,9 +450,15 @@ You have been provided with these additional arguments, that you can access usin
                     ],
                 }
             ]
-            plan_message = self.model(input_messages, stop_sequences=["<end_plan>"])
+            if self.stream_outputs and hasattr(self.model, "generate_stream"):
+                plan_message_content = ""
+                for completion_delta in self.model.generate_stream(input_messages, stop_sequences=["<end_plan>"]):  # type: ignore
+                    plan_message_content += completion_delta.content
+                    yield completion_delta
+            else:
+                plan_message_content = self.model.generate(input_messages, stop_sequences=["<end_plan>"]).content
             plan = textwrap.dedent(
-                f"""Here are the facts I know and the plan of action that I will follow to solve the task:\n```\n{plan_message.content}\n```"""
+                f"""Here are the facts I know and the plan of action that I will follow to solve the task:\n```\n{plan_message_content}\n```"""
             )
         else:
             # Summary mode removes the system prompt and previous planning messages output by the model.
@@ -480,16 +493,22 @@ You have been provided with these additional arguments, that you can access usin
                 ],
             }
             input_messages = [plan_update_pre] + memory_messages + [plan_update_post]
-            plan_message = self.model(input_messages, stop_sequences=["<end_plan>"])
+            if self.stream_outputs and hasattr(self.model, "generate_stream"):
+                plan_message_content = ""
+                for completion_delta in self.model.generate_stream(input_messages, stop_sequences=["<end_plan>"]):  # type: ignore
+                    plan_message_content += completion_delta.content
+                    yield completion_delta
+            else:
+                plan_message_content = self.model.generate(input_messages, stop_sequences=["<end_plan>"]).content
             plan = textwrap.dedent(
-                f"""I still need to solve the task I was given:\n```\n{self.task}\n```\n\nHere are the facts I know and my new/updated plan of action to solve the task:\n```\n{plan_message.content}\n```"""
+                f"""I still need to solve the task I was given:\n```\n{self.task}\n```\n\nHere are the facts I know and my new/updated plan of action to solve the task:\n```\n{plan_message_content}\n```"""
             )
         log_headline = "Initial plan" if is_first_step else "Updated plan"
         self.logger.log(Rule(f"[bold]{log_headline}", style="orange"), Text(plan), level=LogLevel.INFO)
-        return PlanningStep(
+        yield PlanningStep(
             model_input_messages=input_messages,
             plan=plan,
-            model_output_message=plan_message,
+            model_output_message=ChatMessage(role=MessageRole.ASSISTANT, content=plan_message_content),
         )
 
     @property
@@ -522,9 +541,19 @@ You have been provided with these additional arguments, that you can access usin
             messages.extend(memory_step.to_messages(summary_mode=summary_mode))
         return messages
 
-    def visualize(self):
-        """Creates a rich tree visualization of the agent's structure."""
-        self.logger.visualize_agent_tree(self)
+    def _step_stream(self, memory_step: ActionStep) -> Generator[Any]:
+        """
+        Perform one step in the ReAct framework: the agent thinks, acts, and observes the result.
+        Yields either None if the step is not final, or the final answer.
+        """
+        raise NotImplementedError("This method should be implemented in child classes")
+
+    def step(self, memory_step: ActionStep) -> Any:
+        """
+        Perform one step in the ReAct framework: the agent thinks, acts, and observes the result.
+        Returns either None if the step is not final, or the final answer.
+        """
+        return list(self._step_stream(memory_step))[-1]
 
     def extract_action(self, model_output: str, split_token: str) -> tuple[str, str]:
         """
@@ -591,10 +620,9 @@ You have been provided with these additional arguments, that you can access usin
         except Exception as e:
             return f"Error in generating final LLM output:\n{e}"
 
-    @abstractmethod
-    def step(self, memory_step: ActionStep) -> None | Any:
-        """To be implemented in children classes. Should return either None if the step is not final."""
-        pass
+    def visualize(self):
+        """Creates a rich tree visualization of the agent's structure."""
+        self.logger.visualize_agent_tree(self)
 
     def replay(self, detailed: bool = False):
         """Prints a pretty replay of the agent's steps.
@@ -1007,10 +1035,10 @@ class ToolCallingAgent(MultiStepAgent):
         )
         return system_prompt
 
-    def step(self, memory_step: ActionStep) -> None | Any:
+    def _step_stream(self, memory_step: ActionStep) -> Generator[Any]:
         """
         Perform one step in the ReAct framework: the agent thinks, acts, and observes the result.
-        Returns None if the step is not final.
+        Yields either None if the step is not final, or the final answer.
         """
         memory_messages = self.write_memory_to_messages()
 
@@ -1081,7 +1109,7 @@ class ToolCallingAgent(MultiStepAgent):
                 )
 
             memory_step.action_output = final_answer
-            return final_answer
+            yield final_answer
         else:
             if tool_arguments is None:
                 tool_arguments = {}
@@ -1103,7 +1131,7 @@ class ToolCallingAgent(MultiStepAgent):
                 level=LogLevel.INFO,
             )
             memory_step.observations = updated_information
-            return None
+            yield None
 
     def _substitute_state_variables(self, arguments: dict[str, str] | str) -> dict[str, Any] | str:
         """Replace string values in arguments with their corresponding state values if they exist."""
@@ -1271,10 +1299,10 @@ class CodeAgent(MultiStepAgent):
         )
         return system_prompt
 
-    def step(self, memory_step: ActionStep) -> None | Any:
+    def _step_stream(self, memory_step: ActionStep) -> Generator[Any]:
         """
         Perform one step in the ReAct framework: the agent thinks, acts, and observes the result.
-        Returns None if the step is not final.
+        Yields either None if the step is not final, or the final answer.
         """
         memory_messages = self.write_memory_to_messages()
 
@@ -1295,13 +1323,14 @@ class CodeAgent(MultiStepAgent):
                         if event.content is not None:
                             output_text += event.content
                             live.update(Markdown(output_text))
+                        yield event
 
                 model_output = output_text
                 chat_message = ChatMessage(role="assistant", content=model_output)
                 memory_step.model_output_message = chat_message
                 model_output = chat_message.content
             else:
-                chat_message: ChatMessage = self.model(
+                chat_message: ChatMessage = self.model.generate(
                     input_messages,
                     stop_sequences=["<end_code>", "Observation:", "Calling tools:"],
                     **additional_args,
@@ -1381,7 +1410,7 @@ class CodeAgent(MultiStepAgent):
         ]
         self.logger.log(Group(*execution_outputs_console), level=LogLevel.INFO)
         memory_step.action_output = output
-        return output if is_final_answer else None
+        yield output if is_final_answer else None
 
     def to_dict(self) -> dict[str, Any]:
         """Convert the agent to a dictionary representation.
diff --git a/src/smolagents/gradio_ui.py b/src/smolagents/gradio_ui.py
index 83fbaff..ec089b3 100644
--- a/src/smolagents/gradio_ui.py
+++ b/src/smolagents/gradio_ui.py
@@ -16,10 +16,12 @@
 import os
 import re
 import shutil
+from pathlib import Path
 
 from smolagents.agent_types import AgentAudio, AgentImage, AgentText
 from smolagents.agents import MultiStepAgent, PlanningStep
 from smolagents.memory import ActionStep, FinalAnswerStep, MemoryStep
+from smolagents.models import ChatMessageStreamDelta
 from smolagents.utils import _is_package_available
 
 
@@ -36,10 +38,14 @@ def get_step_footnote_content(step_log: MemoryStep, step_name: str) -> str:
     return step_footnote_content
 
 
-def pull_messages_from_step(
-    step_log: MemoryStep,
-):
-    """Extract ChatMessage objects from agent steps with proper nesting"""
+def pull_messages_from_step(step_log: MemoryStep, skip_model_outputs: bool = False):
+    """Extract ChatMessage objects from agent steps with proper nesting.
+
+    Args:
+        step_log: The step log to display as gr.ChatMessage objects.
+        skip_model_outputs: If True, skip the model outputs when creating the gr.ChatMessage objects:
+            This is used for instance when streaming model outputs have already been displayed.
+    """
     if not _is_package_available("gradio"):
         raise ModuleNotFoundError(
             "Please install 'gradio' extra to use the GradioUI: `pip install 'smolagents[gradio]'`"
@@ -49,24 +55,23 @@ def pull_messages_from_step(
     if isinstance(step_log, ActionStep):
         # Output the step number
         step_number = f"Step {step_log.step_number}" if step_log.step_number is not None else "Step"
-        yield gr.ChatMessage(role="assistant", content=f"**{step_number}**")
 
         # First yield the thought/reasoning from the LLM
-        if hasattr(step_log, "model_output") and step_log.model_output is not None:
-            # Clean up the LLM output
+        if not skip_model_outputs:
+            yield gr.ChatMessage(role="assistant", content=f"**{step_number}**", metadata={"status": "done"})
+        elif skip_model_outputs and hasattr(step_log, "model_output") and step_log.model_output is not None:
             model_output = step_log.model_output.strip()
             # Remove any trailing <end_code> and extra backticks, handling multiple possible formats
             model_output = re.sub(r"```\s*<end_code>", "```", model_output)  # handles ```<end_code>
             model_output = re.sub(r"<end_code>\s*```", "```", model_output)  # handles <end_code>```
             model_output = re.sub(r"```\s*\n\s*<end_code>", "```", model_output)  # handles ```\n<end_code>
             model_output = model_output.strip()
-            yield gr.ChatMessage(role="assistant", content=model_output)
+            yield gr.ChatMessage(role="assistant", content=model_output, metadata={"status": "done"})
 
         # For tool calls, create a parent message
         if hasattr(step_log, "tool_calls") and step_log.tool_calls is not None:
             first_tool_call = step_log.tool_calls[0]
             used_code = first_tool_call.name == "python_interpreter"
-            parent_id = f"call_{len(step_log.tool_calls)}"
 
             # Tool call becomes the parent message with timing info
             # First we will handle arguments based on type
@@ -89,7 +94,6 @@ def pull_messages_from_step(
                 content=content,
                 metadata={
                     "title": f"🛠️ Used tool {first_tool_call.name}",
-                    "id": parent_id,
                     "status": "done",
                 },
             )
@@ -128,15 +132,21 @@ def pull_messages_from_step(
 
         # Handle standalone errors but not from tool calls
         if hasattr(step_log, "error") and step_log.error is not None:
-            yield gr.ChatMessage(role="assistant", content=str(step_log.error), metadata={"title": "💥 Error"})
+            yield gr.ChatMessage(
+                role="assistant", content=str(step_log.error), metadata={"title": "💥 Error", "status": "done"}
+            )
 
-        yield gr.ChatMessage(role="assistant", content=get_step_footnote_content(step_log, step_number))
+        yield gr.ChatMessage(
+            role="assistant", content=get_step_footnote_content(step_log, step_number), metadata={"status": "done"}
+        )
         yield gr.ChatMessage(role="assistant", content="-----", metadata={"status": "done"})
 
     elif isinstance(step_log, PlanningStep):
-        yield gr.ChatMessage(role="assistant", content="**Planning step**")
-        yield gr.ChatMessage(role="assistant", content=step_log.plan)
-        yield gr.ChatMessage(role="assistant", content=get_step_footnote_content(step_log, "Planning step"))
+        yield gr.ChatMessage(role="assistant", content="**Planning step**", metadata={"status": "done"})
+        yield gr.ChatMessage(role="assistant", content=step_log.plan, metadata={"status": "done"})
+        yield gr.ChatMessage(
+            role="assistant", content=get_step_footnote_content(step_log, "Planning step"), metadata={"status": "done"}
+        )
         yield gr.ChatMessage(role="assistant", content="-----", metadata={"status": "done"})
 
     elif isinstance(step_log, FinalAnswerStep):
@@ -145,19 +155,24 @@ def pull_messages_from_step(
             yield gr.ChatMessage(
                 role="assistant",
                 content=f"**Final answer:**\n{final_answer.to_string()}\n",
+                metadata={"status": "done"},
             )
         elif isinstance(final_answer, AgentImage):
             yield gr.ChatMessage(
                 role="assistant",
                 content={"path": final_answer.to_string(), "mime_type": "image/png"},
+                metadata={"status": "done"},
             )
         elif isinstance(final_answer, AgentAudio):
             yield gr.ChatMessage(
                 role="assistant",
                 content={"path": final_answer.to_string(), "mime_type": "audio/wav"},
+                metadata={"status": "done"},
             )
         else:
-            yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(final_answer)}")
+            yield gr.ChatMessage(
+                role="assistant", content=f"**Final answer:** {str(final_answer)}", metadata={"status": "done"}
+            )
 
     else:
         raise ValueError(f"Unsupported step type: {type(step_log)}")
@@ -174,6 +189,13 @@ def stream_to_gradio(
     total_input_tokens = 0
     total_output_tokens = 0
 
+    if not _is_package_available("gradio"):
+        raise ModuleNotFoundError(
+            "Please install 'gradio' extra to use the GradioUI: `pip install 'smolagents[gradio]'`"
+        )
+
+    intermediate_text = ""
+
     for step_log in agent.run(
         task, images=task_images, stream=True, reset=reset_agent_memory, additional_args=additional_args
     ):
@@ -185,10 +207,17 @@ def stream_to_gradio(
                 step_log.input_token_count = agent.model.last_input_token_count
                 step_log.output_token_count = agent.model.last_output_token_count
 
-        for message in pull_messages_from_step(
-            step_log,
-        ):
-            yield message
+        if isinstance(step_log, MemoryStep):
+            intermediate_text = ""
+            for message in pull_messages_from_step(
+                step_log,
+                # If we're streaming model outputs, no need to display them twice
+                skip_model_outputs=getattr(agent, "stream_outputs", False),
+            ):
+                yield message
+        elif isinstance(step_log, ChatMessageStreamDelta):
+            intermediate_text += step_log.content or ""
+            yield intermediate_text
 
 
 class GradioUI:
@@ -200,12 +229,12 @@ class GradioUI:
                 "Please install 'gradio' extra to use the GradioUI: `pip install 'smolagents[gradio]'`"
             )
         self.agent = agent
-        self.file_upload_folder = file_upload_folder
+        self.file_upload_folder = Path(file_upload_folder) if file_upload_folder is not None else None
         self.name = getattr(agent, "name") or "Agent interface"
         self.description = getattr(agent, "description", None)
         if self.file_upload_folder is not None:
-            if not os.path.exists(file_upload_folder):
-                os.mkdir(file_upload_folder)
+            if not self.file_upload_folder.exists():
+                self.file_upload_folder.mkdir(parents=True, exist_ok=True)
 
     def interact_with_agent(self, prompt, messages, session_state):
         import gradio as gr
@@ -215,11 +244,22 @@ class GradioUI:
             session_state["agent"] = self.agent
 
         try:
-            messages.append(gr.ChatMessage(role="user", content=prompt))
+            messages.append(gr.ChatMessage(role="user", content=prompt, metadata={"status": "done"}))
             yield messages
 
             for msg in stream_to_gradio(session_state["agent"], task=prompt, reset_agent_memory=False):
-                messages.append(msg)
+                if isinstance(msg, gr.ChatMessage):
+                    messages.append(msg)
+                elif isinstance(msg, str):  # Then it's only a completion delta
+                    try:
+                        if messages[-1].metadata["status"] == "pending":
+                            messages[-1].content = msg
+                        else:
+                            messages.append(
+                                gr.ChatMessage(role="assistant", content=msg, metadata={"status": "pending"})
+                            )
+                    except Exception as e:
+                        raise e
                 yield messages
 
             yield messages
@@ -309,12 +349,9 @@ class GradioUI:
                         [upload_status, file_uploads_log],
                     )
 
-                gr.HTML("<br><br><h4><center>Powered by:</center></h4>")
-                with gr.Row():
-                    gr.HTML("""<div style="display: flex; align-items: center; gap: 8px; font-family: system-ui, -apple-system, sans-serif;">
-            <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png" style="width: 32px; height: 32px; object-fit: contain;" alt="logo">
-            <a target="_blank" href="https://github.com/huggingface/smolagents"><b>huggingface/smolagents</b></a>
-            </div>""")
+                gr.HTML(
+                    "<br><br><h4><center>Powered by <a target='_blank' href='https://github.com/huggingface/smolagents'><b>smolagents</b></a></center></h4>"
+                )
 
             # Main chat interface
             chatbot = gr.Chatbot(


Test patch:
diff --git a/tests/test_agents.py b/tests/test_agents.py
index 826c2f2..31e188a 100644
--- a/tests/test_agents.py
+++ b/tests/test_agents.py
@@ -17,9 +17,9 @@ import os
 import tempfile
 import unittest
 import uuid
+from collections.abc import Generator
 from contextlib import nullcontext as does_not_raise
 from pathlib import Path
-from typing import Any
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -620,8 +620,8 @@ class MockAgent:
 
 
 class DummyMultiStepAgent(MultiStepAgent):
-    def step(self, memory_step: ActionStep) -> None | Any:
-        return super().step(memory_step)
+    def step(self, memory_step: ActionStep) -> Generator[None]:
+        yield None
 
     def initialize_system_prompt(self):
         pass
@@ -682,7 +682,7 @@ class TestMultiStepAgent:
         fake_model.last_input_token_count = 10
         fake_model.last_output_token_count = 20
         max_steps = 2
-        agent = DummyMultiStepAgent(tools=[], model=fake_model, max_steps=max_steps)
+        agent = CodeAgent(tools=[], model=fake_model, max_steps=max_steps)
         assert hasattr(agent, "step_number"), "step_number attribute should be defined"
         assert agent.step_number == 0, "step_number should be initialized to 0"
         agent.run("Test task")
@@ -719,7 +719,8 @@ class TestMultiStepAgent:
             model=fake_model,
         )
         task = "Test task"
-        planning_step = agent._generate_planning_step(task, is_first_step=(step == 1), step=step)
+
+        planning_step = list(agent._generate_planning_step(task, is_first_step=(step == 1), step=step))[-1]
         expected_message_texts = {
             "INITIAL_PLAN_USER_PROMPT": populate_template(
                 agent.prompt_templates["planning"]["initial_plan"],
@@ -764,8 +765,8 @@ class TestMultiStepAgent:
             for content, expected_content in zip(message["content"], expected_message["content"]):
                 assert content == expected_content
         # Test calls to model
-        assert len(fake_model.call_args_list) == 1
-        for call_args, expected_messages in zip(fake_model.call_args_list, expected_messages_list):
+        assert len(fake_model.generate.call_args_list) == 1
+        for call_args, expected_messages in zip(fake_model.generate.call_args_list, expected_messages_list):
             assert len(call_args.args) == 1
             messages = call_args.args[0]
             assert isinstance(messages, list)
diff --git a/tests/test_monitoring.py b/tests/test_monitoring.py
index 41bbc8b..8d2f50c 100644
--- a/tests/test_monitoring.py
+++ b/tests/test_monitoring.py
@@ -137,7 +137,7 @@ class MonitoringTester(unittest.TestCase):
         # Use stream_to_gradio to capture the output
         outputs = list(stream_to_gradio(agent, task="Test task"))
 
-        self.assertEqual(len(outputs), 11)
+        self.assertEqual(len(outputs), 10)
         plan_message = outputs[1]
         self.assertEqual(plan_message.role, "assistant")
         self.assertIn("Code:", plan_message.content)
@@ -161,7 +161,7 @@ class MonitoringTester(unittest.TestCase):
             )
         )
 
-        self.assertEqual(len(outputs), 6)
+        self.assertEqual(len(outputs), 5)
         final_message = outputs[-1]
         self.assertEqual(final_message.role, "assistant")
         self.assertIsInstance(final_message.content, dict)
@@ -182,7 +182,7 @@ class MonitoringTester(unittest.TestCase):
         # Use stream_to_gradio to capture the output
         outputs = list(stream_to_gradio(agent, task="Test task"))
 
-        self.assertEqual(len(outputs), 13)
+        self.assertEqual(len(outputs), 11)
         final_message = outputs[-1]
         self.assertEqual(final_message.role, "assistant")
         self.assertIn("Malformed call", final_message.content)