crfm-helm 0.5.7__py3-none-any.whl → 0.5.8__py3-none-any.whl

helm/clients/audio_language/qwen_omni/qwen2_5_omni_utils/v2_5/vision_process.py

@@ -0,0 +1,380 @@
+from __future__ import annotations
+
+import base64
+import logging
+import math
+import os
+import sys
+import time
+import warnings
+from functools import lru_cache
+from io import BytesIO
+
+import requests
+import torch
+import torchvision
+from packaging import version
+from PIL import Image
+from torchvision import io, transforms
+from torchvision.transforms import InterpolationMode
+from typing import List, Optional, Union
+
+
+logger = logging.getLogger(__name__)
+
+IMAGE_FACTOR = 28
+MIN_PIXELS = 4 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+VIDEO_MIN_PIXELS = 128 * 28 * 28
+VIDEO_MAX_PIXELS = 768 * 28 * 28
+FRAME_FACTOR = 2
+FPS = 2.0
+FPS_MIN_FRAMES = 4
+FPS_MAX_FRAMES = 768
+
+# Set the maximum number of video token inputs.
+# Here, 128K represents the maximum number of input tokens for the VLLM model.
+# Remember to adjust it according to your own configuration.
+VIDEO_TOTAL_PIXELS = int(float(os.environ.get("VIDEO_MAX_PIXELS", 128000 * 28 * 28 * 0.9)))
+logger.info(f"set VIDEO_TOTAL_PIXELS: {VIDEO_TOTAL_PIXELS}")
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(
+    height: int, width: int, factor: int = IMAGE_FACTOR, min_pixels: int = MIN_PIXELS, max_pixels: int = MAX_PIXELS
+) -> tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(int(height / beta), factor)
+        w_bar = floor_by_factor(int(width / beta), factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(int(height * beta), factor)
+        w_bar = ceil_by_factor(int(width * beta), factor)
+    return h_bar, w_bar
+
+
+def to_rgb(pil_image: Image.Image) -> Image.Image:
+    if pil_image.mode == "RGBA":
+        white_background = Image.new("RGB", pil_image.size, (255, 255, 255))
+        white_background.paste(pil_image, mask=pil_image.split()[3])  # Use alpha channel as mask
+        return white_background
+    else:
+        return pil_image.convert("RGB")
+
+
+def fetch_image(ele, size_factor: int = IMAGE_FACTOR) -> Image.Image:
+    if "image" in ele:
+        image = ele["image"]
+    else:
+        image = ele["image_url"]
+    image_obj = None
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        response = requests.get(image, stream=True)
+        image_obj = Image.open(BytesIO(response.content))
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            image_obj = Image.open(BytesIO(data))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
+    image = to_rgb(image_obj)
+    # resize
+    if "resized_height" in ele and "resized_width" in ele:
+        resized_height, resized_width = smart_resize(
+            int(ele["resized_height"]),
+            int(ele["resized_width"]),
+            factor=size_factor,
+        )
+    else:
+        width, height = image.size
+        min_pixels = int(ele.get("min_pixels", MIN_PIXELS))
+        max_pixels = int(ele.get("max_pixels", MAX_PIXELS))
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=size_factor,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+        )
+    image = image.resize((resized_width, resized_height))
+
+    return image
+
+
+def smart_nframes(
+    ele: dict,
+    total_frames: int,
+    video_fps: Union[int, float],
+) -> int:
+    """calculate the number of frames for video used for model inputs.
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+            support either `fps` or `nframes`:
+                - nframes: the number of frames to extract for model inputs.
+                - fps: the fps to extract frames for model inputs.
+                    - min_frames: the minimum number of frames of the video, only used when fps is provided.
+                    - max_frames: the maximum number of frames of the video, only used when fps is provided.
+        total_frames (int): the original total number of frames of the video.
+        video_fps (int | float): the original fps of the video.
+
+    Raises:
+        ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
+
+    Returns:
+        int: the number of frames for video used for model inputs.
+    """
+    assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`"
+    if "nframes" in ele:
+        nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
+    else:
+        fps = ele.get("fps", FPS)
+        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
+        max_frames = floor_by_factor(ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), FRAME_FACTOR)
+        nframes = total_frames / video_fps * fps
+        if nframes > total_frames:
+            logger.warning(f"smart_nframes: nframes[{nframes}] > total_frames[{total_frames}]")
+        nframes = min(min(max(nframes, min_frames), max_frames), total_frames)
+        nframes = floor_by_factor(nframes, FRAME_FACTOR)
+    if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
+        raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
+    return nframes
+
+
+def _read_video_torchvision(
+    ele: dict,
+):
+    """read video using torchvision.io.read_video
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://", "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        torch.Tensor: the video tensor with shape (T, C, H, W).
+    """
+    video_path = ele["video"]
+    if version.parse(torchvision.__version__) < version.parse("0.19.0"):
+        if "http://" in video_path or "https://" in video_path:
+            warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.")
+        if "file://" in video_path:
+            video_path = video_path[7:]
+    st = time.time()
+    video, audio, info = io.read_video(
+        video_path,
+        start_pts=ele.get("video_start", 0.0),
+        end_pts=ele.get("video_end", None),
+        pts_unit="sec",
+        output_format="TCHW",
+    )
+    total_frames, video_fps = video.size(0), info["video_fps"]
+    logger.info(f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(0, total_frames - 1, nframes).round().long()
+    sample_fps = nframes / max(total_frames, 1e-6) * video_fps
+    video = video[idx]
+    return video, sample_fps
+
+
+def is_decord_available() -> bool:
+    import importlib.util
+
+    return importlib.util.find_spec("decord") is not None
+
+
+def _read_video_decord(
+    ele: dict,
+):
+    """read video using decord.VideoReader
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://", "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        torch.Tensor: the video tensor with shape (T, C, H, W).
+    """
+    import decord
+
+    video_path = ele["video"]
+    st = time.time()
+    vr = decord.VideoReader(video_path)
+    # TODO: support start_pts and end_pts
+    if "video_start" in ele or "video_end" in ele:
+        raise NotImplementedError("not support start_pts and end_pts in decord for now.")
+    total_frames, video_fps = len(vr), vr.get_avg_fps()
+    logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
+    video = vr.get_batch(idx).asnumpy()
+    video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
+    sample_fps = nframes / max(total_frames, 1e-6) * video_fps
+    return video, sample_fps
+
+
+VIDEO_READER_BACKENDS = {
+    "decord": _read_video_decord,
+    "torchvision": _read_video_torchvision,
+}
+
+FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
+
+
+@lru_cache(maxsize=1)
+def get_video_reader_backend() -> str:
+    if FORCE_QWENVL_VIDEO_READER is not None:
+        video_reader_backend = FORCE_QWENVL_VIDEO_READER
+    elif is_decord_available():
+        video_reader_backend = "decord"
+    else:
+        video_reader_backend = "torchvision"
+    print(f"qwen-vl-utils using {video_reader_backend} to read video.", file=sys.stderr)
+    return video_reader_backend
+
+
+def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR, return_video_sample_fps: bool = False):
+    if isinstance(ele["video"], str):
+        video_reader_backend = get_video_reader_backend()
+        try:
+            video, sample_fps = VIDEO_READER_BACKENDS[video_reader_backend](ele)
+        except Exception as e:
+            logger.warning(f"video_reader_backend {video_reader_backend} error, use torchvision as default, msg: {e}")
+            video, sample_fps = VIDEO_READER_BACKENDS["torchvision"](ele)
+
+        nframes, _, height, width = video.shape
+        min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
+        total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
+        max_pixels = max(min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), int(min_pixels * 1.05))
+        max_pixels_supposed = ele.get("max_pixels", max_pixels)
+        if max_pixels_supposed > max_pixels:
+            logger.warning(f"The given max_pixels[{max_pixels_supposed}] exceeds limit[{max_pixels}].")
+        max_pixels = min(max_pixels_supposed, max_pixels)
+        if "resized_height" in ele and "resized_width" in ele:
+            resized_height, resized_width = smart_resize(
+                ele["resized_height"],
+                ele["resized_width"],
+                factor=image_factor,
+            )
+        else:
+            resized_height, resized_width = smart_resize(
+                height,
+                width,
+                factor=image_factor,
+                min_pixels=min_pixels,
+                max_pixels=max_pixels,
+            )
+        video = transforms.functional.resize(
+            video,
+            [resized_height, resized_width],
+            interpolation=InterpolationMode.BICUBIC,
+            antialias=True,
+        ).float()
+        if return_video_sample_fps:
+            return video, sample_fps
+        return video
+    else:
+        assert isinstance(ele["video"], (list, tuple))
+        process_info = ele.copy()
+        process_info.pop("type", None)
+        process_info.pop("video", None)
+        images = [
+            fetch_image({"image": video_element, **process_info}, size_factor=image_factor)
+            for video_element in ele["video"]
+        ]
+        nframes = ceil_by_factor(len(images), FRAME_FACTOR)
+        if len(images) < nframes:
+            images.extend([images[-1]] * (nframes - len(images)))
+        if return_video_sample_fps:
+            return images, process_info.pop("fps", 2.0)
+        return images
+
+
+def extract_vision_info(conversations) -> list[dict]:
+    vision_infos = []
+    if isinstance(conversations[0], dict):
+        conversations_p = [conversations]
+    for conversation in conversations_p:
+        for message in conversation:
+            if isinstance(message["content"], list):
+                for ele in message["content"]:
+                    if (
+                        "image" in ele
+                        or "image_url" in ele
+                        or "video" in ele
+                        or ele["type"] in ("image", "image_url", "video")
+                    ):
+                        vision_infos.append(ele)
+    return vision_infos
+
+
+def process_vision_info(
+    conversations: list[dict] | list[list[dict]],
+    return_video_kwargs: bool = False,
+):
+
+    vision_infos = extract_vision_info(conversations)
+    # Read images or videos
+    image_inputs: Optional[List] = []
+    video_inputs: Optional[List] = []
+    video_sample_fps_list = []
+    for vision_info in vision_infos:
+        if "image" in vision_info or "image_url" in vision_info:
+            assert image_inputs is not None
+            image_inputs.append(fetch_image(vision_info))
+        elif "video" in vision_info:
+            assert video_inputs is not None
+            video_input, video_sample_fps = fetch_video(vision_info, return_video_sample_fps=True)
+            video_sample_fps_list.append(video_sample_fps)
+            video_inputs.append(video_input)
+        else:
+            raise ValueError("image, image_url or video should in content.")
+    if image_inputs is not None and len(image_inputs) == 0:
+        image_inputs = None
+    if video_inputs is not None and len(video_inputs) == 0:
+        video_inputs = None
+    if return_video_kwargs:
+        return image_inputs, video_inputs, {"fps": video_sample_fps_list}
+    return image_inputs, video_inputs

helm/clients/image_generation/cogview2/sr_pipeline/dsr_model.py

@@ -184,7 +184,7 @@ def sparse_attention_2d_light(
     attention_dropout=None,
     log_attention_weights=None,
     add_scalar=0,
-    **kwargs
+    **kwargs,
 ):
     """
     q0, k0, v0: [batch_size, 1088, hidden_size]

helm/clients/image_generation/mindalle/models/stage1/layers.py

@@ -141,7 +141,7 @@ class Encoder(nn.Module):
         in_channels: int,
         resolution: int,
         z_channels: int,
-        double_z: Optional[bool] = None
+        double_z: Optional[bool] = None,
     ) -> None:
         super().__init__()
         self.ch = ch
@@ -232,7 +232,7 @@ class Decoder(nn.Module):
         in_channels: int,
         resolution: int,
         z_channels: int,
-        double_z: bool
+        double_z: bool,
     ) -> None:
         super().__init__()
         self.ch = ch

helm/clients/openai_client.py

@@ -33,9 +33,12 @@ class OpenAIClientUtils:
     @classmethod
     def is_reasoning_model(cls, model_engine: str) -> bool:
         # All OpenAI  reasoning models start "o[somenumber]", so we regexp for that to future proof things
-        return bool(re.match(r"^o\d+", model_engine))
+        return bool(re.match(r"^o\d+", model_engine)) or bool(re.match(r"^gpt-5", model_engine))
 
     # Error OpenAI throws when the image in the prompt violates their content policy
+    HARMFUL_INFORMATION_ERROR: str = (
+        "Invalid prompt: we've limited access to this content for safety reasons. This type of information may be used to benefit or to harm people."  # noqa: E501
+    )
     INAPPROPRIATE_IMAGE_ERROR: str = "Your input image may contain content that is not allowed by our safety system"
     INAPPROPRIATE_PROMPT_ERROR: str = "Invalid prompt: your prompt was flagged"
     INAPPROPRIATE_PROMPT_AZURE_ERROR: str = (
@@ -44,12 +47,10 @@ class OpenAIClientUtils:
     INAPPROPRIATE_PROMPT_MICROSOFT_ERROR: str = (
         "The response was filtered due to the prompt triggering Microsoft's content management policy."
     )
-
-    # OpenAI server error
-    OPENAI_SERVER_ERROR: str = (
-        "The server had an error processing your request. Sorry about that! You can retry your request, "
-        "or contact us through our help center at help.openai.com if you keep seeing this error."
-    )
+    # Grok content safety guidelines error message
+    # TODO: Refactor so that this is owned by the Grok client instead.
+    SAFETY_GUIDELINES_GROK_ERROR: str = "Content violates safety guidelines."
+    USAGE_GUIDELINES_GROK_ERROR: str = "Content violates usage guidelines."
 
     # Set the finish reason to this if the prompt violates OpenAI's content policy
     CONTENT_POLICY_VIOLATED_FINISH_REASON: str = (
@@ -74,21 +75,14 @@ class OpenAIClientUtils:
                 completions=[empty_completion] * request.num_completions,
                 embedding=[],
             )
-        elif cls.OPENAI_SERVER_ERROR in str(e):
-            # Handle these errors by returning an empty completion to unblock
-            hwarn(f"OpenAI server error for request: {str(request)}")
-            empty_completion = GeneratedOutput(
-                text="",
-                logprob=0,
-                tokens=[],
-                finish_reason={"reason": cls.OPENAI_SERVER_ERROR},
-            )
+        elif cls.HARMFUL_INFORMATION_ERROR in str(e):
             return RequestResult(
-                success=True,
+                success=False,
                 cached=False,
-                request_time=0,
-                completions=[empty_completion] * request.num_completions,
+                error="Prompt blocked by OpenAI's safety filter",
+                completions=[],
                 embedding=[],
+                error_flags=ErrorFlags(is_retriable=False, is_fatal=False),
             )
         elif cls.INAPPROPRIATE_PROMPT_AZURE_ERROR in str(e) or cls.INAPPROPRIATE_PROMPT_MICROSOFT_ERROR in str(e):
             return RequestResult(
@@ -99,6 +93,24 @@ class OpenAIClientUtils:
                 embedding=[],
                 error_flags=ErrorFlags(is_retriable=False, is_fatal=False),
             )
+        elif cls.SAFETY_GUIDELINES_GROK_ERROR in str(e):
+            return RequestResult(
+                success=False,
+                cached=False,
+                error="Grok API error: Content violates safety guidelines",
+                completions=[],
+                embedding=[],
+                error_flags=ErrorFlags(is_retriable=False, is_fatal=False),
+            )
+        elif cls.USAGE_GUIDELINES_GROK_ERROR in str(e):
+            return RequestResult(
+                success=False,
+                cached=False,
+                error="Grok API error: Content violates usage guidelines",
+                completions=[],
+                embedding=[],
+                error_flags=ErrorFlags(is_retriable=False, is_fatal=False),
+            )
 
         error: str = f"OpenAI error: {e}"
         return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])

helm/clients/openai_responses_client.py

@@ -5,6 +5,7 @@ from typing import Any, Dict, List, Optional, Union
 
 from helm.clients.openai_client import OpenAIClientUtils
 from helm.common.cache import CacheConfig
+from helm.common.hierarchical_logger import hwarn
 from helm.common.media_object import TEXT_TYPE
 from helm.common.request import (
     Thinking,
@@ -60,7 +61,28 @@ class OpenAIResponseClient(CachingClient):
 
     def _make_raw_request(self, request: Request) -> dict[str, Any]:
         input: Union[str, List[Dict[str, Any]]]
-        if request.multimodal_prompt is not None:
+
+        if (
+            (request.prompt and request.messages)
+            or (request.prompt and request.multimodal_prompt)
+            or (request.messages and request.multimodal_prompt)
+        ):
+            raise ValueError(
+                f"More than one of `prompt`, `messages` and `multimodal_prompt` was set in request: {request}"
+            )
+
+        if request.messages is not None:
+            # Checks that all messages have a role and some content
+            for message in request.messages:
+                if not message.get("role") or not message.get("content"):
+                    raise ValueError("All messages must have a role and content")
+            # Checks that the last role is "user"
+            if request.messages[-1]["role"] != "user":
+                raise ValueError("Last message must have role 'user'")
+            if request.prompt != "":
+                hwarn("Since message is set, prompt will be ignored")
+            input = request.messages
+        elif request.multimodal_prompt is not None:
             content = []
             request.validate()
             for media_object in request.multimodal_prompt.media_objects:
@@ -101,6 +123,8 @@ class OpenAIResponseClient(CachingClient):
         # Plus other changes
         model_engine: str = request.model_engine
         if OpenAIClientUtils.is_reasoning_model(model_engine):
+            if "reasoning" not in raw_request:
+                raw_request["reasoning"] = {}
             raw_request["reasoning"]["summary"] = "detailed"
             # Avoid error:
             # "Error code: 400 - {'error': {'message': "Unsupported parameter: 'temperature' is
@@ -150,9 +174,9 @@ class OpenAIResponseClient(CachingClient):
                 ]  # one of "message" or "reasoning" from API observation, but can also include tool calls
 
                 if output_type == "reasoning":
-                    reasoning_output += "\n".join([raw_output["text"] for raw_output in output["summary"]])
+                    reasoning_output += "\n\n".join([raw_output["text"] for raw_output in output["summary"]])
                 elif output_type == "message":
-                    text_output += "\n".join([raw_output["text"] for raw_output in output["content"]])
+                    text_output += "\n\n".join([raw_output["text"] for raw_output in output["content"]])
                 # (Other output types are ignored)
 
             completion = truncate_and_tokenize_response_text(

helm/clients/openrouter_client.py

@@ -0,0 +1,31 @@
+import os
+from typing import Optional
+from helm.clients.openai_client import OpenAIClient
+from helm.common.cache import CacheConfig
+from helm.tokenizers.tokenizer import Tokenizer
+
+
+class OpenRouterClient(OpenAIClient):
+    def __init__(
+        self,
+        tokenizer_name: str,
+        tokenizer: Tokenizer,
+        cache_config: CacheConfig,
+        api_key: Optional[str] = None,
+        model_name: Optional[str] = None,
+        output_processor: Optional[str] = None,
+    ):
+        self.api_key = api_key or os.getenv("OPENROUTER_API_KEY")
+        self.base_url = "https://openrouter.ai/api/v1/"
+        super().__init__(
+            tokenizer,
+            tokenizer_name,
+            cache_config=cache_config,
+            output_processor=output_processor,
+            base_url=self.base_url,
+            api_key=self.api_key,
+        )
+        self.model_name = model_name
+
+    def _get_model_for_request(self, request):
+        return self.model_name or request.model

helm/clients/test_openrouter_client.py

@@ -0,0 +1,69 @@
+import os
+import pytest
+import tempfile
+
+from helm.common.cache import BlackHoleCacheConfig, SqliteCacheConfig
+from helm.common.request import Request
+from helm.clients.openrouter_client import OpenRouterClient
+
+from helm.tokenizers.huggingface_tokenizer import HuggingFaceTokenizer
+
+
+class TestOpenRouterClient:
+    def setup_method(self, method):
+        cache_file = tempfile.NamedTemporaryFile(delete=False)
+        self.cache_path: str = cache_file.name
+        self.tokenizer_name = "mistralai/Mistral-7B-v0.1"
+        self.tokenizer = HuggingFaceTokenizer(
+            cache_config=BlackHoleCacheConfig(),
+            tokenizer_name=self.tokenizer_name,
+        )
+
+    def teardown_method(self, method):
+        os.remove(self.cache_path)
+
+    @pytest.mark.parametrize(
+        "model_name,test_input,expected_model",
+        [
+            (
+                "mistralai/mistral-medium-3.1",
+                Request(
+                    model="mistralai/mistral-medium-3.1",
+                    model_deployment="openrouter/mistral-medium-3.1",
+                ),
+                "mistralai/mistral-medium-3.1",
+            ),
+            (
+                None,
+                Request(model="openai/gpt-oss-20b:free", model_deployment="openrouter/gpt-oss-20b:free"),
+                "openai/gpt-oss-20b:free",
+            ),
+        ],
+    )
+    def test_get_model_for_request(self, model_name, test_input, expected_model):
+        client = OpenRouterClient(
+            tokenizer_name=self.tokenizer_name,
+            tokenizer=self.tokenizer,
+            cache_config=SqliteCacheConfig(self.cache_path),
+            model_name=model_name,
+            api_key="test_key",
+        )
+        assert client._get_model_for_request(test_input) == expected_model
+
+    def test_api_key_env_var(self, monkeypatch):
+        monkeypatch.setenv("OPENROUTER_API_KEY", "test_key")
+        client = OpenRouterClient(
+            tokenizer_name=self.tokenizer_name,
+            tokenizer=self.tokenizer,
+            cache_config=SqliteCacheConfig(self.cache_path),
+        )
+        assert client.api_key == "test_key"
+
+    def test_api_key_argument(self):
+        client = OpenRouterClient(
+            tokenizer_name=self.tokenizer_name,
+            tokenizer=self.tokenizer,
+            cache_config=BlackHoleCacheConfig(),
+            api_key="explicit_key",
+        )
+        assert client.api_key == "explicit_key"