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

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

@@ -0,0 +1,370 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    Generator,
+    Literal,
+    overload,
+)
+
+from inspect_ai._util.hash import mm3_hash
+from inspect_ai._util.path import pretty_path
+from inspect_ai.analysis.beta._dataframe.progress import import_progress
+from inspect_ai.log._file import (
+    list_eval_logs,
+    read_eval_log_sample_summaries,
+    read_eval_log_samples,
+)
+from inspect_ai.log._log import EvalSample, EvalSampleSummary
+from inspect_ai.log._transcript import Event
+from inspect_ai.model._chat_message import ChatMessage
+
+from ..columns import Column, ColumnErrors, ColumnType
+from ..evals.columns import EvalColumn
+from ..evals.table import EVAL_ID, EVAL_SUFFIX, _read_evals_df, ensure_eval_id
+from ..events.columns import EventColumn
+from ..extract import message_as_str
+from ..messages.columns import MessageColumn
+from ..record import import_record, resolve_duplicate_columns
+from ..util import (
+    LogPaths,
+    add_unreferenced_columns,
+    records_to_pandas,
+    resolve_columns,
+    resolve_logs,
+    verify_prerequisites,
+)
+from .columns import SampleColumn, SampleSummary
+from .extract import auto_detail_id, auto_sample_id
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+
+SAMPLE_ID = "sample_id"
+SAMPLE_SUFFIX = "_sample"
+
+
+@overload
+def samples_df(
+    logs: LogPaths = list_eval_logs(),
+    columns: list[Column] = SampleSummary,
+    strict: Literal[True] = True,
+) -> "pd.DataFrame": ...
+
+
+@overload
+def samples_df(
+    logs: LogPaths = list_eval_logs(),
+    columns: list[Column] = SampleSummary,
+    strict: Literal[False] = False,
+) -> tuple["pd.DataFrame", ColumnErrors]: ...
+
+
+def samples_df(
+    logs: LogPaths = list_eval_logs(),
+    columns: list[Column] = SampleSummary,
+    strict: bool = True,
+) -> "pd.DataFrame" | tuple["pd.DataFrame", ColumnErrors]:
+    """Read a dataframe containing samples from a set of evals.
+
+    Args:
+       logs: One or more paths to log files or log directories.
+          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.
+       strict: Raise import errors immediately. Defaults to `True`.
+          If `False` then a tuple of `DataFrame` and errors is returned.
+
+    Returns:
+       For `strict`, a Pandas `DataFrame` with information for the specified logs.
+       For `strict=False`, a tuple of Pandas `DataFrame` and a dictionary of errors
+       encountered (by log file) during import.
+    """
+    return _read_samples_df(logs, columns, strict=strict)
+
+
+@dataclass
+class MessagesDetail:
+    name: str = "message"
+    col_type = MessageColumn
+    filter: Callable[[ChatMessage], bool] = lambda m: True
+
+
+@dataclass
+class EventsDetail:
+    name: str = "event"
+    col_type = EventColumn
+    filter: Callable[[Event], bool] = lambda e: True
+
+
+def _read_samples_df(
+    logs: LogPaths,
+    columns: list[Column],
+    *,
+    strict: bool = True,
+    detail: MessagesDetail | EventsDetail | None = None,
+) -> "pd.DataFrame" | tuple["pd.DataFrame", ColumnErrors]:
+    verify_prerequisites()
+
+    # resolve logs
+    logs = resolve_logs(logs)
+
+    # split columns by type
+    columns_eval: list[Column] = []
+    columns_sample: list[Column] = []
+    columns_detail: list[Column] = []
+    for column in columns:
+        if isinstance(column, EvalColumn):
+            columns_eval.append(column)
+        elif isinstance(column, SampleColumn):
+            columns_sample.append(column)
+            if column._full:
+                require_full_samples = True
+        elif detail and isinstance(column, detail.col_type):
+            columns_detail.append(column)
+        else:
+            raise ValueError(
+                f"Unexpected column type passed to samples_df: {type(column)}"
+            )
+    # resolve duplciates
+    columns_eval = resolve_duplicate_columns(columns_eval)
+    columns_sample = resolve_duplicate_columns(columns_sample)
+    columns_detail = resolve_duplicate_columns(columns_detail)
+
+    # determine if we require full samples
+    require_full_samples = len(columns_detail) > 0 or any(
+        [isinstance(column, SampleColumn) and column._full for column in columns_sample]
+    )
+
+    # make sure eval_id is present
+    ensure_eval_id(columns_eval)
+
+    # determine how we will allocate progress
+    with import_progress("scanning logs", total=len(logs)) as (
+        p,
+        task_id,
+    ):
+
+        def progress() -> None:
+            p.update(task_id, advance=1)
+
+        # read samples from each log
+        sample_records: list[dict[str, ColumnType]] = []
+        detail_records: list[dict[str, ColumnType]] = []
+        all_errors = ColumnErrors()
+
+        # read logs and note total samples
+        evals_table, total_samples = _read_evals_df(
+            logs, columns=columns_eval, strict=True, progress=progress
+        )
+
+        # update progress now that we know the total samples
+        entity = detail.name if detail else "sample"
+        p.reset(
+            task_id, description=f"reading {entity}s", completed=0, total=total_samples
+        )
+
+        # read samples
+        for eval_id, log in zip(evals_table[EVAL_ID].to_list(), logs):
+            # get a generator for the samples (might require reading the full log
+            # or might be fine to just read the summaries)
+            if require_full_samples:
+                samples: Generator[EvalSample | EvalSampleSummary, None, None] = (
+                    read_eval_log_samples(
+                        log, all_samples_required=False, resolve_attachments=True
+                    )
+                )
+            else:
+                samples = (summary for summary in read_eval_log_sample_summaries(log))
+            for sample in samples:
+                if strict:
+                    record = import_record(sample, columns_sample, strict=True)
+                else:
+                    record, errors = import_record(sample, columns_sample, strict=False)
+                    error_key = f"{pretty_path(log)} [{sample.id}, {sample.epoch}]"
+                    all_errors[error_key] = errors
+
+                # inject ids
+                sample_id = sample.uuid or auto_sample_id(eval_id, sample)
+                ids: dict[str, ColumnType] = {
+                    EVAL_ID: eval_id,
+                    SAMPLE_ID: sample_id,
+                }
+
+                # record with ids
+                record = ids | record
+
+                # if there are detail columns then we blow out these records w/ detail
+                if detail is not None:
+                    # filter detail records
+                    assert isinstance(sample, EvalSample)
+                    if isinstance(detail, MessagesDetail):
+                        detail_items: list[ChatMessage] | list[Event] = (
+                            sample_messages_from_events(sample.events, detail.filter)
+                        )
+                    elif isinstance(detail, EventsDetail):
+                        detail_items = [e for e in sample.events if detail.filter(e)]
+                    else:
+                        detail_items = []
+
+                    # read detail records (provide auto-ids)
+                    for index, item in enumerate(detail_items):
+                        if strict:
+                            detail_record = import_record(
+                                item, columns_detail, strict=True
+                            )
+                        else:
+                            detail_record, errors = import_record(
+                                item, columns_detail, strict=False
+                            )
+                            error_key = (
+                                f"{pretty_path(log)} [{sample.id}, {sample.epoch}]"
+                            )
+                            all_errors[error_key] = errors
+
+                        # inject ids
+                        detail_id = detail_record.get(
+                            "id", auto_detail_id(sample_id, detail.name, index)
+                        )
+                        ids = {SAMPLE_ID: sample_id, f"{detail.name}_id": detail_id}
+                        detail_record = ids | detail_record
+
+                        # append detail record
+                        detail_records.append(detail_record)
+
+                # record sample record
+                sample_records.append(record)
+                progress()
+
+    # normalize records and produce samples table
+    samples_table = records_to_pandas(sample_records)
+
+    # if we have detail records then join them into the samples table
+    if detail is not None:
+        details_table = records_to_pandas(detail_records)
+        samples_table = details_table.merge(
+            samples_table,
+            on=SAMPLE_ID,
+            how="left",
+            suffixes=(f"_{detail.name}", SAMPLE_SUFFIX),
+        )
+
+    # join eval_records
+    samples_table = samples_table.merge(
+        evals_table, on=EVAL_ID, how="left", suffixes=(SAMPLE_SUFFIX, EVAL_SUFFIX)
+    )
+
+    # re-order based on original specification
+    samples_table = reorder_samples_df_columns(
+        samples_table,
+        columns_eval,
+        columns_sample,
+        columns_detail,
+        detail.name if detail else "",
+    )
+
+    # return
+    if strict:
+        return samples_table
+    else:
+        return samples_table, all_errors
+
+
+def sample_messages_from_events(
+    events: list[Event], filter: Callable[[ChatMessage], bool]
+) -> list[ChatMessage]:
+    # don't yield the same event twice
+    ids: set[str] = set()
+
+    # we need to look at the full input to every model event and add
+    # messages we haven't seen before
+    messages: list[ChatMessage] = []
+    for event in events:
+        if event.event == "model":
+            event_messages = event.input + (
+                [event.output.message] if not event.output.empty else []
+            )
+            for message in event_messages:
+                id = message.id or message_hash(message_as_str(message))
+                if id not in ids:
+                    messages.append(message)
+                    ids.add(id)
+
+    # then apply the filter
+    return [message for message in messages if filter(message)]
+
+
+@lru_cache(maxsize=100)
+def message_hash(message: str) -> str:
+    return mm3_hash(message)
+
+
+def reorder_samples_df_columns(
+    df: "pd.DataFrame",
+    eval_columns: list[Column],
+    sample_columns: list[Column],
+    detail_columns: list[Column],
+    details_name: str,
+) -> "pd.DataFrame":
+    """Reorder columns in the merged DataFrame.
+
+    Order with:
+    1. sample_id first
+    2. eval_id second
+    3. eval columns
+    4. sample columns
+    5. any remaining columns
+    """
+    actual_columns = list(df.columns)
+    ordered_columns: list[str] = []
+
+    # detail first if we have detail
+    if details_name:
+        ordered_columns.append(f"{details_name}_id")
+
+    # sample_id first
+    if SAMPLE_ID in actual_columns:
+        ordered_columns.append(SAMPLE_ID)
+
+    # eval_id next
+    if EVAL_ID in actual_columns:
+        ordered_columns.append(EVAL_ID)
+
+    # eval columns
+    for column in eval_columns:
+        if column.name == EVAL_ID or column.name == SAMPLE_ID:
+            continue  # Already handled
+
+        ordered_columns.extend(
+            resolve_columns(column.name, EVAL_SUFFIX, actual_columns, ordered_columns)
+        )
+
+    # then sample columns
+    for column in sample_columns:
+        if column.name == EVAL_ID or column.name == SAMPLE_ID:
+            continue  # Already handled
+
+        ordered_columns.extend(
+            resolve_columns(column.name, SAMPLE_SUFFIX, actual_columns, ordered_columns)
+        )
+
+    # then detail columns
+    for column in detail_columns:
+        if column.name == EVAL_ID or column.name == SAMPLE_ID:
+            continue  # Already handled
+
+        ordered_columns.extend(
+            resolve_columns(
+                column.name, f"_{details_name}", actual_columns, ordered_columns
+            )
+        )
+
+    # add any unreferenced columns
+    ordered_columns = add_unreferenced_columns(actual_columns, ordered_columns)
+
+    # reorder the DataFrame
+    return df[ordered_columns]

