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

langfun/core/llms/vertexai_test.py

@@ -11,105 +11,18 @@
 # 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 VertexAI models."""
 
-import base64
 import os
-from typing import Any
 import unittest
 from unittest import mock
 
-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 = (
-    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_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': [
-          {
-              'content': {
-                  'role': 'model',
-                  'parts': [
-                      {
-                          'text': (
-                              f'This is a response to {content} with '
-                              f'temperature={c.temperature}, '
-                              f'top_p={c.topP}, '
-                              f'top_k={c.topK}, '
-                              f'max_tokens={c.maxOutputTokens}, '
-                              f'stop={"".join(c.stopSequences)}.'
-                          )
-                      },
-                  ],
-              },
-          },
-      ],
-      'usageMetadata': {
-          'promptTokenCount': 3,
-          'candidatesTokenCount': 4,
-      }
-  }).encode()
-  return response
 
 
 class VertexAITest(unittest.TestCase):
   """Tests for Vertex model with REST API."""
 
-  def test_content_from_message_text_only(self):
-    text = 'This is a beautiful day'
-    model = vertexai.VertexAIGeminiPro1_5_002()
-    chunks = model._content_from_message(lf.UserMessage(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=image
-    )
-
-    # Non-multimodal model.
-    with self.assertRaisesRegex(lf.ModalityError, 'Unsupported modality'):
-      vertexai.VertexAIGeminiPro1()._content_from_message(message)
-
-    model = vertexai.VertexAIGeminiPro1Vision()
-    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',
-                    }
-                },
-                {'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`'):
@@ -126,87 +39,14 @@ class VertexAITest(unittest.TestCase):
 
     os.environ['VERTEXAI_PROJECT'] = 'abc'
     os.environ['VERTEXAI_LOCATION'] = 'us-central1'
-    self.assertTrue(vertexai.VertexAIGeminiPro1()._api_initialized)
+    model = vertexai.VertexAIGeminiPro1()
+    self.assertTrue(model.model_id.startswith('VertexAI('))
+    self.assertIsNotNone(model.api_endpoint)
+    self.assertTrue(model._api_initialized)
+    self.assertIsNotNone(model._session)
     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',
-    }
-    actual = 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(
-        actual,
-        dict(
-            candidateCount=1,
-            temperature=2.0,
-            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'}
-                },
-                '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(),
-      )
-
-  @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
-
-      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(
-          r.text,
-          (
-              'This is a response to hello with temperature=2.0, '
-              'top_p=1.0, top_k=20, max_tokens=1024, stop=\n.'
-          ),
-      )
-      self.assertEqual(r.metadata.usage.prompt_tokens, 3)
-      self.assertEqual(r.metadata.usage.completion_tokens, 4)
-
 
 if __name__ == '__main__':
   unittest.main()

