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

langfun/core/llms/google_genai.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Langfun Authors
+# 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.
@@ -11,57 +11,21 @@
 # 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 models exposed through Google Generative AI APIs."""
+"""Language models from Google GenAI."""
 
-import abc
-import functools
 import os
-from typing import Annotated, Any, Literal
+from typing import Annotated, Literal
 
 import langfun.core as lf
-from langfun.core import modalities as lf_modalities
-from langfun.core.llms import vertexai
+from langfun.core.llms import gemini
 import pyglove as pg
 
 
-try:
-  import google.generativeai as genai   # pylint: disable=g-import-not-at-top
-  BlobDict = genai.types.BlobDict
-  GenerativeModel = genai.GenerativeModel
-  Completion = getattr(genai.types, 'Completion', Any)
-  ChatResponse = getattr(genai.types, 'ChatResponse', Any)
-  GenerateContentResponse = getattr(genai.types, 'GenerateContentResponse', Any)
-  GenerationConfig = genai.GenerationConfig
-except ImportError:
-  genai = None
-  BlobDict = Any
-  GenerativeModel = Any
-  Completion = Any
-  ChatResponse = Any
-  GenerationConfig = Any
-  GenerateContentResponse = Any
-
-
 @lf.use_init_args(['model'])
-class GenAI(lf.LanguageModel):
+@pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
+class GenAI(gemini.Gemini):
   """Language models provided by Google GenAI."""
 
-  model: Annotated[
-      Literal[
-          'gemini-2.0-flash-thinking-exp-1219',
-          'gemini-2.0-flash-exp',
-          'gemini-exp-1206',
-          'gemini-exp-1114',
-          'gemini-1.5-pro-latest',
-          'gemini-1.5-flash-latest',
-          'gemini-pro',
-          'gemini-pro-vision',
-          'text-bison-001',
-          'chat-bison-001',
-      ],
-      'Model name.',
-  ]
-
   api_key: Annotated[
       str | None,
       (
@@ -71,26 +35,18 @@ class GenAI(lf.LanguageModel):
       ),
   ] = None
 
-  supported_modalities: Annotated[
-      list[str],
-      'A list of MIME types for supported modalities'
-  ] = []
+  api_version: Annotated[
+      Literal['v1beta', 'v1alpha'],
+      'The API version to use.'
+  ] = 'v1beta'
 
-  # Set the default max concurrency to 8 workers.
-  max_concurrency = 8
-
-  def _on_bound(self):
-    super()._on_bound()
-    if genai is None:
-      raise RuntimeError(
-          'Please install "langfun[llm-google-genai]" to use '
-          'Google Generative AI models.'
-      )
-    self.__dict__.pop('_api_initialized', None)
+  @property
+  def model_id(self) -> str:
+    """Returns a string to identify the model."""
+    return f'GenAI({self.model})'
 
-  @functools.cached_property
-  def _api_initialized(self):
-    assert genai is not None
+  @property
+  def api_endpoint(self) -> str:
     api_key = self.api_key or os.environ.get('GOOGLE_API_KEY', None)
     if not api_key:
       raise ValueError(
@@ -100,306 +56,76 @@ class GenAI(lf.LanguageModel):
           'https://cloud.google.com/api-keys/docs/create-manage-api-keys '
           'for more details.'
       )
