inspect-ai 0.3.96__py3-none-any.whl → 0.3.98__py3-none-any.whl

inspect_ai/analysis/beta/_dataframe/events/table.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Callable, Literal, TypeAlias
+from typing import TYPE_CHECKING, Callable, Literal, Sequence, TypeAlias
 
 from inspect_ai.analysis.beta._dataframe.events.columns import EventInfo
 from inspect_ai.log._file import list_eval_logs
@@ -11,61 +11,44 @@ if TYPE_CHECKING:
 
 from typing_extensions import overload
 
-from ..columns import Column, ColumnErrors
+from ..columns import Column, ColumnError
 from ..samples.table import EventsDetail, _read_samples_df
 from ..util import LogPaths, verify_prerequisites
 
-EventFilter: TypeAlias = (
-    list[
-        Literal[
-            "sample_init",
-            "sample_limit",
-            "sandbox",
-            "state",
-            "store",
-            "model",
-            "tool",
-            "sandbox",
-            "approval",
-            "input",
-            "score",
-            "error",
-            "logger",
-            "info",
-            "span_begin",
-            "span_end",
-            "subtask",
-        ]
-    ]
-    | Callable[[Event], bool]
-)
+EventFilter: TypeAlias = Callable[[Event], bool]
 """Filter for `events_df()` rows."""
 
 
 @overload
 def events_df(
     logs: LogPaths = list_eval_logs(),
-    columns: list[Column] = EventInfo,
+    columns: Sequence[Column] = EventInfo,
     filter: EventFilter | None = None,
     strict: Literal[True] = True,
+    parallel: bool | int = False,
+    quiet: bool = False,
 ) -> "pd.DataFrame": ...
 
 
 @overload
 def events_df(
     logs: LogPaths = list_eval_logs(),
-    columns: list[Column] = EventInfo,
+    columns: Sequence[Column] = EventInfo,
     filter: EventFilter | None = None,
     strict: Literal[False] = False,
-) -> tuple["pd.DataFrame", ColumnErrors]: ...
+    parallel: bool | int = False,
+    quiet: bool = False,
+) -> tuple["pd.DataFrame", list[ColumnError]]: ...
 
 
 def events_df(
     logs: LogPaths = list_eval_logs(),
-    columns: list[Column] = EventInfo,
+    columns: Sequence[Column] = EventInfo,
     filter: EventFilter | None = None,
     strict: bool = True,
-) -> "pd.DataFrame" | tuple["pd.DataFrame", ColumnErrors]:
+    parallel: bool | int = False,
+    quiet: bool = False,
+) -> "pd.DataFrame" | tuple["pd.DataFrame", list[ColumnError]]:
     """Read a dataframe containing events from a set of evals.
 
     Args:
@@ -73,9 +56,14 @@ def events_df(
           Defaults to the contents of the currently active log directory
           (e.g. ./logs or INSPECT_LOG_DIR).
        columns: Specification for what columns to read from log files.
-       filter: List of event types to include or callable that performs the filter.
+       filter: Callable that filters event types.
        strict: Raise import errors immediately. Defaults to `True`.
           If `False` then a tuple of `DataFrame` and errors is returned.
+       parallel: If `True`, use `ProcessPoolExecutor` to read logs in parallel
+          (with workers based on `mp.cpu_count()`, capped at 8). If `int`, read
+          in parallel with the specified number of workers. If `False` (the default)
+          do not read in parallel.
+       quiet: If `True` do not print any output or progress (defaults to `False`).
 
     Returns:
        For `strict`, a Pandas `DataFrame` with information for the specified logs.
@@ -85,16 +73,16 @@ def events_df(
     verify_prerequisites()
 
     # resolve filter/detail
-    if filter is None:
-        detail = EventsDetail(filter=lambda e: True)
-    elif callable(filter):
+    if callable(filter):
         detail = EventsDetail(filter=filter)
     else:
-        detail = EventsDetail(filter=lambda e: e.event in filter)
+        detail = EventsDetail()
 
     return _read_samples_df(
         logs=logs,
         columns=columns,
         strict=strict,
         detail=detail,
+        progress=not quiet,
+        parallel=parallel,
     )

inspect_ai/analysis/beta/_dataframe/messages/table.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Callable, Literal, TypeAlias
+from typing import TYPE_CHECKING, Callable, Literal, Sequence, TypeAlias
 
 from inspect_ai.log._file import list_eval_logs
 from inspect_ai.model._chat_message import ChatMessage
@@ -10,41 +10,45 @@ if TYPE_CHECKING:
 
 from typing_extensions import overload
 
-from ..columns import Column, ColumnErrors
+from ..columns import Column, ColumnError
 from ..samples.table import MessagesDetail, _read_samples_df
 from ..util import LogPaths, verify_prerequisites
 from .columns import MessageColumns
 
-MessageFilter: TypeAlias = (
-    list[Literal["system", "user", "assistant", "tool"]] | Callable[[ChatMessage], bool]
-)
+MessageFilter: TypeAlias = Callable[[ChatMessage], bool]
 """Filter for `messages_df()` rows."""
 
 
 @overload
 def messages_df(
     logs: LogPaths = list_eval_logs(),
-    columns: list[Column] = MessageColumns,
+    columns: Sequence[Column] = MessageColumns,
     filter: MessageFilter | None = None,
     strict: Literal[True] = True,
+    parallel: bool | int = False,
+    quiet: bool = False,
 ) -> "pd.DataFrame": ...
 
 
 @overload
 def messages_df(
     logs: LogPaths = list_eval_logs(),
-    columns: list[Column] = MessageColumns,
+    columns: Sequence[Column] = MessageColumns,
     filter: MessageFilter | None = None,
     strict: Literal[False] = False,
-) -> tuple["pd.DataFrame", ColumnErrors]: ...
+    parallel: bool | int = False,
+    quiet: bool = False,
+) -> tuple["pd.DataFrame", list[ColumnError]]: ...
 
 
 def messages_df(
     logs: LogPaths = list_eval_logs(),
-    columns: list[Column] = MessageColumns,
+    columns: Sequence[Column] = MessageColumns,
     filter: MessageFilter | None = None,
     strict: bool = True,
-) -> "pd.DataFrame" | tuple["pd.DataFrame", ColumnErrors]:
+    parallel: bool | int = False,
+    quiet: bool = False,
+) -> "pd.DataFrame" | tuple["pd.DataFrame", list[ColumnError]]:
     """Read a dataframe containing messages from a set of evals.
 
     Args:
