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

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/providers.py

@@ -281,7 +281,7 @@ def none() -> type[ModelAPI]:
 def validate_openai_client(feature: str) -> None:
     FEATURE = feature
     PACKAGE = "openai"
-    MIN_VERSION = "1.75.0"
+    MIN_VERSION = "1.78.0"
 
     # verify we have the package
     try:

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
@@ -217,6 +200,7 @@ def multiple_choice(
     template: str | None = None,
     cot: bool = False,
     multiple_correct: bool = False,
+    max_tokens: int | None = None,
     **kwargs: Unpack[DeprecatedArgs],
 ) -> Solver:
     """Multiple choice question solver. Formats a multiple choice question prompt, then calls `generate()`.
@@ -243,6 +227,8 @@ def multiple_choice(
         squares? A) 3, B) 4, C) 9" has multiple correct answers, B and C. Leave
         as `False` if there's exactly one correct answer from the choices
         available. NOTE: this has no effect if you provide a custom template.
+      max_tokens: Default `None`. Controls the number of tokens generated through the call
+        to generate().
       **kwargs (Any): Deprecated arguments for backward compatibility.
 
     #### Shuffling
@@ -299,7 +285,7 @@ def multiple_choice(
             template=str(template),
         )
 
-        state = await generate(state)
+        state = await generate(state, max_tokens=max_tokens)
 
         answers = parse_answers(state)
         if answers and answers.group(1):

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/_task_state.py

@@ -204,13 +204,17 @@ class TaskState:
         Convenience function for accessing the initial input from the `Sample` as a string.
 
         If the `input` is a `list[ChatMessage]`, this will return the text from
-        the first chat message
+        the last chat message
         """
         if isinstance(self._input, str):
             return self._input
         else:
             input = next(
-                (message.text for message in self._input if message.role == "user"),
+                (
+                    message.text
+                    for message in reversed(self._input)
+                    if message.role == "user"
+                ),
                 None,
             )
             if input:
@@ -231,7 +235,7 @@ class TaskState:
         write access to the user chat prompt. Raises an
         exception if there is no user prompt
         """
-        prompt = next((m for m in self.messages if m.role == "user"), None)
+        prompt = next((m for m in reversed(self.messages) if m.role == "user"), None)
         if prompt:
             return prompt
         else:

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/_context.py

@@ -2,13 +2,11 @@ from contextlib import _AsyncGeneratorContextManager
 from typing import TypeAlias
 
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-from mcp.types import (
-    JSONRPCMessage,
-)
+from mcp.shared.message import SessionMessage
 
 MCPServerContext: TypeAlias = _AsyncGeneratorContextManager[
     tuple[
-        MemoryObjectReceiveStream[JSONRPCMessage | Exception],
-        MemoryObjectSendStream[JSONRPCMessage],
+        MemoryObjectReceiveStream[SessionMessage | Exception],
+        MemoryObjectSendStream[SessionMessage],
     ],
 ]

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/_mcp/server.py

@@ -102,7 +102,7 @@ def mcp_server_sandbox(
 def verfify_mcp_package() -> None:
     FEATURE = "MCP tools"
     PACKAGE = "mcp"
-    MIN_VERSION = "1.6.0"
+    MIN_VERSION = "1.8.0"
 
     # verify we have the package
     try:

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/tool/_tools/_think.py

@@ -41,7 +41,7 @@ def think(
 def think_tool_viewer() -> ToolCallViewer:
     def viewer(tool_call: ToolCall) -> ToolCallView:
         call = ToolCallContent(
-            format="markdown", content=tool_call.arguments["thought"]
+            format="markdown", content=tool_call.arguments.get("thought", "")
         )
         return ToolCallView(call=call)
 

inspect_ai/tool/_tools/_web_search/__init__.py

@@ -0,0 +1,3 @@
+from ._web_search import web_search
+
+__all__ = ["web_search"]

inspect_ai/tool/_tools/{_web_search.py → _web_search/_google.py}

@@ -1,5 +1,5 @@
 import os
-from typing import Literal, Protocol, runtime_checkable
+from typing import Awaitable, Callable
 
 import anyio
 import httpx
@@ -16,8 +16,6 @@ from inspect_ai._util.error import PrerequisiteError
 from inspect_ai._util.httpx import httpx_should_retry, log_httpx_retry_attempt
 from inspect_ai.util._concurrency import concurrency
 
-from .._tool import Tool, ToolResult, tool
-
 DEFAULT_RELEVANCE_PROMPT = """I am trying to answer the following question and need to find the most relevant information on the web. Please let me know if the following content is relevant to the question or not. You should just respond with "yes" or "no".
 
 Question: {question}
