crfm-helm 0.5.2__py3-none-any.whl → 0.5.4__py3-none-any.whl

helm/clients/huggingface_client.py

@@ -9,6 +9,7 @@ from typing import Any, Dict, List, Optional, TypedDict
 
 from helm.common.cache import CacheConfig
 from helm.common.hierarchical_logger import htrack_block, hlog
+from helm.common.optional_dependencies import handle_module_not_found_error
 from helm.common.request import (
     wrap_request_time,
     EMBEDDING_UNAVAILABLE_REQUEST_RESULT,
@@ -58,60 +59,61 @@ class HuggingFaceServer:
         self,
         pretrained_model_name_or_path: str,
         wrapped_tokenizer: WrappedPreTrainedTokenizer,
-        openvino=False,
+        openvino: bool = False,
         **kwargs,
     ):
-        if torch.cuda.is_available():
-            hlog("CUDA is available, initializing with a GPU...")
-            self.device: str = "cuda:0"
+        self.device: Optional[str]
+        if "device_map" in kwargs:
+            try:
+                import accelerate  # noqa: F401
+            except ModuleNotFoundError as e:
+                handle_module_not_found_error(e, ["accelerate"])
+            hlog(f'Hugging Face device_map set to "{kwargs["device_map"]}".')
+            self.device = None
+        elif torch.cuda.is_available():
+            hlog('Hugging Face device set to "cuda:0" because CUDA is available.')
+            self.device = "cuda:0"
         else:
+            hlog('Hugging Face device set to "cpu" because CUDA is unavailable.')
             self.device = "cpu"
+
+        # Security issue: currently we trust remote code by default.
+        # We retain this temporarily to maintain reverse compatibility.
+        # TODO: Delete if-else and don't set trust_remote_code=True
+        if "trust_remote_code" not in kwargs:
+            kwargs["trust_remote_code"] = True
+
         with htrack_block(f"Loading Hugging Face model {pretrained_model_name_or_path}"):
             # WARNING this may fail if your GPU does not have enough memory
             if openvino:
-                """
-                Optimum Intel provides a simple interface to optimize Transformer models and convert them to \
-                OpenVINO™ Intermediate Representation (IR) format to accelerate end-to-end pipelines on \
-                Intel® architectures using OpenVINO™ runtime.
-                """
-                from helm.common.optional_dependencies import handle_module_not_found_error
-
+                # Optimum Intel provides a simple interface to optimize Transformer models and convert them to \
+                # OpenVINO™ Intermediate Representation (IR) format to accelerate end-to-end pipelines on \
+                # Intel® architectures using OpenVINO™ runtime.
                 try:
                     from optimum.intel.openvino import OVModelForCausalLM
                 except ModuleNotFoundError as e:
                     handle_module_not_found_error(e, ["openvino"])
 
                 self.device = "cpu"
-                # Security issue: currently we trust remote code by default.
-                # We retain this temporarily to maintain reverse compatibility.
-                # TODO: Delete if-else and don't set trust_remote_code=True
-                if "trust_remote_code" in kwargs:
-                    self.model = OVModelForCausalLM.from_pretrained(
-                        pretrained_model_name_or_path, export=True, **kwargs
-                    ).to(self.device)
-                else:
-                    self.model = OVModelForCausalLM.from_pretrained(
-                        pretrained_model_name_or_path, export=True, trust_remote_code=True, **kwargs
-                    ).to(self.device)
+                self.model = OVModelForCausalLM.from_pretrained(
+                    pretrained_model_name_or_path, export=True, **kwargs
+                ).to(self.device)
+            elif self.device is None:
+                # kwargs contains device_map=auto
+                # Do not call to() because accelerate will take care of model device placement.
+                self.model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path, **kwargs)
             else:
