inspect-ai 0.3.59__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/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"

inspect_ai/model/_providers/openai.py

@@ -1,4 +1,3 @@
-import json
 import os
 from logging import getLogger
 from typing import Any
@@ -15,51 +14,39 @@ from openai import (
 from openai._types import NOT_GIVEN
 from openai.types.chat import (
     ChatCompletion,
-    ChatCompletionAssistantMessageParam,
-    ChatCompletionContentPartImageParam,
-    ChatCompletionContentPartInputAudioParam,
-    ChatCompletionContentPartParam,
-    ChatCompletionContentPartTextParam,
-    ChatCompletionDeveloperMessageParam,
-    ChatCompletionMessage,
-    ChatCompletionMessageParam,
-    ChatCompletionMessageToolCallParam,
-    ChatCompletionNamedToolChoiceParam,
-    ChatCompletionSystemMessageParam,
-    ChatCompletionToolChoiceOptionParam,
-    ChatCompletionToolMessageParam,
-    ChatCompletionToolParam,
-    ChatCompletionUserMessageParam,
 )
-from openai.types.shared_params.function_definition import FunctionDefinition
 from typing_extensions import override
 
 from inspect_ai._util.constants import DEFAULT_MAX_RETRIES
-from inspect_ai._util.content import Content
 from inspect_ai._util.error import PrerequisiteError
-from inspect_ai._util.images import file_as_data_uri
 from inspect_ai._util.logger import warn_once
-from inspect_ai._util.url import is_http_url
-from inspect_ai.tool import ToolCall, ToolChoice, ToolFunction, ToolInfo
+from inspect_ai.model._openai import chat_choices_from_openai
+from inspect_ai.tool import ToolChoice, ToolInfo
 
-from .._chat_message import ChatMessage, ChatMessageAssistant
+from .._chat_message import ChatMessage
 from .._generate_config import GenerateConfig
 from .._image import image_url_filter
 from .._model import ModelAPI
 from .._model_call import ModelCall
 from .._model_output import (
     ChatCompletionChoice,
-    Logprobs,
     ModelOutput,
     ModelUsage,
     StopReason,
 )
+from .._openai import (
+    is_o1,
+    is_o1_full,
+    is_o1_mini,
+    is_o1_preview,
+    openai_chat_messages,
+    openai_chat_tool_choice,
+    openai_chat_tools,
+)
 from .openai_o1 import generate_o1
 from .util import (
-    as_stop_reason,
     environment_prerequisite_error,
     model_base_url,
-    parse_tool_call,
 )
 
 logger = getLogger(__name__)
@@ -87,20 +74,22 @@ class OpenAIAPI(ModelAPI):
             config=config,
         )
 
-        # pull out azure model_arg
-        AZURE_MODEL_ARG = "azure"
-        is_azure = False
-        if AZURE_MODEL_ARG in model_args:
-            is_azure = model_args.get(AZURE_MODEL_ARG, False)
-            del model_args[AZURE_MODEL_ARG]
+        # extract any service prefix from model name
+        parts = model_name.split("/")
+        if len(parts) > 1:
+            self.service: str | None = parts[0]
+            model_name = "/".join(parts[1:])
+        else:
+            self.service = None
 
         # resolve api_key
         if not self.api_key:
             self.api_key = os.environ.get(
                 AZUREAI_OPENAI_API_KEY, os.environ.get(AZURE_OPENAI_API_KEY, None)
             )
-            if self.api_key:
-                is_azure = True
+            # backward compatibility for when env vars determined service
+            if self.api_key and (os.environ.get(OPENAI_API_KEY, None) is None):
+                self.service = "azure"
             else:
                 self.api_key = os.environ.get(OPENAI_API_KEY, None)
                 if not self.api_key:
@@ -113,7 +102,7 @@ class OpenAIAPI(ModelAPI):
                     )
 
         # azure client
-        if is_azure:
+        if self.is_azure():
             # resolve base_url
             base_url = model_base_url(
                 base_url,
@@ -148,17 +137,20 @@ class OpenAIAPI(ModelAPI):
                 **model_args,
             )
 
+    def is_azure(self) -> bool:
+        return self.service == "azure"
+
     def is_o1(self) -> bool:
-        return self.model_name.startswith("o1")
+        return is_o1(self.model_name)
 
     def is_o1_full(self) -> bool:
