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

langfun/core/llms/google_genai_test.py

@@ -11,223 +11,28 @@
 # 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 models."""
+"""Tests for Google GenAI models."""
 
 import os
 import unittest
-from unittest import mock
-
-from google import generativeai as genai
-import langfun.core as lf
-from langfun.core import modalities as lf_modalities
 from langfun.core.llms import google_genai
-import pyglove as pg
-
-
-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_get_model(model_name, *args, **kwargs):
-  del args, kwargs
-  if 'gemini' in model_name:
-    method = 'generateContent'
-  elif 'chat' in model_name:
-    method = 'generateMessage'
-  else:
-    method = 'generateText'
-  return pg.Dict(supported_generation_methods=[method])
-
-
-def mock_generate_text(*, model, prompt, **kwargs):
-  return pg.Dict(
-      candidates=[pg.Dict(output=f'{prompt} to {model} with {kwargs}')]
-  )
-
-
-def mock_chat(*, model, messages, **kwargs):
-  return pg.Dict(
-      candidates=[pg.Dict(content=f'{messages} to {model} with {kwargs}')]
-  )
-
-
-def mock_generate_content(content, generation_config, **kwargs):
-  del kwargs
-  c = generation_config
-  return genai.types.GenerateContentResponse(
-      done=True,
-      iterator=None,
-      chunks=[],
-      result=pg.Dict(
-          prompt_feedback=pg.Dict(block_reason=None),
-          candidates=[
-              pg.Dict(
-                  content=pg.Dict(
-                      parts=[
-                          pg.Dict(
-                              text=(
-                                  f'This is a response to {content[0]} with '
-                                  f'n={c.candidate_count}, '
-                                  f'temperature={c.temperature}, '
-                                  f'top_p={c.top_p}, '
-                                  f'top_k={c.top_k}, '
-                                  f'max_tokens={c.max_output_tokens}, '
-                                  f'stop={c.stop_sequences}.'
-                              )
-                          )
-                      ]
-                  ),
-              ),
-          ],
-      ),
-  )
 
 
 class GenAITest(unittest.TestCase):
-  """Tests for Google GenAI model."""
-
-  def test_content_from_message_text_only(self):
-    text = 'This is a beautiful day'
-    model = google_genai.GeminiPro()
-    chunks = model._content_from_message(lf.UserMessage(text))
-    self.assertEqual(chunks, [text])
-
-  def test_content_from_message_mm(self):
-    message = lf.UserMessage(
-        'This is an <<[[image]]>>, what is it?',
-        image=lf_modalities.Image.from_bytes(example_image),
-    )
+  """Tests for GenAI model."""
 
