langfun 0.1.1.dev20240812__py3-none-any.whl → 0.1.1.dev20240813__py3-none-any.whl

langfun/core/llms/vertexai.py

@@ -18,18 +18,17 @@ import os
 from typing import Annotated, Any
 
 from google.auth import credentials as credentials_lib
-from google.cloud.aiplatform import aiplatform
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
 import pyglove as pg
 
 
 SUPPORTED_MODELS_AND_SETTINGS = {
-    'gemini-1.5-pro-001': pg.Dict(api='gemini', rpm=5),
-    'gemini-1.5-flash-001': pg.Dict(api='gemini', rpm=5),
-    'gemini-1.5-pro-preview-0514': pg.Dict(api='gemini', rpm=5),
-    'gemini-1.5-pro-preview-0409': pg.Dict(api='gemini', rpm=5),
-    'gemini-1.5-flash-preview-0514': pg.Dict(api='gemini', rpm=5),
+    'gemini-1.5-pro-001': pg.Dict(api='gemini', rpm=50),
+    'gemini-1.5-flash-001': pg.Dict(api='gemini', rpm=200),
+    'gemini-1.5-pro-preview-0514': pg.Dict(api='gemini', rpm=50),
+    'gemini-1.5-pro-preview-0409': pg.Dict(api='gemini', rpm=50),
+    'gemini-1.5-flash-preview-0514': pg.Dict(api='gemini', rpm=200),
     'gemini-1.0-pro': pg.Dict(api='gemini', rpm=300),
     'gemini-1.0-pro-vision': pg.Dict(api='gemini', rpm=100),
     # PaLM APIs.
@@ -136,16 +135,34 @@ class VertexAI(lf.LanguageModel):
     )
 
   def _generation_config(
-      self, options: lf.LMSamplingOptions
+      self, prompt: lf.Message, options: lf.LMSamplingOptions
   ) -> Any:  # generative_models.GenerationConfig
     """Creates generation config from langfun sampling options."""
     from vertexai import generative_models
+    # Users could use `metadata_json_schema` to pass additional
+    # request arguments.
+    json_schema = prompt.metadata.get('json_schema')
+    response_mime_type = None
+    if json_schema is not None:
+      if not isinstance(json_schema, dict):
+        raise ValueError(
+            f'`json_schema` must be a dict, got {json_schema!r}.'
+        )
+      response_mime_type = '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 generative_models.GenerationConfig(
         temperature=options.temperature,
         top_p=options.top_p,
         top_k=options.top_k,
         max_output_tokens=options.max_tokens,
         stop_sequences=options.stop,