-        return self.is_o1() and not self.is_o1_mini() and not self.is_o1_preview()
+        return is_o1_full(self.model_name)
 
     def is_o1_mini(self) -> bool:
-        return self.model_name.startswith("o1-mini")
+        return is_o1_mini(self.model_name)
 
     def is_o1_preview(self) -> bool:
-        return self.model_name.startswith("o1-preview")
+        return is_o1_preview(self.model_name)
 
     async def generate(
         self,
@@ -198,9 +190,11 @@ class OpenAIAPI(ModelAPI):
 
         # prepare request (we do this so we can log the ModelCall)
         request = dict(
-            messages=await as_openai_chat_messages(input, self.is_o1_full()),
-            tools=chat_tools(tools) if len(tools) > 0 else NOT_GIVEN,
-            tool_choice=chat_tool_choice(tool_choice) if len(tools) > 0 else NOT_GIVEN,
+            messages=await openai_chat_messages(input, self.model_name),
+            tools=openai_chat_tools(tools) if len(tools) > 0 else NOT_GIVEN,
+            tool_choice=openai_chat_tool_choice(tool_choice)
+            if len(tools) > 0
+            else NOT_GIVEN,
             **self.completion_params(config, len(tools) > 0),
         )
 
@@ -237,7 +231,7 @@ class OpenAIAPI(ModelAPI):
         self, response: ChatCompletion, tools: list[ToolInfo]
     ) -> list[ChatCompletionChoice]:
         # adding this as a method so we can override from other classes (e.g together)
-        return chat_choices_from_response(response, tools)
+        return chat_choices_from_openai(response, tools)
 
     @override
     def is_rate_limit(self, ex: BaseException) -> bool:
@@ -327,163 +321,3 @@ class OpenAIAPI(ModelAPI):
             )
         else:
             return e