-    genai.configure(api_key=api_key)
-    return True
-
-  @classmethod
-  def dir(cls) -> list[str]:
-    """Lists generative models."""
-    assert genai is not None
-    return [
-        m.name.lstrip('models/')
-        for m in genai.list_models()
-        if (
-            'generateContent' in m.supported_generation_methods
-            or 'generateText' in m.supported_generation_methods
-            or 'generateMessage' in m.supported_generation_methods
-        )
-    ]
-
-  @property
-  def model_id(self) -> str:
-    """Returns a string to identify the model."""
-    return self.model
-
-  @property
-  def resource_id(self) -> str:
-    """Returns a string to identify the resource for rate control."""
-    return self.model_id
-
-  def _generation_config(self, options: lf.LMSamplingOptions) -> dict[str, Any]:
-    """Creates generation config from langfun sampling options."""
-    return GenerationConfig(
-        candidate_count=options.n,
-        temperature=options.temperature,
-        top_p=options.top_p,
-        top_k=options.top_k,
-        max_output_tokens=options.max_tokens,
-        stop_sequences=options.stop,
+    return (
+        f'https://generativelanguage.googleapis.com/{self.api_version}'
+        f'/models/{self.model}:generateContent?'
+        f'key={api_key}'
     )
 
-  def _content_from_message(
-      self, prompt: lf.Message
-  ) -> list[str | BlobDict]:
-    """Gets Evergreen formatted content from langfun message."""
-    formatted = lf.UserMessage(prompt.text)
-    formatted.source = prompt
-
-    chunks = []
-    for lf_chunk in formatted.chunk():
-      if isinstance(lf_chunk, str):
-        chunks.append(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:
-              chunk = modality.to_text()
-            else:
-              chunk = BlobDict(
-                  data=modality.to_bytes(),
-                  mime_type=modality.mime_type
-              )
-            chunks.append(chunk)
-        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 chunks
-
-  def _response_to_result(
-      self, response: GenerateContentResponse | pg.Dict
-  ) -> lf.LMSamplingResult:
-    """Parses generative response into message."""
-    samples = []
-    for candidate in response.candidates:
-      chunks = []
-      for part in candidate.content.parts:
-        # TODO(daiyip): support multi-modal parts when they are available via
-        # Gemini API.
-        if hasattr(part, 'text'):
-          chunks.append(part.text)
-      samples.append(lf.LMSample(lf.AIMessage.from_chunks(chunks), score=0.0))
-    return lf.LMSamplingResult(samples)
-
-  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    assert self._api_initialized, 'Vertex AI API is not initialized.'
-    return self._parallel_execute_with_currency_control(
-        self._sample_single,
-        prompts,
-    )
 
-  def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
-    """Samples a single prompt."""
-    model = _GOOGLE_GENAI_MODEL_HUB.get(self.model)
-    input_content = self._content_from_message(prompt)
-    response = model.generate_content(
-        input_content,
-        generation_config=self._generation_config(self.sampling_options),
-    )
-    return self._response_to_result(response)
-
-
-class _LegacyGenerativeModel(pg.Object):
-  """Base for legacy GenAI generative model."""
-
-  model: str
-
-  def generate_content(
-      self,
-      input_content: list[str | BlobDict],
-      generation_config: GenerationConfig,
-  ) -> pg.Dict:
-    """Generate content."""
-    segments = []
-    for s in input_content:
-      if not isinstance(s, str):
-        raise ValueError(f'Unsupported modality: {s!r}')
-      segments.append(s)
-    return self.generate(' '.join(segments), generation_config)
-
-  @abc.abstractmethod
-  def generate(
-      self, prompt: str, generation_config: GenerationConfig) -> pg.Dict:
-    """Generate response based on prompt."""
-
-
-class _LegacyCompletionModel(_LegacyGenerativeModel):
-  """Legacy GenAI completion model."""
-
-  def generate(
-      self, prompt: str, generation_config: GenerationConfig
-  ) -> pg.Dict:
-    assert genai is not None
-    completion: Completion = genai.generate_text(
-        model=f'models/{self.model}',
-        prompt=prompt,
-        temperature=generation_config.temperature,
-        top_k=generation_config.top_k,
-        top_p=generation_config.top_p,
-        candidate_count=generation_config.candidate_count,
-        max_output_tokens=generation_config.max_output_tokens,
-        stop_sequences=generation_config.stop_sequences,
-    )
-    return pg.Dict(
-        candidates=[
-            pg.Dict(content=pg.Dict(parts=[pg.Dict(text=c['output'])]))
-            for c in completion.candidates
-        ]
-    )
-
-
-class _LegacyChatModel(_LegacyGenerativeModel):
-  """Legacy GenAI chat model."""
-
-  def generate(
-      self, prompt: str, generation_config: GenerationConfig
-  ) -> pg.Dict:
-    assert genai is not None
-    response: ChatResponse = genai.chat(
-        model=f'models/{self.model}',
-        messages=prompt,
-        temperature=generation_config.temperature,
-        top_k=generation_config.top_k,
-        top_p=generation_config.top_p,
-        candidate_count=generation_config.candidate_count,
-    )
-    return pg.Dict(
-        candidates=[
-            pg.Dict(content=pg.Dict(parts=[pg.Dict(text=c['content'])]))
-            for c in response.candidates
-        ]
-    )
-
-
-class _ModelHub:
-  """Google Generative AI model hub."""
-
-  def __init__(self):
-    self._model_cache = {}
-
-  def get(
-      self, model_name: str
-  ) -> GenerativeModel | _LegacyGenerativeModel:
-    """Gets a generative model by model id."""
-    assert genai is not None
-    model = self._model_cache.get(model_name, None)
-    if model is None:
-      model_info = genai.get_model(f'models/{model_name}')
-      if 'generateContent' in model_info.supported_generation_methods:
-        model = genai.GenerativeModel(model_name)
-      elif 'generateText' in model_info.supported_generation_methods:
-        model = _LegacyCompletionModel(model_name)
-      elif 'generateMessage' in model_info.supported_generation_methods:
-        model = _LegacyChatModel(model_name)
-      else:
-        raise ValueError(f'Unsupported model: {model_name!r}')
-      self._model_cache[model_name] = model
-    return model
-
-
-_GOOGLE_GENAI_MODEL_HUB = _ModelHub()
-
-
-#
-# Public Gemini models.
-#
-class GeminiFlash2_0ThinkingExp(GenAI):  # pylint: disable=invalid-name
-  """Gemini 2.0 Flash Thinking Experimental model."""
+class GeminiFlash2_0ThinkingExp_20241219(GenAI):  # pylint: disable=invalid-name
+  """Gemini Flash 2.0 Thinking model launched on 12/19/2024."""
 
+  api_version = 'v1alpha'
   model = 'gemini-2.0-flash-thinking-exp-1219'
-  supported_modalities = (
-      vertexai.DOCUMENT_TYPES
-      + vertexai.IMAGE_TYPES
-      + vertexai.AUDIO_TYPES
-      + vertexai.VIDEO_TYPES
-  )
+  timeout = None
 
 
 class GeminiFlash2_0Exp(GenAI):  # pylint: disable=invalid-name
-  """Gemini Experimental model launched on 12/06/2024."""
+  """Gemini Flash 2.0 model launched on 12/11/2024."""
 
   model = 'gemini-2.0-flash-exp'
-  supported_modalities = (
-      vertexai.DOCUMENT_TYPES
-      + vertexai.IMAGE_TYPES
-      + vertexai.AUDIO_TYPES
-      + vertexai.VIDEO_TYPES
-  )
 
 
 class GeminiExp_20241206(GenAI):  # pylint: disable=invalid-name
   """Gemini Experimental model launched on 12/06/2024."""
 
   model = 'gemini-exp-1206'
-  supported_modalities = (
-      vertexai.DOCUMENT_TYPES
-      + vertexai.IMAGE_TYPES
-      + vertexai.AUDIO_TYPES
-      + vertexai.VIDEO_TYPES
-  )
 
 
 class GeminiExp_20241114(GenAI):  # pylint: disable=invalid-name
   """Gemini Experimental model launched on 11/14/2024."""
 
   model = 'gemini-exp-1114'
-  supported_modalities = (
-      vertexai.DOCUMENT_TYPES
-      + vertexai.IMAGE_TYPES
-      + vertexai.AUDIO_TYPES
-      + vertexai.VIDEO_TYPES
-  )
 
 
 class GeminiPro1_5(GenAI):  # pylint: disable=invalid-name
   """Gemini Pro latest model."""
 
   model = 'gemini-1.5-pro-latest'
-  supported_modalities = (
-      vertexai.DOCUMENT_TYPES
-      + vertexai.IMAGE_TYPES
-      + vertexai.AUDIO_TYPES
-      + vertexai.VIDEO_TYPES
-  )
 
 
-class GeminiFlash1_5(GenAI):  # pylint: disable=invalid-name
-  """Gemini Flash latest model."""
+class GeminiPro1_5_002(GenAI):  # pylint: disable=invalid-name
+  """Gemini Pro latest model."""
+
+  model = 'gemini-1.5-pro-002'
 
-  model = 'gemini-1.5-flash-latest'
-  supported_modalities = (
-      vertexai.DOCUMENT_TYPES
-      + vertexai.IMAGE_TYPES
-      + vertexai.AUDIO_TYPES
-      + vertexai.VIDEO_TYPES
-  )
 
+class GeminiPro1_5_001(GenAI):  # pylint: disable=invalid-name
+  """Gemini Pro latest model."""
 
-class GeminiPro(GenAI):
-  """Gemini Pro model."""
+  model = 'gemini-1.5-pro-001'
 
-  model = 'gemini-pro'
+
+class GeminiFlash1_5(GenAI):  # pylint: disable=invalid-name
+  """Gemini Flash latest model."""
+
+  model = 'gemini-1.5-flash-latest'
 
 
-class GeminiProVision(GenAI):
-  """Gemini Pro vision model."""
+class GeminiFlash1_5_002(GenAI):  # pylint: disable=invalid-name
+  """Gemini Flash 1.5 model stable version 002."""
 
-  model = 'gemini-pro-vision'
-  supported_modalities = vertexai.IMAGE_TYPES + vertexai.VIDEO_TYPES
+  model = 'gemini-1.5-flash-002'
 
 
-class Palm2(GenAI):
-  """PaLM2 model."""
+class GeminiFlash1_5_001(GenAI):  # pylint: disable=invalid-name
+  """Gemini Flash 1.5 model stable version 001."""
 
-  model = 'text-bison-001'
+  model = 'gemini-1.5-flash-001'
 
 
-class Palm2_IT(GenAI):  # pylint: disable=invalid-name
-  """PaLM2 instruction-tuned model."""
+class GeminiPro1(GenAI):  # pylint: disable=invalid-name
+  """Gemini 1.0 Pro model."""
 
-  model = 'chat-bison-001'
+  model = 'gemini-1.0-pro'

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()