langfun 0.0.2.dev20240429__py3-none-any.whl → 0.1.2.dev202501140804__py3-none-any.whl

langfun/core/llms/fake.py

@@ -21,9 +21,13 @@ import langfun.core as lf
 class Fake(lf.LanguageModel):
   """The base class for all fake language models."""
 
-  def _score(self, prompt: lf.Message, completions: list[lf.Message]):
+  def _score(self, prompt: lf.Message| list[lf.Message],
+             completions: list[lf.Message]):
     return [lf.LMScoringResult(score=-i * 1.0) for i in range(len(completions))]
 
+  def _tokenize(self, prompt: lf.Message) -> list[tuple[str | bytes, int]]:
+    return [(w, i) for i, w in enumerate(prompt.text.split(' '))]
+
   def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
     results = []
     for prompt in prompts:
@@ -57,12 +61,12 @@ class StaticResponse(Fake):
   """Language model that always gives the same canned response."""
 
   response: Annotated[
-      str,
+      str | lf.Message,
       'A canned response that will be returned regardless of the prompt.'
   ]
 
   def _response_from(self, prompt: lf.Message) -> lf.Message:
-    return lf.AIMessage(self.response)
+    return lf.AIMessage.from_value(self.response)
 
 
 @lf.use_init_args(['mapping'])
@@ -70,12 +74,12 @@ class StaticMapping(Fake):
   """A static mapping from prompt to response."""
 
   mapping: Annotated[
-      dict[str, str],
+      dict[str, str | lf.Message],
       'A mapping from prompt to response.'
   ]
 
   def _response_from(self, prompt: lf.Message) -> lf.Message:
-    return lf.AIMessage(self.mapping[prompt])
+    return lf.AIMessage.from_value(self.mapping[prompt])
 
 
 @lf.use_init_args(['sequence'])
@@ -83,7 +87,7 @@ class StaticSequence(Fake):
   """A static sequence of responses to use."""
 
   sequence: Annotated[
-      list[str],
+      list[str | lf.Message],
       'A sequence of strings as the response.'
   ]
 
@@ -92,6 +96,6 @@ class StaticSequence(Fake):
     self._pos = 0
 
   def _response_from(self, prompt: lf.Message) -> lf.Message:
-    r = lf.AIMessage(self.sequence[self._pos])
+    r = lf.AIMessage.from_value(self.sequence[self._pos])
     self._pos += 1
     return r

langfun/core/llms/fake_test.py

@@ -34,6 +34,7 @@ class EchoTest(unittest.TestCase):
                             'hi',
                             score=1.0,
                             logprobs=None,
+                            is_cached=False,
                             usage=lf.LMSamplingUsage(2, 2, 4),
                             tags=[lf.Message.TAG_LM_RESPONSE],
                         ),
@@ -62,6 +63,13 @@ class EchoTest(unittest.TestCase):
         [lf.LMScoringResult(0.0), lf.LMScoringResult(-1.0)],
     )
 
+  def test_tokenize(self):
+    lm = fakelm.Echo()
+    self.assertEqual(
+        lm.tokenize('hi'),
+        [('hi', 0)]
+    )
+
 
 class StaticResponseTest(unittest.TestCase):
 
@@ -78,6 +86,7 @@ class StaticResponseTest(unittest.TestCase):
                             canned_response,
                             score=1.0,
                             logprobs=None,
+                            is_cached=False,
                             usage=lf.LMSamplingUsage(2, 38, 40),
                             tags=[lf.Message.TAG_LM_RESPONSE],
                         ),
@@ -99,6 +108,7 @@ class StaticResponseTest(unittest.TestCase):
                             canned_response,
                             score=1.0,
                             logprobs=None,
+                            is_cached=False,
                             usage=lf.LMSamplingUsage(15, 38, 53),
                             tags=[lf.Message.TAG_LM_RESPONSE],
                         ),
@@ -143,6 +153,7 @@ class StaticMappingTest(unittest.TestCase):
                             'Hello',
                             score=1.0,
                             logprobs=None,
+                            is_cached=False,
                             usage=lf.LMSamplingUsage(2, 5, 7),
                             tags=[lf.Message.TAG_LM_RESPONSE],
                         ),
@@ -159,6 +170,7 @@ class StaticMappingTest(unittest.TestCase):
                             'I am fine, how about you?',
                             score=1.0,
                             logprobs=None,
+                            is_cached=False,
                             usage=lf.LMSamplingUsage(12, 25, 37),
                             tags=[lf.Message.TAG_LM_RESPONSE],
                         ),
@@ -192,6 +204,7 @@ class StaticSequenceTest(unittest.TestCase):
                             'Hello',
                             score=1.0,
                             logprobs=None,
+                            is_cached=False,
                             usage=lf.LMSamplingUsage(2, 5, 7),
                             tags=[lf.Message.TAG_LM_RESPONSE],
                         ),
@@ -208,6 +221,7 @@ class StaticSequenceTest(unittest.TestCase):
                             'I am fine, how about you?',
                             score=1.0,
                             logprobs=None,
+                            is_cached=False,
                             usage=lf.LMSamplingUsage(12, 25, 37),
                             tags=[lf.Message.TAG_LM_RESPONSE],
                         ),

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