{langfun-0.1.2.dev202501050804.dist-info → langfun-0.1.2.dev202501060804.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.1.2.dev202501050804
+Version: 0.1.2.dev202501060804
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors
@@ -30,7 +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-generativeai>=0.3.2; extra == "all"
+Requires-Dist: google-auth>=2.16.0; 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"
@@ -39,12 +39,8 @@ Requires-Dist: pandas>=2.0.3; extra == "all"
 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-generativeai>=0.3.2; extra == "llm"
-Provides-Extra: llm-google
-Requires-Dist: google-generativeai>=0.3.2; extra == "llm-google"
-Provides-Extra: llm-google-genai
-Requires-Dist: google-generativeai>=0.3.2; extra == "llm-google-genai"
+Provides-Extra: vertexai
+Requires-Dist: google-auth>=2.16.0; extra == "vertexai"
 Provides-Extra: mime
 Requires-Dist: python-magic>=0.4.27; extra == "mime"
 Requires-Dist: python-docx>=0.8.11; extra == "mime"
@@ -201,9 +197,7 @@ If you want to customize your installation, you can select specific features usi
 | Tag                 |  Description                             |
 | ------------------- | ---------------------------------------- |
 | all                 | All Langfun features.                    |
-| llm                 | All supported LLMs.                      |
-| llm-google          | All supported Google-powered LLMs.       |
-| llm-google-genai    | LLMs powered by Google Generative AI API |
+| vertexai            | VertexAI access.                         |
 | mime                | All MIME supports.                       |
 | mime-auto           | Automatic MIME type detection.           |
 | mime-docx           | DocX format support.                     |
@@ -212,9 +206,9 @@ If you want to customize your installation, you can select specific features usi
 | ui                  | UI enhancements                          |
 
 
-For example, to install a nightly build that includes Google-powered LLMs, full modality support, and UI enhancements, use:
+For example, to install a nightly build that includes VertexAI access, full modality support, and UI enhancements, use:
 ```
-pip install langfun[llm-google,mime,ui] --pre
+pip install langfun[vertexai,mime,ui] --pre
 ```
 
 *Disclaimer: this is not an officially supported Google product.*

{langfun-0.1.2.dev202501050804.dist-info → langfun-0.1.2.dev202501060804.dist-info}/RECORD

@@ -73,19 +73,21 @@ langfun/core/eval/v2/progress.py,sha256=azZgssQgNdv3IgjKEaQBuGI5ucFDNbdi02P4z_nQ
 langfun/core/eval/v2/progress_test.py,sha256=YU7VHzmy5knPZwj9vpBN3rQQH2tukj9eKHkuBCI62h8,2540
 langfun/core/eval/v2/progress_tracking.py,sha256=l9fEkz4oP5McpZzf72Ua7PYm3lAWtRru7gRWNf8H0ms,6083
 langfun/core/eval/v2/progress_tracking_test.py,sha256=fouMVJkFJqHjbhQJngGLGCmA9x3n0dU4USI2dY163mg,2291
-langfun/core/eval/v2/reporting.py,sha256=KF4pE2H1qj3mJcgkv_c5YYFjrU-uJCk_-fuu891Olzs,8061
+langfun/core/eval/v2/reporting.py,sha256=obYPQma5k4K3jxTmV8UXigJnQZvWfxUENBwoa-29JSk,8265
 langfun/core/eval/v2/reporting_test.py,sha256=UmYSAQvD3AIXsSyWQ-WD2uLtEISYpmBeoKY5u5Qwc8E,5696
 langfun/core/eval/v2/runners.py,sha256=DKEmSlGXjOXKWFdBhTpLy7tMsBHZHd1Brl3hWIngsSQ,15931
 langfun/core/eval/v2/runners_test.py,sha256=A37fKK2MvAVTiShsg_laluJzJ9AuAQn52k7HPbfD0Ks,11666
-langfun/core/llms/__init__.py,sha256=6mi0IKTNfq6kymZPlPGA2V7YF1xDLrBCPytojeFMMeA,6716
+langfun/core/llms/__init__.py,sha256=hzddKJKI7NsqGXXyV3Q2jhB95A0TTpBIAhKrACICkk0,6519
 langfun/core/llms/anthropic.py,sha256=a5MmnFsBA0CbfvwzXT1v_0fqLRMrhUNdh1tx6469PQ4,14357
 langfun/core/llms/anthropic_test.py,sha256=-2U4kc_pgBM7wqxu8RuxzyHPGww1EAWqKUvN4PW8Btw,8058
 langfun/core/llms/compositional.py,sha256=csW_FLlgL-tpeyCOTVvfUQkMa_zCN5Y2I-YbSNuK27U,2872
 langfun/core/llms/compositional_test.py,sha256=4eTnOer-DncRKGaIJW2ZQQMLnt5r2R0UIx_DYOvGAQo,2027
 langfun/core/llms/fake.py,sha256=gCHBYBLvBCsC78HI1hpoqXCS-p1FMTgY1P1qh_sGBPk,3070
 langfun/core/llms/fake_test.py,sha256=2h13qkwEz_JR0mtUDPxdAhQo7MueXaFSwsD2DIRDW9g,7653
-langfun/core/llms/google_genai.py,sha256=3iAmLMcBXxkfiiI8BN0S6trKCfyfuajCIHIGpnCrtTg,11973
-langfun/core/llms/google_genai_test.py,sha256=zw14sgWmk0P_irHyb7vpPy1WAuLEE0PmyfiFElu03sA,7686
+langfun/core/llms/gemini.py,sha256=tfM4vrt0WnvnrxRhWXZWh7Gp8dYYfMnSbi9uOstkSak,17399
+langfun/core/llms/gemini_test.py,sha256=2ERhYWCJwnfDTQbCaZHFuB1TdWJFrOBS7yyCBInIdQk,6129
+langfun/core/llms/google_genai.py,sha256=p19jTO3D-wNUIxDcLwypgGpTCjvwQxqxCMBoPjtaetM,3645
+langfun/core/llms/google_genai_test.py,sha256=JZf_cbQ4GGGpwiQCLjFJn7V4jxBBqgZhIx91AzbGKVo,1250
 langfun/core/llms/groq.py,sha256=dCnR3eAECEKuKKAAj-PDTs8NRHl6CQPdf57m1f6a79U,10312
 langfun/core/llms/groq_test.py,sha256=GYF_Qtq5S1H1TrKH38t6_lkdroqT7v-joYLDKnmS9e0,5274
 langfun/core/llms/llama_cpp.py,sha256=9tXQntSCDtjTF3bnyJrAPCr4N6wycy5nXYvp9uduygE,2843
@@ -94,12 +96,12 @@ langfun/core/llms/openai.py,sha256=dLDVBB47nJ30XCwjJpAZMc55ZlZXB__PcfcICCRNuXQ,2
 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=EPPswgaTfPZQ_GGa_dWsqWPV9uRjCmIH2Iwgm1YXOqM,15377
-langfun/core/llms/vertexai_test.py,sha256=ffcA5yPecnQy_rhkuYAw_6o1iLW8AR8FgswmHt6aAys,6725
+langfun/core/llms/vertexai.py,sha256=D9pC-zNRTZ4LDZeV_tjP_c8y3JaGQxTOXpp8VIthwCI,5420
+langfun/core/llms/vertexai_test.py,sha256=KqQkJpjZvpqp-JmDK2yTYK_XyLIyizzb4EeTHfUV8sk,1779
 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
-langfun/core/llms/cache/in_memory_test.py,sha256=taebhky84-VbqwSA-96GZp-uBe5INYITVrM-NCsoE7E,10249
+langfun/core/llms/cache/in_memory.py,sha256=i58oiQL28RDsq37dwqgVpC2mBETJjIEFS20yHiV5MKU,5185
+langfun/core/llms/cache/in_memory_test.py,sha256=V2UPeu5co5vUwSkjekCH1B1iLm9qQKPaacvP6VW3GTg,10388
 langfun/core/memories/__init__.py,sha256=HpghfZ-w1NQqzJXBx8Lz0daRhB2rcy2r9Xm491SBhC4,773
 langfun/core/memories/conversation_history.py,sha256=c9amD8hCxGFiZuVAzkP0dOMWSp8L90uvwkOejjuBqO0,1835
 langfun/core/memories/conversation_history_test.py,sha256=AaW8aNoFjxNusanwJDV0r3384Mg0eAweGmPx5DIkM0Y,2052
@@ -146,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.dev202501050804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202501050804.dist-info/METADATA,sha256=kNngwBqJlzR5sCW7VWtWMwZLA2uItkoRB8q1mXZe5ys,8281
-langfun-0.1.2.dev202501050804.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-langfun-0.1.2.dev202501050804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202501050804.dist-info/RECORD,,
+langfun-0.1.2.dev202501060804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202501060804.dist-info/METADATA,sha256=3MbJlnp-TmGrXXb9jhRPLd7g86sm9b2zDXLWIWMkkZA,7941
+langfun-0.1.2.dev202501060804.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
+langfun-0.1.2.dev202501060804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202501060804.dist-info/RECORD,,

{langfun-0.1.2.dev202501050804.dist-info → langfun-0.1.2.dev202501060804.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.6.0)
+Generator: setuptools (75.7.0)
 Root-Is-Purelib: true
 Tag: py3-none-any