inspect-ai 0.3.81__py3-none-any.whl → 0.3.83__py3-none-any.whl

inspect_ai/log/_recorders/buffer/filestore.py

@@ -0,0 +1,259 @@
+import os
+import tempfile
+from logging import getLogger
+from pathlib import Path
+from typing import Literal
+from zipfile import ZIP_DEFLATED, ZipFile
+
+from pydantic import BaseModel, Field
+from typing_extensions import override
+
+from inspect_ai._display.core.display import TaskDisplayMetric
+from inspect_ai._util.constants import DEFAULT_LOG_SHARED, EVAL_LOG_FORMAT
+from inspect_ai._util.file import FileSystem, basename, dirname, file, filesystem
+from inspect_ai._util.json import to_json_safe, to_json_str_safe
+from inspect_ai.log._file import read_eval_log
+
+from ..types import SampleSummary
+from .types import SampleBuffer, SampleData, Samples
+
+logger = getLogger(__name__)
+
+
+class Segment(BaseModel):
+    id: int
+    last_event_id: int
+    last_attachment_id: int
+
+
+class SegmentFile(BaseModel):
+    id: str | int
+    epoch: int
+    data: SampleData
+
+
+class SampleManifest(BaseModel):
+    summary: SampleSummary
+    segments: list[int] = Field(default_factory=list)
+
+
+class Manifest(BaseModel):
+    metrics: list[TaskDisplayMetric] = Field(default_factory=list)
+    samples: list[SampleManifest] = Field(default_factory=list)
+    segments: list[Segment] = Field(default_factory=list)
+
+
+MANIFEST = "manifest.json"
+
+
+class SampleBufferFilestore(SampleBuffer):
+    def __init__(
+        self,
+        location: str,
+        *,
+        create: bool = True,
+        update_interval: int = DEFAULT_LOG_SHARED,
+    ) -> None:
+        self._fs = filesystem(location)
+        self._dir = f"{sample_buffer_dir(dirname(location), self._fs)}{self._fs.sep}{os.path.splitext(basename(location))[0]}{self._fs.sep}"
+        self.update_interval = update_interval
+
+        if create:
+            self._fs.mkdir(self._dir, exist_ok=True)
+
+            # place a file in the dir to force it to be created
+            self._fs.touch(f"{self._dir}.keep")
+
+    def write_manifest(self, manifest: Manifest) -> None:
+        with file(self._manifest_file(), "wb") as f:
+            f.write(to_json_safe(manifest))
+
+    def write_segment(self, id: int, files: list[SegmentFile]) -> None:
+        # write the file locally
+        with tempfile.NamedTemporaryFile(mode="wb", delete=False) as segment_file:
+            name = segment_file.name
+            with ZipFile(
+                segment_file, mode="w", compression=ZIP_DEFLATED, compresslevel=5
+            ) as zip:
+                for sf in files:
+                    zip.writestr(
+                        segment_file_name(sf.id, sf.epoch),
+                        to_json_str_safe(sf.data),
+                    )
+            segment_file.flush()
+            os.fsync(segment_file.fileno())
+
+        # write then move for atomicity
+        try:
+            with open(name, "rb") as zf:
+                with file(f"{self._dir}{segment_name(id)}", "wb") as f:
+                    f.write(zf.read())
+                    f.flush()
+        finally:
+            os.unlink(name)
+
+    def read_manifest(self) -> Manifest | None:
+        try:
+            with file(self._manifest_file(), "r") as f:
+                contents = f.read()
+                return Manifest.model_validate_json(contents)
+        except FileNotFoundError:
+            return None
+
+    def read_segment_data(
+        self, id: int, sample_id: str | int, epoch_id: int
+    ) -> SampleData:
+        segment_file = f"{self._dir}{segment_name(id)}"
+        with file(segment_file, "rb") as f:
+            with ZipFile(f, mode="r") as zip:
+                with zip.open(segment_file_name(sample_id, epoch_id), "r") as sf:
+                    return SampleData.model_validate_json(sf.read())
+
+    def cleanup(self) -> None:
+        cleanup_sample_buffer_filestore(self._dir, self._fs)
+
+    @classmethod
+    @override
+    def running_tasks(cls, log_dir: str) -> list[str] | None:
+        buffer_dir = Path(sample_buffer_dir(log_dir))
+        if buffer_dir.exists():
+            return [
+                f"{basename(path.name)}.{EVAL_LOG_FORMAT}"
+                for path in buffer_dir.iterdir()
+                if path.is_dir()
+            ]
+        else:
+            return None
+
+    @override
+    def get_samples(
+        self, etag: str | None = None
+    ) -> Samples | Literal["NotModified"] | None:
+        # get the etag on the filestore
+        try:
+            info = self._fs.info(self._manifest_file())
+            fs_etag = info.etag or f"{info.mtime}{info.size}"
+        except FileNotFoundError:
+            return None
+
+        # if the etag matches then return not modified
+        if etag == fs_etag:
+            return "NotModified"
+
+        # read the manifest
+        manifest = self.read_manifest()
+        if manifest is None:
+            return None
+
+        # provide samples + etag from the manifest
+        return Samples(
+            samples=[sm.summary for sm in manifest.samples],
+            metrics=manifest.metrics,
+            refresh=self.update_interval,
+            etag=fs_etag,
+        )
+
+    @override
+    def get_sample_data(
+        self,
+        id: str | int,
+        epoch: int,
+        after_event_id: int | None = None,
+        after_attachment_id: int | None = None,
+    ) -> SampleData | None:
+        # read the manifest
+        manifest = self.read_manifest()
+        if manifest is None:
+            return None
+
+        # find this sample in the manifest
+        sample = next(
+            (
+                sample
+                for sample in manifest.samples
+                if sample.summary.id == id and sample.summary.epoch == epoch
+            ),
+            None,
+        )
+        if sample is None:
+            return None
+
+        # determine which segments we need to return in order to
+        # satisfy the after_event_id and after_attachment_id
+        after_event_id = after_event_id or -1
+        after_attachment_id = after_attachment_id or -1
+        segments = [
+            segment for segment in manifest.segments if segment.id in sample.segments
+        ]
+        segments = [
+            segment
+            for segment in segments
+            if segment.last_event_id > after_event_id
+            or segment.last_attachment_id > after_attachment_id
+        ]
+
+        # collect data from the segments
+        sample_data = SampleData(events=[], attachments=[])
+        for segment in segments:
+            data = self.read_segment_data(segment.id, id, epoch)
+            sample_data.events.extend(data.events)
+            sample_data.attachments.extend(data.attachments)
+
+        return sample_data
+
+    def _manifest_file(self) -> str:
+        return f"{self._dir}{MANIFEST}"
+
+
+def cleanup_sample_buffer_filestores(log_dir: str) -> None:
+    # read log buffer dirs (bail if there is no buffer_dir)
+    fs = filesystem(log_dir)
+    buffer_dir = sample_buffer_dir(log_dir, fs)
+    try:
+        log_buffers = [
+            buffer for buffer in fs.ls(buffer_dir) if buffer.type == "directory"
+        ]
+    except FileNotFoundError:
+        return
+
+    # for each buffer dir, confirm there is a running .eval file
+    # (remove the buffer dir if there is no .eval or the eval is finished)
+    for log_buffer in log_buffers:
+        try:
+            log_file = f"{log_dir}{fs.sep}{basename(log_buffer.name)}.{EVAL_LOG_FORMAT}"
+            log_header = read_eval_log(log_file, header_only=True)
+            if log_header.status != "started":
+                cleanup_sample_buffer_filestore(log_buffer.name, fs)
+
+        except FileNotFoundError:
+            cleanup_sample_buffer_filestore(log_buffer.name, fs)
+
+    # remove the .buffer dir if it's empty
+    try:
+        if len(fs.ls(buffer_dir)) == 0:
+            fs.rm(buffer_dir, recursive=True)
+    except FileNotFoundError:
+        pass
+
+
+def cleanup_sample_buffer_filestore(buffer_dir: str, fs: FileSystem) -> None:
+    try:
+        fs.rm(buffer_dir, recursive=True)
+    except Exception as ex:
+        logger.warning(
+            f"Error cleaning up sample buffer database at {buffer_dir}: {ex}"
+        )
+
+
+def segment_name(id: int) -> str:
+    return f"segment.{id}.zip"
+
+
+def segment_file_name(id: str | int, epoch: int) -> str:
+    return f"{id}_{epoch}.json"
+
+
+def sample_buffer_dir(log_dir: str, fs: FileSystem | None = None) -> str:
+    log_dir = log_dir.rstrip("/\\")
+    fs = fs or filesystem(log_dir)
+    return f"{log_dir}{fs.sep}.buffer"

