inspect-ai 0.3.69__py3-none-any.whl → 0.3.71__py3-none-any.whl

inspect_ai/_eval/task/run.py

@@ -33,6 +33,10 @@ from inspect_ai._util.registry import (
     registry_unqualified_name,
 )
 from inspect_ai._util.timeouts import Timeout, timeout
+from inspect_ai._util.working import (
+    init_sample_working_limit,
+    sample_waiting_time,
+)
 from inspect_ai._view.notify import view_notify_eval
 from inspect_ai.dataset import Dataset, Sample
 from inspect_ai.log import (
@@ -46,16 +50,13 @@ 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,
-    set_active_sample_message_limit,
-    set_active_sample_token_limit,
-)
+from inspect_ai.log._samples import active_sample
 from inspect_ai.log._transcript import (
     ErrorEvent,
     SampleInitEvent,
     SampleLimitEvent,
     ScoreEvent,
+    StepEvent,
     transcript,
 )
 from inspect_ai.model import (
@@ -182,9 +183,9 @@ async def task_run(options: TaskRunOptions) -> EvalLog:
         if isinstance(solver, Plan):
             plan = solver
         elif isinstance(solver, Chain):
-            plan = Plan(list(solver), internal=True)
+            plan = Plan(list(solver), cleanup=task.cleanup, internal=True)
         else:
-            plan = Plan(unroll(solver), internal=True)
+            plan = Plan(unroll(solver), cleanup=task.cleanup, internal=True)
 
         # add setup solver(s) if specified
         if task.setup:
@@ -308,6 +309,7 @@ async def task_run(options: TaskRunOptions) -> EvalLog:
                                 or config.fail_on_error is True
                             ),
                             time_limit=config.time_limit,
+                            working_limit=config.working_limit,
                             semaphore=sample_semaphore,
                         )
                         for (sample, state) in zip(samples, states)
@@ -500,6 +502,7 @@ async def task_run_sample(
     sample_complete: Callable[[dict[str, SampleScore]], None],
     fails_on_error: bool,
     time_limit: int | None,
+    working_limit: int | None,
     semaphore: asyncio.Semaphore | None,
 ) -> dict[str, SampleScore] | None:
     # if there is an existing sample then tick off its progress, log it, and return it
@@ -570,19 +573,37 @@ async def task_run_sample(
             message_limit=state.message_limit,
             token_limit=state.token_limit,
             time_limit=time_limit,
+            working_limit=working_limit,
             fails_on_error=fails_on_error,
             transcript=sample_transcript,
         ) as active,
     ):
+        start_time: float | None = None
         error: EvalError | None = None
         raise_error: BaseException | None = None
         results: dict[str, SampleScore] = {}
         try:
+            # begin init
+            transcript()._event(StepEvent(action="begin", name="init"))
+
+            # 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))
+            )
+
             async with sandboxenv_cm:
                 try:
                     # update active sample wth sandboxes now that we are initialised
                     active.sandboxes = await sandbox_connections()
 