+        response_mime_type=response_mime_type,
+        response_schema=json_schema,
     )
 
   def _content_from_message(
@@ -239,7 +256,9 @@ class VertexAI(lf.LanguageModel):
     input_content = self._content_from_message(prompt)
     response = model.generate_content(
         input_content,
-        generation_config=self._generation_config(self.sampling_options),
+        generation_config=self._generation_config(
+            prompt, self.sampling_options
+        ),
     )
     usage_metadata = response.usage_metadata
     usage = lf.LMSamplingUsage(
@@ -277,7 +296,8 @@ class VertexAI(lf.LanguageModel):
 
   def _sample_endpoint_model(self, prompt: lf.Message) -> lf.LMSamplingResult:
     """Samples a text generation model."""
-    model = aiplatform.Endpoint(self.endpoint_name)
+    from google.cloud.aiplatform import models
+    model = models.Endpoint(self.endpoint_name)
     # TODO(chengrun): Add support for stop_sequences.
     predict_options = dict(
         temperature=self.sampling_options.temperature

langfun/core/llms/vertexai_test.py

@@ -17,7 +17,7 @@ import os
 import unittest
 from unittest import mock
 
-from google.cloud.aiplatform.aiplatform import models as aiplatform_models
+from google.cloud.aiplatform import models as aiplatform_models
 from vertexai import generative_models
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
@@ -175,6 +175,53 @@ class VertexAITest(unittest.TestCase):
     del os.environ['VERTEXAI_PROJECT']
     del os.environ['VERTEXAI_LOCATION']
 
+  def test_generation_config(self):
+    model = vertexai.VertexAIGeminiPro1()
+    json_schema = {
+        'type': 'object',
+        'properties': {
+            'name': {'type': 'string'},
+        },
+        'required': ['name'],
+        'title': 'Person',
+    }
+    config = model._generation_config(
+        lf.UserMessage('hi', json_schema=json_schema),
+        lf.LMSamplingOptions(
+            temperature=2.0,
+            top_p=1.0,
+            top_k=20,
+            max_tokens=1024,
+            stop=['\n'],
+        ),
+    )
+    self.assertEqual(
+        config.to_dict(),
+        dict(
+            temperature=2.0,
+            top_p=1.0,
+            top_k=20.0,
+            max_output_tokens=1024,
+            stop_sequences=['\n'],
+            response_mime_type='application/json',
+            response_schema={
+                'type_': 'OBJECT',
+                'properties': {
+                    'name': {'type_': 'STRING'}
+                },
+                'required': ['name'],
+                'title': 'Person',
+            }
+        ),
+    )
+    with self.assertRaisesRegex(
+        ValueError, '`json_schema` must be a dict, got'
+    ):
+      model._generation_config(
+          lf.UserMessage('hi', json_schema='not a dict'),
+          lf.LMSamplingOptions(),
+      )
+
   def test_call_generative_model(self):
     with mock.patch(
         'vertexai.generative_models.'
@@ -244,27 +291,31 @@ class VertexAITest(unittest.TestCase):
 
   def test_call_endpoint_model(self):
     with mock.patch(
-        'google.cloud.aiplatform.aiplatform.Endpoint.predict'
-    ) as mock_model_predict:
-
-      mock_model_predict.side_effect = mock_endpoint_predict
-      lm = vertexai.VertexAI(
-          'custom',
-          endpoint_name='123',
-          project='abc',
-          location='us-central1',
-      )
-      self.assertEqual(
-          lm(
-              'hello',
-              temperature=2.0,
-              top_p=1.0,
-              top_k=20,
-              max_tokens=50,
-          ),
-          'This is a response to hello with temperature=2.0, top_p=1.0,'
-          ' top_k=20, max_tokens=50.',
-      )
+        'google.cloud.aiplatform.models.Endpoint.__init__'
+    ) as mock_model_init:
+      mock_model_init.side_effect = lambda *args, **kwargs: None
+      with mock.patch(
+          'google.cloud.aiplatform.models.Endpoint.predict'
+      ) as mock_model_predict:
+
+        mock_model_predict.side_effect = mock_endpoint_predict
+        lm = vertexai.VertexAI(
+            'custom',
+            endpoint_name='123',
+            project='abc',
+            location='us-central1',
+        )
+        self.assertEqual(
+            lm(
+                'hello',
+                temperature=2.0,
+                top_p=1.0,
+                top_k=20,
+                max_tokens=50,
+            ),
+            'This is a response to hello with temperature=2.0, top_p=1.0,'
+            ' top_k=20, max_tokens=50.',
+        )
 
 
 if __name__ == '__main__':

{langfun-0.1.1.dev20240812.dist-info → langfun-0.1.1.dev20240813.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.1.1.dev20240812
+Version: 0.1.1.dev20240813
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.1.dev20240812.dist-info → langfun-0.1.1.dev20240813.dist-info}/RECORD

@@ -67,8 +67,8 @@ langfun/core/llms/openai.py,sha256=jILxfFb3vBuyf1u_2-LVfs_wekPF2RVuNFzNVg25pEA,1
 langfun/core/llms/openai_test.py,sha256=3muDTnW7UBOSHq694Fi2bofqhe8Pkj0Tl8IShoLCTOM,15525
 langfun/core/llms/rest.py,sha256=laopuq-zD8V-3Y6eFDngftHEbE66VlUkCD2-rvvRaLU,3388
 langfun/core/llms/rest_test.py,sha256=NZ3Nf0XQVpT9kLP5cBVo_yBHLI7vWTYhWQxYEJVMGs4,3472
-langfun/core/llms/vertexai.py,sha256=-YoEUlG19CWIhJb8S6puPqdX9SoiT5NNAItefwdCfsk,13781
-langfun/core/llms/vertexai_test.py,sha256=N3k4N9_bVjC6_Qtg4WO9jYNv8M9xmv5UdODvIKG2upo,8835
+langfun/core/llms/vertexai.py,sha256=tXAnP357XhcsETTnk6M-hH4xyFi7tk6fsaf3tjzsY6E,14501
+langfun/core/llms/vertexai_test.py,sha256=EPR-mB2hNUpvpf7E8m_k5bh04epdQTVUuYU6hPgZyu8,10321
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=cFfYvOIUae842pncqCAsRvqXCk2AnAsRYVx0mcIoAeY,3338
 langfun/core/llms/cache/in_memory.py,sha256=YfFyJEhLs73cUiB0ZfhMxYpdE8Iuxxw-dvMFwGHTSHw,4742
@@ -117,8 +117,8 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=rBW2Qr8yi-aWYwoTwRR-n1peKyMX9QXPZXURjLgoiRs,2264
-langfun-0.1.1.dev20240812.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.1.dev20240812.dist-info/METADATA,sha256=Kk1dXZKetEZkE1Ycs26AdvbUuJY6TIA7Z75LKGqDlq0,5234
-langfun-0.1.1.dev20240812.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-langfun-0.1.1.dev20240812.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.1.dev20240812.dist-info/RECORD,,
+langfun-0.1.1.dev20240813.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.1.dev20240813.dist-info/METADATA,sha256=pRn2OuwICzmonFmbZuzx6LzEgF_LE44B0ceTIa76fLs,5234
+langfun-0.1.1.dev20240813.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+langfun-0.1.1.dev20240813.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.1.dev20240813.dist-info/RECORD,,