inspect-ai 0.3.59__py3-none-any.whl → 0.3.61__py3-none-any.whl

inspect_ai/model/_providers/anthropic.py

@@ -14,10 +14,13 @@ from anthropic import (
     APIConnectionError,
     AsyncAnthropic,
     AsyncAnthropicBedrock,
+    AsyncAnthropicVertex,
     BadRequestError,
     InternalServerError,
+    NotGiven,
     RateLimitError,
 )
+from anthropic._types import Body
 from anthropic.types import (
     ImageBlockParam,
     Message,
@@ -64,15 +67,25 @@ class AnthropicAPI(ModelAPI):
         base_url: str | None = None,
         api_key: str | None = None,
         config: GenerateConfig = GenerateConfig(),
-        bedrock: bool = False,
         **model_args: Any,
     ):
         # extract any service prefix from model name
         parts = model_name.split("/")
         if len(parts) > 1:
-            service = parts[0]
-            bedrock = service == "bedrock"
+            self.service: str | None = parts[0]
             model_name = "/".join(parts[1:])
+        else:
+            self.service = None
+
+        # collect gemerate model_args (then delete them so we can pass the rest on)
+        def collect_model_arg(name: str) -> Any | None:
+            nonlocal model_args
+            value = model_args.get(name, None)
+            if value is not None:
+                model_args.pop(name)
+            return value
+
+        self.extra_body: Body | None = collect_model_arg("extra_body")
 
         # call super
         super().__init__(
@@ -84,7 +97,7 @@ class AnthropicAPI(ModelAPI):
         )
 
         # create client
-        if bedrock:
+        if self.is_bedrock():
             base_url = model_base_url(
                 base_url, ["ANTHROPIC_BEDROCK_BASE_URL", "BEDROCK_ANTHROPIC_BASE_URL"]
             )
@@ -95,7 +108,9 @@ class AnthropicAPI(ModelAPI):
             if base_region is None:
                 aws_region = os.environ.get("AWS_DEFAULT_REGION", None)
 
