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

inspect_ai/model/_model.py

@@ -19,6 +19,7 @@ from typing import (
     cast,
 )
 
+from pydantic import BaseModel
 from pydantic_core import to_jsonable_python
 from tenacity import (
     RetryCallState,
@@ -402,36 +403,32 @@ class Model:
         start_time = datetime.now()
         working_start = sample_working_time()
         async with self._connection_concurrency(config):
-            from inspect_ai.log._samples import track_active_sample_retries
-
             # generate
-            with track_active_sample_retries():
-                output = await self._generate(
-                    input=input,
-                    tools=tools,
-                    tool_choice=tool_choice,
-                    config=config,
-                    cache=cache,
-                )
+            output, event = await self._generate(
+                input=input,
+                tools=tools,
+                tool_choice=tool_choice,
+                config=config,
+                cache=cache,
+            )
 
             # update the most recent ModelEvent with the actual start/completed
             # times as well as a computation of working time (events are
             # created _after_ the call to _generate, potentially in response
             # to retries, so they need their timestamp updated so it accurately
             # reflects the full start/end time which we know here)
-            from inspect_ai.log._transcript import ModelEvent, transcript
-
-            last_model_event = transcript().find_last_event(ModelEvent)
-            if last_model_event:
-                last_model_event.timestamp = start_time
-                last_model_event.working_start = working_start
-                completed = datetime.now()
-                last_model_event.completed = completed
-                last_model_event.working_time = (
-                    output.time
-                    if output.time is not None
-                    else (completed - start_time).total_seconds()
-                )
+            from inspect_ai.log._transcript import ModelEvent
+
+            assert isinstance(event, ModelEvent)
+            event.timestamp = start_time
+            event.working_start = working_start
+            completed = datetime.now()
+            event.completed = completed
+            event.working_time = (
+                output.time
+                if output.time is not None
+                else (completed - start_time).total_seconds()
+            )
 
             # return output
             return output
@@ -492,9 +489,12 @@ class Model:
         tool_choice: ToolChoice | None,
         config: GenerateConfig,
         cache: bool | CachePolicy = False,
-    ) -> ModelOutput:
+    ) -> tuple[ModelOutput, BaseModel]:
+        from inspect_ai.log._samples import track_active_model_event
+        from inspect_ai.log._transcript import ModelEvent
+
         # default to 'auto' for tool_choice (same as underlying model apis)
-        tool_choice = tool_choice if tool_choice else "auto"
+        tool_choice = tool_choice if tool_choice is not None else "auto"
 
         # resolve top level tool source
         if isinstance(tools, ToolSource):
@@ -581,7 +581,10 @@ class Model:
             stop=stop,
             before_sleep=functools.partial(log_model_retry, self.api.model_name),
         )
-        async def generate() -> ModelOutput:
+        async def generate() -> tuple[ModelOutput, BaseModel]:
+            # type-checker can't see that we made sure tool_choice is not none in the outer frame
+            assert tool_choice is not None
+
             check_sample_interrupt()
 
             cache_entry: CacheEntry | None
@@ -602,7 +605,7 @@ class Model:
                 )
                 existing = cache_fetch(cache_entry)
                 if isinstance(existing, ModelOutput):