-                # Security issue: currently we trust remote code by default.
-                # We retain this temporarily to maintain reverse compatibility.
-                # TODO: Delete if-else and don't set trust_remote_code=True
-                if "trust_remote_code" in kwargs:
-                    self.model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path, **kwargs).to(
-                        self.device
-                    )
-                else:
-                    self.model = AutoModelForCausalLM.from_pretrained(
-                        pretrained_model_name_or_path, trust_remote_code=True, **kwargs
-                    ).to(self.device)
+                self.model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path, **kwargs).to(
+                    self.device
+                )
         self.wrapped_tokenizer = wrapped_tokenizer
 
     def serve_request(self, raw_request: HuggingFaceRequest) -> Dict:
         with self.wrapped_tokenizer as tokenizer:
             encoded_input = tokenizer(raw_request["prompt"], return_tensors="pt", return_token_type_ids=False).to(
-                self.device
+                0 if self.device is None else self.device
             )
+
         stopping_criteria: Optional[StoppingCriteriaList] = None
         optional_args = {}
         if len(raw_request["stop_sequences"]) > 0:
@@ -249,9 +251,8 @@ def _process_huggingface_client_kwargs(raw_kwargs: Dict[str, Any]):
     # e.g. the string "torch.bfloat16" is converted to torch.bfloat16
     torch_dtype = processed_kwargs.get(TORCH_DTYPE_KEY)
     if torch_dtype and isinstance(torch_dtype, str):
-        if not torch_dtype.startswith(TORCH_DTYPE_VALUE_PREFIX):
-            raise ValueError(f'Unknown dtype "{torch_dtype}"; expected a string such as "torch.bfloat16"')
-        processed_kwargs[TORCH_DTYPE_KEY] = getattr(torch, torch_dtype[len(TORCH_DTYPE_VALUE_PREFIX) :])
+        if torch_dtype.startswith(TORCH_DTYPE_VALUE_PREFIX):
+            processed_kwargs[TORCH_DTYPE_KEY] = getattr(torch, torch_dtype[len(TORCH_DTYPE_VALUE_PREFIX) :])
 
     return processed_kwargs
 

helm/clients/nvidia_nim_client.py

@@ -0,0 +1,35 @@
+from typing import Optional
+
+from helm.clients.openai_client import OpenAIClient
+from helm.common.cache import CacheConfig
+from helm.common.request import Request
+from helm.tokenizers.tokenizer import Tokenizer
+
+
+class NvidiaNimClient(OpenAIClient):
+
+    BASE_URL = "https://integrate.api.nvidia.com/v1"
+
+    def __init__(
+        self,
+        tokenizer: Tokenizer,
+        tokenizer_name: str,
+        cache_config: CacheConfig,
+        api_key: Optional[str] = None,
+    ):
+        self.tokenizer = tokenizer
+        self.tokenizer_name = tokenizer_name
+        super().__init__(
+            tokenizer=tokenizer,
+            tokenizer_name=tokenizer_name,
+            cache_config=cache_config,
+            api_key=api_key,
+            org_id=None,
+            base_url=NvidiaNimClient.BASE_URL,
+        )
+
+    def _get_model_for_request(self, request: Request) -> str:
+        return request.model
+
+    def _is_chat_model_engine(self, model_engine: str) -> bool:
+        return True

helm/clients/openai_client.py

@@ -12,8 +12,8 @@ from helm.common.tokenization_request import (
     TokenizationRequest,
     TokenizationRequestResult,
 )
-from helm.tokenizers.tokenizer import Tokenizer
 from .client import CachingClient, truncate_sequence, generate_uid_for_multimodal_prompt
+from helm.tokenizers.tokenizer import Tokenizer
 
 try:
     import openai
@@ -51,7 +51,7 @@ class OpenAIClient(CachingClient):
     def _is_chat_model_engine(self, model_engine: str) -> bool:
         if model_engine == "gpt-3.5-turbo-instruct":
             return False
-        elif model_engine.startswith("gpt-3.5") or model_engine.startswith("gpt-4"):
+        elif model_engine.startswith("gpt-3.5") or model_engine.startswith("gpt-4") or model_engine.startswith("o1"):
             return True
         return False
 
@@ -132,6 +132,7 @@ class OpenAIClient(CachingClient):
             content: Union[str, List[Union[str, Any]]]
             if request.multimodal_prompt is not None:
                 content = []
