langfun 0.1.2.dev202501050804__py3-none-any.whl → 0.1.2.dev202501090804__py3-none-any.whl

langfun/core/llms/gemini.py

@@ -0,0 +1,507 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Gemini REST API (Shared by Google GenAI and Vertex AI)."""
+
+import base64
+from typing import Any
+
+import langfun.core as lf
+from langfun.core import modalities as lf_modalities
+from langfun.core.llms import rest
+import pyglove as pg
+
+# Supported modalities.
+
+IMAGE_TYPES = [
+    'image/png',
+    'image/jpeg',
+    'image/webp',
+    'image/heic',
+    'image/heif',
+]
+
+AUDIO_TYPES = [
+    'audio/aac',
+    'audio/flac',
+    'audio/mp3',
+    'audio/m4a',
+    'audio/mpeg',
+    'audio/mpga',
+    'audio/mp4',
+    'audio/opus',
+    'audio/pcm',
+    'audio/wav',
+    'audio/webm',
+]
+
+VIDEO_TYPES = [
+    'video/mov',
+    'video/mpeg',
+    'video/mpegps',
+    'video/mpg',
+    'video/mp4',
+    'video/webm',
+    'video/wmv',
+    'video/x-flv',
+    'video/3gpp',
+    'video/quicktime',
+]
+
+DOCUMENT_TYPES = [
+    'application/pdf',
+    'text/plain',
+    'text/csv',
+    'text/html',
+    'text/xml',
+    'text/x-script.python',
+    'application/json',
+]
+
+TEXT_ONLY = []
+
+ALL_MODALITIES = (
+    IMAGE_TYPES + AUDIO_TYPES + VIDEO_TYPES + DOCUMENT_TYPES
+)
+
+SUPPORTED_MODELS_AND_SETTINGS = {
+    # For automatically rate control and cost estimation, we explicitly register
+    # supported models here. This may be inconvenient for new models, but it
+    # helps us to keep track of the models and their pricing.
+    # Models and RPM are from
+    # https://ai.google.dev/gemini-api/docs/models/gemini?_gl=1*114hbho*_up*MQ..&gclid=Cj0KCQiAst67BhCEARIsAKKdWOljBY5aQdNQ41zOPkXFCwymUfMNFl_7ukm1veAf75ZTD9qWFrFr11IaApL3EALw_wcB
+    # Pricing in US dollars, from https://ai.google.dev/pricing
+    # as of 2025-01-03.
+    # NOTE: Please update google_genai.py, vertexai.py, __init__.py when
+    # adding new models.
+    # !!! PLEASE KEEP MODELS SORTED BY RELEASE DATE !!!
+    'gemini-2.0-flash-thinking-exp-1219': pg.Dict(
+        latest_update='2024-12-19',
+        experimental=True,
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=10,
+        tpm_free=4_000_000,
+        rpm_paid=0,
+        tpm_paid=0,
+        cost_per_1m_input_tokens_up_to_128k=0,
+        cost_per_1m_output_tokens_up_to_128k=0,
+        cost_per_1m_cached_tokens_up_to_128k=0,
+        cost_per_1m_input_tokens_longer_than_128k=0,
+        cost_per_1m_output_tokens_longer_than_128k=0,
+        cost_per_1m_cached_tokens_longer_than_128k=0,
+    ),
+    'gemini-2.0-flash-exp': pg.Dict(
+        latest_update='2024-12-11',
+        experimental=True,
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=10,
+        tpm_free=4_000_000,
+        rpm_paid=0,
+        tpm_paid=0,
+        cost_per_1m_input_tokens_up_to_128k=0,
+        cost_per_1m_output_tokens_up_to_128k=0,
+        cost_per_1m_cached_tokens_up_to_128k=0,
+        cost_per_1m_input_tokens_longer_than_128k=0,
+        cost_per_1m_output_tokens_longer_than_128k=0,
+        cost_per_1m_cached_tokens_longer_than_128k=0,
+    ),
+    'gemini-exp-1206': pg.Dict(
+        latest_update='2024-12-06',
+        experimental=True,
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=10,
+        tpm_free=4_000_000,
+        rpm_paid=0,
+        tpm_paid=0,
+        cost_per_1m_input_tokens_up_to_128k=0,
+        cost_per_1m_output_tokens_up_to_128k=0,
+        cost_per_1m_cached_tokens_up_to_128k=0,
+        cost_per_1m_input_tokens_longer_than_128k=0,
+        cost_per_1m_output_tokens_longer_than_128k=0,
+        cost_per_1m_cached_tokens_longer_than_128k=0,
+    ),
+    'learnlm-1.5-pro-experimental': pg.Dict(
+        latest_update='2024-11-19',
+        experimental=True,
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=10,
+        tpm_free=4_000_000,
+        rpm_paid=0,
+        tpm_paid=0,
+        cost_per_1m_input_tokens_up_to_128k=0,
+        cost_per_1m_output_tokens_up_to_128k=0,
+        cost_per_1m_cached_tokens_up_to_128k=0,
+        cost_per_1m_input_tokens_longer_than_128k=0,
+        cost_per_1m_output_tokens_longer_than_128k=0,
+        cost_per_1m_cached_tokens_longer_than_128k=0,
+    ),
+    'gemini-exp-1114': pg.Dict(
+        latest_update='2024-11-14',
+        experimental=True,
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=10,
+        tpm_free=4_000_000,
+        rpm_paid=0,
+        tpm_paid=0,
+        cost_per_1m_input_tokens_up_to_128k=0,
+        cost_per_1m_output_tokens_up_to_128k=0,
+        cost_per_1m_cached_tokens_up_to_128k=0,
+        cost_per_1m_input_tokens_longer_than_128k=0,
+        cost_per_1m_output_tokens_longer_than_128k=0,
+        cost_per_1m_cached_tokens_longer_than_128k=0,
+    ),
+    'gemini-1.5-flash-latest': pg.Dict(
+        latest_update='2024-09-30',
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=15,
+        tpm_free=1_000_000,
+        rpm_paid=2000,
+        tpm_paid=4_000_000,
+        cost_per_1m_input_tokens_up_to_128k=0.075,
+        cost_per_1m_output_tokens_up_to_128k=0.3,
+        cost_per_1m_cached_tokens_up_to_128k=0.01875,
+        cost_per_1m_input_tokens_longer_than_128k=0.15,
+        cost_per_1m_output_tokens_longer_than_128k=0.6,
+        cost_per_1m_cached_tokens_longer_than_128k=0.0375,
+    ),
+    'gemini-1.5-flash': pg.Dict(
+        latest_update='2024-09-30',
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=15,
+        tpm_free=1_000_000,
+        rpm_paid=2000,
+        tpm_paid=4_000_000,
+        cost_per_1m_input_tokens_up_to_128k=0.075,
+        cost_per_1m_output_tokens_up_to_128k=0.3,
+        cost_per_1m_cached_tokens_up_to_128k=0.01875,
+        cost_per_1m_input_tokens_longer_than_128k=0.15,
+        cost_per_1m_output_tokens_longer_than_128k=0.6,
+        cost_per_1m_cached_tokens_longer_than_128k=0.0375,
+    ),
+    'gemini-1.5-flash-001': pg.Dict(
+        latest_update='2024-09-30',
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=15,
+        tpm_free=1_000_000,
+        rpm_paid=2000,
+        tpm_paid=4_000_000,
+        cost_per_1m_input_tokens_up_to_128k=0.075,
+        cost_per_1m_output_tokens_up_to_128k=0.3,
+        cost_per_1m_cached_tokens_up_to_128k=0.01875,
+        cost_per_1m_input_tokens_longer_than_128k=0.15,
+        cost_per_1m_output_tokens_longer_than_128k=0.6,
+        cost_per_1m_cached_tokens_longer_than_128k=0.0375,
+    ),
+    'gemini-1.5-flash-002': pg.Dict(
+        latest_update='2024-09-30',
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=15,
+        tpm_free=1_000_000,
+        rpm_paid=2000,
+        tpm_paid=4_000_000,
+        cost_per_1m_input_tokens_up_to_128k=0.075,
+        cost_per_1m_output_tokens_up_to_128k=0.3,
+        cost_per_1m_cached_tokens_up_to_128k=0.01875,
+        cost_per_1m_input_tokens_longer_than_128k=0.15,
+        cost_per_1m_output_tokens_longer_than_128k=0.6,
+        cost_per_1m_cached_tokens_longer_than_128k=0.0375,
+    ),
+    'gemini-1.5-flash-8b': pg.Dict(
+        latest_update='2024-10-30',
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=15,
+        tpm_free=1_000_000,
+        rpm_paid=4000,
+        tpm_paid=4_000_000,
+        cost_per_1m_input_tokens_up_to_128k=0.0375,
+        cost_per_1m_output_tokens_up_to_128k=0.15,
+        cost_per_1m_cached_tokens_up_to_128k=0.01,
+        cost_per_1m_input_tokens_longer_than_128k=0.075,
+        cost_per_1m_output_tokens_longer_than_128k=0.3,
+        cost_per_1m_cached_tokens_longer_than_128k=0.02,
+    ),
+    'gemini-1.5-flash-8b-001': pg.Dict(
+        latest_update='2024-10-30',
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=15,
+        tpm_free=1_000_000,
+        rpm_paid=4000,
+        tpm_paid=4_000_000,
+        cost_per_1m_input_tokens_up_to_128k=0.0375,
+        cost_per_1m_output_tokens_up_to_128k=0.15,
+        cost_per_1m_cached_tokens_up_to_128k=0.01,
+        cost_per_1m_input_tokens_longer_than_128k=0.075,
+        cost_per_1m_output_tokens_longer_than_128k=0.3,
+        cost_per_1m_cached_tokens_longer_than_128k=0.02,
+    ),
+    'gemini-1.5-pro-latest': pg.Dict(
+        latest_update='2024-09-30',
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=2,
+        tpm_free=32_000,
+        rpm_paid=1000,
+        tpm_paid=4_000_000,
+        cost_per_1m_input_tokens_up_to_128k=1.25,
+        cost_per_1m_output_tokens_up_to_128k=5.00,
+        cost_per_1m_cached_tokens_up_to_128k=0.3125,
+        cost_per_1m_input_tokens_longer_than_128k=2.5,
+        cost_per_1m_output_tokens_longer_than_128k=10.00,
+        cost_per_1m_cached_tokens_longer_than_128k=0.625,
+    ),
+    'gemini-1.5-pro': pg.Dict(
+        latest_update='2024-09-30',
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=2,
+        tpm_free=32_000,
+        rpm_paid=1000,
+        tpm_paid=4_000_000,
+        cost_per_1m_input_tokens_up_to_128k=1.25,
+        cost_per_1m_output_tokens_up_to_128k=5.00,
+        cost_per_1m_cached_tokens_up_to_128k=0.3125,
+        cost_per_1m_input_tokens_longer_than_128k=2.5,
+        cost_per_1m_output_tokens_longer_than_128k=10.00,
+        cost_per_1m_cached_tokens_longer_than_128k=0.625,
+    ),
+    'gemini-1.5-pro-001': pg.Dict(
+        latest_update='2024-09-30',
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=2,
+        tpm_free=32_000,
+        rpm_paid=1000,
+        tpm_paid=4_000_000,
+        cost_per_1m_input_tokens_up_to_128k=1.25,
+        cost_per_1m_output_tokens_up_to_128k=5.00,
+        cost_per_1m_cached_tokens_up_to_128k=0.3125,
+        cost_per_1m_input_tokens_longer_than_128k=2.5,
+        cost_per_1m_output_tokens_longer_than_128k=10.00,
+        cost_per_1m_cached_tokens_longer_than_128k=0.625,
+    ),
+    'gemini-1.5-pro-002': pg.Dict(
+        latest_update='2024-09-30',
+        in_service=True,
+        supported_modalities=ALL_MODALITIES,
+        rpm_free=2,
+        tpm_free=32_000,
+        rpm_paid=1000,
+        tpm_paid=4_000_000,
+        cost_per_1m_input_tokens_up_to_128k=1.25,
+        cost_per_1m_output_tokens_up_to_128k=5.00,
+        cost_per_1m_cached_tokens_up_to_128k=0.3125,
+        cost_per_1m_input_tokens_longer_than_128k=2.5,
+        cost_per_1m_output_tokens_longer_than_128k=10.00,
+        cost_per_1m_cached_tokens_longer_than_128k=0.625,
+    ),
+    'gemini-1.0-pro': pg.Dict(
+        in_service=False,
+        supported_modalities=TEXT_ONLY,
+        rpm_free=15,
+        tpm_free=32_000,
+        rpm_paid=360,
+        tpm_paid=120_000,
+        cost_per_1m_input_tokens_up_to_128k=0.5,
+        cost_per_1m_output_tokens_up_to_128k=1.5,
+        cost_per_1m_cached_tokens_up_to_128k=0,
+        cost_per_1m_input_tokens_longer_than_128k=0.5,
+        cost_per_1m_output_tokens_longer_than_128k=1.5,
+        cost_per_1m_cached_tokens_longer_than_128k=0,
+    ),
+}
+
+
+@pg.use_init_args(['model'])
+class Gemini(rest.REST):
+  """Language models provided by Google GenAI."""
+
+  model: pg.typing.Annotated[
+      pg.typing.Enum(
+          pg.MISSING_VALUE, list(SUPPORTED_MODELS_AND_SETTINGS.keys())
+      ),
+      'The name of the model to use.',
+  ]
+
+  @property
+  def supported_modalities(self) -> list[str]:
+    """Returns the list of supported modalities."""
+    return SUPPORTED_MODELS_AND_SETTINGS[self.model].supported_modalities
+
+  @property
+  def max_concurrency(self) -> int:
+    """Returns the maximum number of concurrent requests."""
+    return self.rate_to_max_concurrency(
+        requests_per_min=max(
+            SUPPORTED_MODELS_AND_SETTINGS[self.model].rpm_free,
+            SUPPORTED_MODELS_AND_SETTINGS[self.model].rpm_paid
+        ),
+        tokens_per_min=max(
+            SUPPORTED_MODELS_AND_SETTINGS[self.model].tpm_free,
+            SUPPORTED_MODELS_AND_SETTINGS[self.model].tpm_paid,
+        ),
+    )
+
+  def estimate_cost(
+      self,
+      num_input_tokens: int,
+      num_output_tokens: int
+  ) -> float | None:
+    """Estimate the cost based on usage."""
+    entry = SUPPORTED_MODELS_AND_SETTINGS[self.model]
+    if num_input_tokens < 128_000:
+      cost_per_1m_input_tokens = entry.cost_per_1m_input_tokens_up_to_128k
+      cost_per_1m_output_tokens = entry.cost_per_1m_output_tokens_up_to_128k
+    else:
+      cost_per_1m_input_tokens = entry.cost_per_1m_input_tokens_longer_than_128k
+      cost_per_1m_output_tokens = (
+          entry.cost_per_1m_output_tokens_longer_than_128k
+      )
+    return (
+        cost_per_1m_input_tokens * num_input_tokens
+        + cost_per_1m_output_tokens * num_output_tokens
+    ) / 1000_1000
+
+  @property
+  def model_id(self) -> str:
+    """Returns a string to identify the model."""
+    return self.model
+
+  @classmethod
+  def dir(cls):
+    return [k for k, v in SUPPORTED_MODELS_AND_SETTINGS.items() if v.in_service]
+
+  @property
+  def headers(self):
+    return {
+        'Content-Type': 'application/json; charset=utf-8',
+    }
+
+  def request(
+      self, prompt: lf.Message, sampling_options: lf.LMSamplingOptions
+  ) -> dict[str, Any]:
+    request = dict(
+        generationConfig=self._generation_config(prompt, sampling_options)
+    )
+    request['contents'] = [self._content_from_message(prompt)]
+    return request
+
+  def _generation_config(
+      self, prompt: lf.Message, options: lf.LMSamplingOptions
+  ) -> dict[str, Any]:
+    """Returns a dict as generation config for prompt and LMSamplingOptions."""
+    config = dict(
+        temperature=options.temperature,
+        maxOutputTokens=options.max_tokens,
+        candidateCount=options.n,
+        topK=options.top_k,
+        topP=options.top_p,
+        stopSequences=options.stop,
+        seed=options.random_seed,
+        responseLogprobs=options.logprobs,
+        logprobs=options.top_logprobs,
+    )
+
+    if json_schema := prompt.metadata.get('json_schema'):
+      if not isinstance(json_schema, dict):
+        raise ValueError(
+            f'`json_schema` must be a dict, got {json_schema!r}.'
+        )
+      json_schema = pg.to_json(json_schema)
+      config['responseSchema'] = json_schema
+      config['responseMimeType'] = 'application/json'
+      prompt.metadata.formatted_text = (
+          prompt.text
+          + '\n\n [RESPONSE FORMAT (not part of prompt)]\n'
+          + pg.to_json_str(json_schema, json_indent=2)
+      )
+    return config
+
+  def _content_from_message(self, prompt: lf.Message) -> dict[str, Any]:
+    """Gets generation content from langfun message."""
+    parts = []
+    for lf_chunk in prompt.chunk():
+      if isinstance(lf_chunk, str):
+        parts.append({'text': lf_chunk})
+      elif isinstance(lf_chunk, lf_modalities.Mime):
+        try:
+          modalities = lf_chunk.make_compatible(
+              self.supported_modalities + ['text/plain']
+          )
+          if isinstance(modalities, lf_modalities.Mime):
+            modalities = [modalities]
+          for modality in modalities:
+            if modality.is_text:
+              parts.append({'text': modality.to_text()})
+            else:
+              parts.append({
+                  'inlineData': {
+                      'data': base64.b64encode(modality.to_bytes()).decode(),
+                      'mimeType': modality.mime_type,
+                  }
+              })
+        except lf.ModalityError as e:
+          raise lf.ModalityError(f'Unsupported modality: {lf_chunk!r}') from e
+      else:
+        raise lf.ModalityError(f'Unsupported modality: {lf_chunk!r}')
+    return dict(role='user', parts=parts)
+
+  def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
+    messages = [
+        self._message_from_content_parts(candidate['content']['parts'])
+        for candidate in json['candidates']
+    ]
+    usage = json['usageMetadata']
+    input_tokens = usage['promptTokenCount']
+    output_tokens = usage['candidatesTokenCount']
+    return lf.LMSamplingResult(
+        [lf.LMSample(message) for message in messages],
+        usage=lf.LMSamplingUsage(
+            prompt_tokens=input_tokens,
+            completion_tokens=output_tokens,
+            total_tokens=input_tokens + output_tokens,
+            estimated_cost=self.estimate_cost(
+                num_input_tokens=input_tokens,
+                num_output_tokens=output_tokens,
+            ),
+        ),
+    )
+
+  def _message_from_content_parts(
+      self, parts: list[dict[str, Any]]
+  ) -> lf.Message:
+    """Converts Vertex AI's content parts protocol to message."""
+    chunks = []
+    thought_chunks = []
+    for part in parts:
+      if text_part := part.get('text'):
+        if part.get('thought'):
+          thought_chunks.append(text_part)
+        else:
+          chunks.append(text_part)
+      else:
+        raise ValueError(f'Unsupported part: {part}')
+    message = lf.AIMessage.from_chunks(chunks)
+    if thought_chunks:
+      message.set('thought', lf.AIMessage.from_chunks(thought_chunks))
+    return message

