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

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

@@ -0,0 +1,87 @@
+from datetime import datetime
+from typing import Any, Callable, Mapping, Type
+
+from jsonpath_ng import JSONPath  # type: ignore
+from pydantic import JsonValue
+from typing_extensions import override
+
+from inspect_ai.log._transcript import Event
+
+from ..columns import Column, ColumnType
+from .extract import (
+    completion_as_str,
+    model_event_input_as_str,
+    tool_choice_as_str,
+    tool_view_as_str,
+)
+
+
+class EventColumn(Column):
+    """Column which maps to `Event`."""
+
+    def __init__(
+        self,
+        name: str,
+        *,
+        path: str | JSONPath | Callable[[Event], JsonValue],
+        required: bool = False,
+        default: JsonValue | None = None,
+        type: Type[ColumnType] | None = None,
+        value: Callable[[JsonValue], JsonValue] | None = None,
+    ) -> None:
+        super().__init__(
+            name=name,
+            path=path if not callable(path) else None,
+            required=required,
+            default=default,
+            type=type,
+            value=value,
+        )
+        self._extract_event = path if callable(path) else None
+
+    @override
+    def path_schema(self) -> Mapping[str, Any] | None:
+        return None
+
+
+EventInfo: list[Column] = [
+    EventColumn("event", path="event"),
+    EventColumn("span_id", path="span_id"),
+]
+"""Event basic information columns."""
+
+EventTiming: list[Column] = [
+    EventColumn("timestamp", path="timestamp", type=datetime),
+    EventColumn("completed", path="completed", type=datetime),
+    EventColumn("working_start", path="working_start"),
+    EventColumn("working_time", path="working_time"),
+]
+"""Event timing columns."""
+
+ModelEventColumns: list[Column] = [
+    EventColumn("model_event_model", path="model"),
+    EventColumn("model_event_role", path="role"),
+    EventColumn("model_event_input", path=model_event_input_as_str),
+    EventColumn("model_event_tools", path="tools"),
+    EventColumn("model_event_tool_choice", path=tool_choice_as_str),
+    EventColumn("model_event_config", path="config"),
+    EventColumn("model_event_usage", path="output.usage"),
+    EventColumn("model_event_time", path="output.time"),
+    EventColumn("model_event_completion", path=completion_as_str),
+    EventColumn("model_event_retries", path="retries"),
+    EventColumn("model_event_error", path="error"),
+    EventColumn("model_event_cache", path="cache"),
+    EventColumn("model_event_call", path="call"),
+]
+"""Model event columns."""
+
+ToolEventColumns: list[Column] = [
+    EventColumn("tool_event_function", path="function"),
+    EventColumn("tool_event_arguments", path="arguments"),
+    EventColumn("tool_event_view", path=tool_view_as_str),
+    EventColumn("tool_event_result", path="result"),
+    EventColumn("tool_event_truncated", path="truncated"),
+    EventColumn("tool_event_error_type", path="error.type"),
+    EventColumn("tool_event_error_message", path="error.message"),
+]
+"""Tool event columns."""

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

@@ -0,0 +1,26 @@
+from inspect_ai.log._transcript import ModelEvent, ToolEvent
+
+from ..extract import messages_as_str
+
+
+def model_event_input_as_str(event: ModelEvent) -> str:
+    return messages_as_str(event.input)
+
+
+def tool_choice_as_str(event: ModelEvent) -> str:
+    if isinstance(event.tool_choice, str):
+        return event.tool_choice
+    else:
+        return event.tool_choice.name
+
+
+def completion_as_str(event: ModelEvent) -> str:
+    return event.output.completion
+
+
+def tool_view_as_str(event: ToolEvent) -> str | None:
+    if event.view is not None:
+        title = f"{event.view.title}\n\n" if event.view.title is not None else ""
+        return f"{title}{event.view.content}"
+    else:
+        return None

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

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable, Literal, TypeAlias
+
+from inspect_ai.analysis.beta._dataframe.events.columns import EventInfo
+from inspect_ai.log._file import list_eval_logs
+from inspect_ai.log._transcript import Event
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+from typing_extensions import overload
+
+from ..columns import Column, ColumnErrors
+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]
+)
+"""Filter for `events_df()` rows."""
+
+
+@overload
+def events_df(
+    logs: LogPaths = list_eval_logs(),
+    columns: list[Column] = EventInfo,
+    filter: EventFilter | None = None,
+    strict: Literal[True] = True,
+) -> "pd.DataFrame": ...
+
+
+@overload
+def events_df(
+    logs: LogPaths = list_eval_logs(),
+    columns: list[Column] = EventInfo,
+    filter: EventFilter | None = None,
+    strict: Literal[False] = False,
+) -> tuple["pd.DataFrame", ColumnErrors]: ...
+
+
+def events_df(
+    logs: LogPaths = list_eval_logs(),
+    columns: list[Column] = EventInfo,
+    filter: EventFilter | None = None,
+    strict: bool = True,
+) -> "pd.DataFrame" | tuple["pd.DataFrame", ColumnErrors]:
+    """Read a dataframe containing events 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.
+       filter: List of event types to include or callable that performs the filter.
+       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.
+    """
+    verify_prerequisites()
+
+    # resolve filter/detail
+    if filter is None:
+        detail = EventsDetail(filter=lambda e: True)
+    elif callable(filter):
+        detail = EventsDetail(filter=filter)
+    else:
+        detail = EventsDetail(filter=lambda e: e.event in filter)
+
+    return _read_samples_df(
+        logs=logs,
+        columns=columns,
+        strict=strict,
+        detail=detail,
+    )