inspect_ai/log/_recorders/buffer/types.py

@@ -0,0 +1,84 @@
+import abc
+from typing import Literal, TypeAlias
+
+from pydantic import BaseModel, JsonValue
+
+from inspect_ai._display.core.display import TaskDisplayMetric
+
+from ..types import SampleSummary
+
+JsonData: TypeAlias = dict[str, JsonValue]
+
+
+class Samples(BaseModel):
+    samples: list[SampleSummary]
+    metrics: list[TaskDisplayMetric]
+    refresh: int
+    etag: str
+
+
+class EventData(BaseModel):
+    id: int
+    event_id: str
+    sample_id: str
+    epoch: int
+    event: JsonData
+
+
+class AttachmentData(BaseModel):
+    id: int
+    sample_id: str
+    epoch: int
+    hash: str
+    content: str
+
+
+class SampleData(BaseModel):
+    events: list[EventData]
+    attachments: list[AttachmentData]
+
+
+class SampleBuffer(abc.ABC):
+    @classmethod
+    @abc.abstractmethod
+    def running_tasks(cls, log_dir: str) -> list[str] | None: ...
+
+    @abc.abstractmethod
+    def get_samples(
+        self, etag: str | None = None
+    ) -> Samples | Literal["NotModified"] | None:
+        """Get the manifest of all running samples.
+
+        Args:
+          etag: Optional etag (returned in `Samples`) for checking
+            whether there are any changes in the datatabase.
+
+        Returns:
+          - `Samples` if the database exists and has updates
+          - "NotModifed" if the database exists and has no updates.
+          - None if the database no longer exists
+
+        """
+        ...
+
+    @abc.abstractmethod
+    def get_sample_data(
+        self,
+        id: str | int,
+        epoch: int,
+        after_event_id: int | None = None,
+        after_attachment_id: int | None = None,
+    ) -> SampleData | None:
+        """Get event and attachment data for a sample.
+
+        Args:
+          id: Sample id
+          epoch: Sample epoch
+          after_event_id: Optional. Fetch only event data greater than this id.
+          after_attachment_id: Optioinal. Fetch only attachment data greater than this id.
+
+        Returns:
+          - `SampleData` with event and attachment data.
+          - None if the database no longer exists
+        """
+        ...