+                request.validate()
                 for media_object in request.multimodal_prompt.media_objects:
                     if media_object.is_type("image") and media_object.location:
                         from helm.common.images_utils import encode_base64
@@ -140,8 +141,6 @@ class OpenAIClient(CachingClient):
                         image_object: Dict[str, str] = {"url": f"data:image/jpeg;base64,{base64_image}"}
                         content.append({"type": "image_url", "image_url": image_object})
                     elif media_object.is_type(TEXT_TYPE):
-                        if media_object.text is None:
-                            raise ValueError("MediaObject of text type has missing text field value")
                         content.append({"type": media_object.type, "text": media_object.text})
                     else:
                         raise ValueError(f"Unrecognized MediaObject type {media_object.type}")
@@ -170,6 +169,21 @@ class OpenAIClient(CachingClient):
         if is_vlm(request.model) and raw_request["stop"] is None:
             raw_request.pop("stop")
 
+        # Special handling for o1 models.
+        # Refer to the "Reasoning models" documentation further discussion of o1 model limitations:
+        # https://platform.openai.com/docs/guides/reasoning
+        if request.model_engine.startswith("o1"):
+            # Avoid error:
+            # "Unsupported parameter: 'max_tokens' is not supported with this model. Use 'max_completion_tokens' instead."  # noqa: E501
+            # Note that openai>=1.45 is needed for this
+            if raw_request["max_tokens"]:
+                raw_request["max_completion_tokens"] = raw_request["max_tokens"]
+                raw_request.pop("max_tokens")
+            # Avoid error:
+            # "Invalid type for 'stop': expected an unsupported value, but got null instead."
+            if raw_request["stop"] is None:
+                raw_request.pop("stop")
+
         def do_it() -> Dict[str, Any]:
             return self.client.chat.completions.create(**raw_request).model_dump(mode="json")
 

helm/clients/palmyra_client.py

@@ -3,6 +3,7 @@ import json
 import requests
 from typing import Any, Dict, List
 
+from helm.clients.openai_client import OpenAIClient
 from helm.common.cache import CacheConfig
 from helm.common.hierarchical_logger import hlog
 from helm.common.request import wrap_request_time, Request, RequestResult, GeneratedOutput, Token, ErrorFlags
@@ -142,3 +143,26 @@ class PalmyraClient(CachingClient):
             completions=completions,
             embedding=[],
         )
+
+
+class PalmyraChatClient(OpenAIClient):
+    """Sends request to a Palmyra model using a OpenAI-compatible Chat API."""
+
+    def __init__(
+        self,
+        tokenizer: Tokenizer,
+        tokenizer_name: str,
+        cache_config: CacheConfig,
+        api_key: str,
+    ):
+        super().__init__(
+            tokenizer=tokenizer,
+            tokenizer_name=tokenizer_name,
+            cache_config=cache_config,
+            api_key=api_key,
+            org_id=None,
+            base_url="https://api.writer.com/v1/chat",
+        )
+
+    def _is_chat_model_engine(self, model_engine: str) -> bool:
+        return True

helm/clients/perspective_api_client.py

@@ -4,16 +4,21 @@ from dataclasses import asdict
 from typing import Any, List, Dict, Optional
 
 from dacite import from_dict
-from googleapiclient import discovery
-from googleapiclient.errors import BatchError, HttpError
-from googleapiclient.http import BatchHttpRequest
-from httplib2 import HttpLib2Error
+
 from helm.clients.toxicity_classifier_client import ToxicityClassifierClient
+from helm.common.optional_dependencies import handle_module_not_found_error
 from helm.proxy.retry import NonRetriableException
-
 from helm.common.cache import Cache, CacheConfig
 from helm.common.perspective_api_request import ToxicityAttributes, PerspectiveAPIRequest, PerspectiveAPIRequestResult
