langfun 0.1.2.dev202412010804__py3-none-any.whl → 0.1.2.dev202412030804__py3-none-any.whl

langfun/core/llms/vertexai_test.py

@@ -13,16 +13,17 @@
 # limitations under the License.
 """Tests for Gemini models."""
 
+import base64
 import os
+from typing import Any
 import unittest
 from unittest import mock
 
-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
 from langfun.core.llms import vertexai
 import pyglove as pg
+import requests
 
 
 example_image = (
@@ -37,60 +38,53 @@ example_image = (
 )
 
 
-def mock_generate_content(content, generation_config, **kwargs):
-  del kwargs
-  c = pg.Dict(generation_config.to_dict())
-  return generative_models.GenerationResponse.from_dict({
+def mock_requests_post(url: str, json: dict[str, Any], **kwargs):
+  del url, kwargs
+  c = pg.Dict(json['generationConfig'])
+  content = json['contents'][0]['parts'][0]['text']
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str({
       'candidates': [
           {
-              'index': 0,
               'content': {
                   'role': 'model',
                   'parts': [
                       {
                           'text': (
-                              f'This is a response to {content[0]} with '
+                              f'This is a response to {content} with '
                               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={"".join(c.stop_sequences)}.'
+                              f'top_p={c.topP}, '
+                              f'top_k={c.topK}, '
+                              f'max_tokens={c.maxOutputTokens}, '
+                              f'stop={"".join(c.stopSequences)}.'
                           )
                       },
                   ],
               },
           },
-      ]
-  })
-
-
-def mock_endpoint_predict(instances, **kwargs):
-  del kwargs
-  assert len(instances) == 1
-  return aiplatform_models.Prediction(
-      predictions=[
-          f"This is a response to {instances[0]['prompt']} with"
-          f" temperature={instances[0]['temperature']},"
-          f" top_p={instances[0]['top_p']}, top_k={instances[0]['top_k']},"
-          f" max_tokens={instances[0]['max_tokens']}."
       ],
-      deployed_model_id='',
-  )
+      'usageMetadata': {
+          'promptTokenCount': 3,
+          'candidatesTokenCount': 4,
+      }
+  }).encode()
+  return response
 
 
 class VertexAITest(unittest.TestCase):
-  """Tests for Vertex model."""
+  """Tests for Vertex model with REST API."""
 
   def test_content_from_message_text_only(self):
     text = 'This is a beautiful day'
-    model = vertexai.VertexAIGeminiPro1()
+    model = vertexai.VertexAIGeminiPro1_5_002()
     chunks = model._content_from_message(lf.UserMessage(text))
-    self.assertEqual(chunks, [text])
+    self.assertEqual(chunks, {'role': 'user', 'parts': [{'text': text}]})
 
   def test_content_from_message_mm(self):
+    image = lf_modalities.Image.from_bytes(example_image)
     message = lf.UserMessage(
-        'This is an <<[[image]]>>, what is it?',
-        image=lf_modalities.Image.from_bytes(example_image),
+        'This is an <<[[image]]>>, what is it?', image=image
     )
 
     # Non-multimodal model.
@@ -98,64 +92,25 @@ class VertexAITest(unittest.TestCase):
       vertexai.VertexAIGeminiPro1()._content_from_message(message)
 
     model = vertexai.VertexAIGeminiPro1Vision()
-    chunks = model._content_from_message(message)
-    self.maxDiff = None
-    self.assertEqual([chunks[0], chunks[2]], ['This is an', ', what is it?'])
-    self.assertIsInstance(chunks[1], generative_models.Part)
-
-  def test_generation_response_to_message_text_only(self):
-    response = generative_models.GenerationResponse.from_dict({
-        'candidates': [
-            {
-                'index': 0,
-                'content': {
-                    'role': 'model',
-                    'parts': [
-                        {
-                            'text': 'hello world',
-                        },
-                    ],
+    content = model._content_from_message(message)
+    self.assertEqual(
+        content,
+        {
+            'role': 'user',
+            'parts': [
+                {'text': 'This is an'},
+                {
+                    'inlineData': {
+                        'data': base64.b64encode(example_image).decode(),
+                        'mimeType': 'image/png',
+                    }
                 },
-            },
-        ],
-    })
-    model = vertexai.VertexAIGeminiPro1()
-    message = model._generation_response_to_message(response)
-    self.assertEqual(message, lf.AIMessage('hello world'))
-
-  def test_model_hub(self):
-    with mock.patch(
-        'vertexai.generative_models.'
-        'GenerativeModel.__init__'
-    ) as mock_model_init:
-      mock_model_init.side_effect = lambda *args, **kwargs: None
-      model = vertexai._VERTEXAI_MODEL_HUB.get_generative_model(
-          'gemini-1.0-pro'
-      )
-      self.assertIsNotNone(model)
-      self.assertIs(
-          vertexai._VERTEXAI_MODEL_HUB.get_generative_model('gemini-1.0-pro'),
-          model,
-      )
-
-    with mock.patch(
-        'vertexai.language_models.'
-        'TextGenerationModel.from_pretrained'
-    ) as mock_model_init:
-
-      class TextGenerationModel:
-        pass
-
-      mock_model_init.side_effect = lambda *args, **kw: TextGenerationModel()
-      model = vertexai._VERTEXAI_MODEL_HUB.get_text_generation_model(
-          'text-bison'
-      )
-      self.assertIsNotNone(model)
-      self.assertIs(
-          vertexai._VERTEXAI_MODEL_HUB.get_text_generation_model('text-bison'),
-          model,
-      )
+                {'text': ', what is it?'},
+            ],
+        },
+    )
 
+  @mock.patch.object(vertexai.VertexAI, 'credentials', new=True)
   def test_project_and_location_check(self):
     with self.assertRaisesRegex(ValueError, 'Please specify `project`'):
       _ = vertexai.VertexAIGeminiPro1()._api_initialized
@@ -185,7 +140,7 @@ class VertexAITest(unittest.TestCase):
         'required': ['name'],
         'title': 'Person',
     }
-    config = model._generation_config(
+    actual = model._generation_config(
         lf.UserMessage('hi', json_schema=json_schema),
         lf.LMSamplingOptions(
             temperature=2.0,
@@ -195,29 +150,23 @@ class VertexAITest(unittest.TestCase):
             stop=['\n'],
         ),
     )
-    actual = config.to_dict()
-    # There is a discrepancy between the `property_ordering` in the
-    # Google-internal version and the open-source version.
-    actual['response_schema'].pop('property_ordering', None)
-    if pg.KeyPath.parse('response_schema.type_').get(actual):
-      actual['response_schema']['type'] = actual['response_schema'].pop('type_')
-    if pg.KeyPath.parse('response_schema.properties.name.type_').get(actual):
-      actual['response_schema']['properties']['name']['type'] = actual[
-          'response_schema']['properties']['name'].pop('type_')
-
     self.assertEqual(
         actual,
         dict(
+            candidateCount=1,
             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',
+            topP=1.0,
+            topK=20,
+            maxOutputTokens=1024,
+            stopSequences=['\n'],
+            responseLogprobs=False,
+            logprobs=None,
+            seed=None,
+            responseMimeType='application/json',
+            responseSchema={
+                'type': 'object',
                 'properties': {
-                    'name': {'type': 'STRING'}
+                    'name': {'type': 'string'}
                 },
                 'required': ['name'],
                 'title': 'Person',
@@ -232,100 +181,31 @@ class VertexAITest(unittest.TestCase):
           lf.LMSamplingOptions(),
       )
 
-  def test_call_generative_model(self):
-    with mock.patch(
-        'vertexai.generative_models.'
-        'GenerativeModel.__init__'
-    ) as mock_model_init:
-      mock_model_init.side_effect = lambda *args, **kwargs: None
-
-      with mock.patch(
-          'vertexai.generative_models.'
-          'GenerativeModel.generate_content'
-      ) as mock_generate:
-        mock_generate.side_effect = mock_generate_content
-
-        lm = vertexai.VertexAIGeminiPro1(project='abc', location='us-central1')
-        self.assertEqual(
-            lm(
-                'hello',
-                temperature=2.0,
-                top_p=1.0,
-                top_k=20,
-                max_tokens=1024,
-                stop='\n',
-            ).text,
-            (
-                'This is a response to hello with temperature=2.0, '
-                'top_p=1.0, top_k=20.0, max_tokens=1024, stop=\n.'
-            ),
-        )
-
-  def test_call_text_generation_model(self):
-    with mock.patch(
-        'vertexai.language_models.'
-        'TextGenerationModel.from_pretrained'
-    ) as mock_model_init:
+  @mock.patch.object(vertexai.VertexAI, 'credentials', new=True)
+  def test_call_model(self):
+    with mock.patch('requests.Session.post') as mock_generate:
+      mock_generate.side_effect = mock_requests_post
 
-      class TextGenerationModel:
-
-        def predict(self, prompt, **kwargs):
-          c = pg.Dict(kwargs)
-          return pg.Dict(
-              text=(
-                  f'This is a response to {prompt} with '
-                  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={"".join(c.stop_sequences)}.'
-              )
-          )
-
-      mock_model_init.side_effect = lambda *args, **kw: TextGenerationModel()
-      lm = vertexai.VertexAIPalm2(project='abc', location='us-central1')
+      lm = vertexai.VertexAIGeminiPro1_5_002(
+          project='abc', location='us-central1'
+      )
+      r = lm(
+          'hello',
+          temperature=2.0,
+          top_p=1.0,
+          top_k=20,
+          max_tokens=1024,
+          stop='\n',
+      )
       self.assertEqual(
-          lm(
-              'hello',
-              temperature=2.0,
-              top_p=1.0,
-              top_k=20,
-              max_tokens=1024,
-              stop='\n',
-          ).text,
+          r.text,
           (
               'This is a response to hello with temperature=2.0, '
               'top_p=1.0, top_k=20, max_tokens=1024, stop=\n.'
           ),
       )
-
-  def test_call_endpoint_model(self):
-    with mock.patch(
-        '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.',
-        )
+      self.assertEqual(r.metadata.usage.prompt_tokens, 3)
+      self.assertEqual(r.metadata.usage.completion_tokens, 4)
 
 
 if __name__ == '__main__':

{langfun-0.1.2.dev202412010804.dist-info → langfun-0.1.2.dev202412030804.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.1.2.dev202412010804
+Version: 0.1.2.dev202412030804
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors
@@ -30,9 +30,7 @@ Requires-Dist: jinja2>=3.1.2; extra == "all"
 Requires-Dist: requests>=2.31.0; extra == "all"
 Requires-Dist: termcolor==1.1.0; extra == "all"
 Requires-Dist: tqdm>=4.64.1; extra == "all"
-Requires-Dist: google-cloud-aiplatform>=1.5.0; extra == "all"
 Requires-Dist: google-generativeai>=0.3.2; extra == "all"
-Requires-Dist: openai>=0.27.2; extra == "all"
 Requires-Dist: python-magic>=0.4.27; extra == "all"
 Requires-Dist: python-docx>=0.8.11; extra == "all"
 Requires-Dist: pillow>=10.0.0; extra == "all"
@@ -42,18 +40,11 @@ Provides-Extra: ui
 Requires-Dist: termcolor==1.1.0; extra == "ui"
 Requires-Dist: tqdm>=4.64.1; extra == "ui"
 Provides-Extra: llm
-Requires-Dist: google-cloud-aiplatform>=1.5.0; extra == "llm"
 Requires-Dist: google-generativeai>=0.3.2; extra == "llm"
-Requires-Dist: openai>=0.27.2; extra == "llm"
 Provides-Extra: llm-google
-Requires-Dist: google-cloud-aiplatform>=1.5.0; extra == "llm-google"
 Requires-Dist: google-generativeai>=0.3.2; extra == "llm-google"
-Provides-Extra: llm-google-vertex
-Requires-Dist: google-cloud-aiplatform>=1.5.0; extra == "llm-google-vertex"
 Provides-Extra: llm-google-genai
 Requires-Dist: google-generativeai>=0.3.2; extra == "llm-google-genai"
-Provides-Extra: llm-openai
-Requires-Dist: openai>=0.27.2; extra == "llm-openai"
 Provides-Extra: mime
 Requires-Dist: python-magic>=0.4.27; extra == "mime"
 Requires-Dist: python-docx>=0.8.11; extra == "mime"
@@ -212,9 +203,7 @@ If you want to customize your installation, you can select specific features usi
 | all                 | All Langfun features.                    |
 | llm                 | All supported LLMs.                      |
 | llm-google          | All supported Google-powered LLMs.       |
-| llm-google-vertexai | LLMs powered by Google Cloud VertexAI    |
 | llm-google-genai    | LLMs powered by Google Generative AI API |
-| llm-openai          | LLMs powered by OpenAI                   |
 | mime                | All MIME supports.                       |
 | mime-auto           | Automatic MIME type detection.           |
 | mime-docx           | DocX format support.                     |

{langfun-0.1.2.dev202412010804.dist-info → langfun-0.1.2.dev202412030804.dist-info}/RECORD

@@ -72,14 +72,14 @@ langfun/core/eval/v2/metrics.py,sha256=bl8i6u-ZHRBz4hAc3LzsZ2Dc7ZRQcuTYeUhhH-Gxf
 langfun/core/eval/v2/metrics_test.py,sha256=p4FzLJsE8XAzAQuyP9hfEf9YeKWZ__PO_ue8a9P0-cc,6082
 langfun/core/eval/v2/progress.py,sha256=azZgssQgNdv3IgjKEaQBuGI5ucFDNbdi02P4z_nQ8GE,10292
 langfun/core/eval/v2/progress_test.py,sha256=YU7VHzmy5knPZwj9vpBN3rQQH2tukj9eKHkuBCI62h8,2540
-langfun/core/eval/v2/progress_tracking.py,sha256=1imwSbllxHWG3zYrzo2NvytBZsVtjqum6bmXGGsvT1E,5987
-langfun/core/eval/v2/progress_tracking_test.py,sha256=eY2HvZeEXDA5Zyfi2m5NDWO_9kSfQsaAOEcIhkSbWCY,1874
+langfun/core/eval/v2/progress_tracking.py,sha256=l9fEkz4oP5McpZzf72Ua7PYm3lAWtRru7gRWNf8H0ms,6083
+langfun/core/eval/v2/progress_tracking_test.py,sha256=iO-DslCJWncU7-27XaMKxDeKrsGbwdk_tKfoRk3KboE,2271
 langfun/core/eval/v2/reporting.py,sha256=TGkli1IDwqfqsCJ_WslOMGk_24JDg7oRRTGXlAJlWpc,4361
 langfun/core/eval/v2/reporting_test.py,sha256=JxffbUPWInUyLjo-AQVFrllga884Mdfm05R86FtxSss,1482
 langfun/core/eval/v2/runners.py,sha256=zJmu-amUiYv1g0Ek4c3mXkBgp-AFvSF7WpXVZCCf7Y4,14245
 langfun/core/eval/v2/runners_test.py,sha256=UeiUNygux_U6iGVG18rhp68ZE4hoWeoT6XsXvSjxNQg,11620
 langfun/core/eval/v2/test_helper.py,sha256=pDpZTBnWRR5xjJv3Uy3NWEzArqlL8FTMOgeR4C53F5M,2348
-langfun/core/llms/__init__.py,sha256=tM9LjVQDFFG7Et2vjUr45GlYNEu73LAcWsd_1DO0KcQ,6460
+langfun/core/llms/__init__.py,sha256=C4hyLflqOQT841nMfclcxcnOhdP83zR0GGW29PnA-vU,6216
 langfun/core/llms/anthropic.py,sha256=uJXVgaFONL8okOSVQ4VGMGht_VZ30m1hoLzmDbIjmks,13990
 langfun/core/llms/anthropic_test.py,sha256=-2U4kc_pgBM7wqxu8RuxzyHPGww1EAWqKUvN4PW8Btw,8058
 langfun/core/llms/compositional.py,sha256=csW_FLlgL-tpeyCOTVvfUQkMa_zCN5Y2I-YbSNuK27U,2872
@@ -92,12 +92,12 @@ langfun/core/llms/groq.py,sha256=dCnR3eAECEKuKKAAj-PDTs8NRHl6CQPdf57m1f6a79U,103
 langfun/core/llms/groq_test.py,sha256=GYF_Qtq5S1H1TrKH38t6_lkdroqT7v-joYLDKnmS9e0,5274
 langfun/core/llms/llama_cpp.py,sha256=9tXQntSCDtjTF3bnyJrAPCr4N6wycy5nXYvp9uduygE,2843
 langfun/core/llms/llama_cpp_test.py,sha256=MWO_qaOeKjRniGjcaWPDScd7HPaIJemqUZoslrt4FPs,1806
-langfun/core/llms/openai.py,sha256=HVJct35hWwer7iq_1ZfTy0F20ZXSB2X-H4CVwStR0rg,23054
-langfun/core/llms/openai_test.py,sha256=_8cd3VRNEUfE0-Ko1RiM6MlC5hjalRj7nYTJNhG1p3E,18907
+langfun/core/llms/openai.py,sha256=l49v6RubfInvV0iG114AymTKNogTX4u4N-UFCeSgIxw,20963
+langfun/core/llms/openai_test.py,sha256=kOWa1nf-nJvtYY10REUw5wojh3ZgfU8tRaCZ8wUgJbA,16623
 langfun/core/llms/rest.py,sha256=sWbYUV8S3SuOg9giq7xwD-xDRfaF7NP_ig7bI52-Rj4,3442
 langfun/core/llms/rest_test.py,sha256=NZ3Nf0XQVpT9kLP5cBVo_yBHLI7vWTYhWQxYEJVMGs4,3472
-langfun/core/llms/vertexai.py,sha256=sl8TIhTuz8ZZFSGo97K6zqXcJQNgF2CVAOLTEiMuiDA,19006
-langfun/core/llms/vertexai_test.py,sha256=I8gEHLRXZZGq_d2VDtJAkAIzf-lNSCoB8y2lwFckY-w,10885
+langfun/core/llms/vertexai.py,sha256=48GFuf04NOeInvwHeLhs-iEyxL03VCjtw05BGcKio0s,14172
+langfun/core/llms/vertexai_test.py,sha256=ffcA5yPecnQy_rhkuYAw_6o1iLW8AR8FgswmHt6aAys,6725
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=rt3zwmyw0y9jsSGW-ZbV1vAfLxQ7_3AVk0l2EySlse4,3918
 langfun/core/llms/cache/in_memory.py,sha256=l6b-iU9OTfTRo9Zmg4VrQIuArs4cCJDOpXiEpvNocjo,5004
@@ -148,8 +148,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=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
-langfun-0.1.2.dev202412010804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202412010804.dist-info/METADATA,sha256=HuT85enj3_pC21Mt-MaOLQHZ6DKwkHd7AesNlxcqlu0,8890
-langfun-0.1.2.dev202412010804.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-langfun-0.1.2.dev202412010804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202412010804.dist-info/RECORD,,
+langfun-0.1.2.dev202412030804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202412030804.dist-info/METADATA,sha256=BG5zoGi1sstPc97R4aTE51aiVfa02LLq-BRHRb8R_3Q,8281
+langfun-0.1.2.dev202412030804.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+langfun-0.1.2.dev202412030804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202412030804.dist-info/RECORD,,