langfun 0.1.2.dev202502110804__py3-none-any.whl → 0.1.2.dev202502120804__py3-none-any.whl

langfun/core/llms/gemini_test.py

@@ -78,6 +78,33 @@ def mock_requests_post(url: str, json: dict[str, Any], **kwargs):
 class GeminiTest(unittest.TestCase):
   """Tests for Vertex model with REST API."""
 
+  def test_dir(self):
+    self.assertIn('gemini-1.5-pro', gemini.Gemini.dir())
+
+  def test_estimate_cost(self):
+    model = gemini.Gemini('gemini-1.5-pro', api_endpoint='')
+    self.assertEqual(
+        model.estimate_cost(
+            lf.LMSamplingUsage(
+                prompt_tokens=100_000,
+                completion_tokens=1000,
+                total_tokens=101_000,
+            )
+        ),
+        0.13
+    )
+    # Prompt length greater than 128K.
+    self.assertEqual(
+        model.estimate_cost(
+            lf.LMSamplingUsage(
+                prompt_tokens=200_000,
+                completion_tokens=1000,
+                total_tokens=201_000,
+            )
+        ),
+        0.51
+    )
+
   def test_content_from_message_text_only(self):
     text = 'This is a beautiful day'
     model = gemini.Gemini('gemini-1.5-pro', api_endpoint='')