-from google.auth.exceptions import DefaultCredentialsError
+
+try:
+    from googleapiclient import discovery
+    from googleapiclient.errors import BatchError, HttpError
+    from googleapiclient.http import BatchHttpRequest
+    from httplib2 import HttpLib2Error
+    from google.auth.exceptions import DefaultCredentialsError
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["metrics"])
 
 
 class PerspectiveAPIClientCredentialsError(NonRetriableException):

helm/clients/test_client.py

@@ -23,30 +23,28 @@ def test_truncate_sequence():
     # echo_prompt = True, nothing gets truncated
     truncate_sequence_helper(
         ["a", "b", "c"],
-        Request(
-            model="openai/text-davinci-002", model_deployment="openai/text-davinci-002", prompt="abc", echo_prompt=True
-        ),
+        Request(model="openai/gpt2", model_deployment="huggingface/gpt2", prompt="abc", echo_prompt=True),
         ["a", "b", "c"],
     )
 
     # Nothing gets truncated
     truncate_sequence_helper(
         ["hello", " world"],
-        Request(model="openai/text-davinci-002", model_deployment="openai/text-davinci-002", stop_sequences=["#"]),
+        Request(model="openai/gpt2", model_deployment="huggingface/gpt2", stop_sequences=["#"]),
         ["hello", " world"],
     )
 
     # Truncate using stop sequences
     truncate_sequence_helper(
         ["hello", " world", "\n", "what"],
-        Request(model="openai/text-davinci-002", model_deployment="openai/text-davinci-002", stop_sequences=["\n"]),
+        Request(model="openai/gpt2", model_deployment="huggingface/gpt2", stop_sequences=["\n"]),
         ["hello", " world"],
     )
 
     # Truncate using max tokens
     truncate_sequence_helper(
         ["a", "b", "c"],
-        Request(model="openai/text-davinci-002", model_deployment="openai/text-davinci-002", max_tokens=2),
+        Request(model="openai/gpt2", model_deployment="huggingface/gpt2", max_tokens=2),
         ["a", "b"],
     )
 

helm/clients/together_client.py

@@ -7,6 +7,7 @@ import requests
 from retrying import retry
 
 from helm.common.cache import CacheConfig
+from helm.common.media_object import IMAGE_TYPE, TEXT_TYPE
 from helm.common.optional_dependencies import handle_module_not_found_error
 from helm.common.request import wrap_request_time, Request, RequestResult, GeneratedOutput, Token
 from helm.clients.client import CachingClient, truncate_sequence, cleanup_str
@@ -323,8 +324,29 @@ class TogetherChatClient(CachingClient):
         self._together_model = together_model
 
     def convert_to_raw_chat_request(self, request: Request) -> TogetherRawChatRequest:
+        request.validate()
+        messages: List[Dict[str, Any]]
         if request.messages:
             messages = request.messages
+        elif request.multimodal_prompt:
+            message_contents = []
+            for media_object in request.multimodal_prompt.media_objects:
+                if media_object.is_type(IMAGE_TYPE) and media_object.location:
+                    assert media_object.location
+                    if media_object.is_local_file:
+                        from helm.common.images_utils import encode_base64
+
+                        base64_image: str = encode_base64(media_object.location)
+                        image_url = f"data:image/jpeg;base64,{base64_image}"
+                    else:
+                        image_url = media_object.location
+                    message_contents.append({"type": "image_url", "image_url": {"url": image_url}})
+                elif media_object.is_type(TEXT_TYPE):
+                    assert media_object.text
+                    message_contents.append({"type": "text", "text": media_object.text})
+                else:
+                    raise ValueError(f"Unrecognized MediaObject type {media_object.type}")
+            messages = [{"role": "user", "content": message_contents}]
         else:
             messages = [{"role": "user", "content": request.prompt}]
         if self._together_model is not None:

helm/clients/vision_language/open_flamingo_client.py

@@ -82,13 +82,12 @@ class OpenFlamingoClient(CachingClient):
         # Build the prompt
         prompt_text: str = ""
         images: List[Image.Image] = []