@@ -31,59 +29,35 @@ class SearchLink:
         self.snippet = snippet
 
 
-@runtime_checkable
-class SearchProvider(Protocol):
-    async def __call__(self, query: str, start_idx: int) -> list[SearchLink]: ...
-
-
-@tool
-def web_search(
-    provider: Literal["google"] = "google",
-    num_results: int = 3,
-    max_provider_calls: int = 3,
-    max_connections: int = 10,
-    model: str | None = None,
-) -> Tool:
-    """Web search tool.
-
-    A tool that can be registered for use by models to search the web. Use
-    the `use_tools()` solver to make the tool available (e.g. `use_tools(web_search())`))
-
-    A web search is conducted using the specified provider, the results are parsed for relevance
-    using the specified model, and the top 'num_results' relevant pages are returned.
-
-    See further documentation at <https://inspect.aisi.org.uk/tools-standard.html#sec-web-search>.
-
-    Args:
-      provider: Search provider (defaults to "google", currently
-        the only provider). Possible future providers include "brave" and "bing".
-      num_results: Number of web search result pages to return to the model.
-      max_provider_calls: Maximum number of search calls to make to the search provider.
-      max_connections: Maximum number of concurrent connections to API
-        endpoint of search provider.
-      model: Model used to parse web pages for relevance.
+def maybe_get_google_api_keys() -> tuple[str, str] | None:
+    """
+    Get Google API keys from environment variables.
 
     Returns:
-       A tool that can be registered for use by models to search the web.
+        tuple: A tuple containing the Google API key and the Google CSE ID.
     """
-    # get search client
-    client = httpx.AsyncClient()
+    google_api_key = os.environ.get("GOOGLE_CSE_API_KEY", None)
+    google_cse_id = os.environ.get("GOOGLE_CSE_ID", None)
+    return (google_api_key, google_cse_id) if google_api_key and google_cse_id else None
 
-    if provider == "google":
-        search_provider = google_search_provider(client)
-    else:
-        raise ValueError(
-            f"Provider {provider} not supported. Only 'google' is supported."
+
+def google_search_provider(
+    num_results: int,
+    max_provider_calls: int,
+    max_connections: int,
+    model: str | None,
+) -> Callable[[str], Awaitable[str | None]]:
+    keys = maybe_get_google_api_keys()
+    if not keys:
+        raise PrerequisiteError(
+            "GOOGLE_CSE_ID and/or GOOGLE_CSE_API_KEY not set in the environment. Please ensure these variables are defined to use Google Custom Search with the web_search tool.\n\nLearn more about the Google web search provider at https://inspect.aisi.org.uk/tools.html#google-provider"
         )
+    google_api_key, google_cse_id = keys
 
-    # resolve provider (only google for now)
-    async def execute(query: str) -> ToolResult:
-        """
-        Use the web_search tool to perform keyword searches of the web.
+    # Create the client within the provider
+    client = httpx.AsyncClient()
 