-            self.client: AsyncAnthropic | AsyncAnthropicBedrock = AsyncAnthropicBedrock(
+            self.client: (
+                AsyncAnthropic | AsyncAnthropicBedrock | AsyncAnthropicVertex
+            ) = AsyncAnthropicBedrock(
                 base_url=base_url,
                 max_retries=(
                     config.max_retries if config.max_retries else DEFAULT_MAX_RETRIES
@@ -103,6 +118,21 @@ class AnthropicAPI(ModelAPI):
                 aws_region=aws_region,
                 **model_args,
             )
+        elif self.is_vertex():
+            base_url = model_base_url(
+                base_url, ["ANTHROPIC_VERTEX_BASE_URL", "VERTEX_ANTHROPIC_BASE_URL"]
+            )
+            region = os.environ.get("ANTHROPIC_VERTEX_REGION", NotGiven())
+            project_id = os.environ.get("ANTHROPIC_VERTEX_PROJECT_ID", NotGiven())
+            self.client = AsyncAnthropicVertex(
+                region=region,
+                project_id=project_id,
+                base_url=base_url,
+                max_retries=(
+                    config.max_retries if config.max_retries else DEFAULT_MAX_RETRIES
+                ),
+                **model_args,
+            )
         else:
             # resolve api_key
             if not self.api_key:
@@ -119,6 +149,12 @@ class AnthropicAPI(ModelAPI):
                 **model_args,
             )
 
+    def is_bedrock(self) -> bool:
+        return self.service == "bedrock"
+
+    def is_vertex(self) -> bool:
+        return self.service == "vertex"
+
     async def generate(
         self,
         input: list[ChatMessage],
@@ -163,6 +199,10 @@ class AnthropicAPI(ModelAPI):
             if computer_use:
                 request["extra_headers"] = {"anthropic-beta": "computer-use-2024-10-22"}
 
+            # extra_body
+            if self.extra_body is not None:
+                request["extra_body"] = self.extra_body
+
             # make request
             message = await self.client.messages.create(**request, stream=False)
 
@@ -251,9 +291,6 @@ class AnthropicAPI(ModelAPI):
         elif "content filtering" in error:
             content = "Sorry, but I am unable to help with that request."
             stop_reason = "content_filter"
-        else:
-            content = error
-            stop_reason = "unknown"
 
         if content and stop_reason:
             return ModelOutput.from_content(
@@ -466,6 +503,12 @@ def message_tool_choice(tool_choice: ToolChoice) -> message_create_params.ToolCh
 
 
 async def message_param(message: ChatMessage) -> MessageParam:
+    # if content is empty that is going to result in an error when we replay
+    # this message to claude, so in that case insert a NO_CONTENT message
+    if isinstance(message.content, list) and len(message.content) == 0:
+        message = message.model_copy()
+        message.content = [ContentText(text=NO_CONTENT)]
+
     # no system role for anthropic (this is more like an assertion,
     # as these should have already been filtered out)
     if message.role == "system":
@@ -507,7 +550,7 @@ async def message_param(message: ChatMessage) -> MessageParam:
     elif message.role == "assistant" and message.tool_calls:
         # first include content (claude <thinking>)
         tools_content: list[TextBlockParam | ImageBlockParam | ToolUseBlockParam] = (
-            [TextBlockParam(type="text", text=message.content)]
+            [TextBlockParam(type="text", text=message.content or NO_CONTENT)]
             if isinstance(message.content, str)
             else (
                 [(await message_param_content(content)) for content in message.content]
@@ -576,11 +619,6 @@ def model_output_from_message(message: Message, tools: list[ToolInfo]) -> ModelO
                 )
             )
 
-    # if content is empty that is going to result in an error when we replay
-    # this message to claude, so in that case insert a NO_CONTENT message
-    if len(content) == 0:
-        content = [ContentText(text=NO_CONTENT)]
-
     # resolve choice
     choice = ChatCompletionChoice(
         message=ChatMessageAssistant(

inspect_ai/model/_providers/azureai.py

@@ -37,6 +37,7 @@ from inspect_ai.tool import ToolChoice, ToolInfo
 from inspect_ai.tool._tool_call import ToolCall
 from inspect_ai.tool._tool_choice import ToolFunction
 
+from .._call_tools import parse_tool_call
 from .._chat_message import (
     ChatMessage,
     ChatMessageAssistant,
@@ -60,7 +61,6 @@ from .util import (
 )
 from .util.chatapi import ChatAPIHandler
 from .util.llama31 import Llama31Handler
-from .util.util import parse_tool_call
 
 AZUREAI_API_KEY = "AZUREAI_API_KEY"
 AZUREAI_ENDPOINT_KEY = "AZUREAI_ENDPOINT_KEY"

inspect_ai/model/_providers/goodfire.py

@@ -0,0 +1,248 @@
+import os
+from typing import Any, List, Literal, get_args
+
+from goodfire import AsyncClient
+from goodfire.api.chat.interfaces import ChatMessage as GoodfireChatMessage
+from goodfire.api.exceptions import InvalidRequestException, RateLimitException
+from goodfire.variants.variants import SUPPORTED_MODELS, Variant
+from typing_extensions import override
+
+from inspect_ai.tool._tool_choice import ToolChoice
+from inspect_ai.tool._tool_info import ToolInfo
+
+from .._chat_message import (
+    ChatMessage,
+    ChatMessageAssistant,
+    ChatMessageSystem,
+    ChatMessageTool,
+    ChatMessageUser,
+)
+from .._generate_config import GenerateConfig
+from .._model import ModelAPI
+from .._model_call import ModelCall
+from .._model_output import (
+    ChatCompletionChoice,
+    ModelOutput,
+    ModelUsage,
+)
+from .util import environment_prerequisite_error, model_base_url
+
+# Constants
+GOODFIRE_API_KEY = "GOODFIRE_API_KEY"
+DEFAULT_BASE_URL = "https://api.goodfire.ai"
+DEFAULT_MAX_TOKENS = 4096
+DEFAULT_TEMPERATURE = 1.0  # Standard sampling temperature (baseline)
+DEFAULT_TOP_P = 1.0  # No nucleus sampling truncation (baseline)
+
+
+class GoodfireAPI(ModelAPI):
+    """Goodfire API provider.
+
+    This provider implements the Goodfire API for LLM inference. It supports:
+    - Chat completions with standard message formats
+    - Basic parameter controls (temperature, top_p, etc.)
+    - Usage statistics tracking
+    - Stop reason handling
+
+    Does not currently support:
+    - Tool calls
+    - Feature analysis
+    - Streaming responses
+
+    Known limitations:
+    - Limited role support (system/user/assistant only)
+    - Tool messages converted to user messages
+    """
+
+    client: AsyncClient
+    variant: Variant
+    model_args: dict[str, Any]
+
+    def __init__(
+        self,
+        model_name: str,
+        base_url: str | None = None,
+        api_key: str | None = None,
+        config: GenerateConfig = GenerateConfig(),
+        **model_args: Any,
+    ) -> None:
+        """Initialize the Goodfire API provider.
+
+        Args:
+            model_name: Name of the model to use
+            base_url: Optional custom API base URL
+            api_key: Optional API key (will check env vars if not provided)
+            config: Generation config options
+            **model_args: Additional arguments passed to the API
+        """
+        super().__init__(
+            model_name=model_name,
+            base_url=base_url,
+            api_key=api_key,
+            api_key_vars=[GOODFIRE_API_KEY],
+            config=config,
+        )
+
+        # resolve api_key
+        if not self.api_key:
+            self.api_key = os.environ.get(GOODFIRE_API_KEY)
+            if not self.api_key:
+                raise environment_prerequisite_error("Goodfire", GOODFIRE_API_KEY)
+
+        # Validate model name against supported models
+        supported_models = list(get_args(SUPPORTED_MODELS))
+        if self.model_name not in supported_models:
+            raise ValueError(
+                f"Model {self.model_name} not supported. Supported models: {supported_models}"
+            )
+
+        # Initialize client with minimal configuration
+        base_url_val = model_base_url(base_url, "GOODFIRE_BASE_URL")
+        assert isinstance(base_url_val, str) or base_url_val is None
+
+        # Store model args for use in generate
+        self.model_args = model_args
+
+        self.client = AsyncClient(
+            api_key=self.api_key,
+            base_url=base_url_val or DEFAULT_BASE_URL,
+        )
+
+        # Initialize variant directly with model name
+        self.variant = Variant(self.model_name)  # type: ignore
+
+    def _to_goodfire_message(self, message: ChatMessage) -> GoodfireChatMessage:
+        """Convert an Inspect message to a Goodfire message format.
+
+        Args:
+            message: The message to convert
+
+        Returns:
+            The converted message in Goodfire format
+
+        Raises:
+            ValueError: If the message type is unknown
+        """
+        role: Literal["system", "user", "assistant"] = "user"
+        if isinstance(message, ChatMessageSystem):
+            role = "system"
+        elif isinstance(message, ChatMessageUser):
+            role = "user"
+        elif isinstance(message, ChatMessageAssistant):
+            role = "assistant"
+        elif isinstance(message, ChatMessageTool):
+            role = "user"  # Convert tool messages to user messages
+        else:
+            raise ValueError(f"Unknown message type: {type(message)}")
+
+        content = str(message.content)
+        if isinstance(message, ChatMessageTool):
+            content = f"Tool {message.function}: {content}"
+
+        return GoodfireChatMessage(role=role, content=content)
+
+    def handle_error(self, ex: Exception) -> ModelOutput | Exception:
+        """Handle only errors that need special treatment for retry logic or model limits."""
+        # Handle token/context length errors
+        if isinstance(ex, InvalidRequestException):
+            error_msg = str(ex).lower()
+            if "context length" in error_msg or "max tokens" in error_msg:
+                return ModelOutput.from_content(
+                    model=self.model_name,
+                    content=str(ex),
+                    stop_reason="model_length",
+                    error=error_msg,
+                )
+
+        # Let all other errors propagate
+        return ex
+
+    @override
+    def is_rate_limit(self, ex: BaseException) -> bool:
+        """Check if exception is due to rate limiting."""
+        return isinstance(ex, RateLimitException)
+
+    @override
+    def connection_key(self) -> str:
+        """Return key for connection pooling."""
+        return f"goodfire:{self.api_key}"
+
+    @override
+    def max_tokens(self) -> int | None:
+        """Return maximum tokens supported by model."""
+        return DEFAULT_MAX_TOKENS  # Let Goodfire's Variant handle model-specific limits
+
+    async def generate(
+        self,
+        input: List[ChatMessage],
+        tools: List[ToolInfo],
+        tool_choice: ToolChoice,
+        config: GenerateConfig,
+        *,
+        cache: bool = True,
+    ) -> tuple[ModelOutput | Exception, ModelCall]:
+        """Generate output from the model."""
+        # Convert messages and prepare request params
+        messages = [self._to_goodfire_message(msg) for msg in input]
+        # Build request parameters with type hints
+        params: dict[str, Any] = {
+            "model": self.variant.base_model,  # Use base_model instead of stringifying the Variant
+            "messages": messages,
+            "max_completion_tokens": int(config.max_tokens)
+            if config.max_tokens
+            else DEFAULT_MAX_TOKENS,
+            "stream": False,
+        }
+
+        # Add generation parameters from config if not in model_args
+        if "temperature" not in self.model_args and config.temperature is not None:
+            params["temperature"] = float(config.temperature)
+        elif "temperature" not in self.model_args:
+            params["temperature"] = DEFAULT_TEMPERATURE
+
+        if "top_p" not in self.model_args and config.top_p is not None:
+            params["top_p"] = float(config.top_p)
+        elif "top_p" not in self.model_args:
+            params["top_p"] = DEFAULT_TOP_P
+
+        # Add any additional model args (highest priority)
+        api_params = {
+            k: v
+            for k, v in self.model_args.items()
+            if k not in ["api_key", "base_url", "model_args"]
+        }
+        params.update(api_params)
+
+        try:
+            # Use native async client
+            response = await self.client.chat.completions.create(**params)
+            response_dict = response.model_dump()
+
+            output = ModelOutput(
+                model=self.model_name,
+                choices=[
+                    ChatCompletionChoice(
+                        message=ChatMessageAssistant(
+                            content=response_dict["choices"][0]["message"]["content"]
+                        ),
+                        stop_reason="stop",
+                    )
+                ],
+                usage=ModelUsage(**response_dict["usage"])
+                if "usage" in response_dict
+                else None,
+            )
+            model_call = ModelCall.create(request=params, response=response_dict)
+            return (output, model_call)
+        except Exception as ex:
+            result = self.handle_error(ex)
+            model_call = ModelCall.create(
+                request=params,
+                response={},  # Empty response for error case
+            )
+            return (result, model_call)
+
+    @property
+    def name(self) -> str:
+        """Get provider name."""
+        return "goodfire"

inspect_ai/model/_providers/groq.py

@@ -27,6 +27,7 @@ from inspect_ai._util.images import file_as_data_uri
 from inspect_ai._util.url import is_http_url
 from inspect_ai.tool import ToolCall, ToolChoice, ToolFunction, ToolInfo
 
+from .._call_tools import parse_tool_call
 from .._chat_message import (
     ChatMessage,
     ChatMessageAssistant,
@@ -37,12 +38,15 @@ from .._chat_message import (
 from .._generate_config import GenerateConfig
 from .._model import ModelAPI
 from .._model_call import ModelCall
-from .._model_output import ChatCompletionChoice, ModelOutput, ModelUsage
-from .util import (
+from .._model_output import (
+    ChatCompletionChoice,
+    ModelOutput,
+    ModelUsage,
     as_stop_reason,
+)
+from .util import (
     environment_prerequisite_error,
     model_base_url,
-    parse_tool_call,
 )
 
 GROQ_API_KEY = "GROQ_API_KEY"

inspect_ai/model/_providers/hf.py

@@ -150,6 +150,12 @@ class HuggingFaceAPI(ModelAPI):
             kwargs["output_logits"] = config.logprobs
         if "return_dict_in_generate" in kwargs:
             assert kwargs["return_dict_in_generate"]
+        if config.stop_seqs is not None:
+            from transformers.generation import StopStringCriteria  # type: ignore
+
+            stopping_criteria = [StopStringCriteria(self.tokenizer, config.stop_seqs)]
+            kwargs["stopping_criteria"] = stopping_criteria
+
         kwargs["return_dict_in_generate"] = True
         generator = functools.partial(self.model.generate, **kwargs)
 

inspect_ai/model/_providers/mistral.py

@@ -46,6 +46,7 @@ from inspect_ai._util.content import Content, ContentImage, ContentText
 from inspect_ai._util.images import file_as_data_uri
 from inspect_ai.tool import ToolCall, ToolChoice, ToolFunction, ToolInfo
 
+from .._call_tools import parse_tool_call
 from .._chat_message import (
     ChatMessage,
     ChatMessageAssistant,
@@ -59,7 +60,7 @@ from .._model_output import (
     ModelUsage,
     StopReason,
 )
-from .util import environment_prerequisite_error, model_base_url, parse_tool_call
+from .util import environment_prerequisite_error, model_base_url
 
 AZURE_MISTRAL_API_KEY = "AZURE_MISTRAL_API_KEY"
 AZUREAI_MISTRAL_API_KEY = "AZUREAI_MISTRAL_API_KEY"