+        request.validate()
         for media_object in request.multimodal_prompt.media_objects:
             if media_object.is_type("image") and media_object.location:
                 images.append(open_image(media_object.location))
                 prompt_text += self.IMAGE_TOKEN
             elif media_object.is_type(TEXT_TYPE):
-                if media_object.text is None:
-                    raise ValueError("MediaObject of text type has missing text field value")
                 prompt_text += media_object.text
             else:
                 raise ValueError(f"Unrecognized MediaObject type {media_object.type}")

helm/clients/vision_language/palmyra_vision_client.py

@@ -6,13 +6,19 @@ import requests
 from helm.common.cache import CacheConfig
 from helm.common.images_utils import encode_base64
 from helm.common.media_object import TEXT_TYPE
-from helm.common.request import Request, RequestResult, GeneratedOutput
+from helm.common.request import Request, RequestResult, GeneratedOutput, ErrorFlags
 from helm.common.request import wrap_request_time
 from helm.clients.client import CachingClient, generate_uid_for_multimodal_prompt, truncate_and_tokenize_response_text
 from helm.tokenizers.tokenizer import Tokenizer
 
 
+class PalmyraVisionContentBlockedError(Exception):
+    pass
+
+
 class PalmyraVisionClient(CachingClient):
+    CONTENT_BLOCKED_ERROR: str = "fail.input.content.moderation"
+
     def __init__(self, tokenizer: Tokenizer, tokenizer_name: str, endpoint: str, cache_config: CacheConfig):
         super().__init__(cache_config)
         self.tokenizer: Tokenizer = tokenizer
@@ -49,17 +55,19 @@ class PalmyraVisionClient(CachingClient):
                 response = requests.post(
                     self.endpoint, headers={"Content-Type": "application/json"}, data=json.dumps({"parts": prompt})
                 )
-                if response.status_code != 200:
-                    curl_command: str = (
-                        f"curl --location '{self.endpoint}' --header 'Content-Type: application/json' "
-                        f"--data '{json.dumps({'parts': prompt})}'"
-                    )
-                    assert False, f"Got status code {response.status_code}. Try {curl_command}"
-
                 json_response = json.loads(response.text)
-                assert (
-                    "choices" in json_response and "errors" not in json_response
-                ), f"Invalid response: {response.text}"
+
+                # Check for content blocked error
+                if (
+                    "errors" in json_response
+                    and "tpe" in json_response
+                    and json_response["tpe"] == self.CONTENT_BLOCKED_ERROR
+                ):
+                    raise PalmyraVisionContentBlockedError(json_response["errors"])
+
+                # Hard fail if the `choices` is missing from the response
+                assert "choices" in json_response, f"Invalid response: {response.text}"
+
                 return json_response
 
             cache_key = CachingClient.make_cache_key(
@@ -67,8 +75,15 @@ class PalmyraVisionClient(CachingClient):
                 request=request,
             )
             result, cached = self.cache.get(cache_key, wrap_request_time(do_it))