-        Args:
-            query (str): Search query.
-        """
+    async def search(query: str) -> str | None:
         # limit number of concurrent searches
         page_contents: list[str] = []
         urls: list[str] = []
@@ -92,8 +66,8 @@ def web_search(
 
         # Paginate through search results until we have successfully extracted num_results pages or we have reached max_provider_calls
         while len(page_contents) < num_results and search_calls < max_provider_calls:
-            async with concurrency(f"{provider}_web_search", max_connections):
-                links = await search_provider(query, start_idx=search_calls * 10)
+            async with concurrency("google_web_search", max_connections):
+                links = await _search(query, start_idx=search_calls * 10)
 
             async with anyio.create_task_group() as tg:
 
@@ -114,19 +88,39 @@ def web_search(
             search_calls += 1
 
         all_page_contents = "\n\n".join(page_contents)
-        if all_page_contents == "":
-            response: ToolResult = (
-                "I'm sorry, I couldn't find any relevant information on the web."
-            )
-        else:
-            response = (
-                "Here are your web search results. Please read them carefully as they may be useful later! "
-                + all_page_contents
-            )
+        return None if all_page_contents == "" else all_page_contents
 
-        return response
+    async def _search(query: str, start_idx: int) -> list[SearchLink]:
+        # List of allowed parameters can be found https://developers.google.com/custom-search/v1/reference/rest/v1/cse/list
+        search_params = {
+            "q": query,
+            "key": google_api_key,
+            "cx": google_cse_id,
+            "start": start_idx,
+        }
+        search_url = "https://www.googleapis.com/customsearch/v1?" + "&".join(
+            [f"{key}={value}" for key, value in search_params.items()]
+        )
 
-    return execute
+        # retry up to 5 times over a period of up to 1 minute
+        @retry(
+            wait=wait_exponential_jitter(),
+            stop=stop_after_attempt(5) | stop_after_delay(60),
+            retry=retry_if_exception(httpx_should_retry),
+            before_sleep=log_httpx_retry_attempt(search_url),
+        )
+        async def execute_search() -> httpx.Response:
+            return await client.get(search_url)
+
+        result = await execute_search()
+        data = result.json()
+
+        if "items" in data:
+            return [SearchLink(item["link"], item["snippet"]) for item in data["items"]]
+        else:
+            return []
+
+    return search
 
 
 async def page_if_relevant(
@@ -183,44 +177,3 @@ async def page_if_relevant(
         return full_text
     else:
         return None
-
-
-def google_search_provider(client: httpx.AsyncClient) -> SearchProvider:
-    google_api_key = os.environ.get("GOOGLE_CSE_API_KEY", None)
-    google_cse_id = os.environ.get("GOOGLE_CSE_ID", None)
-    if not google_api_key or not google_cse_id:
-        raise PrerequisiteError(
-            "GOOGLE_CSE_ID and/or GOOGLE_CSE_API_KEY not set in the environment. Please ensure these variables are defined to use Google Custom Search with the web_search tool.\n\nLearn more about the Google web search provider at https://inspect.aisi.org.uk/tools.html#google-provider"
-        )
-
-    async def search(query: str, start_idx: int) -> list[SearchLink]:
-        # List of allowed parameters can be found https://developers.google.com/custom-search/v1/reference/rest/v1/cse/list
-        search_params = {
-            "q": query,
-            "key": google_api_key,
-            "cx": google_cse_id,
-            "start": start_idx,
-        }
-        search_url = "https://www.googleapis.com/customsearch/v1?" + "&".join(
-            [f"{key}={value}" for key, value in search_params.items()]
-        )
-
-        # retry up to 5 times over a period of up to 1 minute
-        @retry(
-            wait=wait_exponential_jitter(),
-            stop=stop_after_attempt(5) | stop_after_delay(60),
-            retry=retry_if_exception(httpx_should_retry),
-            before_sleep=log_httpx_retry_attempt(search_url),
-        )
-        async def execute_search() -> httpx.Response:
-            return await client.get(search_url)
-
-        result = await execute_search()
-        data = result.json()
-
-        if "items" in data:
-            return [SearchLink(item["link"], item["snippet"]) for item in data["items"]]
-        else:
-            return []
-
-    return search