-                    self._record_model_interaction(
+                    _, event = self._record_model_interaction(
                         input=input,
                         tools=tools_info,
                         tool_choice=tool_choice,
@@ -611,7 +614,7 @@ class Model:
                         output=existing,
                         call=None,
                     )
-                    return existing
+                    return existing, event
             else:
                 cache_entry = None
 
@@ -620,7 +623,7 @@ class Model:
 
             # record the interaction before the call to generate
             # (we'll update it with the results once we have them)
-            complete = self._record_model_interaction(
+            complete, event = self._record_model_interaction(
                 input=input,
                 tools=tools_info,
                 tool_choice=tool_choice,
@@ -631,12 +634,14 @@ class Model:
             with trace_action(logger, "Model", f"generate ({str(self)})"):
                 time_start = time.monotonic()
                 try:
-                    result = await self.api.generate(
-                        input=input,
-                        tools=tools_info,
-                        tool_choice=tool_choice,
-                        config=config,
-                    )
+                    assert isinstance(event, ModelEvent)
+                    with track_active_model_event(event):
+                        result = await self.api.generate(
+                            input=input,
+                            tools=tools_info,
+                            tool_choice=tool_choice,
+                            config=config,
+                        )
                 finally:
                     time_elapsed = time.monotonic() - time_start
 
@@ -686,18 +691,18 @@ class Model:
             if cache and cache_entry:
                 cache_store(entry=cache_entry, output=output)
 
-            return output
+            return output, event
 
         # call the model (this will so retries, etc., so report waiting time
         # as elapsed time - actual time for successful model call)
         time_start = time.monotonic()
-        model_output = await generate()
+        model_output, event = await generate()
         total_time = time.monotonic() - time_start
         if model_output.time:
             report_sample_waiting_time(total_time - model_output.time)
 
         # return results
-        return model_output
+        return model_output, event
 
     def should_retry(self, ex: BaseException) -> bool:
         if isinstance(ex, Exception):
@@ -769,7 +774,7 @@ class Model:
         cache: Literal["read", "write"] | None,
         output: ModelOutput | None = None,
         call: ModelCall | None = None,
-    ) -> Callable[[ModelOutput | Exception, ModelCall | None], None]:
+    ) -> tuple[Callable[[ModelOutput | Exception, ModelCall | None], None], BaseModel]:
         from inspect_ai.log._transcript import ModelEvent, transcript
 
         # create event and add it to the transcript
@@ -809,7 +814,7 @@ class Model:
         if output:
             complete(output, call)
 
-        return complete
+        return complete, event
 
 
 class ModelName:

inspect_ai/model/_providers/hf.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import concurrent
 import concurrent.futures
 import copy
@@ -26,7 +28,12 @@ from transformers import (  # type: ignore
 from typing_extensions import override
 
 from inspect_ai._util.constants import DEFAULT_MAX_TOKENS
-from inspect_ai._util.content import ContentText
+from inspect_ai._util.content import (
+    ContentAudio,
+    ContentImage,
+    ContentText,
+    ContentVideo,
+)
 from inspect_ai._util.trace import trace_action
 from inspect_ai.tool import ToolChoice, ToolInfo
 
@@ -85,6 +92,7 @@ class HuggingFaceAPI(ModelAPI):
         self.batch_size = collect_model_arg("batch_size")
         self.chat_template = collect_model_arg("chat_template")
         self.tokenizer_call_args = collect_model_arg("tokenizer_call_args")
+        self.enable_thinking = collect_model_arg("enable_thinking")
         if self.tokenizer_call_args is None:
             self.tokenizer_call_args = {}
 
@@ -263,6 +271,7 @@ class HuggingFaceAPI(ModelAPI):
             elif "qwen" in self.model_name.lower():
                 hf_messages = inspect_tools_to_string(hf_messages)
 
+        hf_messages = message_content_to_string(hf_messages)
         # apply chat template
         if self.tokenizer.chat_template is not None:
             chat = self.tokenizer.apply_chat_template(
@@ -270,6 +279,7 @@ class HuggingFaceAPI(ModelAPI):
                 add_generation_prompt=True,
                 tokenize=False,
                 tools=tools_list if len(tools_list) > 0 else None,
+                enable_thinking=self.enable_thinking,  # not all models use this, check if it is supported
             )
         else:
             chat = ""
@@ -279,6 +289,22 @@ class HuggingFaceAPI(ModelAPI):
         return cast(str, chat)
 
 
+def message_content_to_string(messages: list[ChatMessage]) -> list[ChatMessage]:
+    """Convert list of content in `ChatMessageAssistant`, `ChatMessageUser` or `ChatMessageSystem` to a string."""
+    for message in messages:
+        if isinstance(message.content, list):
+            is_multimodal = any(
+                isinstance(item, ContentAudio | ContentImage | ContentVideo)
+                for item in message.content
+            )
+            if is_multimodal:
+                raise NotImplementedError(
+                    "HuggingFace provider does not support multimodal content, please provide text inputs only."
+                )
+            message.content = message.text
+    return messages
+
+
 def shorten_tool_id(messages: list[ChatMessage]) -> list[ChatMessage]:
     """Shorten the tool_call_id in the messages to the last 9 characters for Mistral."""
     for i, message in enumerate(messages):

inspect_ai/model/_providers/sglang.py

@@ -71,6 +71,7 @@ class SGLangAPI(OpenAICompatibleAPI):
             SGLANG_DEFAULT_SERVER_ARGS, server_args, logger
         )
 
+        self.server_found = True
         try:
             # Try to initialize with existing server
             super().__init__(
@@ -83,7 +84,9 @@ class SGLangAPI(OpenAICompatibleAPI):
             )
             logger.info(f"Using existing SGLang server at {self.base_url}")
         except PrerequisiteError:
-            # No existing server found, start a new one
+            self.server_found = False
+
+        if not self.server_found:
             logger.warning(
                 f"Existing SGLang server not found. Starting new server for {model_name}."
             )
@@ -125,7 +128,9 @@ class SGLangAPI(OpenAICompatibleAPI):
             api_key = "inspectai"  # Create a default API key if not provided
 
         # Handle device configuration
-        self.server_args = configure_devices(self.server_args, parallel_size_param="tp")
+        self.server_args, env_vars = configure_devices(
+            self.server_args, parallel_size_param="tp"
+        )
 
         timeout = self.server_args.pop("timeout", None)
         host = self.server_args.pop("host", "0.0.0.0")
@@ -149,6 +154,7 @@ class SGLangAPI(OpenAICompatibleAPI):
             server_type="SGLang",
             timeout=timeout,
             server_args=self.server_args,
+            env=env_vars,
         )
 
         # Register cleanup function to run when Python exits