-        except RuntimeError as ex:
-            return RequestResult(success=False, cached=False, error=str(ex), completions=[], embedding=[])
+        except PalmyraVisionContentBlockedError as ex:
+            return RequestResult(
+                success=False,
+                cached=False,
+                error=f"Content blocked: {str(ex)}",
+                completions=[],
+                embedding=[],
+                error_flags=ErrorFlags(is_retriable=False, is_fatal=False),
+            )
 
         # The internal endpoint doesn't support any other parameters, so we have to truncate ourselves
         completions: List[GeneratedOutput] = [

helm/common/cache.py

@@ -1,6 +1,6 @@
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Dict, Callable, Generator, Mapping, Optional, Tuple
+from typing import Dict, Callable, Generator, Mapping, Tuple
 import json
 import threading
 
@@ -38,6 +38,12 @@ class CacheConfig:
 class KeyValueStoreCacheConfig(CacheConfig):
     """Configuration for a cache backed by a key-value store."""
 
+    # This was originally to distinguish between "primitive" cache configs
+    # and "compound" cache configs. But we don't have any "compound" cache configs currently.
+    # Hypthetical "compound" example: ReadOnlyCacheConfig(SqliteCacheConfig("path"))
+    # TODO: Maybe remove this eventually?
+    pass
+
 
 @dataclass(frozen=True)
 class SqliteCacheConfig(KeyValueStoreCacheConfig):
@@ -78,24 +84,6 @@ class MongoCacheConfig(KeyValueStoreCacheConfig):
         return f"{self.uri}/{self.collection_name}"
 
 
-@dataclass(frozen=True)
-class WithFollowerCacheConfig(CacheConfig):
-    """Configuration of a cache backed by a main cache and a follower cache."""
-
-    # Configuration for the main cache.
-    # Responses will be written to and served out of this cache.
-    main: KeyValueStoreCacheConfig
-
-    # Configuration for the follower cache.
-    # The follower cache is a write-only cache. Responses will be written to this cache,
-    # but not served out of this cache.
-    follower: KeyValueStoreCacheConfig
-
-    @property
-    def cache_stats_key(self) -> str:
-        return self.main.cache_stats_key
-
-
 def get_all_from_sqlite(path: str) -> Generator[Tuple[Dict, Dict], None, None]:
     """Yields all decoded key, value pairs from the SQLite cache.
 
@@ -126,7 +114,7 @@ def create_key_value_store(config: KeyValueStoreCacheConfig) -> KeyValueStore:
     elif isinstance(config, BlackHoleCacheConfig):
         return BlackHoleKeyValueStore()
     else:
-        raise ValueError(f"KeyValueStoreCacheConfig with unknown type: {config}")
+        raise ValueError(f"CacheConfig with unknown type: {config}")
 
 
 @retry
@@ -189,14 +177,8 @@ class Cache(object):
 
     def __init__(self, config: CacheConfig):
         hlog(f"Created cache with config: {config}")
-        self.config: KeyValueStoreCacheConfig
-        self.follower_config: Optional[KeyValueStoreCacheConfig]
         if isinstance(config, KeyValueStoreCacheConfig):
             self.config = config
-            self.follower_config = None
-        elif isinstance(config, WithFollowerCacheConfig):
-            self.config = config.main
-            self.follower_config = config.follower
         else:
             raise ValueError(f"CacheConfig with unknown type: {config}")
 
@@ -216,8 +198,4 @@ class Cache(object):
                 response = compute()
 
                 write_to_key_value_store(key_value_store, request, response)
-        if self.follower_config is not None:
-            # TODO: Initialize follower_key_value_store in constructor
-            with create_key_value_store(self.follower_config) as follower_key_value_store:
-                write_to_key_value_store(follower_key_value_store, request, response)
         return response, cached

helm/common/images_utils.py

@@ -1,3 +1,4 @@
+from hashlib import md5
 import base64
 import io
 import os
@@ -44,6 +45,11 @@ def encode_base64(image_location: str, format="JPEG") -> str:
     return base64.b64encode(image_file.getvalue()).decode("ascii")
 
 
+def generate_hash(image: Image.Image) -> str:
+    """Generates a hash for the image."""
+    return md5(image.tobytes()).hexdigest()
+
+
 def copy_image(src: str, dest: str, width: Optional[int] = None, height: Optional[int] = None) -> None:
     """
     Copies the image file from `src` path to `dest` path. If dimensions `width` and `height`

helm/common/key_value_store.py

@@ -15,11 +15,11 @@ class KeyValueStore(contextlib.AbstractContextManager):
     """Key value store that persists writes."""
 
     @abstractmethod
-    def contains(self, key: Dict) -> bool:
+    def contains(self, key: Mapping) -> bool:
         pass
 
     @abstractmethod
-    def get(self, key: Dict) -> Optional[Dict]:
+    def get(self, key: Mapping) -> Optional[Dict]:
         pass
 
     @abstractmethod
@@ -35,7 +35,7 @@ class KeyValueStore(contextlib.AbstractContextManager):
         pass
 
     @abstractmethod
-    def remove(self, key: Dict) -> None:
+    def remove(self, key: Mapping) -> None:
         pass
 
 
@@ -53,10 +53,10 @@ class SqliteKeyValueStore(KeyValueStore):
     def __exit__(self, exc_type, exc_value, traceback) -> None:
         self._sqlite_dict.__exit__(exc_type, exc_value, traceback)
 
-    def contains(self, key: Dict) -> bool:
+    def contains(self, key: Mapping) -> bool:
         return request_to_key(key) in self._sqlite_dict
 
-    def get(self, key: Dict) -> Optional[Dict]:
+    def get(self, key: Mapping) -> Optional[Dict]:
         key_string = request_to_key(key)
         result = self._sqlite_dict.get(key_string)
         if result is not None:
@@ -77,7 +77,7 @@ class SqliteKeyValueStore(KeyValueStore):
         for key, value in pairs:
             self.put(key, value)
 
-    def remove(self, key: Dict) -> None:
+    def remove(self, key: Mapping) -> None:
         del self._sqlite_dict[key]
         self._sqlite_dict.commit()
 
@@ -91,10 +91,10 @@ class BlackHoleKeyValueStore(KeyValueStore):
     def __exit__(self, exc_type, exc_value, traceback) -> None:
         pass
 
-    def contains(self, key: Dict) -> bool:
+    def contains(self, key: Mapping) -> bool:
         return False
 
-    def get(self, key: Dict) -> Optional[Dict]:
+    def get(self, key: Mapping) -> Optional[Dict]:
         return None
 
     def get_all(self) -> Generator[Tuple[Dict, Dict], None, None]:
@@ -109,5 +109,5 @@ class BlackHoleKeyValueStore(KeyValueStore):
     def multi_put(self, pairs: Iterable[Tuple[Dict, Dict]]) -> None:
         return None
 
-    def remove(self, key: Dict) -> None:
+    def remove(self, key: Mapping) -> None:
         return None

helm/common/mongo_key_value_store.py

@@ -39,11 +39,11 @@ class MongoKeyValueStore(KeyValueStore):
         serialized = json.dumps(key, sort_keys=True)
         return json.loads(serialized, object_pairs_hook=SON)
 
-    def contains(self, key: Dict) -> bool:
+    def contains(self, key: Mapping) -> bool:
         query = {self._REQUEST_KEY: self._canonicalize_key(key)}
         return self._collection.find_one(query) is not None
 
-    def get(self, key: Dict) -> Optional[Dict]:
+    def get(self, key: Mapping) -> Optional[Dict]:
         query = {self._REQUEST_KEY: self._canonicalize_key(key)}
         document = self._collection.find_one(query)
         if document is not None:
@@ -84,5 +84,6 @@ class MongoKeyValueStore(KeyValueStore):
         # Note: unlike put, multi_put does not support documents with null bytes in keys.
         self._collection.bulk_write(operations)
 
-    def remove(self, key: Dict) -> None:
-        self._collection.delete_one(key)
+    def remove(self, key: Mapping) -> None:
+        query = {self._REQUEST_KEY: self._canonicalize_key(key)}
+        self._collection.delete_one(query)

helm/common/request.py

@@ -72,6 +72,22 @@ class Request:
     image_generation_parameters: Optional[ImageGenerationParameters] = None
     """Parameters for image generation."""
 
+    def validate(self):
+        if (
+            (self.messages and self.prompt)
+            or (self.messages and self.multimodal_prompt)
+            or (self.prompt and self.multimodal_prompt)
+        ):
+            raise ValueError("Exactly one of the messages, prompt, multimodal_prompt fields should be set")
+
+        if self.multimodal_prompt:
+            for media_object in self.multimodal_prompt.media_objects:
+                if media_object.content_type == "text" and media_object.text is None:
+                    raise ValueError("Media object with text content type must have text set")
+
+                if media_object.content_type == "image" and media_object.location is None:
+                    raise ValueError("Media object with image content type must have location set")
+
     @property
     def model_host(self) -> str:
         """Returns the model host (referring to the deployment).