inspect-ai 0.3.53__py3-none-any.whl → 0.3.55__py3-none-any.whl

inspect_ai/_util/constants.py

@@ -14,12 +14,12 @@ DEFAULT_VIEW_PORT = 7575
 DEFAULT_SERVER_HOST = "127.0.0.1"
 HTTP = 15
 HTTP_LOG_LEVEL = "HTTP"
-SANDBOX = 17
-SANDBOX_LOG_LEVEL = "SANDBOX"
+TRACE = 13
+TRACE_LOG_LEVEL = "TRACE"
 ALL_LOG_LEVELS = [
     "DEBUG",
+    TRACE_LOG_LEVEL,
     HTTP_LOG_LEVEL,
-    SANDBOX_LOG_LEVEL,
     "INFO",
     "WARNING",
     "ERROR",

inspect_ai/_util/display.py

@@ -14,6 +14,7 @@ _display_type: DisplayType | None = None
 
 def init_display_type(display: str | None = None) -> DisplayType:
     global _display_type
+    global _display_metrics
     display = (
         display or os.environ.get("INSPECT_DISPLAY", DEFAULT_DISPLAY).lower().strip()
     )

inspect_ai/_util/logger.py

@@ -11,6 +11,7 @@ from logging import (
     getLevelName,
     getLogger,
 )
+from pathlib import Path
 
 import rich
 from rich.console import ConsoleRenderable
@@ -18,17 +19,20 @@ from rich.logging import RichHandler
 from rich.text import Text
 from typing_extensions import override
 
-from inspect_ai._util.constants import (
+from .constants import (
     ALL_LOG_LEVELS,
     DEFAULT_LOG_LEVEL,
     DEFAULT_LOG_LEVEL_TRANSCRIPT,
     HTTP,
     HTTP_LOG_LEVEL,
     PKG_NAME,
-    SANDBOX,
-    SANDBOX_LOG_LEVEL,
+    TRACE,
+    TRACE_LOG_LEVEL,
 )
-from inspect_ai._util.error import PrerequisiteError
+from .error import PrerequisiteError
+from .trace import TraceFileHandler, TraceFormatter, inspect_trace_dir
+
+TRACE_FILE_NAME = "trace.log"
 
 
 # log handler that filters messages to stderr and the log file
@@ -52,6 +56,24 @@ class LogHandler(RichHandler):
         else:
             self.file_logger_level = 0
 
+        # add a trace handler
+        default_trace_file = inspect_trace_dir() / TRACE_FILE_NAME
+        have_existing_trace_file = default_trace_file.exists()
+        env_trace_file = os.environ.get("INSPECT_TRACE_FILE", None)
+        trace_file = Path(env_trace_file) if env_trace_file else default_trace_file
+        trace_total_files = 10
+        self.trace_logger = TraceFileHandler(
+            trace_file.as_posix(),
+            backupCount=trace_total_files - 1,  # exclude the current file (10 total)
+        )
+        self.trace_logger.setFormatter(TraceFormatter())
+        if have_existing_trace_file:
+            self.trace_logger.doRollover()
+
+        # set trace level
+        trace_level = os.environ.get("INSPECT_TRACE_LEVEL", TRACE_LOG_LEVEL)
+        self.trace_logger_level = int(getLevelName(trace_level.upper()))
+
     @override
     def emit(self, record: LogRecord) -> None:
         # demote httpx and return notifications to log_level http
@@ -79,6 +101,10 @@ class LogHandler(RichHandler):
         ):
             self.file_logger.emit(record)
 
+        # write to trace if the trace level matches.
+        if self.trace_logger and record.levelno >= self.trace_logger_level:
+            self.trace_logger.emit(record)
+
         # eval log always gets info level and higher records
         # eval log only gets debug or http if we opt-in
         write = record.levelno >= self.transcript_levelno
@@ -95,12 +121,12 @@ def init_logger(
     log_level: str | None = None, log_level_transcript: str | None = None
 ) -> None:
     # backwards compatibility for 'tools'
-    if log_level == "tools":
-        log_level = "sandbox"
+    if log_level == "sandbox" or log_level == "tools":
+        log_level = "trace"
 
     # register http and tools levels
     addLevelName(HTTP, HTTP_LOG_LEVEL)
-    addLevelName(SANDBOX, SANDBOX_LOG_LEVEL)
+    addLevelName(TRACE, TRACE_LOG_LEVEL)
 
     def validate_level(option: str, level: str) -> None:
         if level not in ALL_LOG_LEVELS:
@@ -134,7 +160,7 @@ def init_logger(
         getLogger().addHandler(_logHandler)
 
     # establish default capture level
-    capture_level = min(HTTP, levelno)
+    capture_level = min(TRACE, levelno)
 
     # see all the messages (we won't actually display/write all of them)
     getLogger().setLevel(capture_level)

inspect_ai/_util/trace.py

@@ -0,0 +1,275 @@
+import asyncio
+import datetime
+import gzip
+import json
+import logging
+import os
+import shutil
+import time
+import traceback
+from contextlib import contextmanager
+from logging import Logger
+from logging.handlers import RotatingFileHandler
+from pathlib import Path
+from typing import Any, Generator, Literal, TextIO
+
+import jsonlines
+from pydantic import BaseModel, Field, JsonValue
+from shortuuid import uuid
+
+from .appdirs import inspect_data_dir
+from .constants import TRACE
+
+
+def inspect_trace_dir() -> Path:
+    return inspect_data_dir("traces")
+
+
+@contextmanager
+def trace_action(
+    logger: Logger, action: str, message: str, *args: Any, **kwargs: Any
+) -> Generator[None, None, None]:
+    trace_id = uuid()
+    start_monotonic = time.monotonic()
+    start_wall = time.time()
+    pid = os.getpid()
+    detail = message % args if args else message % kwargs if kwargs else message
+
+    def trace_message(event: str) -> str:
+        return f"{action}: {detail} ({event})"
+
+    logger.log(
+        TRACE,
+        trace_message("enter"),
+        extra={
+            "action": action,
+            "detail": detail,
+            "event": "enter",
+            "trace_id": str(trace_id),
+            "start_time": start_wall,
+            "pid": pid,
+        },
+    )
+
+    try:
+        yield
+        duration = time.monotonic() - start_monotonic
+        logger.log(
+            TRACE,
+            trace_message("exit"),
+            extra={
+                "action": action,
+                "detail": detail,
+                "event": "exit",
+                "trace_id": str(trace_id),
+                "duration": duration,
+                "pid": pid,
+            },
+        )
+    except (KeyboardInterrupt, asyncio.CancelledError):
+        duration = time.monotonic() - start_monotonic
+        logger.log(
+            TRACE,
+            trace_message("cancel"),
+            extra={
+                "action": action,
+                "detail": detail,
+                "event": "cancel",
+                "trace_id": str(trace_id),
+                "duration": duration,
+                "pid": pid,
+            },
+        )
+        raise
+    except TimeoutError:
+        duration = time.monotonic() - start_monotonic
+        logger.log(
+            TRACE,
+            trace_message("timeout"),
+            extra={
+                "action": action,
+                "detail": detail,
+                "event": "timeout",
+                "trace_id": str(trace_id),
+                "duration": duration,
+                "pid": pid,
+            },
+        )
+        raise
+    except Exception as ex:
+        duration = time.monotonic() - start_monotonic
+        logger.log(
+            TRACE,
+            trace_message("error"),
+            extra={
+                "action": action,
+                "detail": detail,
+                "event": "error",
+                "trace_id": str(trace_id),
+                "duration": duration,
+                "error": getattr(ex, "message", str(ex)) or repr(ex),
+                "error_type": type(ex).__name__,
+                "stacktrace": traceback.format_exc(),
+                "pid": pid,
+            },
+        )
+        raise
+
+
+def trace_message(
+    logger: Logger, category: str, message: str, *args: Any, **kwargs: Any
+) -> None:
+    logger.log(TRACE, f"[{category}] {message}", *args, **kwargs)
+
+
+class TraceFormatter(logging.Formatter):
+    def format(self, record: logging.LogRecord) -> str:
+        # Base log entry with standard fields
+        output: dict[str, JsonValue] = {
+            "timestamp": self.formatTime(record),
+            "level": record.levelname,
+            "message": record.getMessage(),  # This handles the % formatting of the message
+        }
+
+        # Add basic context its not a TRACE message
+        if record.levelname != "TRACE":
+            if hasattr(record, "module"):
+                output["module"] = record.module
+            if hasattr(record, "funcName"):
+                output["function"] = record.funcName
+            if hasattr(record, "lineno"):
+                output["line"] = record.lineno
+
+        # Add any structured fields from extra
+        elif hasattr(record, "action"):
+            # This is a trace_action log
+            for key in [
+                "action",
+                "detail",
+                "event",
+                "trace_id",
+                "start_time",
+                "duration",
+                "error",
+                "error_type",
+                "stacktrace",
+                "pid",
+            ]:
+                if hasattr(record, key):
+                    output[key] = getattr(record, key)
+
+        # Handle any unexpected extra attributes
+        for key, value in record.__dict__.items():
+            if key not in output and key not in (
+                "args",
+                "lineno",
+                "funcName",
+                "module",
+                "asctime",
+                "created",
+                "exc_info",
+                "exc_text",
+                "filename",
+                "levelno",
+                "levelname",
+                "msecs",
+                "msg",
+                "name",
+                "pathname",
+                "process",
+                "processName",
+                "relativeCreated",
+                "stack_info",
+                "thread",
+                "threadName",
+            ):
+                output[key] = value
+
+        return json.dumps(
+            output, default=str
+        )  # default=str handles non-serializable objects
+
+    def formatTime(self, record: logging.LogRecord, datefmt: str | None = None) -> str:
+        # ISO format with timezone
+        dt = datetime.datetime.fromtimestamp(record.created)
+        return dt.isoformat()
+
+
+class TraceRecord(BaseModel):
+    timestamp: str
+    level: str
+    message: str
+
+
+class SimpleTraceRecord(TraceRecord):
+    action: None = Field(default=None)
+
+
+class ActionTraceRecord(TraceRecord):
+    action: str
+    event: Literal["enter", "cancel", "error", "timeout", "exit"]
+    trace_id: str
+    detail: str = Field(default="")
+    start_time: float | None = Field(default=None)
+    duration: float | None = Field(default=None)
+    error: str | None = Field(default=None)
+    error_type: str | None = Field(default=None)
+    stacktrace: str | None = Field(default=None)
+    pid: int | None = Field(default=None)
+
+
+def read_trace_file(file: Path) -> list[TraceRecord]:
+    def read_file(f: TextIO) -> list[TraceRecord]:
+        jsonlines_reader = jsonlines.Reader(f)
+        trace_records: list[TraceRecord] = []
+        for trace in jsonlines_reader.iter(type=dict):
+            if "action" in trace:
+                trace_records.append(ActionTraceRecord(**trace))
+            else:
+                trace_records.append(SimpleTraceRecord(**trace))
+        return trace_records
+
+    if file.name.endswith(".gz"):
+        with gzip.open(file, "rt") as f:
+            return read_file(f)
+    else:
+        with open(file, "r") as f:
+            return read_file(f)
+
+
+class TraceFileHandler(RotatingFileHandler):
+    def __init__(
+        self,
+        filename: str,
+        mode: str = "a",
+        maxBytes: int = 0,
+        backupCount: int = 0,
+        encoding: str | None = None,
+        delay: bool = False,
+    ) -> None:
+        super().__init__(filename, mode, maxBytes, backupCount, encoding, delay)
+
+    def rotation_filename(self, default_name: str) -> str:
+        """
+        Returns the name of the rotated file.
+
+        Args:
+            default_name: The default name that would be used for rotation
+
+        Returns:
+            The modified filename with .gz extension
+        """
+        return default_name + ".gz"
+
+    def rotate(self, source: str, dest: str) -> None:
+        """
+        Compresses the source file and moves it to destination.
+
+        Args:
+            source: The source file to be compressed
+            dest: The destination path for the compressed file
+        """
+        with open(source, "rb") as f_in:
+            with gzip.open(dest, "wb") as f_out:
+                shutil.copyfileobj(f_in, f_out)
+        os.remove(source)

inspect_ai/log/_log.py

@@ -94,6 +94,9 @@ class EvalConfig(BaseModel):
     log_buffer: int | None = Field(default=None)
     """Number of samples to buffer before writing log file."""
 
+    score_display: bool | None = Field(default=None)
+    """Display scoring metrics realtime."""
+
     @property
     def max_messages(self) -> int | None:
         """Deprecated max_messages property."""

inspect_ai/log/_message.py

@@ -64,7 +64,7 @@ class LoggingMessage(BaseModel):
     ) -> dict[str, Any]:
         if "level" in values:
             level = values["level"]
-            if level == "tools":
-                values["level"] = "sandbox"
+            if level == "tools" or level == "sandbox":
+                values["level"] = "trace"
 
         return values

inspect_ai/log/_recorders/eval.py

@@ -17,6 +17,7 @@ from inspect_ai._util.content import ContentImage, ContentText
 from inspect_ai._util.error import EvalError
 from inspect_ai._util.file import FileSystem, async_fileystem, dirname, file, filesystem
 from inspect_ai._util.json import jsonable_python
+from inspect_ai._util.trace import trace_action
 from inspect_ai.model._chat_message import ChatMessage
 from inspect_ai.scorer._metric import Score
 
@@ -351,25 +352,13 @@ class ZipLogFile:
             self._temp_file.seek(0)
             log_bytes = self._temp_file.read()
 
-            # attempt async write
-            written = False
-            try:
-                if self._async_fs:
-                    await self._async_fs._pipe_file(self._file, log_bytes)
-                    written = True
-            except Exception as ex:
-                logger.warning(
-                    f"Error occurred during async write to {self._file}: {ex}. Falling back to sync write."
-                )
-
-            try:
-                # write sync if we need to
-                if not written:
+            with trace_action(logger, "Log Write", self._file):
+                try:
                     with file(self._file, "wb") as f:
                         f.write(log_bytes)
-            finally:
-                # re-open zip file w/ self.temp_file pointer at end
-                self._open()
+                finally:
+                    # re-open zip file w/ self.temp_file pointer at end
+                    self._open()
 
     async def close(self) -> EvalLog:
         async with self._lock:

inspect_ai/log/_recorders/json.py

@@ -15,6 +15,7 @@ from inspect_ai._util.file import (
     file,
     filesystem,
 )
+from inspect_ai._util.trace import trace_action
 
 from .._log import (
     EvalLog,
@@ -181,23 +182,24 @@ class JSONRecorder(FileRecorder):
         # get log as bytes
         log_bytes = eval_log_json(log)
 
-        # try to write async for async filesystems
-        written = False
-        try:
-            fs = filesystem(location)
-            if fs.is_async():
-                async with async_fileystem(location) as async_fs:
-                    await async_fs._pipe_file(location, log_bytes)
-                    written = True
-        except Exception as ex:
-            logger.warning(
-                f"Error occurred during async write to {location}: {ex}. Falling back to sync write."
-            )
-
-        # otherwise use sync
-        if not written:
-            with file(location, "wb") as f:
-                f.write(log_bytes)
+        with trace_action(logger, "Log Write", location):
+            # try to write async for async filesystems
+            written = False
+            try:
+                fs = filesystem(location)
+                if fs.is_async():
+                    async with async_fileystem(location) as async_fs:
+                        await async_fs._pipe_file(location, log_bytes)
+                        written = True
+            except Exception as ex:
+                logger.warning(
+                    f"Error occurred during async write to {location}: {ex}. Falling back to sync write."
+                )
+
+            # otherwise use sync
+            if not written:
+                with file(location, "wb") as f:
+                    f.write(log_bytes)
 
 
 def _validate_version(ver: int) -> None:

inspect_ai/model/_cache.py

@@ -6,10 +6,12 @@ from datetime import datetime, timezone
 from hashlib import md5
 from pathlib import Path
 from shutil import rmtree
+from typing import Any
 
 from dateutil.relativedelta import relativedelta
 
 from inspect_ai._util.appdirs import inspect_cache_dir
+from inspect_ai._util.trace import trace_message
 from inspect_ai.tool import ToolChoice, ToolInfo
 
 from ._chat_message import ChatMessage
@@ -19,6 +21,10 @@ from ._model_output import ModelOutput
 logger = logging.getLogger(__name__)
 
 
+def trace(msg: str, *args: Any) -> None:
+    trace_message(logger, "Cache", msg, *args)
+
+
 def _path_is_in_cache(path: Path | str) -> bool:
     """This ensures the path is in our cache directory, just in case the `model` is ../../../home/ubuntu/maliciousness"""
     if isinstance(path, str):
@@ -153,7 +159,7 @@ def _cache_key(entry: CacheEntry) -> str:
 
     base_string = "|".join([str(component) for component in components])
 
-    logger.debug(_cache_key_debug_string([str(component) for component in components]))
+    trace(_cache_key_debug_string([str(component) for component in components]))
 
     return md5(base_string.encode("utf-8")).hexdigest()
 
@@ -192,11 +198,11 @@ def cache_store(
 
         with open(filename, "wb") as f:
             expiry = _cache_expiry(entry.policy)
-            logger.debug("Storing in cache: %s (expires: %s)", filename, expiry)
+            trace("Storing in cache: %s (expires: %s)", filename, expiry)
             pickle.dump((expiry, output), f)
         return True
     except Exception as e:
-        logger.debug(f"Failed to cache {filename}: {e}")
+        trace(f"Failed to cache {filename}: {e}")
         return False
 
 
@@ -204,12 +210,12 @@ def cache_fetch(entry: CacheEntry) -> ModelOutput | None:
     """Fetch a value from the cache directory."""
     filename = cache_path(model=entry.model) / _cache_key(entry)
     try:
-        logger.debug("Fetching from cache: %s", filename)
+        trace("Fetching from cache: %s", filename)
 
         with open(filename, "rb") as f:
             expiry, output = pickle.load(f)
             if not isinstance(output, ModelOutput):
-                logger.debug(
+                trace(
                     "Unexpected cached type, can only fetch ModelOutput: %s (%s)",
                     type(output),
                     filename,
@@ -217,7 +223,7 @@ def cache_fetch(entry: CacheEntry) -> ModelOutput | None:
                 return None
 
             if _is_expired(expiry):
-                logger.debug("Cache expired for %s (%s)", filename, expiry)
+                trace("Cache expired for %s (%s)", filename, expiry)
                 # If it's expired, no point keeping it as we'll never access it
                 # successfully again.
                 filename.unlink(missing_ok=True)
@@ -225,7 +231,7 @@ def cache_fetch(entry: CacheEntry) -> ModelOutput | None:
 
             return output
     except Exception as e:
-        logger.debug(f"Failed to fetch from cache {filename}: {e}")
+        trace(f"Failed to fetch from cache {filename}: {e}")
         return None
 
 
@@ -235,7 +241,7 @@ def cache_clear(model: str = "") -> bool:
         path = cache_path(model)
 
         if (model == "" or _path_is_in_cache(path)) and path.exists():
-            logger.debug("Clearing cache: %s", path)
+            trace("Clearing cache: %s", path)
             rmtree(path)
             return True
 
@@ -351,24 +357,24 @@ def cache_list_expired(filter_by: list[str] = []) -> list[Path]:
         # "../../foo/bar") but we don't want to search the entire cache
         return []
 
-    logger.debug("Filtering by paths: %s", filter_by_paths)
+    trace("Filtering by paths: %s", filter_by_paths)
     for dirpath, _dirnames, filenames in os.walk(cache_path()):
         if filter_by_paths and Path(dirpath) not in filter_by_paths:
-            logger.debug("Skipping path %s", dirpath)
+            trace("Skipping path %s", dirpath)
             continue
 
-        logger.debug("Checking dirpath %s", dirpath)
+        trace("Checking dirpath %s", dirpath)
         for filename in filenames:
             path = Path(dirpath) / filename
-            logger.debug("Checking path %s", path)
+            trace("Checking path %s", path)
             try:
                 with open(path, "rb") as f:
                     expiry, _cache_entry = pickle.load(f)
                     if _is_expired(expiry):
-                        logger.debug("Expired cache entry found: %s (%s)", path, expiry)
+                        trace("Expired cache entry found: %s (%s)", path, expiry)
                         expired_cache_entries.append(path)
             except Exception as e:
-                logger.debug("Failed to load cached item %s: %s", path, e)
+                trace("Failed to load cached item %s: %s", path, e)
                 continue
 
     return expired_cache_entries
@@ -389,8 +395,8 @@ def cache_prune(files: list[Path] = []) -> None:
             with open(file, "rb") as f:
                 expiry, _cache_entry = pickle.load(f)
                 if _is_expired(expiry):
-                    logger.debug("Pruning expired cache: %s", file)
+                    trace("Pruning expired cache: %s", file)
                     file.unlink(missing_ok=True)
         except Exception as e:
-            logger.debug("Failed to prune cache %s: %s", file, e)
+            trace("Failed to prune cache %s: %s", file, e)
             continue

inspect_ai/model/_call_tools.py

@@ -1,6 +1,7 @@
 import asyncio
 import inspect
 from dataclasses import is_dataclass
+from logging import getLogger
 from textwrap import dedent
 from typing import (
     Any,
@@ -19,7 +20,9 @@ from jsonschema import Draft7Validator
 from pydantic import BaseModel
 
 from inspect_ai._util.content import Content, ContentImage, ContentText
+from inspect_ai._util.format import format_function_call
 from inspect_ai._util.text import truncate_string_to_bytes
+from inspect_ai._util.trace import trace_action
 from inspect_ai.model._trace import trace_tool_mesage
 from inspect_ai.tool import Tool, ToolCall, ToolError, ToolInfo
 from inspect_ai.tool._tool import (
@@ -35,6 +38,8 @@ from inspect_ai.util import OutputLimitExceededError
 from ._chat_message import ChatMessageAssistant, ChatMessageTool
 from ._generate_config import active_generate_config
 
+logger = getLogger(__name__)
+
 
 async def call_tools(
     message: ChatMessageAssistant,
@@ -215,7 +220,10 @@ async def call_tool(tools: list[ToolDef], message: str, call: ToolCall) -> Any:
     arguments = tool_params(call.arguments, tool_def.tool)
 
     # call the tool
-    result = await tool_def.tool(**arguments)
+    with trace_action(
+        logger, "Tool Call", format_function_call(tool_def.name, arguments, width=1000)
+    ):
+        result = await tool_def.tool(**arguments)
 
     # return result
     return result