inspect_ai/analysis/beta/_dataframe/extract.py

@@ -0,0 +1,73 @@
+import hashlib
+import uuid
+from typing import Any, cast
+
+import shortuuid
+from pydantic import BaseModel, JsonValue
+
+from inspect_ai.model._chat_message import (
+    ChatMessage,
+    ChatMessageAssistant,
+    ChatMessageTool,
+    ChatMessageUser,
+)
+
+
+def model_to_record(model: BaseModel) -> dict[str, JsonValue]:
+    return cast(dict[str, JsonValue], model.model_dump(mode="json", exclude_none=True))
+
+
+def list_as_str(x: JsonValue) -> str:
+    return ",".join([str(e) for e in (x if isinstance(x, list) else [x])])
+
+
+def score_values(x: JsonValue) -> dict[str, JsonValue]:
+    scores = cast(dict[str, Any], x)
+    return {k: v["value"] for k, v in scores.items()}
+
+
+def auto_id(base: str, index: str) -> str:
+    seed = f"{base}_{index}"
+    hash_bytes = hashlib.md5(seed.encode("utf-8")).digest()
+    long_uuid = uuid.UUID(bytes=hash_bytes)
+    return shortuuid.encode(long_uuid)
+
+
+def messages_as_str(messages: str | list[ChatMessage]) -> str:
+    if isinstance(messages, str):
+        messages = [ChatMessageUser(content=messages)]
+    return "\n\n".join([message_as_str(message) for message in messages])
+
+
+def message_as_str(message: ChatMessage) -> str:
+    transcript: list[str] = []
+    role = message.role
+    content = message.text.strip() if message.text else ""
+
+    # assistant messages with tool calls
+    if isinstance(message, ChatMessageAssistant) and message.tool_calls is not None:
+        entry = f"{role}:\n{content}\n"
+
+        for tool in message.tool_calls:
+            func_name = tool.function
+            args = tool.arguments
+
+            if isinstance(args, dict):
+                args_text = "\n".join(f"{k}: {v}" for k, v in args.items())
+                entry += f"\nTool Call: {func_name}\nArguments:\n{args_text}"
+            else:
+                entry += f"\nTool Call: {func_name}\nArguments: {args}"
+
+        transcript.append(entry)
+
+    # tool responses with errors
+    elif isinstance(message, ChatMessageTool) and message.error is not None:
+        func_name = message.function or "unknown"
+        entry = f"{role}:\n{content}\n\nError in tool call '{func_name}':\n{message.error.message}\n"
+        transcript.append(entry)
+
+    # normal messages
+    else:
+        transcript.append(f"{role}:\n{content}\n")
+
+    return "\n".join(transcript)

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