inspect_ai/analysis/beta/_dataframe/util.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import re
+from os import PathLike
+from pathlib import Path
+from re import Pattern
+from typing import TYPE_CHECKING, Sequence, TypeAlias
+
+from inspect_ai._util.error import pip_dependency_error
+from inspect_ai._util.file import FileInfo, filesystem
+from inspect_ai._util.version import verify_required_version
+from inspect_ai.log._file import EvalLogInfo, log_files_from_ls
+
+if TYPE_CHECKING:
+    import pandas as pd
+    import pyarrow as pa
+
+from .columns import ColumnType
+
+LogPaths: TypeAlias = (
+    PathLike[str] | str | EvalLogInfo | Sequence[PathLike[str] | str | EvalLogInfo]
+)
+
+
+def verify_prerequisites() -> None:
+    # ensure we have all of the optional packages we need
+    required_packages: list[str] = []
+    try:
+        import pandas  # noqa: F401
+    except ImportError:
+        required_packages.append("pandas")
+
+    try:
+        import pyarrow  # noqa: F401
+    except ImportError:
+        required_packages.append("pyarrow")
+
+    if len(required_packages) > 0:
+        raise pip_dependency_error("inspect_ai.analysis", required_packages)
+
+    # enforce version constraints
+    verify_required_version("inspect_ai.analysis", "pandas", "2.0.0")
+    verify_required_version("inspect_ai.analysis", "pyarrow", "10.0.1")
+
+
+def resolve_logs(logs: LogPaths) -> list[str]:
+    # normalize to list of str
+    logs = [logs] if isinstance(logs, str | PathLike | EvalLogInfo) else logs
+    logs_str = [
+        Path(log).as_posix()
+        if isinstance(log, PathLike)
+        else log.name
+        if isinstance(log, EvalLogInfo)
+        else log
+        for log in logs
+    ]
+
+    # expand directories
+    log_paths: list[FileInfo] = []
+    for log_str in logs_str:
+        fs = filesystem(log_str)
+        info = fs.info(log_str)
+        if info.type == "directory":
+            log_paths.extend(
+                [fi for fi in fs.ls(info.name, recursive=True) if fi.type == "file"]
+            )
+        else:
+            log_paths.append(info)
+
+    log_files = log_files_from_ls(log_paths, sort=False)
+    return [log_file.name for log_file in log_files]
+
+
+def normalize_records(
+    records: list[dict[str, ColumnType]],
+) -> list[dict[str, ColumnType]]:
+    all_keys: set[str] = set()
+    for record in records:
+        all_keys.update(record.keys())
+    normalized_records = []
+    for record in records:
+        normalized_record = {key: record.get(key, None) for key in all_keys}
+        normalized_records.append(normalized_record)
+    return normalized_records
+
+
+def resolve_columns(
+    col_pattern: str, suffix: str, columns: list[str], processed_columns: list[str]
+) -> list[str]:
+    resolved_columns: list[str] = []
+
+    if "*" not in col_pattern:
+        # Regular column - check with suffix
+        col_with_suffix = f"{col_pattern}{suffix}"
+        if col_with_suffix in columns and col_with_suffix not in processed_columns:
+            resolved_columns.append(col_with_suffix)
+        # Then without suffix
+        elif col_pattern in columns and col_pattern not in processed_columns:
+            resolved_columns.append(col_pattern)
+    else:
+        # Wildcard pattern - check both with and without suffix
+        suffix_pattern = col_pattern + suffix
+        matching_with_suffix = match_col_pattern(
+            suffix_pattern, columns, processed_columns
+        )
+        matching_without_suffix = match_col_pattern(
+            col_pattern, columns, processed_columns
+        )
+
+        # Add all matches
+        matched_columns = sorted(set(matching_with_suffix + matching_without_suffix))
+        resolved_columns.extend(matched_columns)
+
+    return resolved_columns
+
+
+def match_col_pattern(
+    pattern: str, columns: list[str], processed_columns: list[str]
+) -> list[str]:
+    regex = _col_pattern_to_regex(pattern)
+    return [c for c in columns if regex.match(c) and c not in processed_columns]
+
+
+def _col_pattern_to_regex(pattern: str) -> Pattern[str]:
+    parts = []
+    for part in re.split(r"(\*)", pattern):
+        if part == "*":
+            parts.append(".*")
+        else:
+            parts.append(re.escape(part))
+    return re.compile("^" + "".join(parts) + "$")
+
+
+def add_unreferenced_columns(
+    columns: list[str], referenced_columns: list[str]
+) -> list[str]:
+    unreferenced_columns = sorted([c for c in columns if c not in referenced_columns])
+    return referenced_columns + unreferenced_columns
+
+
+def records_to_pandas(records: list[dict[str, ColumnType]]) -> "pd.DataFrame":
+    import pyarrow as pa
+
+    records = normalize_records(records)
+    table = pa.Table.from_pylist(records).to_pandas(types_mapper=arrow_types_mapper)
+    return table
+
+
+def arrow_types_mapper(
+    arrow_type: "pa.DataType",
+) -> "pd.api.extensions.ExtensionDtype" | None:
+    import pandas as pd
+    import pyarrow as pa
+
+    # convert str => str
+    if pa.types.is_string(arrow_type):
+        return pd.StringDtype()
+    # default conversion for other types
+    else:
+        return None