-    # Non-multimodal model.
-    with self.assertRaisesRegex(lf.ModalityError, 'Unsupported modality'):
-      google_genai.GeminiPro()._content_from_message(message)
-
-    model = google_genai.GeminiProVision()
-    chunks = model._content_from_message(message)
-    self.maxDiff = None
-    self.assertEqual(
-        chunks,
-        [
-            'This is an',
-            genai.types.BlobDict(mime_type='image/png', data=example_image),
-            ', what is it?',
-        ],
-    )
-
-  def test_response_to_result_text_only(self):
-    response = genai.types.GenerateContentResponse(
-        done=True,
-        iterator=None,
-        chunks=[],
-        result=pg.Dict(
-            prompt_feedback=pg.Dict(block_reason=None),
-            candidates=[
-                pg.Dict(
-                    content=pg.Dict(
-                        parts=[pg.Dict(text='This is response 1.')]
-                    ),
-                ),
-                pg.Dict(
-                    content=pg.Dict(parts=[pg.Dict(text='This is response 2.')])
-                ),
-            ],
-        ),
-    )
-    model = google_genai.GeminiProVision()
-    result = model._response_to_result(response)
-    self.assertEqual(
-        result,
-        lf.LMSamplingResult([
-            lf.LMSample(lf.AIMessage('This is response 1.'), score=0.0),
-            lf.LMSample(lf.AIMessage('This is response 2.'), score=0.0),
-        ]),
-    )
-
-  def test_model_hub(self):
-    orig_get_model = genai.get_model
-    genai.get_model = mock_get_model
-
-    model = google_genai._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro')
-    self.assertIsNotNone(model)
-    self.assertIs(google_genai._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro'), model)
-
-    genai.get_model = orig_get_model
-
-  def test_api_key_check(self):
+  def test_basics(self):
     with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
-      _ = google_genai.GeminiPro()._api_initialized
+      _ = google_genai.GeminiPro1_5().api_endpoint
+
+    self.assertIsNotNone(google_genai.GeminiPro1_5(api_key='abc').api_endpoint)
 
-    self.assertTrue(google_genai.GeminiPro(api_key='abc')._api_initialized)
     os.environ['GOOGLE_API_KEY'] = 'abc'
-    self.assertTrue(google_genai.GeminiPro()._api_initialized)
+    lm = google_genai.GeminiPro1_5()
+    self.assertIsNotNone(lm.api_endpoint)
+    self.assertTrue(lm.model_id.startswith('GenAI('))
     del os.environ['GOOGLE_API_KEY']
 
-  def test_call(self):
-    with mock.patch(
-        'google.generativeai.GenerativeModel.generate_content',
-    ) as mock_generate:
-      orig_get_model = genai.get_model
-      genai.get_model = mock_get_model
-      mock_generate.side_effect = mock_generate_content
-
-      lm = google_genai.GeminiPro(api_key='test_key')
-      self.maxDiff = None
-      self.assertEqual(
-          lm('hello', temperature=2.0, top_k=20, max_tokens=1024).text,
-          (
-              'This is a response to hello with n=1, temperature=2.0, '
-              'top_p=None, top_k=20, max_tokens=1024, stop=None.'
-          ),
-      )
-      genai.get_model = orig_get_model
-
-  def test_call_with_legacy_completion_model(self):
-    orig_get_model = genai.get_model
-    genai.get_model = mock_get_model
-    orig_generate_text = getattr(genai, 'generate_text', None)
-    if orig_generate_text is not None:
-      genai.generate_text = mock_generate_text
-
-      lm = google_genai.Palm2(api_key='test_key')
-      self.maxDiff = None
-      self.assertEqual(
-          lm('hello', temperature=2.0, top_k=20).text,
-          (
-              "hello to models/text-bison-001 with {'temperature': 2.0, "
-              "'top_k': 20, 'top_p': None, 'candidate_count': 1, "
-              "'max_output_tokens': None, 'stop_sequences': None}"
-          ),
-      )
-      genai.generate_text = orig_generate_text
-    genai.get_model = orig_get_model
-
-  def test_call_with_legacy_chat_model(self):
-    orig_get_model = genai.get_model
-    genai.get_model = mock_get_model
-    orig_chat = getattr(genai, 'chat', None)
-    if orig_chat is not None:
-      genai.chat = mock_chat
-
-      lm = google_genai.Palm2_IT(api_key='test_key')
-      self.maxDiff = None
-      self.assertEqual(
-          lm('hello', temperature=2.0, top_k=20).text,
-          (
-              "hello to models/chat-bison-001 with {'temperature': 2.0, "
-              "'top_k': 20, 'top_p': None, 'candidate_count': 1}"
-          ),
-      )
-      genai.chat = orig_chat
-    genai.get_model = orig_get_model
-
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/llms/vertexai.py

@@ -13,14 +13,12 @@
 # limitations under the License.
 """Vertex AI generative models."""
 
-import base64
 import functools
 import os
 from typing import Annotated, Any
 
 import langfun.core as lf
-from langfun.core import modalities as lf_modalities
-from langfun.core.llms import rest
+from langfun.core.llms import gemini
 import pyglove as pg
 
 try:
@@ -38,114 +36,11 @@ except ImportError:
   Credentials = Any
 
 
-# https://cloud.google.com/vertex-ai/generative-ai/pricing
-# describes that the average number of characters per token is about 4.
-AVGERAGE_CHARS_PER_TOKEN = 4
-
-
-# Price in US dollars,
-# from https://cloud.google.com/vertex-ai/generative-ai/pricing
-# as of 2024-10-10.
-SUPPORTED_MODELS_AND_SETTINGS = {
-    'gemini-1.5-pro-001': pg.Dict(
-        rpm=100,
-        cost_per_1k_input_chars=0.0003125,
-        cost_per_1k_output_chars=0.00125,
-    ),
-    'gemini-1.5-pro-002': pg.Dict(
-        rpm=100,
-        cost_per_1k_input_chars=0.0003125,
-        cost_per_1k_output_chars=0.00125,
-    ),
-    'gemini-1.5-flash-002': pg.Dict(
-        rpm=500,
-        cost_per_1k_input_chars=0.00001875,
-        cost_per_1k_output_chars=0.000075,
-    ),
-    'gemini-1.5-flash-001': pg.Dict(
-        rpm=500,
-        cost_per_1k_input_chars=0.00001875,
-        cost_per_1k_output_chars=0.000075,
-    ),
-    'gemini-1.5-pro': pg.Dict(
-        rpm=100,
-        cost_per_1k_input_chars=0.0003125,
-        cost_per_1k_output_chars=0.00125,
-    ),
-    'gemini-1.5-flash': pg.Dict(
-        rpm=500,
-        cost_per_1k_input_chars=0.00001875,
-        cost_per_1k_output_chars=0.000075,
-    ),
-    'gemini-1.5-pro-preview-0514': pg.Dict(
-        rpm=50,
-        cost_per_1k_input_chars=0.0003125,
-        cost_per_1k_output_chars=0.00125,
-    ),
-    'gemini-1.5-pro-preview-0409': pg.Dict(
-        rpm=50,
-        cost_per_1k_input_chars=0.0003125,
-        cost_per_1k_output_chars=0.00125,
-    ),
-    'gemini-1.5-flash-preview-0514': pg.Dict(
-        rpm=200,
-        cost_per_1k_input_chars=0.00001875,
-        cost_per_1k_output_chars=0.000075,
-    ),
-    'gemini-1.0-pro': pg.Dict(
-        rpm=300,
-        cost_per_1k_input_chars=0.000125,
-        cost_per_1k_output_chars=0.000375,
-    ),
-    'gemini-1.0-pro-vision': pg.Dict(
-        rpm=100,
-        cost_per_1k_input_chars=0.000125,
-        cost_per_1k_output_chars=0.000375,
-    ),
-    # TODO(sharatsharat): Update costs when published
-    'gemini-exp-1206': pg.Dict(
-        rpm=20,
-        cost_per_1k_input_chars=0.000,
-        cost_per_1k_output_chars=0.000,
-    ),
-    # TODO(sharatsharat): Update costs when published
-    'gemini-2.0-flash-exp': pg.Dict(
-        rpm=10,
-        cost_per_1k_input_chars=0.000,
-        cost_per_1k_output_chars=0.000,
-    ),
-    # TODO(yifenglu): Update costs when published
-    'gemini-2.0-flash-thinking-exp-1219': pg.Dict(
-        rpm=10,
-        cost_per_1k_input_chars=0.000,
-        cost_per_1k_output_chars=0.000,
-    ),
-    # TODO(chengrun): Set a more appropriate rpm for endpoint.
-    'vertexai-endpoint': pg.Dict(
-        rpm=20,
-        cost_per_1k_input_chars=0.0000125,
-        cost_per_1k_output_chars=0.0000375,
-    ),
-}
-
-
 @lf.use_init_args(['model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
-class VertexAI(rest.REST):
+class VertexAI(gemini.Gemini):
   """Language model served on VertexAI with REST API."""
 
-  model: pg.typing.Annotated[
-      pg.typing.Enum(
-          pg.MISSING_VALUE, list(SUPPORTED_MODELS_AND_SETTINGS.keys())
-      ),
-      (
-          'Vertex AI model name with REST API support. See '
-          'https://cloud.google.com/vertex-ai/generative-ai/docs/'
-          'model-reference/inference#supported-models'
-          ' for details.'
-      ),
-  ]
-
   project: Annotated[
       str | None,
       (
@@ -170,11 +65,6 @@ class VertexAI(rest.REST):
       ),
   ] = None
 
-  supported_modalities: Annotated[
-      list[str],
-      'A list of MIME types for supported modalities'
-  ] = []
-
   def _on_bound(self):
     super()._on_bound()
     if google_auth is None:
@@ -209,31 +99,9 @@ class VertexAI(rest.REST):
     self._credentials = credentials
 
   @property
-  def max_concurrency(self) -> int:
-    """Returns the maximum number of concurrent requests."""
-    return self.rate_to_max_concurrency(
-        requests_per_min=SUPPORTED_MODELS_AND_SETTINGS[self.model].rpm,
-        tokens_per_min=0,
-    )
-
-  def estimate_cost(
-      self,
-      num_input_tokens: int,
-      num_output_tokens: int
-  ) -> float | None:
-    """Estimate the cost based on usage."""
-    cost_per_1k_input_chars = SUPPORTED_MODELS_AND_SETTINGS[self.model].get(
-        'cost_per_1k_input_chars', None
-    )
-    cost_per_1k_output_chars = SUPPORTED_MODELS_AND_SETTINGS[self.model].get(
-        'cost_per_1k_output_chars', None
-    )
-    if cost_per_1k_output_chars is None or cost_per_1k_input_chars is None:
-      return None
-    return (
-        cost_per_1k_input_chars * num_input_tokens
-        + cost_per_1k_output_chars * num_output_tokens
-    ) * AVGERAGE_CHARS_PER_TOKEN / 1000
+  def model_id(self) -> str:
+    """Returns a string to identify the model."""
+    return f'VertexAI({self.model})'
 
   @functools.cached_property
   def _session(self):
@@ -244,12 +112,6 @@ class VertexAI(rest.REST):
     s.headers.update(self.headers or {})
     return s
 
-  @property
-  def headers(self):
-    return {
-        'Content-Type': 'application/json; charset=utf-8',
-    }
-
   @property
   def api_endpoint(self) -> str:
     return (
@@ -258,263 +120,69 @@ class VertexAI(rest.REST):
         f'models/{self.model}:generateContent'
     )
 
-  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,
-            ),
-        ),
-    )
+class VertexAIGeminiFlash2_0ThinkingExp_20241219(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini Flash 2.0 Thinking model launched on 12/19/2024."""
 
-  def _message_from_content_parts(
-      self, parts: list[dict[str, Any]]
-  ) -> lf.Message:
-    """Converts Vertex AI's content parts protocol to message."""
-    chunks = []
-    for part in parts:
-      if text_part := part.get('text'):
-        chunks.append(text_part)
-      else:
-        raise ValueError(f'Unsupported part: {part}')
-    return lf.AIMessage.from_chunks(chunks)
-
-
-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',
-]
-
-
-class VertexAIGemini2_0(VertexAI):  # pylint: disable=invalid-name
-  """Vertex AI Gemini 2.0 model."""
-
-  supported_modalities: pg.typing.List(str).freeze(  # pytype: disable=invalid-annotation
-      DOCUMENT_TYPES + IMAGE_TYPES + AUDIO_TYPES + VIDEO_TYPES
-  )
-
-
-class VertexAIGeminiFlash2_0Exp(VertexAIGemini2_0):  # pylint: disable=invalid-name
+  api_version = 'v1alpha'
+  model = 'gemini-2.0-flash-thinking-exp-1219'
+
+
+class VertexAIGeminiFlash2_0Exp(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 2.0 Flash model."""
 
   model = 'gemini-2.0-flash-exp'
 
 
-class VertexAIGeminiFlash2_0ThinkingExp(VertexAIGemini2_0):  # pylint: disable=invalid-name
-  """Vertex AI Gemini 2.0 Flash model."""
+class VertexAIGeminiExp_20241206(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini Experimental model launched on 12/06/2024."""
 
-  model = 'gemini-2.0-flash-thinking-exp-1219'
+  model = 'gemini-exp-1206'
 
 
-class VertexAIGemini1_5(VertexAI):  # pylint: disable=invalid-name
-  """Vertex AI Gemini 1.5 model."""
+class VertexAIGeminiExp_20241114(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini Experimental model launched on 11/14/2024."""
 
-  supported_modalities: pg.typing.List(str).freeze(  # pytype: disable=invalid-annotation
-      DOCUMENT_TYPES + IMAGE_TYPES + AUDIO_TYPES + VIDEO_TYPES
-  )
+  model = 'gemini-exp-1114'
 
 
-class VertexAIGeminiPro1_5(VertexAIGemini1_5):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1_5(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
-  model = 'gemini-1.5-pro'
+  model = 'gemini-1.5-pro-latest'
 
 
-class VertexAIGeminiPro1_5_002(VertexAIGemini1_5):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1_5_002(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-002'
 
 
-class VertexAIGeminiPro1_5_001(VertexAIGemini1_5):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1_5_001(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-001'
 
 
-class VertexAIGeminiPro1_5_0514(VertexAIGemini1_5):  # pylint: disable=invalid-name
-  """Vertex AI Gemini 1.5 Pro preview model."""
-
-  model = 'gemini-1.5-pro-preview-0514'
-
-
-class VertexAIGeminiPro1_5_0409(VertexAIGemini1_5):  # pylint: disable=invalid-name
-  """Vertex AI Gemini 1.5 Pro preview model."""
-
-  model = 'gemini-1.5-pro-preview-0409'
-
-
-class VertexAIGeminiFlash1_5(VertexAIGemini1_5):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash1_5(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash'
 
 
-class VertexAIGeminiFlash1_5_002(VertexAIGemini1_5):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash1_5_002(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-002'
 
 
-class VertexAIGeminiFlash1_5_001(VertexAIGemini1_5):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash1_5_001(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-001'
 
 
-class VertexAIGeminiFlash1_5_0514(VertexAIGemini1_5):  # pylint: disable=invalid-name
-  """Vertex AI Gemini 1.5 Flash preview model."""
-
-  model = 'gemini-1.5-flash-preview-0514'
-
-
 class VertexAIGeminiPro1(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.0 Pro model."""
 
   model = 'gemini-1.0-pro'
-
-
-class VertexAIGeminiPro1Vision(VertexAI):  # pylint: disable=invalid-name
-  """Vertex AI Gemini 1.0 Pro Vision model."""
-
-  model = 'gemini-1.0-pro-vision'
-  supported_modalities: pg.typing.List(str).freeze(  # pytype: disable=invalid-annotation
-      IMAGE_TYPES + VIDEO_TYPES
-  )
-
-
-class VertexAIEndpoint(VertexAI):  # pylint: disable=invalid-name
-  """Vertex AI Endpoint model."""
-
-  model = 'vertexai-endpoint'
-
-  endpoint: Annotated[str, 'Vertex AI Endpoint ID.']
-
-  @property
-  def api_endpoint(self) -> str:
-    return (
-        f'https://{self.location}-aiplatform.googleapis.com/v1/projects/'
-        f'{self.project}/locations/{self.location}/'
-        f'endpoints/{self.endpoint}:generateContent'
-    )