inspect_ai/model/_providers/vllm.py

@@ -76,6 +76,7 @@ class VLLMAPI(OpenAICompatibleAPI):
             VLLM_DEFAULT_SERVER_ARGS, server_args, logger
         )
 
+        self.server_found = True
         try:
             # Try to initialize with existing server
             super().__init__(
@@ -88,7 +89,9 @@ class VLLMAPI(OpenAICompatibleAPI):
             )
             logger.info(f"Using existing vLLM server at {self.base_url}")
         except PrerequisiteError:
-            # No existing server found, start a new one
+            self.server_found = False
+
+        if not self.server_found:
             logger.warning(
                 f"Existing vLLM server not found. Starting new server for {model_name}."
             )
@@ -131,7 +134,7 @@ class VLLMAPI(OpenAICompatibleAPI):
             raise pip_dependency_error("vLLM Server", ["vllm"])
 
         # Handle device configuration
-        self.server_args = configure_devices(
+        self.server_args, env_vars = configure_devices(
             self.server_args, parallel_size_param="tensor_parallel_size"
         )
 
@@ -152,6 +155,7 @@ class VLLMAPI(OpenAICompatibleAPI):
             server_type="vLLM",
             timeout=timeout,
             server_args=self.server_args,
+            env=env_vars,
         )
 
         # Register cleanup function to run when Python exits

inspect_ai/scorer/_choice.py

@@ -1,6 +1,5 @@
+from inspect_ai._util.answer import answer_character, answer_index
 from inspect_ai.solver._multiple_choice import (
-    answer_character,
-    answer_index,
     answer_options,
     unshuffle_choices,
 )

inspect_ai/solver/_chain.py

@@ -82,7 +82,7 @@ class Chain(Sequence[Solver], Solver):
         from ._transcript import solver_transcript
 
         for slv in self._solvers:
-            with solver_transcript(slv, state) as st:
+            async with solver_transcript(slv, state) as st:
                 state = await slv(state, generate)
                 st.complete(state)
             if state.completed:

inspect_ai/solver/_fork.py

@@ -73,7 +73,7 @@ async def solver_subtask(state: TaskState, solver: Solver) -> TaskState:
     @subtask(name=name, store=state.store, type="fork", input=input)  # type: ignore
     async def solve() -> TaskState:
         if not isinstance(solver, Chain):
-            with solver_transcript(solver, state) as st:
+            async with solver_transcript(solver, state) as st:
                 new_state = await solver(state, generate)
                 st.complete(new_state)
             return new_state

inspect_ai/solver/_multiple_choice.py

@@ -6,6 +6,7 @@ from typing import Match, TypedDict
 
 from typing_extensions import Unpack
 
+from inspect_ai._util.answer import answer_character, answer_index
 from inspect_ai._util.logger import warn_once
 from inspect_ai.util import resource
 
@@ -64,31 +65,13 @@ def answer_options(choices: Choices) -> str:
     indexes = list(range(len(choices)))
 
     return "\n".join(
-        [f"{chr(65 + i)}) {choices[j].value}" for i, j in enumerate(indexes)]
+        [f"{answer_character(i)}) {choices[j].value}" for i, j in enumerate(indexes)]
     )
 
 
-def answer_character(index: int) -> str:
-    r"""
-    Helper to go from array index to char, for example:
-
-        0 -> 'A', 1 -> 'B', etc
-    """
-    return chr(ord("A") + index)
-
-
-def answer_index(char: str) -> int:
-    r"""
-    Helper to go from char to array index, for example:
-
-        'A' -> 0, 'B' -> 1, etc
-    """
-    return ord(char.upper()) - ord("A")
-
-
 def prompt(question: str, choices: Choices, template: str) -> str:
     choices_text = answer_options(choices)