@@ -52,9 +56,14 @@ def messages_df(
           Defaults to the contents of the currently active log directory
           (e.g. ./logs or INSPECT_LOG_DIR).
        columns: Specification for what columns to read from log files.
-       filter: List of message role types to include or callable that performs the filter.
+       filter: Callable that filters messages
        strict: Raise import errors immediately. Defaults to `True`.
           If `False` then a tuple of `DataFrame` and errors is returned.
+       parallel: If `True`, use `ProcessPoolExecutor` to read logs in parallel
+          (with workers based on `mp.cpu_count()`, capped at 8). If `int`, read
+          in parallel with the specified number of workers. If `False` (the default)
+          do not read in parallel.
+       quiet: If `True` do not print any output or progress (defaults to `False`).
 
     Returns:
        For `strict`, a Pandas `DataFrame` with information for the specified logs.
@@ -64,16 +73,16 @@ def messages_df(
     verify_prerequisites()
 
     # resolve filter/detail
-    if filter is None:
-        detail = MessagesDetail(filter=lambda m: True)
-    elif callable(filter):
+    if callable(filter):
         detail = MessagesDetail(filter=filter)
     else:
-        detail = MessagesDetail(filter=lambda m: m.role in filter)
+        detail = MessagesDetail()
 
     return _read_samples_df(
         logs=logs,
         columns=columns,
         strict=strict,
         detail=detail,
+        parallel=parallel,
+        progress=not quiet,
     )

inspect_ai/analysis/beta/_dataframe/progress.py

@@ -1,5 +1,5 @@
 from contextlib import contextmanager
-from typing import Iterator
+from typing import Iterator, Protocol
 
 from rich.progress import (
     BarColumn,
@@ -11,10 +11,40 @@ from rich.progress import (
 )
 
 
+class ImportProgress(Protocol):
+    def update(self) -> None: ...
+    def reset(self, description: str, completed: int, total: int) -> None: ...
+
+
+class NoProgress(ImportProgress):
+    def update(self) -> None:
+        pass
+
+    def reset(self, description: str, completed: int, total: int) -> None:
+        pass
+
+
+class RichImportProgress(ImportProgress):
+    def __init__(self, progress: Progress, task_id: TaskID) -> None:
+        self._progress = progress
+        self._task_id = task_id
+
+    def update(self) -> None:
+        self._progress.update(self._task_id, advance=1)
+
+    def reset(self, description: str, completed: int, total: int) -> None:
+        self._progress.reset(
+            self._task_id, description=description, completed=completed, total=total
+        )
+
+
+@contextmanager
+def no_progress() -> Iterator[ImportProgress]:
+    yield NoProgress()
+
+
 @contextmanager
-def import_progress(
-    description: str, total: float | None
-) -> Iterator[tuple[Progress, TaskID]]:
+def import_progress(description: str, total: float | None) -> Iterator[ImportProgress]:
     with Progress(
         TextColumn("[progress.description]{task.description:<18}"),
         BarColumn(),
@@ -23,4 +53,4 @@ def import_progress(
         transient=True,
     ) as progress:
         task_id = progress.add_task(description, total=total)
-        yield progress, task_id
+        yield RichImportProgress(progress, task_id)

inspect_ai/analysis/beta/_dataframe/record.py

@@ -1,6 +1,6 @@
 import json
 from datetime import date, datetime, time, timezone
-from typing import Any, Callable, Literal, Type, cast, overload
+from typing import Any, Callable, Literal, Sequence, Type, cast, overload
 
 import yaml
 from jsonpath_ng import JSONPath  # type: ignore
@@ -20,38 +20,41 @@ from .extract import model_to_record
 
 @overload
 def import_record(
+    log: EvalLog,
     record: EvalLog
     | EvalSampleSummary
     | EvalSample
     | ChatMessage
     | Event
     | dict[str, JsonValue],
-    columns: list[Column],
+    columns: Sequence[Column],
     strict: Literal[True] = True,
 ) -> dict[str, ColumnType]: ...
 
 
 @overload
 def import_record(
+    log: EvalLog,
     record: EvalLog
     | EvalSampleSummary
     | EvalSample
     | ChatMessage
     | Event
     | dict[str, JsonValue],
-    columns: list[Column],
+    columns: Sequence[Column],
     strict: Literal[False],
 ) -> tuple[dict[str, ColumnType], list[ColumnError]]: ...
 
 
 def import_record(
+    log: EvalLog,
     record: EvalLog
     | EvalSampleSummary
     | EvalSample
     | ChatMessage
     | Event
     | dict[str, JsonValue],
-    columns: list[Column],
+    columns: Sequence[Column],
     strict: bool = True,
 ) -> dict[str, ColumnType] | tuple[dict[str, ColumnType], list[ColumnError]]:
     # resolve the record BaseModel into a dict (and optionally a summary dict).
@@ -80,7 +83,7 @@ def import_record(
         try:
             result[name] = _resolve_value(value, column.type)
         except ValueError as ex:
-            error = ColumnError(name, path=column.path, message=str(ex))
+            error = ColumnError(name, path=column.path, error=ex, log=log)
             if strict:
                 raise ValueError(str(error))
             else:
@@ -90,10 +93,10 @@ def import_record(
     def field_not_found(
         name: str, path: JSONPath | None, required_type: str | None = None
     ) -> None:
-        message = (
+        ex = ValueError(
             f"field not of type {required_type}" if required_type else "field not found"
         )
-        error = ColumnError(name, path=path, message=f"{message}")
+        error = ColumnError(name, path=path, error=ex, log=log)
         if strict:
             raise ValueError(str(error))
         else:
@@ -157,7 +160,8 @@ def import_record(
             error = ColumnError(
                 column.name,
                 path=str(column.path) if column.path else None,
-                message=str(ex),
+                error=ex,
+                log=log,
             )
             if strict:
                 raise ValueError(str(error))
@@ -190,7 +194,7 @@ def import_record(
         return result, errors
 
 
-def resolve_duplicate_columns(columns: list[Column]) -> list[Column]:
+def resolve_duplicate_columns(columns: Sequence[Column]) -> list[Column]:
     """Remove duplicate columns (with the later columns winning)"""
     seen = set[str]()
     deduped: list[Column] = []

inspect_ai/analysis/beta/_dataframe/samples/columns.py

@@ -65,7 +65,7 @@ SampleSummary: list[Column] = [
     SampleColumn("model_usage", path="model_usage"),
     SampleColumn("total_time", path="total_time"),
     SampleColumn("working_time", path="total_time"),
-    SampleColumn("error", path="error"),
+    SampleColumn("error", path="error", default=""),
     SampleColumn("limit", path="limit"),
     SampleColumn("retries", path="retries"),
 ]