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

inspect_ai/_eval/task/log.py

@@ -4,6 +4,7 @@ from typing import Any, Iterator, Literal, cast
 
 from shortuuid import uuid
 
+from inspect_ai._display.core.display import TaskDisplayMetric
 from inspect_ai._eval.task.util import slice_dataset
 from inspect_ai._util.constants import PKG_NAME
 from inspect_ai._util.datetime import iso_now
@@ -34,6 +35,9 @@ from inspect_ai.log._log import (
     eval_config_defaults,
 )
 from inspect_ai.log._recorders import Recorder
+from inspect_ai.log._recorders.buffer import SampleBufferDatabase
+from inspect_ai.log._recorders.types import SampleEvent, SampleSummary
+from inspect_ai.log._transcript import Event
 from inspect_ai.model import (
     GenerateConfig,
     Model,
@@ -53,6 +57,7 @@ class TaskLogger:
         task_name: str,
         task_version: int,
         task_file: str | None,
+        task_registry_name: str | None,
         task_id: str | None,
         run_id: str,
         solver: SolverSpec | None,
@@ -127,6 +132,7 @@ class TaskLogger:
             task_id=task_id if task_id else uuid(),
             task_version=task_version,
             task_file=task_file,
+            task_registry_name=task_registry_name,
             task_attribs=task_attribs,
             task_args=task_args,
             solver=solver.solver if solver else None,
@@ -159,10 +165,15 @@ class TaskLogger:
 
         # size of flush buffer (how many samples we buffer before hitting storage)
         self.flush_buffer = eval_config.log_buffer or recorder.default_log_buffer()
-        self.flush_pending = 0
+        self.flush_pending: list[tuple[str | int, int]] = []
 
     async def init(self) -> None:
         self._location = await self.recorder.log_init(self.eval)
+        self._buffer_db = SampleBufferDatabase(
+            location=self._location,
+            log_images=self.eval.config.log_images is not False,
+            log_shared=self.eval.config.log_shared,
+        )
 
     @property
     def location(self) -> str:
@@ -174,22 +185,53 @@ class TaskLogger:
 
     async def log_start(self, plan: EvalPlan) -> None:
         await self.recorder.log_start(self.eval, plan)
+        await self.recorder.flush(self.eval)
+
+    async def start_sample(self, sample: SampleSummary) -> None:
+        self._buffer_db.start_sample(sample)
+
+    def log_sample_event(self, id: str | int, epoch: int, event: Event) -> None:
+        # log the sample event
+        self._buffer_db.log_events([SampleEvent(id=id, epoch=epoch, event=event)])
 
-    async def log_sample(self, sample: EvalSample, *, flush: bool) -> None:
+    async def complete_sample(self, sample: EvalSample, *, flush: bool) -> None:
         # log the sample
         await self.recorder.log_sample(self.eval, sample)
 
+        # mark complete
+        self._buffer_db.complete_sample(
+            SampleSummary(
+                id=sample.id,
+                epoch=sample.epoch,
+                input=sample.input,
+                target=sample.target,
+                completed=True,
+                scores=sample.scores,
+                error=sample.error.message if sample.error is not None else None,
+                limit=f"{sample.limit.type}" if sample.limit is not None else None,
+            )
+        )
+
         # flush if requested
         if flush:
-            self.flush_pending += 1
-            if self.flush_pending >= self.flush_buffer:
+            self.flush_pending.append((sample.id, sample.epoch))
+            if len(self.flush_pending) >= self.flush_buffer:
+                # flush to disk
                 await self.recorder.flush(self.eval)
-                self.flush_pending = 0
+
+                # notify the event db it can remove these
+                self._buffer_db.remove_samples(self.flush_pending)
+
+                # Clear
+                self.flush_pending.clear()
 
         # track sucessful samples logged
         if sample.error is None:
             self._samples_completed += 1
 
+    def update_metrics(self, metrics: list[TaskDisplayMetric]) -> None:
+        self._buffer_db.update_metrics(metrics)
+
     async def log_finish(
         self,
         status: Literal["success", "cancelled", "error"],
@@ -198,10 +240,17 @@ class TaskLogger:
         reductions: list[EvalSampleReductions] | None = None,
         error: EvalError | None = None,
     ) -> EvalLog:
-        return await self.recorder.log_finish(
+        # finish and get log
+        log = await self.recorder.log_finish(
             self.eval, status, stats, results, reductions, error
         )
 
+        # cleanup the events db
+        self._buffer_db.cleanup()
+
+        # return log
+        return log
+
 
 async def log_start(
     logger: TaskLogger,

inspect_ai/_eval/task/run.py

@@ -19,7 +19,7 @@ from inspect_ai._display import (
     TaskSuccess,
     display,
 )
-from inspect_ai._display.core.display import TaskDisplay, TaskDisplayMetric
+from inspect_ai._display.core.display import TaskDisplayMetric
 from inspect_ai._util._async import tg_collect
 from inspect_ai._util.constants import (
     DEFAULT_EPOCHS,
@@ -29,6 +29,7 @@ from inspect_ai._util.constants import (
 from inspect_ai._util.datetime import iso_now
 from inspect_ai._util.error import exception_message
 from inspect_ai._util.hooks import send_telemetry
+from inspect_ai._util.json import to_json_str_safe
 from inspect_ai._util.registry import (
     is_registry_object,
     registry_log_name,
@@ -51,13 +52,17 @@ from inspect_ai.log import (
 from inspect_ai.log._condense import condense_sample
 from inspect_ai.log._file import eval_log_json_str
 from inspect_ai.log._log import EvalSampleLimit, EvalSampleReductions, eval_error
-from inspect_ai.log._samples import active_sample
+from inspect_ai.log._recorders.types import SampleSummary
+from inspect_ai.log._samples import (
+    active_sample,
+)
 from inspect_ai.log._transcript import (
     ErrorEvent,
     SampleInitEvent,
     SampleLimitEvent,
     ScoreEvent,
     StepEvent,
+    Transcript,
     transcript,
 )
 from inspect_ai.model import (
@@ -264,8 +269,13 @@ async def task_run(options: TaskRunOptions) -> EvalLog:
 
                     # track when samples complete and update progress as we go
                     progress_results: list[dict[str, SampleScore]] = []
+
+                    def update_metrics(metrics: list[TaskDisplayMetric]) -> None:
+                        td.update_metrics(metrics)
+                        logger.update_metrics(metrics)
+
                     update_metrics_display = update_metrics_display_fn(
-                        td,
+                        update_metrics,
                         display_metrics=profile.eval_config.score_display is not False,
                     )
 
@@ -423,7 +433,7 @@ async def task_run(options: TaskRunOptions) -> EvalLog:
 
 
 def update_metrics_display_fn(
-    td: TaskDisplay,
+    update_fn: Callable[[list[TaskDisplayMetric]], None],
     initial_interval: float = 0,
     min_interval: float = 0.9,
     display_metrics: bool = True,
@@ -463,7 +473,7 @@ def update_metrics_display_fn(
             )
 
             # Name, reducer, value
-            task_metrics = []
+            task_metrics: list[TaskDisplayMetric] = []
             if len(results.scores) > 0:
                 for score in results.scores:
                     for key, metric in score.metrics.items():
@@ -475,7 +485,7 @@ def update_metrics_display_fn(
                                 reducer=score.reducer,
                             )
                         )
-                td.update_metrics(task_metrics)
+                update_fn(task_metrics)
 
             # determine how long to wait before recomputing metrics
             time_end = time.perf_counter()
@@ -516,7 +526,7 @@ async def task_run_sample(
 
             # log if requested
             if logger:
-                await logger.log_sample(previous_sample, flush=False)
+                await logger.complete_sample(previous_sample, flush=False)
 
             # return score
             sample_scores = (
@@ -539,10 +549,19 @@ async def task_run_sample(
         semaphore if semaphore else contextlib.nullcontext()
     )
 
+    # validate that we have sample_id (mostly for the typechecker)
+    sample_id = sample.id
+    if sample_id is None:
+        raise ValueError("sample must have id to run")
+
     # initialise subtask and scoring context
     init_sample_model_usage()
     set_sample_state(state)
-    sample_transcript = init_subtask(SAMPLE_SUBTASK, state.store)
+    sample_transcript: Transcript = init_subtask(SAMPLE_SUBTASK, state.store)
+    if logger:
+        sample_transcript._subscribe(
+            lambda event: logger.log_sample_event(sample_id, state.epoch, event)
+        )
     if scorers:
         init_scoring_context(scorers, Target(sample.target))
 
@@ -626,6 +645,28 @@ async def task_run_sample(
                         # mark started
                         active.started = datetime.now().timestamp()
 
+                        if logger is not None:
+                            await logger.start_sample(
+                                SampleSummary(
+                                    id=sample_id,
+                                    epoch=state.epoch,
+                                    input=sample.input,
+                                    target=sample.target,
+                                )
+                            )
+
+                        # sample init event (remove file bodies as they have content or absolute paths)
+                        event_sample = sample.model_copy(
+                            update=dict(files={k: "" for k in sample.files.keys()})
+                            if sample.files
+                            else None
+                        )
+                        transcript()._event(
+                            SampleInitEvent(
+                                sample=event_sample, state=state_jsonable(state)
+                            )
+                        )
+
                         # set progress for plan then run it
                         state = await plan(state, generate)
 
@@ -824,7 +865,7 @@ async def log_sample(
     id = sample.id
     if id is None:
         raise ValueError(
-            f"Samples without IDs cannot be logged: {sample.model_dump_json()}"
+            f"Samples without IDs cannot be logged: {to_json_str_safe(sample)}"
         )
 
     # construct sample for logging
@@ -866,7 +907,7 @@ async def log_sample(
         limit=limit,
     )
 
-    await logger.log_sample(condense_sample(eval_sample, log_images), flush=True)
+    await logger.complete_sample(condense_sample(eval_sample, log_images), flush=True)
 
 
 async def resolve_dataset(

inspect_ai/_eval/task/task.py

@@ -1,4 +1,3 @@
-from copy import deepcopy
 from dataclasses import dataclass
 from logging import getLogger
 from typing import Any, Awaitable, Callable, Sequence, cast
@@ -9,6 +8,8 @@ from typing_extensions import TypedDict, Unpack
 from inspect_ai._util.logger import warn_once
 from inspect_ai._util.notgiven import NOT_GIVEN, NotGiven
 from inspect_ai._util.registry import is_registry_object, registry_info
+from inspect_ai.agent._agent import Agent, is_agent
+from inspect_ai.agent._as_solver import as_solver
 from inspect_ai.approval._policy import ApprovalPolicy, approval_policies_from_config
 from inspect_ai.dataset import Dataset, MemoryDataset, Sample
 from inspect_ai.log import EvalLog
@@ -47,7 +48,7 @@ class Task:
         self,
         dataset: Dataset | Sequence[Sample] | None = None,
         setup: Solver | list[Solver] | None = None,
-        solver: Solver | list[Solver] = generate(),
+        solver: Solver | Agent | list[Solver] = generate(),
         cleanup: Callable[[TaskState], Awaitable[None]] | None = None,
         scorer: Scorer | list[Scorer] | None = None,
         metrics: list[Metric] | dict[str, list[Metric]] | None = None,
@@ -158,6 +159,13 @@ class Task:
         else:
             return "task"
 
+    @property
+    def registry_name(self) -> str | None:
+        if is_registry_object(self):
+            return registry_info(self).name
+        else:
+            return None
+
     @property
     def attribs(self) -> dict[str, Any]:
         if is_registry_object(self):
@@ -191,8 +199,12 @@ def task_with(
 ) -> Task:
     """Task adapted with alternate values for one or more options.
 
+    This function modifies the passed task in place and returns it.
+    If you want to create multiple variations of a single task using
+    `task_with()` you should create the underlying task multiple times.
+
     Args:
-        task: Task to adapt (it is deep copied prior to mutating options)
+        task: Task to adapt
         dataset: Dataset to evaluate
         setup: Setup step (always run even when the main `solver` is replaced).
         solver: Solver or list of solvers. Defaults to generate(), a normal call to the model.
@@ -227,11 +239,8 @@ def task_with(
         metadata:  Additional metadata to associate with the task.
 
     Returns:
-        Task: Task adapted with alternate options.
+        Task: Passed `task` with modifications.
     """
-    # deep copy the task
-    task = deepcopy(task)
-
     if not isinstance(dataset, NotGiven):
         task.dataset = resolve_dataset(dataset)
     if not isinstance(setup, NotGiven):
@@ -340,8 +349,13 @@ def resolve_dataset(dataset: Dataset | Sequence[Sample] | None) -> Dataset:
     return dataset if isinstance(dataset, Dataset) else MemoryDataset(list(dataset))
 
 
-def resolve_solver(solver: Solver | list[Solver]) -> Solver:
-    return chain(solver) if isinstance(solver, list) else solver
+def resolve_solver(solver: Solver | Agent | list[Solver]) -> Solver:
+    if isinstance(solver, list):
+        return chain(solver)
+    elif is_agent(solver):
+        return as_solver(solver)
+    else:
+        return cast(Solver, solver)
 
 
 def resolve_model(model: str | Model | None) -> Model | None:

inspect_ai/_util/constants.py

@@ -25,8 +25,10 @@ ALL_LOG_LEVELS = [
 ]
 DEFAULT_LOG_LEVEL = "warning"
 DEFAULT_LOG_LEVEL_TRANSCRIPT = "info"
+DEFAULT_LOG_SHARED = 10
 ALL_LOG_FORMATS = ["eval", "json"]
 DEFAULT_LOG_FORMAT: Literal["eval", "json"] = "eval"
+JSON_LOG_FORMAT = "json"
 EVAL_LOG_FORMAT = "eval"
 DEFAULT_DISPLAY = "full"
 LOG_SCHEMA_VERSION = 2

inspect_ai/_util/file.py

@@ -13,7 +13,7 @@ from urllib.parse import urlparse
 import fsspec  # type: ignore  # type: ignore
 from fsspec.core import split_protocol  # type: ignore  # type: ignore
 from fsspec.implementations.local import make_path_posix  # type: ignore
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from s3fs import S3FileSystem  # type: ignore
 from shortuuid import uuid
 
@@ -158,6 +158,9 @@ class FileInfo(BaseModel):
     mtime: float | None
     """File modification time (None if the file is a directory on S3)."""
 
+    etag: str | None = Field(default=None)
+    """Etag (provided by some remote filesystems)"""
+
 
 class FileSystem:
     def __init__(self, fs: Any) -> None:
@@ -178,6 +181,9 @@ class FileSystem:
     ) -> None:
         self.fs.rm(path, recursive=recursive, maxdepth=maxdepth)
 
+    def mv(self, lpath: str, rpath: str) -> None:
+        self.fs.mv(lpath, rpath)
+
     def mkdir(self, path: str, exist_ok: bool = False) -> None:
         if self.is_s3():
             # try to avoid calling create_bucket on s3 filesystems (as that requires distinct
@@ -199,6 +205,9 @@ class FileSystem:
     def info(self, path: str, **kwargs: dict[str, Any]) -> FileInfo:
         return self._file_info(self.fs.info(path, **kwargs))
 
+    def path_as_uri(self, path: str) -> str:
+        return str(self.fs.unstrip_protocol(path))
+
     def ls(
         self, path: str, recursive: bool = False, **kwargs: dict[str, Any]
     ) -> list[FileInfo]:
@@ -267,11 +276,18 @@ class FileSystem:
         else:
             file["mtime"] = None
 
+        # S3 filesystems provided an ETag
+        if "ETag" in file.keys():
+            etag: str | None = file["ETag"].strip('"')
+        else:
+            etag = None
+
         return FileInfo(
             name=file["name"],
             type=file["type"],
             size=file["size"],
             mtime=file["mtime"],
+            etag=etag,
         )
 
 
@@ -306,6 +322,19 @@ def absolute_file_path(file: str) -> str:
     return file
 
 
+def to_uri(path_or_uri: str) -> str:
+    # Check if it's already a URI
+    parsed = urlparse(path_or_uri)
+
+    if parsed.scheme:
+        # Already has a scheme, return as is
+        return path_or_uri
+
+    # It's a file path, convert to URI
+    path_obj = Path(path_or_uri).absolute()
+    return path_obj.as_uri()
+
+
 def default_fs_options(file: str) -> dict[str, Any]:
     scheme = urlparse(file).scheme
     if (

inspect_ai/_util/json.py

@@ -6,7 +6,10 @@ from typing import (
 
 import jsonpatch
 from pydantic import BaseModel, Field, JsonValue
-from pydantic_core import to_jsonable_python
+from pydantic_core import to_json, to_jsonable_python
+
+JSONType = Literal["string", "integer", "number", "boolean", "array", "object", "null"]
+"""Valid types within JSON schema."""
 
 
 def jsonable_python(x: Any) -> Any:
@@ -23,6 +26,39 @@ def jsonable_dict(x: Any) -> dict[str, JsonValue]:
         )
 
 
+def to_json_safe(x: Any) -> bytes:
+    return to_json(value=x, indent=2, exclude_none=True, fallback=lambda _x: None)
+
+
+def to_json_str_safe(x: Any) -> str:
+    return to_json_safe(x).decode()
+
+
+def python_type_to_json_type(python_type: str | None) -> JSONType:
+    match python_type:
+        case "str":
+            return "string"
+        case "int":
+            return "integer"
+        case "float":
+            return "number"
+        case "bool":
+            return "boolean"
+        case "list":
+            return "array"
+        case "dict":
+            return "object"
+        case "None":
+            return "null"
+        # treat 'unknown' as string as anything can be converted to string
+        case None:
+            return "string"
+        case _:
+            raise ValueError(
+                f"Unsupported type: {python_type} for Python to JSON conversion."
+            )
+
+
 class JsonChange(BaseModel):
     """Describes a change to data using JSON Patch format."""
 

inspect_ai/_util/registry.py

@@ -21,6 +21,7 @@ RegistryType = Literal[
     "scorer",
     "metric",
     "tool",
+    "agent",
     "sandboxenv",
     "score_reducer",
     "approver",

inspect_ai/_util/vscode.py

@@ -1,13 +1,19 @@
 import os
+from logging import getLogger
 from pathlib import Path
 from typing import Any
 
 from pydantic import BaseModel, Field
 from pydantic_core import to_json
+from semver import Version
 from shortuuid import uuid
 
 from .appdirs import inspect_data_dir
 
+logger = getLogger(__name__)
+
+EXTENSION_COMMAND_VERSIONS = {"inspect.openLogViewer": Version(0, 3, 61)}
+
 
 class VSCodeCommand(BaseModel):
     command: str
@@ -34,6 +40,25 @@ def can_execute_vscode_commands() -> bool:
     return vs_code_commands_dir() is not None
 
 
+def can_execute_vscode_command(command: str) -> bool:
+    if not can_execute_vscode_commands():
+        return False
+
+    required_version = EXTENSION_COMMAND_VERSIONS.get(command)
+    if required_version is None:
+        return True
+    else:
+        return has_vscode_version(required_version)
+
+
+def has_vscode_version(required_version: Version) -> bool:
+    current_version = vscode_extension_version()
+    if current_version is None:
+        return False
+    else:
+        return current_version.is_compatible(required_version)
+
+
 def vs_code_commands_dir() -> Path | None:
     workspace_id = vscode_workspace_id()
     if workspace_id:
@@ -49,3 +74,15 @@ def vs_code_commands_dir() -> Path | None:
 
 def vscode_workspace_id() -> str | None:
     return os.environ.get("INSPECT_WORKSPACE_ID", None)
+
+
+def vscode_extension_version() -> Version | None:
+    version = os.environ.get("INSPECT_VSCODE_EXT_VERSION", None)
+    if version is not None:
+        try:
+            return Version.parse(version)
+        except Exception:
+            logger.warning(f"Invalid Inspect vscode extension version: {version}")
+            return None
+    else:
+        return None