-    letters = ",".join(chr(65 + i) for i in range(len(choices)))
+    letters = ",".join(answer_character(i) for i in range(len(choices)))
 
     return template.format(
         choices=choices_text,
@@ -112,7 +95,7 @@ def parse_answers(state: TaskState) -> Match[str] | None:
     # In this case, we're looking for a single line which contains the expected
     # ANSWER: B,C string with only whitespace after it
     match = re.search(
-        r"(?i)^ANSWER\s*:\s*([A-Za-z ,]+)\s*(?:$|\n)",
+        r"(?i)^ANSWER\s*:\s*([A-Za-z\d ,]+)\s*(?:$|\n)",
         state.output.completion,
         flags=re.MULTILINE,
     )
@@ -121,7 +104,7 @@ def parse_answers(state: TaskState) -> Match[str] | None:
     # version for backward compatibility
     if match is None:
         return re.search(
-            r"(?i)ANSWER\s*:\s*([A-Za-z ,]+)(?:[^\w]|\n|$)", state.output.completion
+            r"(?i)ANSWER\s*:\s*([A-Za-z\d ,]+)(?:[^\w]|\n|$)", state.output.completion
         )
     else:
         return match

inspect_ai/solver/_plan.py

@@ -102,7 +102,7 @@ class Plan(Solver):
             # execute steps
             for index, solver in enumerate(self.steps):
                 # run solver
-                with solver_transcript(solver, state) as st:
+                async with solver_transcript(solver, state) as st:
                     state = await solver(state, generate)
                     st.complete(state)
 
@@ -113,7 +113,7 @@ class Plan(Solver):
 
             # execute finish
             if self.finish:
-                with solver_transcript(self.finish, state) as st:
+                async with solver_transcript(self.finish, state) as st:
                     state = await self.finish(state, generate)
                     st.complete(state)
                 check_sample_interrupt()

inspect_ai/solver/_transcript.py

@@ -1,8 +1,9 @@
 import contextlib
-from typing import Iterator
+from typing import AsyncIterator
 
 from inspect_ai._util.json import json_changes
 from inspect_ai._util.registry import registry_log_name
+from inspect_ai.util._span import span
 
 from ._solver import Solver
 from ._task_state import TaskState, state_jsonable
@@ -22,12 +23,10 @@ class SolverTranscript:
             transcript()._event(StateEvent(changes=changes))
 
 
-@contextlib.contextmanager
-def solver_transcript(
+@contextlib.asynccontextmanager
+async def solver_transcript(
     solver: Solver, state: TaskState, name: str | None = None
-) -> Iterator[SolverTranscript]:
-    from inspect_ai.log._transcript import transcript
-
+) -> AsyncIterator[SolverTranscript]:
     name = registry_log_name(name or solver)
-    with transcript().step(name=name, type="solver"):
+    async with span(name=name, type="solver"):
         yield SolverTranscript(name, state)

inspect_ai/tool/_mcp/_mcp.py

@@ -61,16 +61,17 @@ class MCPServerImpl(MCPServer):
     ) -> list[Tool]:
         return await self._task_session()._list_tools(tools)
 
-    # create a separate MCPServer session per async task
-    _task_sessions: dict[int, "MCPServerSession"] = {}
+    # create a separate MCPServer session per async task / server name
+    _task_sessions: dict[str, "MCPServerSession"] = {}
 
     def _task_session(self) -> "MCPServerSession":
         task_id = anyio.get_current_task().id
