inspect-ai 0.3.81__py3-none-any.whl → 0.3.83__py3-none-any.whl

inspect_ai/_eval/eval.py

@@ -2,9 +2,11 @@ import logging
 import os
 import sys
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 from inspect_ai._util.notgiven import NOT_GIVEN, NotGiven
+from inspect_ai.agent._agent import Agent, is_agent
+from inspect_ai.agent._as_solver import as_solver
 
 if sys.version_info < (3, 11):
     from exceptiongroup import ExceptionGroup
@@ -15,7 +17,11 @@ from typing_extensions import Unpack
 from inspect_ai._cli.util import parse_cli_args
 from inspect_ai._display.core.active import display as task_display
 from inspect_ai._util.config import resolve_args
-from inspect_ai._util.constants import DEFAULT_LOG_FORMAT
+from inspect_ai._util.constants import (
+    DEFAULT_LOG_FORMAT,
+    DEFAULT_LOG_SHARED,
+    JSON_LOG_FORMAT,
+)
 from inspect_ai._util.error import PrerequisiteError
 from inspect_ai._util.file import absolute_file_path
 from inspect_ai._util.logger import warn_once
@@ -31,6 +37,7 @@ from inspect_ai.approval._policy import (
 from inspect_ai.log import EvalConfig, EvalLog, EvalLogInfo
 from inspect_ai.log._file import read_eval_log_async
 from inspect_ai.log._recorders import create_recorder_for_format
+from inspect_ai.log._recorders.buffer import cleanup_sample_buffers
 from inspect_ai.model import (
     GenerateConfig,
     GenerateConfigArgs,
@@ -66,7 +73,7 @@ def eval(
     task_args: dict[str, Any] | str = dict(),
     sandbox: SandboxEnvironmentType | None = None,
     sandbox_cleanup: bool | None = None,
-    solver: Solver | list[Solver] | SolverSpec | None = None,
+    solver: Solver | SolverSpec | Agent | list[Solver] | None = None,
     tags: list[str] | None = None,
     metadata: dict[str, Any] | None = None,
     trace: bool | None = None,
@@ -92,6 +99,7 @@ def eval(
     log_samples: bool | None = None,
     log_images: bool | None = None,
     log_buffer: int | None = None,
+    log_shared: bool | int | None = None,
     score: bool = True,
     score_display: bool | None = None,
     **kwargs: Unpack[GenerateConfigArgs],
@@ -161,6 +169,9 @@ def eval(
         log_buffer: Number of samples to buffer before writing log file.
             If not specified, an appropriate default for the format and filesystem is
             chosen (10 for most all cases, 100 for JSON logs on remote filesystems).
+        log_shared: Sync sample events to log directory so that users on other systems
+            can see log updates in realtime (defaults to no syncing). Specify `True`
+            to sync every 10 seconds, otherwise an integer to sync every `n` seconds.
         score: Score output (defaults to True)
         score_display: Show scoring metrics in realtime (defaults to True)
         **kwargs: Model generation options.
@@ -210,6 +221,7 @@ def eval(
                 log_samples=log_samples,
                 log_images=log_images,
                 log_buffer=log_buffer,
+                log_shared=log_shared,
                 score=score,
                 score_display=score_display,
                 **kwargs,
@@ -236,7 +248,7 @@ async def eval_async(
     task_args: dict[str, Any] | str = dict(),
     sandbox: SandboxEnvironmentType | None = None,
     sandbox_cleanup: bool | None = None,
-    solver: Solver | list[Solver] | SolverSpec | None = None,
+    solver: Solver | SolverSpec | Agent | list[Solver] | None = None,
     tags: list[str] | None = None,
     metadata: dict[str, Any] | None = None,
     approval: str | list[ApprovalPolicy] | ApprovalPolicyConfig | None = None,
@@ -260,6 +272,7 @@ async def eval_async(
     log_samples: bool | None = None,
     log_images: bool | None = None,
     log_buffer: int | None = None,
+    log_shared: bool | int | None = None,
     score: bool = True,
     score_display: bool | None = None,
     **kwargs: Unpack[GenerateConfigArgs],
@@ -312,6 +325,7 @@ async def eval_async(
         log_buffer: Number of samples to buffer before writing log file.
            If not specified, an appropriate default for the format and filesystem is
            chosen (10 for most all cases, 100 for JSON logs on remote filesystems).
+        log_shared: Indicate that the log directory is shared, which results in additional syncing of realtime log data for Inspect View.
         score: Score output (defaults to True)
         score_display: Show scoring metrics in realtime (defaults to True)
         **kwargs: Model generation options.
@@ -341,13 +355,10 @@ async def eval_async(
 
     try:
         # intialise eval
-        model, approval, resolved_tasks = eval_init(
-            tasks=tasks,
+        model, approval = eval_init(
             model=model,
             model_base_url=model_base_url,
             model_args=model_args,
-            task_args=task_args,
-            sandbox=sandbox,
             approval=approval,
             max_subprocesses=max_subprocesses,
             log_level=log_level,
@@ -355,6 +366,11 @@ async def eval_async(
             **kwargs,
         )
 
+        # resolve tasks
+        resolved_tasks = eval_resolve_tasks(
+            tasks, task_args, model, GenerateConfig(**kwargs), sandbox
+        )
+
         # warn and return empty string if we resolved no tasks
         if len(resolved_tasks) == 0:
             log.warning("No inspect tasks were found at the specified paths.")
@@ -390,8 +406,22 @@ async def eval_async(
                 f"ERROR: You do not have write permission for the log_dir '{log_dir}'"
             )
 
+        # resolve log_shared
+        log_shared = DEFAULT_LOG_SHARED if log_shared is True else log_shared
+
+        # validate that --log-shared can't use used with 'json' format
+        if log_shared and log_format == JSON_LOG_FORMAT:
+            raise PrerequisiteError(
+                "ERROR: --log-shared is not compatible with the json log format."
+            )
+
         # resolve solver
-        solver = chain(solver) if isinstance(solver, list) else solver
+        if isinstance(solver, list):
+            solver = chain(solver)
+        elif is_agent(solver):
+            solver = as_solver(solver)
+        else:
+            solver = cast(Solver | SolverSpec | None, solver)
 
         # ensure consistency of limit and sample_id
         if sample_id is not None and limit is not None:
@@ -426,6 +456,7 @@ async def eval_async(
             log_samples=log_samples,
             log_images=log_images,
             log_buffer=log_buffer,
+            log_shared=log_shared,
             score_display=score_display,
         )
 
@@ -485,6 +516,9 @@ async def eval_async(
             )
             logs = EvalLogs(results)
 
+        # cleanup sample buffers if required
+        cleanup_sample_buffers(log_dir)
+
     finally:
         _eval_async_running = False
 
@@ -510,6 +544,7 @@ def eval_retry(
     log_samples: bool | None = None,
     log_images: bool | None = None,
     log_buffer: int | None = None,
+    log_shared: bool | int | None = None,
     score: bool = True,
     score_display: bool | None = None,
     max_retries: int | None = None,
@@ -551,6 +586,9 @@ def eval_retry(
         log_buffer: Number of samples to buffer before writing log file.
             If not specified, an appropriate default for the format and filesystem is
             chosen (10 for most all cases, 100 for JSON logs on remote filesystems).
+        log_shared: Sync sample events to log directory so that users on other systems
+            can see log updates in realtime (defaults to no syncing). Specify `True`
+            to sync every 10 seconds, otherwise an integer to sync every `n` seconds.
         score: Score output (defaults to True)
         score_display: Show scoring metrics in realtime (defaults to True)
         max_retries:
@@ -586,6 +624,7 @@ def eval_retry(
             log_samples=log_samples,
             log_images=log_images,
             log_buffer=log_buffer,
+            log_shared=log_shared,
             score=score,
             score_display=score_display,
             max_retries=max_retries,
@@ -612,6 +651,7 @@ async def eval_retry_async(
     log_samples: bool | None = None,
     log_images: bool | None = None,
     log_buffer: int | None = None,
+    log_shared: bool | int | None = None,
     score: bool = True,
     score_display: bool | None = None,
     max_retries: int | None = None,
@@ -651,6 +691,8 @@ async def eval_retry_async(
         log_buffer: (int | None): Number of samples to buffer before writing log file.
            If not specified, an appropriate default for the format and filesystem is
            chosen (10 for most all cases, 100 for JSON logs on remote filesystems).
+        log_shared: Indicate that the log directory is shared, which results in
+            additional syncing of realtime log data for Inspect View.
         score (bool): Score output (defaults to True)
         score_display (bool | None): Show scoring metrics in realtime (defaults to True)
         max_retries (int | None):
@@ -691,7 +733,7 @@ async def eval_retry_async(
         # context to reconstruct ephemeral Task instances)
         task: str | None
         task_id = eval_log.eval.task_id
-        task_name = eval_log.eval.task
+        task_name = eval_log.eval.task_registry_name or eval_log.eval.task
         task_file = eval_log.eval.task_file
         if task_file:
             if not Path(task_file).exists():
@@ -750,6 +792,9 @@ async def eval_retry_async(
         log_buffer = (
             log_buffer if log_buffer is not None else eval_log.eval.config.log_buffer
         )
+        log_shared = (
+            log_shared if log_shared is not None else eval_log.eval.config.log_shared
+        )
         score_display = (
             score_display
             if score_display is not None
@@ -796,6 +841,7 @@ async def eval_retry_async(
                 log_samples=log_samples,
                 log_images=log_images,
                 log_buffer=log_buffer,
+                log_shared=log_shared,
                 score=score,
                 score_display=score_display,
                 **dict(config),
@@ -809,24 +855,20 @@ async def eval_retry_async(
 
 
 def eval_init(
-    tasks: Tasks,
     model: str | Model | list[str] | list[Model] | None | NotGiven = NOT_GIVEN,
     model_base_url: str | None = None,
     model_args: dict[str, Any] | str = dict(),
-    task_args: dict[str, Any] | str = dict(),
-    sandbox: SandboxEnvironmentType | None = None,
     approval: str | list[ApprovalPolicy] | ApprovalPolicyConfig | None = None,
     max_subprocesses: int | None = None,
     log_level: str | None = None,
     log_level_transcript: str | None = None,
     **kwargs: Unpack[GenerateConfigArgs],
-) -> tuple[list[Model], list[ApprovalPolicy] | None, list[ResolvedTask]]:
+) -> tuple[list[Model], list[ApprovalPolicy] | None]:
     # init eval context
     init_eval_context(log_level, log_level_transcript, max_subprocesses)
 
     # resolve model and task args
     model_args = resolve_args(model_args)
-    task_args = resolve_args(task_args)
 
     # resolve model args from environment if not specified
     if len(model_args) == 0:
@@ -839,21 +881,28 @@ def eval_init(
     generate_config = GenerateConfig(**kwargs)
     models = resolve_models(model, model_base_url, model_args, generate_config)
 
-    # resolve tasks (set active model to resolve uses of the
-    # 'default' model in tools, solvers, and scorers)
-
-    with task_display().suspend_task_app():
-        resolved_tasks: list[ResolvedTask] = []
-        for m in models:
-            init_active_model(m, generate_config)
-            resolved_tasks.extend(resolve_tasks(tasks, task_args, m, sandbox))
-
     # resolve approval
     if isinstance(approval, str | ApprovalPolicyConfig):
         approval = approval_policies_from_config(approval)
     init_tool_approval(approval)
 
-    return models, approval, resolved_tasks
+    return models, approval
+
+
+def eval_resolve_tasks(
+    tasks: Tasks,
+    task_args: dict[str, Any] | str,
+    models: list[Model],
+    config: GenerateConfig,
+    sandbox: SandboxEnvironmentType | None,
+) -> list[ResolvedTask]:
+    task_args = resolve_args(task_args)
+    with task_display().suspend_task_app():
+        resolved_tasks: list[ResolvedTask] = []
+        for m in models:
+            init_active_model(m, config)
+            resolved_tasks.extend(resolve_tasks(tasks, task_args, m, sandbox))
+        return resolved_tasks
 
 
 def init_eval_display(

inspect_ai/_eval/evalset.py

@@ -1,6 +1,5 @@
 import hashlib
 import logging
-from copy import deepcopy
 from typing import Any, Literal, NamedTuple, Set, cast
 
 import rich
@@ -18,6 +17,7 @@ from typing_extensions import Unpack
 from inspect_ai._util.error import PrerequisiteError
 from inspect_ai._util.file import basename, filesystem
 from inspect_ai._util.notgiven import NOT_GIVEN, NotGiven
+from inspect_ai.agent._agent import Agent
 from inspect_ai.approval._policy import ApprovalPolicy
 from inspect_ai.log import EvalLog
 from inspect_ai.log._bundle import bundle_log_dir
@@ -37,7 +37,7 @@ from inspect_ai.solver._solver import Solver, SolverSpec
 from inspect_ai.util import DisplayType, SandboxEnvironmentType
 from inspect_ai.util._display import display_type_initialized, init_display_type
 
-from .eval import eval, eval_init
+from .eval import eval, eval_init, eval_resolve_tasks
 from .loader import resolve_task_args
 from .task import Epochs
 from .task.resolved import ResolvedTask
@@ -66,7 +66,7 @@ def eval_set(
     task_args: dict[str, Any] | str = dict(),
     sandbox: SandboxEnvironmentType | None = None,
     sandbox_cleanup: bool | None = None,
-    solver: Solver | list[Solver] | SolverSpec | None = None,
+    solver: Solver | SolverSpec | Agent | list[Solver] | None = None,
     tags: list[str] | None = None,
     metadata: dict[str, Any] | None = None,
     trace: bool | None = None,
@@ -92,6 +92,7 @@ def eval_set(
     log_samples: bool | None = None,
     log_images: bool | None = None,
     log_buffer: int | None = None,
+    log_shared: bool | int | None = None,
     bundle_dir: str | None = None,
     bundle_overwrite: bool = False,
     **kwargs: Unpack[GenerateConfigArgs],
@@ -171,6 +172,9 @@ def eval_set(
         log_buffer: Number of samples to buffer before writing log file.
             If not specified, an appropriate default for the format and filesystem is
             chosen (10 for most all cases, 100 for JSON logs on remote filesystems).
+        log_shared: Sync sample events to log directory so that users on other systems
+            can see log updates in realtime (defaults to no syncing). Specify `True`
+            to sync every 10 seconds, otherwise an integer to sync every `n` seconds.
         bundle_dir: If specified, the log viewer and logs generated
             by this eval set will be bundled into this directory.
         bundle_overwrite: Whether to overwrite files in the bundle_dir.
@@ -219,6 +223,7 @@ def eval_set(
             log_samples=log_samples,
             log_images=log_images,
             log_buffer=log_buffer,
+            log_shared=log_shared,
             score=score,
             **kwargs,
         )
@@ -242,29 +247,21 @@ def eval_set(
     if display == "conversation":
         raise RuntimeError("eval_set cannot be used with conversation display.")
 
-    # resolve tasks
-    models, _, resolved_tasks = eval_init(
-        tasks=tasks,
+    # initialize eval
+    models, _ = eval_init(
         model=model,
         model_base_url=model_base_url,
         model_args=model_args,
-        task_args=task_args,
-        sandbox=sandbox,
         max_subprocesses=max_subprocesses,
         log_level=log_level,
         log_level_transcript=log_level_transcript,
         **kwargs,
     )
 
-    # ensure log_dir and list all logs
+    # ensure log_dir
     fs = filesystem(log_dir)
     fs.mkdir(log_dir, exist_ok=True)
 
-    # validate that:
-    #  (1) All tasks have a unique identifier
-    #  (2) All logs have identifiers that map to tasks
-    validate_eval_set_prerequisites(resolved_tasks, list_all_eval_logs(log_dir))
-
     # resolve some parameters
     retry_connections = retry_connections or 0.5
     retry_cleanup = retry_cleanup is not False
@@ -305,11 +302,21 @@ def eval_set(
     #   - tasks with a successful log (they'll just be returned)
     #   - tasks with failed logs (they'll be retried)
     def try_eval() -> list[EvalLog]:
+        # resolve tasks
+        resolved_tasks = eval_resolve_tasks(
+            tasks, task_args, models, GenerateConfig(**kwargs), sandbox
+        )
+
         # list all logs currently in the log directory (update manifest if there are some)
         all_logs = list_all_eval_logs(log_dir)
         if len(all_logs) > 0:
             write_log_dir_manifest(log_dir)
 
+        # validate that:
+        #  (1) All tasks have a unique identifier
+        #  (2) All logs have identifiers that map to tasks
+        validate_eval_set_prerequisites(resolved_tasks, all_logs)
+
         # see which tasks are yet to run (to complete successfully we need
         # a successful eval for every [task_file/]task_name/model combination)
         # for those that haven't run, schedule them into models => tasks groups
@@ -414,13 +421,10 @@ def as_previous_tasks(
         # want to bring this back but we'd need to resolve the
         # directory issues.
 
-        # deepcopy so the same instance is not run twice
-        prev_task = deepcopy(task.task)
-
         previous_tasks.append(
             PreviousTask(
                 id=log.header.eval.task_id,
-                task=prev_task,
+                task=task.task,
                 task_args=resolve_task_args(task.task),
                 model=task.model,
                 log=read_eval_log(log.info),

inspect_ai/_eval/loader.py

@@ -26,6 +26,8 @@ from inspect_ai._util.registry import (
     registry_lookup,
     registry_params,
 )
+from inspect_ai.agent._agent import Agent
+from inspect_ai.agent._as_solver import as_solver
 from inspect_ai.model import Model
 from inspect_ai.scorer._scorer import Scorer, ScorerSpec, scorer_create
 from inspect_ai.solver._bridge import bridge
@@ -421,20 +423,32 @@ def solver_from_spec(spec: SolverSpec) -> Solver:
         if solver_file is None:
             if solver_name is None:
                 raise ValueError(f"Unable to resolve solver name from {spec.solver}")
-            return cast(Solver, registry_create("solver", solver_name, **spec.args))
+            elif registry_lookup("solver", solver_name) is not None:
+                return cast(Solver, registry_create("solver", solver_name, **spec.args))
+            elif registry_lookup("agent", solver_name) is not None:
+                agent = cast(Agent, registry_create("agent", solver_name, **spec.args))
+                return as_solver(agent)
+            else:
+                raise ValueError(
+                    f"Unkonwn solver {solver_name} (not registered as a @solver or @agent)"
+                )
 
         # we do have a solver file
         else:
             # load the module and parse decorators
             solver_module = load_module(solver_file)
-            decorators = parse_decorators(solver_file, "solver")
+            solver_decorators = parse_decorators(solver_file, "solver")
+            agent_decorators = parse_decorators(solver_file, "agent")
 
             # if there is no solver_name see if we can discover it
             if solver_name is None:
-                if len(decorators) == 1:
+                if len(solver_decorators) == 1:
                     # decorator based solver
-                    solver_name = decorators[0][0]
-                elif len(decorators) == 0:
+                    solver_name = solver_decorators[0][0]
+                elif len(agent_decorators) == 1:
+                    # decorator based agent
+                    solver_name = agent_decorators[0][0]
+                elif len(solver_decorators) == 0 and len(agent_decorators) == 0:
                     # see if we can find an agent based solver
                     functions = [
                         function
@@ -454,26 +468,35 @@ def solver_from_spec(spec: SolverSpec) -> Solver:
 
                     elif len(agent_functions) == 0:
                         raise PrerequisiteError(
-                            f"The source file {pretty_solver_file} does not contain any @solver functions or agent functions."
+                            f"The source file {pretty_solver_file} does not contain any @solver, @agent or bridged agent functions."
                         )
                     else:
                         raise PrerequisiteError(
-                            f"The source file {pretty_solver_file} has more than one agent function (qualify which agent using e.g. '{solver_file.name}@agent_fn')"
+                            f"The source file {pretty_solver_file} has more than one bridged agent function (qualify which agent using e.g. '{solver_file.name}@agent_fn')"
                         )
-                else:
+                elif len(solver_decorators) > 1:
                     raise PrerequisiteError(
                         f"The source file {pretty_solver_file} has more than one @solver function (qualify which solver using e.g. '{solver_file.name}y@solver_fn')"
                     )
+                else:
+                    raise PrerequisiteError(
+                        f"The source file {pretty_solver_file} has more than one @agent function (qualify which agent using e.g. '{solver_file.name}y@agent_fn')"
+                    )
 
             # create decorator based solvers using the registry
-            if any(solver[0] == solver_name for solver in decorators):
+            if any(solver[0] == solver_name for solver in solver_decorators):
                 return cast(Solver, registry_create("solver", solver_name, **spec.args))
 
-            # create agent based solvers by calling the function and wrapping it in bridge()
+            # create decorator based agents using the registry
+            elif any(agent[0] == solver_name for agent in agent_decorators):
+                agent = cast(Agent, registry_create("agent", solver_name, **spec.args))
+                return as_solver(agent)
+
+            # create bridge based solvers by calling the function and wrapping it in bridge()
             else:
                 agent_fn = getattr(solver_module, solver_name, None)
                 if inspect.isfunction(agent_fn):
-                    return bridge.bridge(agent_fn(**spec.args))
+                    return bridge(agent_fn(**spec.args))
                 elif agent_fn is not None:
                     raise PrerequisiteError(
                         f"The object {solver_name} in file {pretty_solver_file} is not a Python function."

inspect_ai/_eval/run.py

@@ -1,4 +1,3 @@
-import functools
 import logging
 import os
 import sys
@@ -20,7 +19,6 @@ from inspect_ai._display.core.active import (
     init_task_screen,
 )
 from inspect_ai._display.core.display import TaskSpec
-from inspect_ai._util._async import tg_collect
 from inspect_ai._util.error import PrerequisiteError, exception_message
 from inspect_ai._util.path import chdir
 from inspect_ai._util.registry import registry_unqualified_name
@@ -195,6 +193,7 @@ async def eval_run(
                     task_name=task.name,
                     task_version=task.version,
                     task_file=resolved_task.task_file,
+                    task_registry_name=resolved_task.task.registry_name,
                     task_id=resolved_task.id if resolved_task.id else uuid(),
                     run_id=run_id,
                     solver=eval_solver_spec,
@@ -359,17 +358,13 @@ async def run_multiple(tasks: list[TaskRunOptions], parallel: int) -> list[EvalL
                         "Run Task",
                         f"task: {task_options.task.name} ({task_options.model})",
                     ):
-                        tg_results = await tg_collect(
-                            [functools.partial(task_run, task_options)]
-                        )
-                    # check for empty results list (indicates cancellation)
-                    if len(tg_results) == 0:
-                        # task was cancelled, break out of the worker loop
-                        result = None
-
-                    else:
-                        result = tg_results[0]
-                        results.append(result)
+                        async with anyio.create_task_group() as tg:
+
+                            async def run_task() -> None:
+                                result = await task_run(task_options)
+                                results.append(result)
+
+                            tg.start_soon(run_task)
 
                 except Exception as ex:
                     # errors generally don't escape from tasks (the exception being if an error
@@ -407,12 +402,15 @@ async def run_multiple(tasks: list[TaskRunOptions], parallel: int) -> list[EvalL
         # Use anyio task group instead of manual task management
         try:
             async with anyio.create_task_group() as tg:
+                # computer number of workers (never more than total_tasks)
+                num_workers = min(parallel, total_tasks)
+
                 # start worker tasks
-                for _ in range(parallel):
+                for _ in range(num_workers):
                     tg.start_soon(worker)
 
                 # enqueue initial set of tasks
-                for _ in range(min(parallel, total_tasks)):
+                for _ in range(num_workers):
                     await enque_next_task()
         except anyio.get_cancelled_exc_class():
             pass

inspect_ai/_eval/score.py

@@ -7,8 +7,8 @@ import anyio
 
 from inspect_ai._display import display
 from inspect_ai._eval.loader import scorer_from_spec
-from inspect_ai._util._async import tg_collect
-from inspect_ai._util.platform import platform_init
+from inspect_ai._util._async import configured_async_backend, run_coroutine, tg_collect
+from inspect_ai._util.platform import platform_init, running_in_notebook
 from inspect_ai._util.registry import registry_create, registry_unqualified_name
 from inspect_ai.log import (
     EvalLog,
@@ -56,7 +56,17 @@ def score(
     # resolve scorers into a list
     scorers = [scorers] if isinstance(scorers, Scorer) else scorers
 
-    return anyio.run(score_async, log, scorers, epochs_reducer, action)
+    if running_in_notebook():
+        return run_coroutine(score_async(log, scorers, epochs_reducer, action))
+    else:
+        return anyio.run(
+            score_async,
+            log,
+            scorers,
+            epochs_reducer,
+            action,
+            backend=configured_async_backend(),
+        )
 
 
 async def score_async(

inspect_ai/_eval/task/generate.py

@@ -1,12 +1,8 @@
 from typing import Literal
 
-from inspect_ai.model import (
-    CachePolicy,
-    GenerateConfig,
-    Model,
-    call_tools,
-)
+from inspect_ai.model import CachePolicy, GenerateConfig, Model
 from inspect_ai.model._cache import epoch
+from inspect_ai.model._call_tools import execute_tools
 from inspect_ai.solver import TaskState
 from inspect_ai.solver._limit import SampleLimitExceededError
 from inspect_ai.tool import ToolFunction
@@ -48,10 +44,13 @@ async def task_generate(
 
             # resolve tool calls if necessary
             if tool_calls != "none" and message.tool_calls:
-                # call tools and append messages to state
-                state.messages.extend(
-                    await call_tools(message, state.tools, config.max_tool_output)
+                # call tools and update messages and output
+                messages, output = await execute_tools(
+                    state.messages, state.tools, config.max_tool_output
                 )
+                state.messages.extend(messages)
+                if output is not None:
+                    state.output = output
 
                 # check for completed or only executing a single tool call
                 if state.completed or tool_calls == "single":