inspect_ai/log/_recorders/eval.py

@@ -23,7 +23,6 @@ from inspect_ai._util.file import FileSystem, 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
 
 from .._log import (
     EvalLog,
@@ -36,20 +35,11 @@ from .._log import (
     sort_samples,
 )
 from .file import FileRecorder
+from .types import SampleSummary
 
 logger = getLogger(__name__)
 
 
-class SampleSummary(BaseModel):
-    id: int | str
-    epoch: int
-    input: str | list[ChatMessage]
-    target: str | list[str]
-    scores: dict[str, Score] | None = Field(default=None)
-    error: str | None = Field(default=None)
-    limit: str | None = Field(default=None)
-
-
 class LogStart(BaseModel):
     version: int
     eval: EvalSpec
@@ -331,6 +321,7 @@ class ZipLogFile:
                         epoch=sample.epoch,
                         input=text_inputs(sample.input),
                         target=sample.target,
+                        completed=True,
                         scores=sample.scores,
                         error=sample.error.message
                         if sample.error is not None

inspect_ai/log/_recorders/types.py

@@ -0,0 +1,30 @@
+from pydantic import BaseModel, Field, model_validator
+
+from inspect_ai.log._transcript import Event
+from inspect_ai.model._chat_message import ChatMessage
+from inspect_ai.scorer._metric import Score
+
+
+class SampleEvent(BaseModel):
+    id: str | int
+    epoch: int
+    event: Event
+
+
+class SampleSummary(BaseModel):
+    id: int | str
+    epoch: int
+    input: str | list[ChatMessage]
+    target: str | list[str]
+    completed: bool = Field(default=False)
+    scores: dict[str, Score] | None = Field(default=None)
+    error: str | None = Field(default=None)
+    limit: str | None = Field(default=None)
+
+    @model_validator(mode="after")
+    def thin_scores(self) -> "SampleSummary":
+        if self.scores is not None:
+            self.scores = {
+                key: Score(value=score.value) for key, score in self.scores.items()
+            }
+        return self

inspect_ai/log/_transcript.py

@@ -15,6 +15,7 @@ from typing import (
 )
 
 from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_serializer
+from shortuuid import uuid
 
 from inspect_ai._util.constants import SAMPLE_SUBTASK
 from inspect_ai._util.error import EvalError
@@ -43,6 +44,13 @@ logger = getLogger(__name__)
 
 
 class BaseEvent(BaseModel):
+    model_config = {
+        "json_schema_extra": lambda schema: schema.get("properties", {}).pop(
+            "id_", None
+        )
+    }
+    id_: str = Field(default_factory=lambda: str(uuid()), exclude=True)
+
     timestamp: datetime = Field(default_factory=datetime.now)
     """Clock time at which event occurred."""
 
@@ -170,8 +178,8 @@ class ToolEvent(BaseEvent):
     arguments: dict[str, JsonValue]
     """Arguments to function."""
 
-    internal_name: str | None = Field(default=None)
-    """Internal name for tool (if any)."""
+    internal: JsonValue | None = Field(default=None)
+    """Model provider specific payload - typically used to aid transformation back to model types."""
 
     view: ToolCallContent | None = Field(default=None)
     """Custom view of tool call input."""
@@ -194,6 +202,12 @@ class ToolEvent(BaseEvent):
     working_time: float | None = Field(default=None)
     """Working time for tool call (i.e. time not spent waiting on semaphores)."""
 
+    agent: str | None = Field(default=None)
+    """Name of agent if the tool call was an agent handoff."""
+
+    failed: bool | None = Field(default=None)
+    """Did the tool call fail with a hard error?."""
+
     def _set_result(
         self,
         result: ToolResult,
@@ -201,6 +215,8 @@ class ToolEvent(BaseEvent):
         error: ToolCallError | None,
         events: list["Event"],
         waiting_time: float,
+        agent: str | None,
+        failed: bool | None,
     ) -> None:
         self.result = result
         self.truncated = truncated
@@ -210,6 +226,8 @@ class ToolEvent(BaseEvent):
         completed = datetime.now()
         self.completed = completed
         self.working_time = (completed - self.timestamp).total_seconds() - waiting_time
+        self.agent = agent
+        self.failed = failed
 
     # mechanism for operator to cancel the tool call
 
@@ -451,8 +469,11 @@ ET = TypeVar("ET", bound=BaseEvent)
 class Transcript:
     """Transcript of events."""
 
+    _event_logger: Callable[[Event], None] | None
+
     def __init__(self, name: str = "") -> None:
         self.name = name
+        self._event_logger = None
         self._events: list[Event] = []
 
     def info(self, data: JsonValue, *, source: str | None = None) -> None:
@@ -493,8 +514,17 @@ class Transcript:
         return None
 
     def _event(self, event: Event) -> None:
+        if self._event_logger:
+            self._event_logger(event)
         self._events.append(event)
 
+    def _event_updated(self, event: Event) -> None:
+        if self._event_logger:
+            self._event_logger(event)
+
+    def _subscribe(self, event_logger: Callable[[Event], None]) -> None:
+        self._event_logger = event_logger
+
 
 def transcript() -> Transcript:
     """Get the current `Transcript`."""

inspect_ai/model/__init__.py

@@ -18,7 +18,7 @@ from ._cache import (
     cache_prune,
     cache_size,
 )
-from ._call_tools import call_tools
+from ._call_tools import ExecuteToolsResult, call_tools, execute_tools
 from ._chat_message import (
     ChatMessage,
     ChatMessageAssistant,
@@ -27,6 +27,7 @@ from ._chat_message import (
     ChatMessageTool,
     ChatMessageUser,
 )
+from ._conversation import ModelConversation
 from ._generate_config import GenerateConfig, GenerateConfigArgs, ResponseSchema
 from ._model import (
     Model,
@@ -34,6 +35,7 @@ from ._model import (
     ModelName,
     get_model,
 )
+from ._model_call import ModelCall
 from ._model_output import (
     ChatCompletionChoice,
     Logprob,
@@ -64,7 +66,9 @@ __all__ = [
     "ChatMessageAssistant",
     "ChatMessageTool",
     "ChatCompletionChoice",
+    "ModelCall",
     "ModelOutput",
+    "ModelConversation",
     "Logprobs",
     "Logprob",
     "TopLogprob",
@@ -74,6 +78,8 @@ __all__ = [
     "ModelUsage",
     "StopReason",
     "call_tools",
+    "execute_tools",
+    "ExecuteToolsResult",
     "cache_clear",
     "cache_list_expired",
     "cache_path",