-
-
-async def as_openai_chat_messages(
-    messages: list[ChatMessage], o1_full: bool
-) -> list[ChatCompletionMessageParam]:
-    return [await openai_chat_message(message, o1_full) for message in messages]
-
-
-async def openai_chat_message(
-    message: ChatMessage, o1_full: bool
-) -> ChatCompletionMessageParam:
-    if message.role == "system":
-        if o1_full:
-            return ChatCompletionDeveloperMessageParam(
-                role="developer", content=message.text
-            )
-        else:
-            return ChatCompletionSystemMessageParam(
-                role=message.role, content=message.text
-            )
-    elif message.role == "user":
-        return ChatCompletionUserMessageParam(
-            role=message.role,
-            content=(
-                message.content
-                if isinstance(message.content, str)
-                else [
-                    await as_chat_completion_part(content)
-                    for content in message.content
-                ]
-            ),
-        )
-    elif message.role == "assistant":
-        if message.tool_calls:
-            return ChatCompletionAssistantMessageParam(
-                role=message.role,
-                content=message.text,
-                tool_calls=[chat_tool_call(call) for call in message.tool_calls],
-            )
-        else:
-            return ChatCompletionAssistantMessageParam(
-                role=message.role, content=message.text
-            )
-    elif message.role == "tool":
-        return ChatCompletionToolMessageParam(
-            role=message.role,
-            content=(
-                f"Error: {message.error.message}" if message.error else message.text
-            ),
-            tool_call_id=str(message.tool_call_id),
-        )
-    else:
-        raise ValueError(f"Unexpected message role {message.role}")
-
-
-def chat_tool_call(tool_call: ToolCall) -> ChatCompletionMessageToolCallParam:
-    return ChatCompletionMessageToolCallParam(
-        id=tool_call.id,
-        function=dict(
-            name=tool_call.function, arguments=json.dumps(tool_call.arguments)
-        ),
-        type=tool_call.type,
-    )
-
-
-def chat_tools(tools: list[ToolInfo]) -> list[ChatCompletionToolParam]:
-    return [chat_tool_param(tool) for tool in tools]
-
-
-def chat_tool_param(tool: ToolInfo) -> ChatCompletionToolParam:
-    function = FunctionDefinition(
-        name=tool.name,
-        description=tool.description,
-        parameters=tool.parameters.model_dump(exclude_none=True),
-    )
-    return ChatCompletionToolParam(type="function", function=function)
-
-
-def chat_tool_choice(tool_choice: ToolChoice) -> ChatCompletionToolChoiceOptionParam:
-    if isinstance(tool_choice, ToolFunction):
-        return ChatCompletionNamedToolChoiceParam(
-            type="function", function=dict(name=tool_choice.name)
-        )
-    # openai supports 'any' via the 'required' keyword
-    elif tool_choice == "any":
-        return "required"
-    else:
-        return tool_choice
-
-
-def chat_tool_calls(
-    message: ChatCompletionMessage, tools: list[ToolInfo]
-) -> list[ToolCall] | None:
-    if message.tool_calls:
-        return [
-            parse_tool_call(call.id, call.function.name, call.function.arguments, tools)
-            for call in message.tool_calls
-        ]
-    else:
-        return None
-
-
-def chat_choices_from_response(
-    response: ChatCompletion, tools: list[ToolInfo]
-) -> list[ChatCompletionChoice]:
-    choices = list(response.choices)
-    choices.sort(key=lambda c: c.index)
-    return [
-        ChatCompletionChoice(
-            message=chat_message_assistant(choice.message, tools),
-            stop_reason=as_stop_reason(choice.finish_reason),
-            logprobs=(
-                Logprobs(**choice.logprobs.model_dump())
-                if choice.logprobs is not None
-                else None
-            ),
-        )
-        for choice in choices
-    ]
-
-
-def chat_message_assistant(
-    message: ChatCompletionMessage, tools: list[ToolInfo]
-) -> ChatMessageAssistant:
-    return ChatMessageAssistant(
-        content=message.content or "",
-        source="generate",
-        tool_calls=chat_tool_calls(message, tools),
-    )
-
-
-async def as_chat_completion_part(
-    content: Content,
-) -> ChatCompletionContentPartParam:
-    if content.type == "text":
-        return ChatCompletionContentPartTextParam(type="text", text=content.text)
-    elif content.type == "image":
-        # API takes URL or base64 encoded file. If it's a remote file or
-        # data URL leave it alone, otherwise encode it
-        image_url = content.image
-        detail = content.detail
-
-        if not is_http_url(image_url):
-            image_url = await file_as_data_uri(image_url)
-
-        return ChatCompletionContentPartImageParam(
-            type="image_url",
-            image_url=dict(url=image_url, detail=detail),
-        )
-    elif content.type == "audio":
-        audio_data = await file_as_data_uri(content.audio)
-
-        return ChatCompletionContentPartInputAudioParam(
-            type="input_audio", input_audio=dict(data=audio_data, format=content.format)
-        )
-
-    else:
-        raise RuntimeError(
-            "Video content is not currently supported by Open AI chat models."
-        )

inspect_ai/model/_providers/openai_o1.py

@@ -24,15 +24,13 @@ from inspect_ai.model import (
 )
 from inspect_ai.tool import ToolCall, ToolInfo
 
+from .._call_tools import parse_tool_call, tool_parse_error_message
 from .._model_call import ModelCall
-from .._model_output import ModelUsage, StopReason
+from .._model_output import ModelUsage, StopReason, as_stop_reason
 from .._providers.util import (
     ChatAPIHandler,
     ChatAPIMessage,
-    as_stop_reason,
     chat_api_input,
-    parse_tool_call,
-    tool_parse_error_message,
 )
 
 logger = getLogger(__name__)

inspect_ai/model/_providers/providers.py

@@ -239,6 +239,28 @@ def mockllm() -> type[ModelAPI]:
     return MockLLM
 
 
+@modelapi("goodfire")
+def goodfire() -> type[ModelAPI]:
+    """Get the Goodfire API provider."""
+    FEATURE = "Goodfire API"
+    PACKAGE = "goodfire"
+    MIN_VERSION = "0.3.4"  # Support for newer Llama models and OpenAI compatibility
+
+    # verify we have the package
+    try:
+        import goodfire  # noqa: F401
+    except ImportError:
+        raise pip_dependency_error(FEATURE, [PACKAGE])
+
+    # verify version
+    verify_required_version(FEATURE, PACKAGE, MIN_VERSION)
+
+    # in the clear
+    from .goodfire import GoodfireAPI
+
+    return GoodfireAPI
+
+
 def validate_openai_client(feature: str) -> None:
     FEATURE = feature
     PACKAGE = "openai"