inspect_ai/analysis/beta/_dataframe/validate.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+from logging import getLogger
+from typing import Any, Iterator, Mapping, Type
+
+import jsonref  # type: ignore
+from jsonpath_ng import Fields, Index, JSONPath, Slice, Where, WhereNot  # type: ignore
+from jsonpath_ng.ext.filter import Filter  # type: ignore
+from pydantic import BaseModel
+
+logger = getLogger(__name__)
+
+Schema = Mapping[str, Any]
+
+
+def resolved_schema(model: Type[BaseModel]) -> Schema:
+    schema_dict = model.model_json_schema()
+    base = "file:///memory/inspect_schema.json"
+    schema: Schema = jsonref.replace_refs(
+        schema_dict, base_uri=base, jsonschema=True, proxies=False
+    )
+    return schema
+
+
+def jsonpath_in_schema(expr: JSONPath, schema: Schema) -> bool:
+    # don't validate unsupported constructs
+    if find_unsupported(expr):
+        return True
+
+    def descend(sch: Schema, tok: str | int | None) -> list[Schema]:
+        # First, branch through anyOf/oneOf/allOf
+        outs: list[Schema] = []
+        for branch in _expand_union(sch):
+            outs.extend(descend_concrete(branch, tok))
+        return outs
+
+    def descend_concrete(sch: Schema, tok: str | int | None) -> list[Schema]:
+        # totally open object – accept any child
+        if sch == {}:
+            return [{}]  # stay alive, accept any key
+
+        outs: list[Schema] = []
+
+        def open_dict(node: Schema) -> None:
+            """Append the schema that governs unknown keys.
+
+            - None / missing  -> open object  ->   {}
+            - True            -> open object  ->   {}
+            - Mapping         -> that mapping (could be {} or a real subschema)
+            - False           -> closed object ->   (do nothing)
+            """
+            if "additionalProperties" not in node:
+                if not node.get("properties"):
+                    outs.append({})
+            else:
+                ap = node["additionalProperties"]
+                if ap is True:
+                    outs.append({})
+                elif isinstance(ap, Mapping):  # {} or {...}
+                    outs.append(ap)
+                # ap is False  -> closed dict  ->  ignore
+
+        # Wildcard -----------------------------------------------------------
+        if tok is None:
+            if "properties" in sch:
+                outs.extend(sch["properties"].values())
+            if "object" in _types(sch):
+                open_dict(sch)
+            if "array" in _types(sch) and "items" in sch:
+                outs.extend(_normalize_items(sch["items"]))
+            return outs
+
+        # Property access ----------------------------------------------------
+        if isinstance(tok, str):
+            if "properties" in sch and tok in sch["properties"]:
+                outs.append(sch["properties"][tok])
+            elif "additionalProperties" in sch:  # PRESENCE, not truthiness
+                open_dict(sch)
+            elif "object" in _types(sch):
+                open_dict(sch)
+
+        # Array index --------------------------------------------------------
+        else:  # tok is int or None from an Index node
+            if "array" in _types(sch) and "items" in sch:
+                outs.extend(_normalize_items(sch["items"], index=tok))
+
+        return outs
+
+    def _types(sch: Schema) -> set[str]:
+        t = sch.get("type")
+        return set(t) if isinstance(t, list) else {t} if t else set()
+
+    def _normalize_items(items: Any, index: int | None = None) -> list[Schema]:
+        if isinstance(items, list):
+            if index is None:  # wildcard/slice
+                return items
+            if 0 <= index < len(items):
+                return [items[index]]
+            return []
+        if isinstance(items, Mapping):
+            return [items]
+        return []
+
+    states = [schema]
+    for tok in iter_tokens(expr):
+        next_states: list[Schema] = []
+        for st in states:
+            next_states.extend(descend(st, tok))
+        if not next_states:  # nothing matched this segment
+            return False
+        states = next_states
+    return True  # every segment found at least one schema
+
+
+def iter_tokens(node: JSONPath) -> Iterator[str | int | None]:
+    """Linearise a jsonpath-ng AST into a stream of tokens we care about."""
+    if hasattr(node, "left"):  # Child, Descendants, etc.
+        yield from iter_tokens(node.left)
+        yield from iter_tokens(node.right)
+    elif isinstance(node, Fields):
+        yield from node.fields  # e.g. ["foo"]
+    elif isinstance(node, Index):
+        yield node.index  # 0  /  -1  /  None for wildcard
+    elif isinstance(node, Slice):
+        yield None  # treat any slice as wildcard
+
+
+COMBINATORS = ("anyOf", "oneOf", "allOf")
+
+
+def _expand_union(sch: Schema) -> list[Schema]:
+    """Return sch itself or the list of subschemas if it is a combinator."""
+    for key in COMBINATORS:
+        if key in sch:
+            subs: list[Schema] = []
+            for sub in sch[key]:
+                # a sub-schema might itself be an anyOf/oneOf/allOf
+                subs.extend(_expand_union(sub))
+            return subs
+    return [sch]
+
+
+UNSUPPORTED: tuple[type[JSONPath], ...] = (
+    Filter,  # [?foo > 0]
+    Where,  # .foo[(@.bar < 42)]
+    WhereNot,
+    Slice,  # [1:5]  (wildcard “[*]” is Index/None, not Slice)
+)
+
+
+def find_unsupported(node: JSONPath) -> list[type[JSONPath]]:
+    """Return a list of node types present in `node` that we do not validate."""
+    bad: list[type[JSONPath]] = []
+    stack: list[JSONPath] = [node]
+    while stack:
+        n = stack.pop()
+        if isinstance(n, UNSUPPORTED):
+            bad.append(type(n))
+        # Drill into children (jsonpath-ng uses .left / .right / .child attributes)
+        for attr in ("left", "right", "child", "expression"):
+            stack.extend(
+                [getattr(n, attr)]
+                if hasattr(n, attr) and isinstance(getattr(n, attr), JSONPath)
+                else []
+            )
+        # handle containers like Fields(fields=[...]) and Index(index=[...])
+        if hasattr(n, "__dict__"):
+            for v in n.__dict__.values():
+                if isinstance(v, list):
+                    stack.extend(x for x in v if isinstance(x, JSONPath))
+    return bad