langfun 0.1.2.dev202501010804__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/openai.py

@@ -32,6 +32,13 @@ SUPPORTED_MODELS_AND_SETTINGS = {
     # o1 (preview) models.
     # Pricing in US dollars, from https://openai.com/api/pricing/
     # as of 2024-10-10.
+    'o1': pg.Dict(
+        in_service=True,
+        rpm=10000,
+        tpm=5000000,
+        cost_per_1k_input_tokens=0.015,
+        cost_per_1k_output_tokens=0.06,
+    ),
     'o1-preview': pg.Dict(
         in_service=True,
         rpm=10000,
@@ -255,25 +262,17 @@ SUPPORTED_MODELS_AND_SETTINGS = {
     ),
     # GPT-3.5 models
     'text-davinci-003': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
+        in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM
     ),
     'text-davinci-002': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
+        in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM
     ),
     'code-davinci-002': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
+        in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM
     ),
     # GPT-3 instruction-tuned models (Deprecated)
     'text-curie-001': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
+        in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM
     ),
     'text-babbage-001': pg.Dict(
         in_service=False,
@@ -290,32 +289,12 @@ SUPPORTED_MODELS_AND_SETTINGS = {
         rpm=_DEFAULT_RPM,
         tpm=_DEFAULT_TPM,
     ),
-    'curie': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
-    ),
-    'babbage': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
-    ),
-    'ada': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
-    ),
+    'curie': pg.Dict(in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'babbage': pg.Dict(in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'ada': pg.Dict(in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
     # GPT-3 base models that are still in service.
-    'babbage-002': pg.Dict(
-        in_service=True,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
-    ),
-    'davinci-002': pg.Dict(
-        in_service=True,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
-    ),
+    'babbage-002': pg.Dict(in_service=True, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'davinci-002': pg.Dict(in_service=True, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
 }
 
 
@@ -569,6 +548,13 @@ class OpenAI(rest.REST):
     )
 
 
+class GptO1(OpenAI):
+  """GPT-O1."""
+
+  model = 'o1'
+  multimodal = True
+
+
 class GptO1Preview(OpenAI):
   """GPT-O1."""
   model = 'o1-preview'