inspect-ai 0.3.103__py3-none-any.whl → 0.3.104__py3-none-any.whl

inspect_ai/model/_providers/openai_compatible.py

@@ -61,7 +61,8 @@ class OpenAICompatibleAPI(ModelAPI):
             self.service = service
 
         # compute api key
-        api_key_var = f"{self.service.upper()}_API_KEY"
+        service_env_name = self.service.upper().replace("-", "_")
+        api_key_var = f"{service_env_name}_API_KEY"
 
         super().__init__(
             model_name=model_name,
@@ -82,7 +83,7 @@ class OpenAICompatibleAPI(ModelAPI):
 
         # use service prefix to lookup base_url
         if not self.base_url:
-            base_url_var = f"{self.service.upper()}_BASE_URL"
+            base_url_var = f"{service_env_name}_BASE_URL"
             self.base_url = model_base_url(base_url, [base_url_var]) or service_base_url
             if not self.base_url:
                 raise environment_prerequisite_error(

inspect_ai/model/_providers/openai_responses.py

@@ -40,7 +40,6 @@ async def generate_responses(
     tool_choice: ToolChoice,
     config: GenerateConfig,
     service_tier: str | None,
-    store: bool,
 ) -> ModelOutput | tuple[ModelOutput | Exception, ModelCall]:
     # allocate request_id (so we can see it from ModelCall)
     request_id = http_hooks.start_request()
@@ -65,7 +64,7 @@ async def generate_responses(
         else NOT_GIVEN
     )
     request = dict(
-        input=await openai_responses_inputs(input, model_name, store),
+        input=await openai_responses_inputs(input, model_name),
         tools=tool_params,
         tool_choice=openai_responses_tool_choice(tool_choice, tool_params)
         if isinstance(tool_params, list) and tool_choice != "auto"
@@ -77,7 +76,6 @@ async def generate_responses(
             config=config,
             service_tier=service_tier,
             tools=len(tools) > 0,
-            store=store,
         ),
     )
 
@@ -125,7 +123,6 @@ def completion_params_responses(
     config: GenerateConfig,
     service_tier: str | None,
     tools: bool,
-    store: bool,
 ) -> dict[str, Any]:
     # TODO: we'll need a computer_use_preview bool for the 'include'
     # and 'reasoning' parameters
@@ -135,7 +132,7 @@ def completion_params_responses(
             f"OpenAI Responses API does not support the '{param}' parameter.",
         )
 
-    params: dict[str, Any] = dict(model=model_name, store=store)
+    params: dict[str, Any] = dict(model=model_name)
     if service_tier is not None:
         params["service_tier"] = service_tier
     if config.max_tokens is not None:

inspect_ai/model/_providers/perplexity.py

@@ -0,0 +1,123 @@
+from typing import Any, cast
+
+from openai.types.chat import ChatCompletion
+
+from inspect_ai._util.citation import UrlCitation
+from inspect_ai._util.content import ContentText
+from inspect_ai.model._generate_config import GenerateConfig
+from inspect_ai.model._model_output import ModelOutput, ModelUsage
+from inspect_ai.model._openai import chat_choices_from_openai
+from inspect_ai.model._providers.openai_compatible import OpenAICompatibleAPI
+from inspect_ai.tool import ToolChoice, ToolInfo
+
+from .._chat_message import ChatMessage
+from .._model_call import ModelCall
+from .._model_output import ChatCompletionChoice
+
+
+class PerplexityAPI(OpenAICompatibleAPI):
+    """Model provider for Perplexity AI."""
+
+    def __init__(
+        self,
+        model_name: str,
+        base_url: str | None = None,
+        api_key: str | None = None,
+        config: GenerateConfig = GenerateConfig(),
+        **model_args: Any,
+    ) -> None:
+        super().__init__(
+            model_name=model_name,
+            base_url=base_url,
+            api_key=api_key,
+            config=config,
+            service="Perplexity",
+            service_base_url="https://api.perplexity.ai",
+            **model_args,
+        )
+
+        self._response: dict[str, Any] | None = None
+
+    def on_response(self, response: dict[str, Any]) -> None:
+        """Capture the raw response for post-processing."""
+        self._response = response
+
+    async def generate(
+        self,
+        input: list["ChatMessage"],
+        tools: list["ToolInfo"],
+        tool_choice: "ToolChoice",
+        config: GenerateConfig,
+    ) -> tuple[ModelOutput | Exception, "ModelCall"]:
+        result = await super().generate(input, tools, tool_choice, config)
+        output, call = cast(tuple[ModelOutput, "ModelCall"], result)
+
+        if self._response:
+            response = self._response
+
+            # attach citations if search results are returned
+            search_results = response.get("search_results")
+            if isinstance(search_results, list):
+                citations = [
+                    UrlCitation(title=sr.get("title"), url=sr.get("url", ""))
+                    for sr in search_results
+                    if isinstance(sr, dict) and sr.get("url") is not None
+                ]
+                if citations:
+                    for choice in output.choices:
+                        msg = choice.message
+                        if isinstance(msg.content, str):
+                            msg.content = [
+                                ContentText(text=msg.content, citations=citations)
+                            ]
+                        else:
+                            added = False
+                            for content in msg.content:
+                                if (
+                                    isinstance(content, ContentText)
+                                    and getattr(content, "citations", None) is None
+                                ):
+                                    content.citations = citations
+                                    added = True
+                                    break
+                            if not added:
+                                msg.content.append(
+                                    ContentText(text="", citations=citations)
+                                )
+
+            # update usage with additional metrics
+            usage_data = response.get("usage")
+            if isinstance(usage_data, dict):
+                extra_usage = {
+                    k: usage_data.get(k)
+                    for k in [
+                        "search_context_size",
+                        "citation_tokens",
+                        "num_search_queries",
+                    ]
+                    if k in usage_data
+                }
+                if output.usage:
+                    output.usage.reasoning_tokens = usage_data.get("reasoning_tokens")
+                else:
+                    output.usage = ModelUsage(
+                        input_tokens=usage_data.get("prompt_tokens", 0),
+                        output_tokens=usage_data.get("completion_tokens", 0),
+                        total_tokens=usage_data.get("total_tokens", 0),
+                        reasoning_tokens=usage_data.get("reasoning_tokens"),
+                    )
+                if extra_usage:
+                    output.metadata = output.metadata or {}
+                    output.metadata.update(extra_usage)
+
+            # keep search_results for reference
+            if search_results:
+                output.metadata = output.metadata or {}
+                output.metadata["search_results"] = search_results
+
+        return output, call
+
+    def chat_choices_from_completion(
+        self, completion: ChatCompletion, tools: list[ToolInfo]
+    ) -> list[ChatCompletionChoice]:
+        return chat_choices_from_openai(completion, tools)

inspect_ai/model/_providers/providers.py

@@ -59,7 +59,7 @@ def openai_api() -> type[ModelAPI]:
 def anthropic() -> type[ModelAPI]:
     FEATURE = "Anthropic API"
     PACKAGE = "anthropic"
-    MIN_VERSION = "0.49.0"
+    MIN_VERSION = "0.52.0"
 
     # verify we have the package
     try:
@@ -157,7 +157,7 @@ def cf() -> type[ModelAPI]:
 def mistral() -> type[ModelAPI]:
     FEATURE = "Mistral API"
     PACKAGE = "mistralai"
-    MIN_VERSION = "1.6.0"
+    MIN_VERSION = "1.8.2"
 
     # verify we have the package
     try:
@@ -218,6 +218,17 @@ def openrouter() -> type[ModelAPI]:
     return OpenRouterAPI
 
 
+@modelapi(name="perplexity")
+def perplexity() -> type[ModelAPI]:
+    # validate
+    validate_openai_client("Perplexity API")
+
+    # in the clear
+    from .perplexity import PerplexityAPI
+
+    return PerplexityAPI
+
+
 @modelapi(name="llama-cpp-python")
 def llama_cpp_python() -> type[ModelAPI]:
     # validate

inspect_ai/model/_providers/vertex.py

@@ -33,6 +33,7 @@ from inspect_ai._util.constants import BASE_64_DATA_REMOVED, NO_CONTENT
 from inspect_ai._util.content import (
     Content,
     ContentAudio,
+    ContentData,
     ContentImage,
     ContentReasoning,
     ContentText,
@@ -338,6 +339,8 @@ async def content_part(content: Content | str) -> Part:
     else:
         if isinstance(content, ContentAudio):
             file = content.audio
+        elif isinstance(content, ContentData):
+            assert False, "Vertex provider should never encounter ContentData"
         else:
             # it's ContentVideo
             file = content.video

inspect_ai/model/_trim.py

@@ -13,6 +13,7 @@ def trim_messages(
     - Retaining the 'input' messages from the sample.
     - Preserving a proportion of the remaining messages (`preserve=0.7` by default).
     - Ensuring that all assistant tool calls have corresponding tool messages.
+    - Ensuring that the sequence of messages doesn't end with an assistant message.
 
     Args:
         messages: List of messages to trim.
@@ -49,6 +50,10 @@ def trim_messages(
             active_tool_ids = set()
             conversation_messages.append(message)
 
+    # it's possible that we end with an assistant message w/ if so, remove it
+    if len(conversation_messages) and conversation_messages[-1].role == "assistant":
+        conversation_messages.pop()
+
     # return trimmed messages
     return partitioned.system + partitioned.input + conversation_messages
 

inspect_ai/tool/__init__.py

@@ -1,6 +1,14 @@
+from inspect_ai._util.citation import (
+    Citation,
+    CitationBase,
+    ContentCitation,
+    DocumentCitation,
+    UrlCitation,
+)
 from inspect_ai._util.content import (
     Content,
     ContentAudio,
+    ContentData,
     ContentImage,
     ContentReasoning,
     ContentText,
@@ -62,6 +70,7 @@ __all__ = [
     "MCPServer",
     "Content",
     "ContentAudio",
+    "ContentData",
     "ContentImage",
     "ContentReasoning",
     "ContentText",
@@ -77,6 +86,11 @@ __all__ = [
     "ToolInfo",
     "ToolParam",
     "ToolParams",
+    "Citation",
+    "CitationBase",
+    "DocumentCitation",
+    "ContentCitation",
+    "UrlCitation",
 ]
 
 _UTIL_MODULE_VERSION = "0.3.19"

inspect_ai/tool/_mcp/_mcp.py

@@ -12,6 +12,7 @@ from mcp.client.session import ClientSession, SamplingFnT
 from mcp.client.sse import sse_client
 from mcp.client.stdio import StdioServerParameters, stdio_client
 from mcp.types import (
+    AudioContent,
     EmbeddedResource,
     ImageContent,
     TextContent,
@@ -282,14 +283,16 @@ def create_server_sandbox(
 
 
 def tool_result_as_text(
-    content: list[TextContent | ImageContent | EmbeddedResource],
+    content: list[TextContent | ImageContent | AudioContent | EmbeddedResource],
 ) -> str:
     content_list: list[str] = []
     for c in content:
         if isinstance(c, TextContent):
             content_list.append(c.text)
         elif isinstance(c, ImageContent):
-            content_list.append("(base64 encoded image ommitted)")
+            content_list.append("(base64 encoded image omitted)")
+        elif isinstance(c, AudioContent):
+            content_list.append("(base64 encoded audio omitted)")
         elif isinstance(c.resource, TextResourceContents):
             content_list.append(c.resource.text)
 

inspect_ai/tool/_mcp/sampling.py

@@ -1,9 +1,10 @@
-from typing import Any
+from typing import Any, Literal
 
 from mcp.client.session import ClientSession
 from mcp.shared.context import RequestContext
 from mcp.types import (
     INTERNAL_ERROR,
+    AudioContent,
     CreateMessageRequestParams,
     CreateMessageResult,
     EmbeddedResource,
@@ -16,7 +17,7 @@ from mcp.types import (
     StopReason as MCPStopReason,
 )
 
-from inspect_ai._util.content import Content, ContentImage, ContentText
+from inspect_ai._util.content import Content, ContentAudio, ContentImage, ContentText
 from inspect_ai._util.error import exception_message
 from inspect_ai._util.url import data_uri_mime_type, data_uri_to_base64
 
@@ -93,7 +94,7 @@ async def sampling_fn(
 
 
 def as_inspect_content(
-    content: TextContent | ImageContent | EmbeddedResource,
+    content: TextContent | ImageContent | AudioContent | EmbeddedResource,
 ) -> Content:
     if isinstance(content, TextContent):
         return ContentText(text=content.text)
@@ -101,6 +102,11 @@ def as_inspect_content(
         return ContentImage(
             image=f"data:image/{content.mimeType};base64,{content.data}"
         )
+    elif isinstance(content, AudioContent):
+        return ContentAudio(
+            audio=f"data:audio/{content.mimeType};base64,{content.data}",
+            format=_get_audio_format(content.mimeType),
+        )
     elif isinstance(content.resource, TextResourceContents):
         return ContentText(text=content.resource.text)
     else:
@@ -116,3 +122,13 @@ def as_mcp_content(content: ContentText | ContentImage) -> TextContent | ImageCo
             mimeType=data_uri_mime_type(content.image) or "image/png",
             data=data_uri_to_base64(content.image),
         )
+
+
+def _get_audio_format(mime_type: str) -> Literal["wav", "mp3"]:
+    """Helper function to determine audio format from MIME type."""
+    if mime_type in ("audio/wav", "audio/x-wav"):
+        return "wav"
+    elif mime_type == "audio/mpeg":
+        return "mp3"
+    else:
+        raise ValueError(f"Unsupported audio mime type: {mime_type}")

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.8.0"
+    MIN_VERSION = "1.9.4"
 
     # verify we have the package
     try:

inspect_ai/tool/_tool.py

@@ -13,6 +13,7 @@ from typing import (
 
 from inspect_ai._util.content import (
     ContentAudio,
+    ContentData,
     ContentImage,
     ContentReasoning,
     ContentText,
@@ -41,7 +42,15 @@ ToolResult = (
     | ContentImage
     | ContentAudio
     | ContentVideo
-    | list[ContentText | ContentReasoning | ContentImage | ContentAudio | ContentVideo]
+    | ContentData
+    | list[
+        ContentText
+        | ContentReasoning
+        | ContentImage
+        | ContentAudio
+        | ContentVideo
+        | ContentData
+    ]
 )
 """Valid types for results from tool calls."""
 

inspect_ai/tool/_tools/_web_search/_base_http_provider.py

@@ -0,0 +1,104 @@
+import os
+from abc import ABC, abstractmethod
+from typing import Any
+
+import httpx
+from tenacity import (
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    stop_after_delay,
+    wait_exponential_jitter,
+)
+
+from inspect_ai._util.content import ContentText
+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
+
+
+class BaseHttpProvider(ABC):
+    """Base class for HTTP-based web search providers (Exa, Tavily, etc.)."""
+
+    def __init__(
+        self,
+        env_key_name: str,
+        api_endpoint: str,
+        provider_name: str,
+        concurrency_key: str,
+        options: dict[str, Any] | None = None,
+    ):
+        self.env_key_name = env_key_name
+        self.api_endpoint = api_endpoint
+        self.provider_name = provider_name
+        self.concurrency_key = concurrency_key
+
+        self.max_connections = self._extract_max_connections(options)
+        self.api_options = self._prepare_api_options(options)
+        self.api_key = self._validate_api_key()
+        self.client = httpx.AsyncClient(timeout=30)
+
+    @abstractmethod
+    def prepare_headers(self, api_key: str) -> dict[str, str]:
+        """Prepare HTTP headers for the request."""
+        pass
+
+    @abstractmethod
+    def parse_response(self, response_data: dict[str, Any]) -> ContentText | None:
+        """Parse the API response and extract content with citations."""
+        pass
+
+    @abstractmethod
+    def set_default_options(self, options: dict[str, Any]) -> dict[str, Any]:
+        """Set provider-specific default options."""
+        pass
+
+    def _extract_max_connections(self, options: dict[str, Any] | None) -> int:
+        """Extract max_connections from options, defaulting to 10."""
+        if not options:
+            return 10
+        max_conn = options.get("max_connections", 10)
+        return int(max_conn) if max_conn is not None else 10
+
+    def _prepare_api_options(self, options: dict[str, Any] | None) -> dict[str, Any]:
+        """Prepare API options by removing max_connections and setting defaults."""
+        if not options:
+            api_options = {}
+        else:
+            # Remove max_connections as it's not an API option
+            api_options = {k: v for k, v in options.items() if k != "max_connections"}
+
+        # Apply provider-specific defaults
+        return self.set_default_options(api_options)
+
+    def _validate_api_key(self) -> str:
+        """Validate that the required API key is set in environment."""
+        api_key = os.environ.get(self.env_key_name)
+        if not api_key:
+            raise PrerequisiteError(
+                f"{self.env_key_name} not set in the environment. Please ensure this variable is defined to use {self.provider_name} with the web_search tool.\n\nLearn more about the {self.provider_name} web search provider at https://inspect.aisi.org.uk/tools.html#{self.provider_name.lower()}-provider"
+            )
+        return api_key
+
+    async def search(self, query: str) -> ContentText | None:
+        """Execute a search query using the provider's API."""
+
+        # Common retry logic for all HTTP providers
+        @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(self.api_endpoint),
+        )
+        async def _search() -> httpx.Response:
+            response = await self.client.post(
+                self.api_endpoint,
+                headers=self.prepare_headers(self.api_key),
+                json={"query": query, **self.api_options},
+            )
+            response.raise_for_status()
+            return response
+
+        async with concurrency(self.concurrency_key, self.max_connections):
+            response_data = (await _search()).json()
+            return self.parse_response(response_data)

inspect_ai/tool/_tools/_web_search/_exa.py

@@ -0,0 +1,78 @@
+from typing import Any, Literal
+
+from pydantic import BaseModel
+
+from inspect_ai._util.citation import UrlCitation
+from inspect_ai._util.content import ContentText
+
+from ._base_http_provider import BaseHttpProvider
+from ._web_search_provider import SearchProvider
+
+
+class ExaOptions(BaseModel):
+    # See https://docs.exa.ai/reference/answer
+    text: bool | None = None
+    """Whether to include text content in citations"""
+    model: Literal["exa", "exa-pro"] | None = None
+    """LLM model to use for generating the answer"""
+    max_connections: int | None = None
+    """max_connections is not an Exa API option, but an inspect option"""
+
+
+class ExaCitation(BaseModel):
+    id: str
+    url: str
+    title: str
+    author: str | None = None
+    publishedDate: str | None = None
+    text: str
+
+
+class ExaSearchResponse(BaseModel):
+    answer: str
+    citations: list[ExaCitation]
+
+
+class ExaSearchProvider(BaseHttpProvider):
+    def __init__(self, options: dict[str, Any] | None = None):
+        super().__init__(
+            env_key_name="EXA_API_KEY",
+            api_endpoint="https://api.exa.ai/answer",
+            provider_name="Exa",
+            concurrency_key="exa_web_search",
+            options=options,
+        )
+
+    def prepare_headers(self, api_key: str) -> dict[str, str]:
+        return {
+            "x-api-key": api_key,
+            "Content-Type": "application/json",
+        }
+
+    def set_default_options(self, options: dict[str, Any]) -> dict[str, Any]:
+        return options
+
+    def parse_response(self, response_data: dict[str, Any]) -> ContentText | None:
+        exa_search_response = ExaSearchResponse.model_validate(response_data)
+
+        if not exa_search_response.answer and not exa_search_response.citations:
+            return None
+
+        return ContentText(
+            text=exa_search_response.answer,
+            citations=[
+                UrlCitation(
+                    cited_text=citation.text, title=citation.title, url=citation.url
+                )
+                for citation in exa_search_response.citations
+            ],
+        )
+
+
+def exa_search_provider(
+    in_options: dict[str, object] | None = None,
+) -> SearchProvider:
+    options = ExaOptions.model_validate(in_options) if in_options else None
+    return ExaSearchProvider(
+        options.model_dump(exclude_none=True) if options else None
+    ).search