@@ -0,0 +1,60 @@
+from typing import Any, Callable, Mapping, Type
+
+from jsonpath_ng import JSONPath  # type: ignore
+from pydantic import JsonValue
+from typing_extensions import override
+
+from inspect_ai.model._chat_message import ChatMessage
+
+from ..columns import Column, ColumnType
+from .extract import (
+    message_text,
+    message_tool_calls,
+)
+
+
+class MessageColumn(Column):
+    """Column which maps to `ChatMessage`."""
+
+    def __init__(
+        self,
+        name: str,
+        *,
+        path: str | JSONPath | Callable[[ChatMessage], JsonValue],
+        required: bool = False,
+        default: JsonValue | None = None,
+        type: Type[ColumnType] | None = None,
+        value: Callable[[JsonValue], JsonValue] | None = None,
+    ) -> None:
+        super().__init__(
+            name=name,
+            path=path if not callable(path) else None,
+            required=required,
+            default=default,
+            type=type,
+            value=value,
+        )
+        self._extract_message = path if callable(path) else None
+
+    @override
+    def path_schema(self) -> Mapping[str, Any] | None:
+        return None
+
+
+MessageContent: list[Column] = [
+    MessageColumn("role", path="role", required=True),
+    MessageColumn("source", path="source"),
+    MessageColumn("content", path=message_text),
+]
+"""Message content columns."""
+
+MessageToolCalls: list[Column] = [
+    MessageColumn("tool_calls", path=message_tool_calls),
+    MessageColumn("tool_call_id", path="tool_call_id"),
+    MessageColumn("tool_call_function", path="function"),
+    MessageColumn("tool_call_error", path="error.message"),
+]
+"""Message tool call columns."""
+
+MessageColumns: list[Column] = MessageContent + MessageToolCalls
+"""Chat message columns."""

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

@@ -0,0 +1,21 @@
+from inspect_ai._util.format import format_function_call
+from inspect_ai.model._chat_message import ChatMessage, ChatMessageAssistant
+
+
+def message_text(message: ChatMessage) -> str:
+    return message.text
+
+
+def message_tool_calls(message: ChatMessage) -> str | None:
+    if isinstance(message, ChatMessageAssistant) and message.tool_calls is not None:
+        tool_calls = "\n".join(
+            [
+                format_function_call(
+                    tool_call.function, tool_call.arguments, width=1000
+                )
+                for tool_call in message.tool_calls
+            ]
+        )
+        return tool_calls
+    else:
+        return None

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

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable, Literal, TypeAlias
+
+from inspect_ai.log._file import list_eval_logs
+from inspect_ai.model._chat_message import ChatMessage
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+from typing_extensions import overload
+
+from ..columns import Column, ColumnErrors
+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]
+)
+"""Filter for `messages_df()` rows."""
+
+
+@overload
+def messages_df(
+    logs: LogPaths = list_eval_logs(),
+    columns: list[Column] = MessageColumns,
+    filter: MessageFilter | None = None,
+    strict: Literal[True] = True,
+) -> "pd.DataFrame": ...
+
+
+@overload
+def messages_df(
+    logs: LogPaths = list_eval_logs(),
+    columns: list[Column] = MessageColumns,
+    filter: MessageFilter | None = None,
+    strict: Literal[False] = False,
+) -> tuple["pd.DataFrame", ColumnErrors]: ...
+
+
+def messages_df(
+    logs: LogPaths = list_eval_logs(),
+    columns: list[Column] = MessageColumns,
+    filter: MessageFilter | None = None,
+    strict: bool = True,
+) -> "pd.DataFrame" | tuple["pd.DataFrame", ColumnErrors]:
+    """Read a dataframe containing messages 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.
+       filter: List of message role types to include or callable that performs the filter.
+       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.
+    """
+    verify_prerequisites()
+
+    # resolve filter/detail
+    if filter is None:
+        detail = MessagesDetail(filter=lambda m: True)
+    elif callable(filter):
+        detail = MessagesDetail(filter=filter)
+    else:
+        detail = MessagesDetail(filter=lambda m: m.role in filter)
+
+    return _read_samples_df(
+        logs=logs,
+        columns=columns,
+        strict=strict,
+        detail=detail,
+    )

inspect_ai/analysis/beta/_dataframe/progress.py

@@ -0,0 +1,26 @@
+from contextlib import contextmanager
+from typing import Iterator
+
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TaskID,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
+
+
+@contextmanager
+def import_progress(
+    description: str, total: float | None
+) -> Iterator[tuple[Progress, TaskID]]:
+    with Progress(
+        TextColumn("[progress.description]{task.description:<18}"),
+        BarColumn(),
+        TaskProgressColumn(),
+        TimeElapsedColumn(),
+        transient=True,
+    ) as progress:
+        task_id = progress.add_task(description, total=total)
+        yield progress, task_id