langfun/core/llms/gemini_test.py

@@ -0,0 +1,195 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for Gemini API."""
+
+import base64
+from typing import Any
+import unittest
+from unittest import mock
+
+import langfun.core as lf
+from langfun.core import modalities as lf_modalities
+from langfun.core.llms import gemini
+import pyglove as pg
+import requests
+
+
+example_image = (
+    b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x18\x00\x00\x00\x18\x04'
+    b'\x03\x00\x00\x00\x12Y \xcb\x00\x00\x00\x18PLTE\x00\x00'
+    b'\x00fff_chaag_cg_ch^ci_ciC\xedb\x94\x00\x00\x00\x08tRNS'
+    b'\x00\n\x9f*\xd4\xff_\xf4\xe4\x8b\xf3a\x00\x00\x00>IDATx'
+    b'\x01c \x05\x08)"\xd8\xcc\xae!\x06pNz\x88k\x19\\Q\xa8"\x10'
+    b'\xc1\x14\x95\x01%\xc1\n\xa143Ta\xa8"D-\x84\x03QM\x98\xc3'
+    b'\x1a\x1a\x1a@5\x0e\x04\xa0q\x88\x05\x00\x07\xf8\x18\xf9'
+    b'\xdao\xd0|\x00\x00\x00\x00IEND\xaeB`\x82'
+)
+
+
+def mock_requests_post(url: str, json: dict[str, Any], **kwargs):
+  del url, kwargs
+  c = pg.Dict(json['generationConfig'])
+  content = json['contents'][0]['parts'][0]['text']
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str({
+      'candidates': [
+          {
+              'content': {
+                  'role': 'model',
+                  'parts': [
+                      {
+                          'text': (
+                              f'This is a response to {content} with '
+                              f'temperature={c.temperature}, '
+                              f'top_p={c.topP}, '
+                              f'top_k={c.topK}, '
+                              f'max_tokens={c.maxOutputTokens}, '
+                              f'stop={"".join(c.stopSequences)}.'
+                          ),
+                      },
+                      {
+                          'text': 'This is the thought.',
+                          'thought': True,
+                      }
+                  ],
+              },
+          },
+      ],
+      'usageMetadata': {
+          'promptTokenCount': 3,
+          'candidatesTokenCount': 4,
+      }
+  }).encode()
+  return response
+
+
+class GeminiTest(unittest.TestCase):
+  """Tests for Vertex model with REST API."""
+
+  def test_content_from_message_text_only(self):
+    text = 'This is a beautiful day'
+    model = gemini.Gemini('gemini-1.5-pro', api_endpoint='')
+    chunks = model._content_from_message(lf.UserMessage(text))
+    self.assertEqual(chunks, {'role': 'user', 'parts': [{'text': text}]})
+
+  def test_content_from_message_mm(self):
+    image = lf_modalities.Image.from_bytes(example_image)
+    message = lf.UserMessage(
+        'This is an <<[[image]]>>, what is it?', image=image
+    )
+
+    # Non-multimodal model.
+    with self.assertRaisesRegex(lf.ModalityError, 'Unsupported modality'):
+      gemini.Gemini(
+          'gemini-1.0-pro', api_endpoint=''
+      )._content_from_message(message)
+
+    model = gemini.Gemini('gemini-1.5-pro', api_endpoint='')
+    content = model._content_from_message(message)
+    self.assertEqual(
+        content,
+        {
+            'role': 'user',
+            'parts': [
+                {'text': 'This is an'},
+                {
+                    'inlineData': {
+                        'data': base64.b64encode(example_image).decode(),
+                        'mimeType': 'image/png',
+                    }
+                },
+                {'text': ', what is it?'},
+            ],
+        },
+    )
+
+  def test_generation_config(self):
+    model = gemini.Gemini('gemini-1.5-pro', api_endpoint='')
+    json_schema = {
+        'type': 'object',
+        'properties': {
+            'name': {'type': 'string'},
+        },
+        'required': ['name'],
+        'title': 'Person',
+    }
+    actual = model._generation_config(
+        lf.UserMessage('hi', json_schema=json_schema),
+        lf.LMSamplingOptions(
+            temperature=2.0,
+            top_p=1.0,
+            top_k=20,
+            max_tokens=1024,
+            stop=['\n'],
+        ),
+    )
+    self.assertEqual(
+        actual,
+        dict(
+            candidateCount=1,
+            temperature=2.0,
+            topP=1.0,
+            topK=20,
+            maxOutputTokens=1024,
+            stopSequences=['\n'],
+            responseLogprobs=False,
+            logprobs=None,
+            seed=None,
+            responseMimeType='application/json',
+            responseSchema={
+                'type': 'object',
+                'properties': {
+                    'name': {'type': 'string'}
+                },
+                'required': ['name'],
+                'title': 'Person',
+            }
+        ),
+    )
+    with self.assertRaisesRegex(
+        ValueError, '`json_schema` must be a dict, got'
+    ):
+      model._generation_config(
+          lf.UserMessage('hi', json_schema='not a dict'),
+          lf.LMSamplingOptions(),
+      )
+
+  def test_call_model(self):
+    with mock.patch('requests.Session.post') as mock_generate:
+      mock_generate.side_effect = mock_requests_post
+
+      lm = gemini.Gemini('gemini-1.5-pro', api_endpoint='')
+      r = lm(
+          'hello',
+          temperature=2.0,
+          top_p=1.0,
+          top_k=20,
+          max_tokens=1024,
+          stop='\n',
+      )
+      self.assertEqual(
+          r.text,
+          (
+              'This is a response to hello with temperature=2.0, '
+              'top_p=1.0, top_k=20, max_tokens=1024, stop=\n.'
+          ),
+      )
+      self.assertEqual(r.metadata.thought, 'This is the thought.')
+      self.assertEqual(r.metadata.usage.prompt_tokens, 3)
+      self.assertEqual(r.metadata.usage.completion_tokens, 4)
+
+
+if __name__ == '__main__':
+  unittest.main()