+                    # end init
+                    transcript()._event(StepEvent(action="end", name="init"))
+
                     # initialise timeout context manager
                     timeout_cm = (
                         timeout(time_limit)
@@ -590,23 +611,15 @@ async def task_run_sample(
                         else contextlib.nullcontext()
                     )
 
+                    # record start time
+                    start_time = time.monotonic()
+                    init_sample_working_limit(start_time, working_limit)
+
                     # run sample w/ optional timeout
                     async with timeout_cm:
                         # mark started
                         active.started = datetime.now().timestamp()
 
-                        # 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)
 
@@ -661,11 +674,13 @@ async def task_run_sample(
 
                     # capture most recent state for scoring
                     state = ex.state or sample_state() or state
-                    state.completed = True
 
                 except BaseException as ex:
                     error, raise_error = handle_error(ex)
 
+                # mark completed
+                state.completed = True
+
                 # set timeout for scoring. if the original timeout was hit we still
                 # want to provide opportunity for scoring, but we don't necessarily
                 # want to wait the full timeout again (especially in the case where
@@ -676,9 +691,10 @@ async def task_run_sample(
                     assert time_limit
                     timeout_cm = timeout(time_limit / 2)
 
-                # turn off sample limits
-                set_active_sample_token_limit(None)
-                set_active_sample_message_limit(None)
+                # turn off message and token limits
+                state.message_limit = None
+                state.token_limit = None
+                set_sample_state(state)
 
                 # scoring
                 try:
@@ -768,6 +784,7 @@ async def task_run_sample(
 
             # log the sample
             await log_sample(
+                start_time=start_time,
                 logger=logger,
                 sample=sample,
                 state=state,
@@ -788,6 +805,7 @@ async def task_run_sample(
 
 
 async def log_sample(
+    start_time: float | None,
     logger: TaskLogger,
     sample: Sample,
     state: TaskState,
@@ -804,6 +822,9 @@ async def log_sample(
 
     # construct sample for logging
 
+    # compute total time if we can
+    total_time = time.monotonic() - start_time if start_time is not None else None
+
     # if a limit was hit, note that in the Eval Sample
     limit = None
     for e in transcript().events:
@@ -827,8 +848,13 @@ async def log_sample(
         output=state.output,
         scores={k: v.score for k, v in scores.items()},
         store=dict(state.store.items()),
+        uuid=state.uuid,
         events=list(transcript().events),
         model_usage=sample_model_usage(),
+        total_time=round(total_time, 3) if total_time is not None else None,
+        working_time=round(total_time - sample_waiting_time(), 3)
+        if total_time is not None
+        else None,
         error=error,
         limit=limit,
     )

inspect_ai/_eval/task/task.py

@@ -1,7 +1,7 @@
 from copy import deepcopy
 from dataclasses import dataclass
 from logging import getLogger
-from typing import Any, Callable, Sequence, cast
+from typing import Any, Awaitable, Callable, Sequence, cast
 
 from pydantic import BaseModel
 from typing_extensions import TypedDict, Unpack
@@ -17,6 +17,7 @@ from inspect_ai.scorer import Metric, Scorer
 from inspect_ai.scorer._reducer import ScoreReducers, create_reducers
 from inspect_ai.solver import Plan, Solver, generate
 from inspect_ai.solver._chain import chain
+from inspect_ai.solver._task_state import TaskState
 from inspect_ai.util._sandbox.environment import (
     SandboxEnvironmentSpec,
     SandboxEnvironmentType,
@@ -46,6 +47,7 @@ class Task:
         dataset: Dataset | Sequence[Sample] | None = None,
         setup: Solver | list[Solver] | None = None,
         solver: Solver | 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,
         config: GenerateConfig = GenerateConfig(),
@@ -56,6 +58,7 @@ class Task:
         message_limit: int | None = None,
         token_limit: int | None = None,
         time_limit: int | None = None,
+        working_limit: int | None = None,
         name: str | None = None,
         version: int = 0,
         metadata: dict[str, Any] | None = None,
@@ -69,6 +72,9 @@ class Task:
                 even when the main `solver` is replaced).
             solver: (Solver | list[Solver]): Solver or list of solvers.
                 Defaults to generate(), a normal call to the model.
+            cleanup: Optional cleanup function for task. Called after
+                all solvers have run for each sample (including if an
+                exception occurs during the run)
             scorer: (Scorer | list[Scorer] | None): Scorer used to evaluate model output.
             metrics (list[Metric] | dict[str, list[Metric]] | None):
                 Alternative metrics (overrides the metrics provided by the specified scorer).
@@ -86,7 +92,10 @@ class Task:
                 eval if a count of samples fails.
             message_limit (int | None): Limit on total messages used for each sample.
             token_limit (int | None): Limit on total tokens used for each sample.
-            time_limit (int | None): Limit on time (in seconds) for execution of each sample.
+            time_limit: Limit on clock time (in seconds) for samples.
+            working_limit: Limit on working time (in seconds) for sample. Working
+                time includes model generation, tool calls, etc. but does not include
+                time spent waiting on retries or shared resources.
             name: (str | None): Task name. If not specified is automatically
                 determined based on the name of the task directory (or "task")
                 if its anonymous task (e.g. created in a notebook and passed to
@@ -123,6 +132,7 @@ class Task:
         self.dataset = resolve_dataset(dataset)
         self.setup = setup
         self.solver = resolve_solver(solver)
+        self.cleanup = cleanup
         self.scorer = resolve_scorer(scorer)
         self.metrics = metrics
         self.config = config
@@ -135,6 +145,7 @@ class Task:
         self.message_limit = message_limit
         self.token_limit = token_limit
         self.time_limit = time_limit
+        self.working_limit = working_limit
         self.version = version
         self._name = name
         self.metadata = metadata
@@ -162,6 +173,7 @@ def task_with(
     dataset: Dataset | Sequence[Sample] | None | NotGiven = NOT_GIVEN,
     setup: Solver | list[Solver] | None | NotGiven = NOT_GIVEN,
     solver: Solver | list[Solver] | NotGiven = NOT_GIVEN,
+    cleanup: Callable[[TaskState], Awaitable[None]] | None | NotGiven = NOT_GIVEN,
     scorer: Scorer | list[Scorer] | None | NotGiven = NOT_GIVEN,
     metrics: list[Metric] | dict[str, list[Metric]] | None | NotGiven = NOT_GIVEN,
     config: GenerateConfig | NotGiven = NOT_GIVEN,
@@ -172,6 +184,7 @@ def task_with(
     message_limit: int | None | NotGiven = NOT_GIVEN,
     token_limit: int | None | NotGiven = NOT_GIVEN,
     time_limit: int | None | NotGiven = NOT_GIVEN,
+    working_limit: int | None | NotGiven = NOT_GIVEN,
     name: str | None | NotGiven = NOT_GIVEN,
     version: int | NotGiven = NOT_GIVEN,
     metadata: dict[str, Any] | None | NotGiven = NOT_GIVEN,
@@ -185,6 +198,9 @@ def task_with(
             even when the main `solver` is replaced).
         solver: (Solver | list[Solver]): Solver or list of solvers.
             Defaults to generate(), a normal call to the model.
+        cleanup: Optional cleanup function for task. Called after
+            all solvers have run for each sample (including if an
+            exception occurs during the run)
         scorer: (Scorer | list[Scorer] | None): Scorer used to evaluate model output.
         metrics (list[Metric] | dict[str, list[Metric]] | None):
             Alternative metrics (overrides the metrics provided by the specified scorer).
@@ -202,7 +218,10 @@ def task_with(
             eval if a count of samples fails.
         message_limit (int | None): Limit on total messages used for each sample.
         token_limit (int | None): Limit on total tokens used for each sample.
-        time_limit (int | None): Limit on time (in seconds) for execution of each sample.
+        time_limit: Limit on clock time (in seconds) for samples.
+        working_limit: Limit on execution time (in seconds) for sample. Execution
+            time includes model generation, tool calls, etc. but does not include
+            time spent waiting on retries or shared resources.
         name: (str | None): Task name. If not specified is automatically
             determined based on the name of the task directory (or "task")
             if its anonymous task (e.g. created in a notebook and passed to
@@ -223,6 +242,8 @@ def task_with(
         task.setup = setup
     if not isinstance(solver, NotGiven):
         task.solver = resolve_solver(solver)
+    if not isinstance(cleanup, NotGiven):
+        task.cleanup = cleanup
     if not isinstance(scorer, NotGiven):
         task.scorer = resolve_scorer(scorer)
     if not isinstance(metrics, NotGiven):
@@ -245,6 +266,8 @@ def task_with(
         task.token_limit = token_limit
     if not isinstance(time_limit, NotGiven):
         task.time_limit = time_limit
+    if not isinstance(working_limit, NotGiven):
+        task.working_limit = working_limit
     if not isinstance(version, NotGiven):
         task.version = version
     if not isinstance(name, NotGiven):

inspect_ai/_util/content.py

@@ -13,6 +13,25 @@ class ContentText(BaseModel):
     """Text content."""
 
 
+class ContentReasoning(BaseModel):
+    """Reasoning content.
+
+    See the specification for [thinking blocks](https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking#understanding-thinking-blocks) for Claude models.
+    """
+
+    type: Literal["reasoning"] = Field(default="reasoning")
+    """Type."""
+
+    reasoning: str
+    """Reasoning content."""
+
+    signature: str | None = Field(default=None)
+    """Signature for reasoning content (used by some models to ensure that reasoning content is not modified for replay)"""
+
+    redacted: bool = Field(default=False)
+    """Indicates that the explicit content of this reasoning block has been redacted."""
+
+
 class ContentImage(BaseModel):
     """Image content."""
 
@@ -55,5 +74,5 @@ class ContentVideo(BaseModel):
     """Format of video data ('mp4', 'mpeg', or 'mov')"""
 
 
-Content = Union[ContentText, ContentImage, ContentAudio, ContentVideo]
+Content = Union[ContentText, ContentReasoning, ContentImage, ContentAudio, ContentVideo]
 """Content sent to or received from a model."""

inspect_ai/_util/interrupt.py

@@ -1,9 +1,15 @@
 import asyncio
 
+from .working import check_sample_working_limit
+
 
 def check_sample_interrupt() -> None:
     from inspect_ai.log._samples import sample_active
 
+    # check for user interrupt
     sample = sample_active()
     if sample and sample.interrupt_action:
         raise asyncio.CancelledError()
+
+    # check for working_limit
+    check_sample_working_limit()

inspect_ai/_util/logger.py

@@ -160,7 +160,9 @@ def init_logger(
 
     # init logging handler on demand
     global _logHandler
+    removed_root_handlers = False
     if not _logHandler:
+        removed_root_handlers = remove_non_pytest_root_logger_handlers()
         _logHandler = LogHandler(min(DEBUG, levelno), transcript_levelno)
         getLogger().addHandler(_logHandler)
 
@@ -173,6 +175,11 @@ def init_logger(
     getLogger("httpx").setLevel(capture_level)
     getLogger("botocore").setLevel(DEBUG)
 
+    if removed_root_handlers:
+        getLogger(PKG_NAME).warning(
+            "Inspect removed pre-existing root logger handlers and replaced them with its own handler."
+        )
+
     # set the levelno on the global handler
     _logHandler.display_level = levelno
 
@@ -180,6 +187,18 @@ def init_logger(
 _logHandler: LogHandler | None = None
 
 
+def remove_non_pytest_root_logger_handlers() -> bool:
+    root_logger = getLogger()
+    non_pytest_handlers = [
+        handler
+        for handler in root_logger.handlers
+        if handler.__module__ != "_pytest.logging"
+    ]
+    for handler in non_pytest_handlers:
+        root_logger.removeHandler(handler)
+    return len(non_pytest_handlers) > 0
+
+
 def notify_logger_record(record: LogRecord, write: bool) -> None:
     from inspect_ai.log._message import LoggingMessage
     from inspect_ai.log._transcript import LoggerEvent, transcript

inspect_ai/_util/rich.py

@@ -2,23 +2,22 @@ from rich.console import RenderableType
 from rich.style import Style
 from rich.text import Text
 
+from inspect_ai._util.text import truncate_lines
+
 
 def lines_display(
     text: str, max_lines: int = 100, style: str | Style = ""
 ) -> list[RenderableType]:
-    lines = text.splitlines()
-    if len(lines) > max_lines:
-        content: list[RenderableType] = [
-            Text("\n".join(lines[0:max_lines]), style=style)
-        ]
+    lines, truncated = truncate_lines(text, max_lines)
+
+    content: list[RenderableType] = [Text(lines, style=style)]
+    if truncated is not None:
         content.append(Text())
         content.append(
             Text.from_markup(
-                f"[italic]Output truncated ({len(lines) - max_lines} additional lines)...[/italic]",
+                f"[italic]Output truncated ({truncated} additional lines)...[/italic]",
                 style=style,
             )
         )
-    else:
-        content = [Text(text, style=style)]
 
     return content

inspect_ai/_util/text.py

@@ -134,6 +134,19 @@ def truncate(text: str, length: int, overflow: str = "...", pad: bool = True) ->
     return truncated
 
 
+def truncate_lines(
+    text: str, max_lines: int = 100, max_characters: int | None = 100 * 100
+) -> tuple[str, int | None]:
+    if max_characters is not None:
+        text = truncate(text, max_characters)
+    lines = text.splitlines()
+    if len(lines) > max_lines:
+        output = "\n".join(lines[0:max_lines])
+        return output, len(lines) - max_lines
+    else:
+        return text, None
+
+
 def generate_large_text(target_tokens: int) -> str:
     """Generate a large amount of text with approximately the target number of tokens"""
     generated_text = []

inspect_ai/_util/transcript.py

@@ -10,6 +10,8 @@ from rich.panel import Panel
 from rich.rule import Rule
 from rich.text import Text
 
+from inspect_ai._util.content import ContentReasoning
+
 from .format import format_function_call
 
 
@@ -111,19 +113,31 @@ def transcript_panel(
     )
 
 
-def transcript_reasoning(reasoning: str) -> list[RenderableType]:
+def transcript_reasoning(reasoning: ContentReasoning) -> list[RenderableType]:
     content: list[RenderableType] = []
+    text = (
+        reasoning.reasoning
+        if not reasoning.redacted
+        else "Reasoning encrypted by model provider."
+    )
+
     content.append(
-        transcript_markdown(
-            f"**<think>**  \n{reasoning}  \n**</think>**\n\n", escape=True
-        )
+        transcript_markdown(f"**<think>**  \n{text}  \n**</think>**\n\n", escape=True)
     )
     content.append(Text())
     return content
 
 
-def transcript_separator(title: str, color: str) -> RenderableType:
-    return Rule(title=title, style=f"{color} bold", align="center", end="\n\n")
+def transcript_separator(
+    title: str, color: str, characters: str = "─"
+) -> RenderableType:
+    return Rule(
+        title=title,
+        characters=characters,
+        style=f"{color} bold",
+        align="center",
+        end="\n\n",
+    )
 
 
 def transcript_function(function: str, arguments: dict[str, Any]) -> RenderableType:

inspect_ai/_util/working.py

@@ -0,0 +1,50 @@
+import time
+from contextvars import ContextVar
+
+
+def init_sample_working_limit(start_time: float, working_limit: float | None) -> None:
+    _sample_working_limit.set(working_limit)
+    _sample_start_time.set(start_time)
+    _sample_waiting_time.set(0)
+
+
+def sample_waiting_time() -> float:
+    return _sample_waiting_time.get()
+
+
+def sample_working_time() -> float:
+    return time.monotonic() - _sample_start_time.get() - sample_waiting_time()
+
+
+def report_sample_waiting_time(waiting_time: float) -> None:
+    _sample_waiting_time.set(_sample_waiting_time.get() + waiting_time)
+    check_sample_working_limit()
+
+
+def check_sample_working_limit() -> None:
+    # no check if we don't have a limit
+    working_limit = _sample_working_limit.get()
+    if working_limit is None:
+        return
+
+    # are we over the limit?
+    running_time = time.monotonic() - _sample_start_time.get()
+    working_time = running_time - sample_waiting_time()
+    if working_time > working_limit:
+        from inspect_ai.solver._limit import SampleLimitExceededError
+
+        raise SampleLimitExceededError(
+            type="working",
+            value=int(working_time),
+            limit=int(working_limit),
+            message=f"Exceeded working time limit ({working_limit:,} seconds)",
+        )
+
+
+_sample_working_limit: ContextVar[float | None] = ContextVar(
+    "sample_working_limit", default=None
+)
+
+_sample_start_time: ContextVar[float] = ContextVar("sample_start_time", default=0)
+
+_sample_waiting_time: ContextVar[float] = ContextVar("sample_waiting_time", default=0)

inspect_ai/_view/www/App.css

@@ -805,15 +805,21 @@ table.table.table-sm td {
   overflow: unset;
 }
 
+.markdown-content pre[class*="language-"],
 pre[class*="language-"].tool-output,
 .tool-output {
   background-color: #f8f8f8;
 }
+
+.vscode-dark .model-call pre[class*="language-"],
+.vscode-dark .markdown-content pre[class*="language-"],
 .vscode-dark pre[class*="language-"].tool-output,
 .vscode-dark .tool-output {
   background-color: #333333;
 }
 
+.model-call pre[class*="language-"],
+.markdown-content pre[class*="language-"],
 pre[class*="language-"].tool-output {
   border: none !important;
   box-shadow: none !important;