@@ -89,13 +116,6 @@ class GeminiTest(unittest.TestCase):
     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(

langfun/core/llms/google_genai.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Language models from Google GenAI."""
 
+import functools
 import os
 from typing import Annotated, Literal
 
@@ -26,6 +27,20 @@ import pyglove as pg
 class GenAI(gemini.Gemini):
   """Language models provided by Google GenAI."""
 
+  model: pg.typing.Annotated[
+      pg.typing.Enum(
+          pg.MISSING_VALUE,
+          [
+              m.model_id for m in gemini.SUPPORTED_MODELS
+              if m.provider == 'Google GenAI' or (
+                  isinstance(m.provider, pg.hyper.OneOf)
+                  and 'Google GenAI' in m.provider.candidates
+              )
+          ]
+      ),
+      'The name of the model to use.',
+  ]
+
   api_key: Annotated[
       str | None,
       (
@@ -40,10 +55,11 @@ class GenAI(gemini.Gemini):
       'The API version to use.'
   ] = 'v1beta'
 
-  @property
-  def model_id(self) -> str:
-    """Returns a string to identify the model."""
-    return f'GenAI({self.model})'
+  @functools.cached_property
+  def model_info(self) -> lf.ModelInfo:
+    return super().model_info.clone(
+        override=dict(provider='Google GenAI')
+    )
 
   @property
   def api_endpoint(self) -> str:
@@ -63,91 +79,105 @@ class GenAI(gemini.Gemini):
     )
 
 
-class Gemini2Flash(GenAI):  # pylint: disable=invalid-name
-  """Gemini Flash 2.0 model launched on 02/05/2025."""
-
-  api_version = 'v1beta'
-  model = 'gemini-2.0-flash'
+# pylint: disable=invalid-name
 
+#
+# Experimental models.
+#
 
-class Gemini2ProExp_20250205(GenAI):  # pylint: disable=invalid-name
-  """Gemini Flash 2.0 Pro model launched on 02/05/2025."""
 
-  api_version = 'v1beta'
+class Gemini2ProExp_20250205(GenAI):
+  """Gemini 2.0 Pro experimental model launched on 02/05/2025."""
   model = 'gemini-2.0-pro-exp-02-05'
 
 
-class Gemini2FlashThinkingExp_20250121(GenAI):  # pylint: disable=invalid-name
-  """Gemini Flash 2.0 Thinking model launched on 01/21/2025."""
-
+class Gemini2FlashThinkingExp_20250121(GenAI):
+  """Gemini 2.0 Flash Thinking model launched on 01/21/2025."""
   api_version = 'v1beta'
   model = 'gemini-2.0-flash-thinking-exp-01-21'
   timeout = None
 
 
-class GeminiFlash2_0ThinkingExp_20241219(GenAI):  # pylint: disable=invalid-name
-  """Gemini Flash 2.0 Thinking model launched on 12/19/2024."""
+class GeminiExp_20241206(GenAI):
+  """Gemini Experimental model launched on 12/06/2024."""
+  model = 'gemini-exp-1206'
 
-  api_version = 'v1beta'
-  model = 'gemini-2.0-flash-thinking-exp-1219'
-  timeout = None
 
+#
+# Production models.
+#
 
-class GeminiFlash2_0Exp(GenAI):  # pylint: disable=invalid-name
-  """Gemini Flash 2.0 model launched on 12/11/2024."""
 
-  model = 'gemini-2.0-flash-exp'
+class Gemini2Flash(GenAI):
+  """Gemini 2.0 Flash model (latest stable)."""
+  model = 'gemini-2.0-flash'
 
 
-class GeminiExp_20241206(GenAI):  # pylint: disable=invalid-name
-  """Gemini Experimental model launched on 12/06/2024."""
+class Gemini2Flash_001(GenAI):
+  """Gemini 2.0 Flash model launched on 02/05/2025."""
+  model = 'gemini-2.0-flash-001'
 
-  model = 'gemini-exp-1206'
+
+class Gemini2FlashLitePreview_20250205(GenAI):
+  """Gemini 2.0 Flash lite preview model launched on 02/05/2025."""
+  model = 'gemini-2.0-flash-lite-preview-02-05'
 
 
-class GeminiExp_20241114(GenAI):  # pylint: disable=invalid-name
-  """Gemini Experimental model launched on 11/14/2024."""
+class Gemini15Pro(GenAI):
+  """Gemini 1.5 Pro latest stable model."""
+  model = 'gemini-1.5-pro'
 
-  model = 'gemini-exp-1114'
 
+class Gemini15Pro_002(GenAI):
+  """Gemini 1.5 Pro stable version 002."""
+  model = 'gemini-1.5-pro-002'
 
-class GeminiPro1_5(GenAI):  # pylint: disable=invalid-name
-  """Gemini Pro latest model."""
 
-  model = 'gemini-1.5-pro-latest'
+class Gemini15Pro_001(GenAI):
+  """Gemini 1.5 Pro stable version 001."""
+  model = 'gemini-1.5-pro-001'
 
 
-class GeminiPro1_5_002(GenAI):  # pylint: disable=invalid-name
-  """Gemini Pro latest model."""
+class Gemini15Flash(GenAI):
+  """Gemini 1.5 Flash latest model."""
+  model = 'gemini-1.5-flash'
 
-  model = 'gemini-1.5-pro-002'
 
+class Gemini15Flash_002(GenAI):
+  """Gemini 1.5 Flash model stable version 002."""
+  model = 'gemini-1.5-flash-002'
 
-class GeminiPro1_5_001(GenAI):  # pylint: disable=invalid-name
-  """Gemini Pro latest model."""
 
-  model = 'gemini-1.5-pro-001'
+class Gemini15Flash_001(GenAI):
+  """Gemini 1.5 Flash model stable version 001."""
+  model = 'gemini-1.5-flash-001'
 
 
-class GeminiFlash1_5(GenAI):  # pylint: disable=invalid-name
-  """Gemini Flash latest model."""
+class Gemini15Flash8B(GenAI):
+  """Gemini 1.5 Flash 8B modle (latest stable)."""
+  model = 'gemini-1.5-flash-8b'
 
-  model = 'gemini-1.5-flash-latest'
 
+class Gemini15Flash8B_001(GenAI):
+  """Gemini 1.5 Flash 8B model (version 001)."""
+  model = 'gemini-1.5-flash-8b-001'
 
-class GeminiFlash1_5_002(GenAI):  # pylint: disable=invalid-name
-  """Gemini Flash 1.5 model stable version 002."""
 
-  model = 'gemini-1.5-flash-002'
+# For backward compatibility.
+GeminiPro1_5 = Gemini15Pro
+GeminiFlash1_5 = Gemini15Flash
 
+# pylint: enable=invalid-name
 
-class GeminiFlash1_5_001(GenAI):  # pylint: disable=invalid-name
-  """Gemini Flash 1.5 model stable version 001."""
 
-  model = 'gemini-1.5-flash-001'
+def _genai_model(model: str, *args, **kwargs) -> GenAI:
+  model = model.removeprefix('google_genai://')
+  return GenAI(model=model, *args, **kwargs)
 
 
-class GeminiPro1(GenAI):  # pylint: disable=invalid-name
-  """Gemini 1.0 Pro model."""
+def _register_genai_models():
+  """Register GenAI models."""
+  for m in gemini.SUPPORTED_MODELS:
+    lf.LanguageModel.register('google_genai://' + m.model_id, _genai_model)
 
-  model = 'gemini-1.0-pro'
+_register_genai_models()

langfun/core/llms/google_genai_test.py

@@ -15,6 +15,7 @@
 
 import os
 import unittest
+import langfun.core as lf
 from langfun.core.llms import google_genai
 
 
@@ -23,16 +24,22 @@ class GenAITest(unittest.TestCase):
 
   def test_basics(self):
     with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
-      _ = google_genai.GeminiPro1_5().api_endpoint
+      _ = google_genai.Gemini15Pro().api_endpoint
 
-    self.assertIsNotNone(google_genai.GeminiPro1_5(api_key='abc').api_endpoint)
+    self.assertIsNotNone(google_genai.Gemini15Pro(api_key='abc').api_endpoint)
 
     os.environ['GOOGLE_API_KEY'] = 'abc'
-    lm = google_genai.GeminiPro1_5()
+    lm = google_genai.Gemini15Pro_001()
     self.assertIsNotNone(lm.api_endpoint)
-    self.assertTrue(lm.model_id.startswith('GenAI('))
+    self.assertEqual(lm.model_id, 'gemini-1.5-pro-001')
+    self.assertEqual(lm.resource_id, 'google_genai://gemini-1.5-pro-001')
     del os.environ['GOOGLE_API_KEY']
 
+  def test_lm_get(self):
+    self.assertIsInstance(
+        lf.LanguageModel.get('google_genai://gemini-1.5-pro'),
+        google_genai.GenAI,
+    )
 
 if __name__ == '__main__':
   unittest.main()