-        if task_id not in self._task_sessions:
-            MCPServerImpl._task_sessions[task_id] = MCPServerSession(
+        session_key = f"{task_id}_{self._name}"
+        if session_key not in self._task_sessions:
+            MCPServerImpl._task_sessions[session_key] = MCPServerSession(
                 self._client, name=self._name, events=self._events
             )
-        return MCPServerImpl._task_sessions[task_id]
+        return MCPServerImpl._task_sessions[session_key]
 
 
 class MCPServerSession(MCPServer):

inspect_ai/tool/_tools/_execute.py

@@ -96,7 +96,10 @@ def python(
           The output of the Python code.
         """
         result = await sandbox_env(sandbox).exec(
-            cmd=["python3"], input=code, timeout=timeout, user=user
+            cmd=["bash", "--login", "-c", "python3 -"],
+            input=code,
+            timeout=timeout,
+            user=user,
         )
         # return output (including stderr if any)
         output = ""

inspect_ai/util/__init__.py

@@ -8,6 +8,7 @@ from inspect_ai.util._limit import (
     token_limit,
 )
 
+from ._collect import collect
 from ._concurrency import concurrency
 from ._console import input_screen
 from ._display import DisplayType, display_counter, display_type
@@ -28,6 +29,7 @@ from ._sandbox import (
     sandbox_with,
     sandboxenv,
 )
+from ._span import span
 from ._store import Store, store
 from ._store_model import StoreModel, store_as
 from ._subprocess import (
@@ -71,6 +73,8 @@ __all__ = [
     "store",
     "StoreModel",
     "store_as",
+    "span",
+    "collect",
     "Subtask",
     "subtask",
     "throttle",

inspect_ai/util/_anyio.py

@@ -1,6 +1,10 @@
 import itertools
 import sys
 
+import anyio
+
+from inspect_ai._util._async import current_async_backend
+
 if sys.version_info < (3, 11):
     from exceptiongroup import ExceptionGroup
 
@@ -36,3 +40,10 @@ def _flatten_exception(exc: Exception) -> list[Exception]:
     ]
 
     return maybe_this_exception + other_exceptions
+
+
+def safe_current_task_id() -> int | None:
+    if current_async_backend() is not None:
+        return anyio.get_current_task().id
+    else:
+        return None

inspect_ai/util/_collect.py

@@ -0,0 +1,50 @@
+import sys
+from typing import Awaitable, TypeVar, cast
+
+import anyio
+
+from ._span import span
+
+if sys.version_info < (3, 11):
+    from exceptiongroup import ExceptionGroup
+
+
+T = TypeVar("T")
+
+
+async def collect(*tasks: Awaitable[T]) -> list[T]:
+    """Run and collect the results of one or more async coroutines.
+
+    Similar to [`asyncio.gather()`](https://docs.python.org/3/library/asyncio-task.html#asyncio.gather),
+    but also works when [Trio](https://trio.readthedocs.io/en/stable/) is the async backend.
+
+    Automatically includes each task in a `span()`, which
+    ensures that its events are grouped together in the transcript.
+
+    Using `collect()` in preference to `asyncio.gather()` is highly recommended
+    for both Trio compatibility and more legible transcript output.
+
+    Args:
+        *tasks: Tasks to run
+
+    Returns:
+        List of task results.
+    """
+    results: list[None | T] = [None] * len(tasks)
+
+    try:
+        async with anyio.create_task_group() as tg:
+
+            async def run_task(index: int, task: Awaitable[T]) -> None:
+                async with span(f"task-{index + 1}", type="task"):
+                    results[index] = await task
+
+            for i, task in enumerate(tasks):
+                tg.start_soon(run_task, i, task)
+    except ExceptionGroup as ex:
+        if len(ex.exceptions) == 1:
+            raise ex.exceptions[0] from None
+        else:
+            raise
+
+    return cast(list[T], results)

inspect_ai/util/_span.py

@@ -0,0 +1,58 @@
+import contextlib
+from contextvars import ContextVar
+from typing import AsyncIterator
+from uuid import uuid4
+
+
+@contextlib.asynccontextmanager
+async def span(name: str, *, type: str | None = None) -> AsyncIterator[None]:
+    """Context manager for establishing a transcript span.
+
+    Args:
+        name (str): Step name.
+        type (str | None): Optional span type.
+    """
+    from inspect_ai.log._transcript import (
+        SpanBeginEvent,
+        SpanEndEvent,
+        track_store_changes,
+        transcript,
+    )
+
+    # span id
+    id = uuid4().hex
+
+    # capture parent id
+    parent_id = _current_span_id.get()
+
+    # set new current span (reset at the end)
+    token = _current_span_id.set(id)
+
+    # run the span
+    try:
+        # span begin event
+        transcript()._event(
+            SpanBeginEvent(
+                id=id,
+                parent_id=parent_id,
+                type=type,
+                name=name,
+            )
+        )
+
+        # run span w/ store change events
+        with track_store_changes():
+            yield
+
+    finally:
+        # send end event
+        transcript()._event(SpanEndEvent(id=id))
+
+        _current_span_id.reset(token)
+
+
+def current_span_id() -> str | None:
+    return _current_span_id.get()
+
+
+_current_span_id: ContextVar[str | None] = ContextVar("_current_span_id", default=None)