inspect-ai 0.3.58__py3-none-any.whl → 0.3.60__py3-none-any.whl

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

@@ -11,7 +11,6 @@ import proto  # type: ignore
 from google.ai.generativelanguage import (
     Blob,
     Candidate,
-    File,
     FunctionCall,
     FunctionCallingConfig,
     FunctionDeclaration,
@@ -29,29 +28,29 @@ from google.api_core.exceptions import (
     TooManyRequests,
 )
 from google.api_core.retry.retry_base import if_transient_error
-from google.generativeai import (  # type: ignore
-    GenerationConfig,
-    GenerativeModel,
-    configure,
-    get_file,
-    upload_file,
-)
-from google.generativeai.types import (  # type: ignore
-    AsyncGenerateContentResponse,
+from google.generativeai.client import configure
+from google.generativeai.files import get_file, upload_file
+from google.generativeai.generative_models import GenerativeModel
+from google.generativeai.types import (
     ContentDict,
-    HarmBlockThreshold,
-    HarmCategory,
+    GenerationConfig,
     PartDict,
     PartType,
-    SafetySettingDict,
     Tool,
 )
+from google.generativeai.types.file_types import File
+from google.generativeai.types.generation_types import AsyncGenerateContentResponse
+from google.generativeai.types.safety_types import (
+    EasySafetySettingDict,
+    HarmBlockThreshold,
+    HarmCategory,
+)
 from google.protobuf.json_format import MessageToDict, ParseDict
 from google.protobuf.struct_pb2 import Struct
 from pydantic import JsonValue
 from typing_extensions import override
 
-from inspect_ai._util.constants import BASE_64_DATA_REMOVED
+from inspect_ai._util.constants import BASE_64_DATA_REMOVED, NO_CONTENT
 from inspect_ai._util.content import (
     Content,
     ContentAudio,
@@ -89,7 +88,7 @@ logger = getLogger(__name__)
 
 SAFETY_SETTINGS = "safety_settings"
 
-DEFAULT_SAFETY_SETTINGS: SafetySettingDict = {
+DEFAULT_SAFETY_SETTINGS: EasySafetySettingDict = {
     HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
     HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
     HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
@@ -141,7 +140,7 @@ class GoogleAPI(ModelAPI):
         tools: list[ToolInfo],
         tool_choice: ToolChoice,
         config: GenerateConfig,
-    ) -> ModelOutput | tuple[ModelOutput, ModelCall]:
+    ) -> ModelOutput | tuple[ModelOutput | Exception, ModelCall]:
         parameters = GenerationConfig(
             temperature=config.temperature,
             top_p=config.top_p,
@@ -149,11 +148,8 @@ class GoogleAPI(ModelAPI):
             max_output_tokens=config.max_tokens,
             stop_sequences=config.stop_seqs,
             candidate_count=config.num_choices,
-            seed=config.seed,
             presence_penalty=config.presence_penalty,
             frequency_penalty=config.frequency_penalty,
-            response_logprobs=config.logprobs,
-            logprobs=config.top_logprobs,
         )
 
         # google-native messages
@@ -176,18 +172,15 @@ class GoogleAPI(ModelAPI):
                 response=response,
             )
 
-        # cast to AsyncGenerateContentResponse since we passed stream=False
         try:
-            response = cast(
-                AsyncGenerateContentResponse,
-                await self.model.generate_content_async(
-                    contents=contents,
-                    safety_settings=self.safety_settings,
-                    generation_config=parameters,
-                    tools=gemini_tools,
-                    tool_config=gemini_tool_config,
-                ),
+            response = await self.model.generate_content_async(
+                contents=contents,
+                safety_settings=self.safety_settings,
+                generation_config=parameters,
+                tools=gemini_tools,
+                tool_config=gemini_tool_config,
             )
+
         except InvalidArgument as ex:
             return self.handle_invalid_argument(ex), model_call()
 
@@ -205,15 +198,13 @@ class GoogleAPI(ModelAPI):
         # return
         return output, model_call()
 
-    def handle_invalid_argument(self, ex: InvalidArgument) -> ModelOutput:
+    def handle_invalid_argument(self, ex: InvalidArgument) -> ModelOutput | Exception:
         if "size exceeds the limit" in ex.message.lower():
             return ModelOutput.from_content(
                 model=self.model_name, content=ex.message, stop_reason="model_length"
             )
         else:
-            return ModelOutput.from_content(
-                model=self.model_name, content=ex.message, stop_reason="unknown"
-            )
+            return ex
 
     @override
     def is_rate_limit(self, ex: BaseException) -> bool:
@@ -231,7 +222,7 @@ class GoogleAPI(ModelAPI):
 def build_model_call(
     contents: list[ContentDict],
     generation_config: GenerationConfig,
-    safety_settings: SafetySettingDict,
+    safety_settings: EasySafetySettingDict,
     tools: list[Tool] | None,
     tool_config: ToolConfig | None,
     response: AsyncGenerateContentResponse | None,
@@ -248,7 +239,7 @@ def build_model_call(
             if tool_config is not None
             else None,
         ),
-        response=response.to_dict() if response is not None else {},
+        response=response.to_dict() if response is not None else {},  # type: ignore[no-untyped-call]
         filter=model_call_filter,
     )
 
@@ -269,12 +260,12 @@ def model_call_content(content: ContentDict) -> ContentDict:
 
 def model_call_part(part: PartType) -> PartType:
     if isinstance(part, proto.Message):
-        return MessageToDict(part._pb)
+        return cast(PartDict, MessageToDict(part._pb))
     elif isinstance(part, dict):
         part = part.copy()
         keys = list(part.keys())
         for key in keys:
-            part[key] = model_call_part(part[key])
+            part[key] = model_call_part(part[key])  # type: ignore[literal-required]
         return part
     else:
         return part
@@ -316,9 +307,6 @@ def consective_tool_message_reducer(
     return messages
 
 
-NO_CONTENT = "(no content)"
-
-
 async def content_dict(
     message: ChatMessageUser | ChatMessageAssistant | ChatMessageTool,
 ) -> ContentDict:
@@ -326,13 +314,13 @@ async def content_dict(
         return ContentDict(
             role="user",
             parts=(
-                [PartDict(text=message.content or NO_CONTENT)]
+                [message.content or NO_CONTENT]
                 if isinstance(message.content, str)
                 else [await content_part(content) for content in message.content]
             ),
         )
     elif isinstance(message, ChatMessageAssistant):
-        content_parts: list[Part] = []
+        content_parts: list[PartType] = []
         # tool call parts
         if message.tool_calls is not None:
             content_parts.extend(
@@ -383,9 +371,9 @@ def dict_to_struct(x: dict[str, Any]) -> Struct:
 
 async def content_part(content: Content | str) -> PartType:
     if isinstance(content, str):
-        return PartDict(text=content or NO_CONTENT)
+        return content or NO_CONTENT
     elif isinstance(content, ContentText):
-        return PartDict(text=content.text or NO_CONTENT)
+        return content.text or NO_CONTENT
     else:
         return await chat_content_to_part(content)
 
@@ -404,7 +392,9 @@ def prepend_system_messages(
     messages: list[ContentDict], system_messages: list[ChatMessageSystem]
 ) -> None:
     # create system_parts
-    system_parts = [Part(text=message.content) for message in system_messages]
+    system_parts: list[PartType] = [
+        Part(text=message.content) for message in system_messages
+    ]
 
     # we want the system messages to be prepended to the first user message
     # (if there is no first user message then prepend one)
@@ -476,6 +466,8 @@ def schema_from_param(param: ToolParam | ToolParams, nullable: bool = False) ->
             return schema_from_param(param.anyOf[0], nullable=True)
         else:
             return Schema(type=Type.TYPE_UNSPECIFIED)
+    elif param.enum:
+        return Schema(type=Type.STRING, format="enum", enum=param.enum)
     else:
         return Schema(type=Type.TYPE_UNSPECIFIED)
 
@@ -600,14 +592,14 @@ def gapi_should_retry(ex: BaseException) -> bool:
 
 def parse_safety_settings(
     safety_settings: Any,
-) -> dict[HarmCategory, HarmBlockThreshold]:
+) -> EasySafetySettingDict:
     # ensure we have a dict
     if isinstance(safety_settings, str):
         safety_settings = json.loads(safety_settings)
     if not isinstance(safety_settings, dict):
         raise ValueError(f"{SAFETY_SETTINGS} must be dictionary.")
 
-    parsed_settings: dict[HarmCategory, HarmBlockThreshold] = {}
+    parsed_settings: EasySafetySettingDict = {}
     for key, value in safety_settings.items():
         if isinstance(key, str):
             key = str_to_harm_category(key)
@@ -623,23 +615,23 @@ def parse_safety_settings(
     return parsed_settings
 
 
-def str_to_harm_category(category: str) -> HarmCategory:
+def str_to_harm_category(category: str) -> int:
     category = category.upper()
     if "HARASSMENT" in category:
-        return HarmCategory.HARM_CATEGORY_HARASSMENT
+        return cast(int, HarmCategory.HARM_CATEGORY_HARASSMENT)
     elif "HATE_SPEECH" in category:
-        return HarmCategory.HARM_CATEGORY_HATE_SPEECH
+        return cast(int, HarmCategory.HARM_CATEGORY_HATE_SPEECH)
     elif "SEXUALLY_EXPLICIT" in category:
-        return HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT
+        return cast(int, HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT)
     elif "DANGEROUS_CONTENT" in category:
-        return HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT
+        return cast(int, HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT)
     else:
         # NOTE: Although there is an "UNSPECIFIED" category, in the
         # documentation, the API does not accept it.
         raise ValueError(f"Unknown HarmCategory: {category}")
 
 
-def str_to_harm_block_threshold(threshold: str) -> HarmBlockThreshold:
+def str_to_harm_block_threshold(threshold: str) -> int:
     threshold = threshold.upper()
     if "LOW" in threshold:
         return HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
@@ -673,7 +665,7 @@ async def file_for_content(content: ContentAudio | ContentVideo) -> File:
         uploaded_file = files_db.get(content_sha256)
         if uploaded_file:
             try:
-                upload = cast(File, get_file(uploaded_file))
+                upload = get_file(uploaded_file)
                 if upload.state.name == "ACTIVE":
                     trace(f"Using uploaded file: {uploaded_file}")
                     return upload

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

@@ -40,11 +40,13 @@ from typing_extensions import override
 # https://github.com/mistralai/client-python/blob/main/MIGRATION.md
 from inspect_ai._util.constants import (
     DEFAULT_TIMEOUT,
+    NO_CONTENT,
 )
 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,
@@ -58,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"
@@ -122,7 +124,7 @@ class MistralAPI(ModelAPI):
         tools: list[ToolInfo],
         tool_choice: ToolChoice,
         config: GenerateConfig,
-    ) -> ModelOutput | tuple[ModelOutput, ModelCall]:
+    ) -> ModelOutput | tuple[ModelOutput | Exception, ModelCall]:
         # build request
         request: dict[str, Any] = dict(
             model=self.model_name,
@@ -146,7 +148,7 @@ class MistralAPI(ModelAPI):
             response = await self.client.chat.complete_async(**request)
         except SDKError as ex:
             if ex.status_code == 400:
-                return self.handle_bad_request(ex)
+                return self.handle_bad_request(ex), mistral_model_call(request, None)
             else:
                 raise ex
 
@@ -181,25 +183,27 @@ class MistralAPI(ModelAPI):
     def connection_key(self) -> str:
         return str(self.api_key)
 
-    def handle_bad_request(self, ex: SDKError) -> ModelOutput:
+    def handle_bad_request(self, ex: SDKError) -> ModelOutput | Exception:
+        body = json.loads(ex.body)
+        content = body.get("message", ex.body)
         if "maximum context length" in ex.body:
-            body = json.loads(ex.body)
-            content = body.get("message", ex.body)
             return ModelOutput.from_content(
                 model=self.model_name, content=content, stop_reason="model_length"
             )
         else:
-            raise ex
+            return ex
 
 
 def mistral_model_call(
-    request: dict[str, Any], response: MistralChatCompletionResponse
+    request: dict[str, Any], response: MistralChatCompletionResponse | None
 ) -> ModelCall:
     request = request.copy()
     request.update(messages=[message.model_dump() for message in request["messages"]])
     if request.get("tools", None) is not None:
         request["tools"] = [tool.model_dump() for tool in request["tools"]]
-    return ModelCall(request=request, response=response.model_dump())
+    return ModelCall(
+        request=request, response=response.model_dump() if response else {}
+    )
 
 
 def mistral_chat_tools(tools: list[ToolInfo]) -> list[MistralTool]:
@@ -326,9 +330,6 @@ async def mistral_chat_message(
         )
 
 
-NO_CONTENT = "(no content)"
-
-
 async def mistral_message_content(
     content: str | list[Content],
 ) -> str | list[ContentChunk]: