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

langfun/core/llms/vertexai_test.py

@@ -20,6 +20,7 @@ from unittest import mock
 from google.auth import exceptions
 import langfun.core as lf
 from langfun.core.llms import vertexai
+import pyglove as pg
 
 
 class VertexAITest(unittest.TestCase):
@@ -28,21 +29,22 @@ class VertexAITest(unittest.TestCase):
   @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
+      _ = vertexai.VertexAIGemini15Pro()._api_initialized
 
     with self.assertRaisesRegex(ValueError, 'Please specify `location`'):
-      _ = vertexai.VertexAIGeminiPro1(project='abc')._api_initialized
+      _ = vertexai.VertexAIGemini15Pro(
+          project='abc', location=None)._api_initialized
 
     self.assertTrue(
-        vertexai.VertexAIGeminiPro1(
+        vertexai.VertexAIGemini15Pro(
             project='abc', location='us-central1'
         )._api_initialized
     )
 
     os.environ['VERTEXAI_PROJECT'] = 'abc'
     os.environ['VERTEXAI_LOCATION'] = 'us-central1'
-    model = vertexai.VertexAIGeminiPro1()
-    self.assertTrue(model.model_id.startswith('VertexAI('))
+    model = vertexai.VertexAIGemini15Pro(location=pg.MISSING_VALUE)
+    self.assertEqual(model.resource_id, 'vertexai://gemini-1.5-pro-002')
     self.assertIn('us-central1', model.api_endpoint)
     self.assertTrue(model._api_initialized)
     self.assertIsNotNone(model.session())
@@ -55,11 +57,18 @@ class VertexAIAnthropicTest(unittest.TestCase):
 
   def test_basics(self):
     with self.assertRaisesRegex(ValueError, 'Please specify `project`'):
-      lm = vertexai.VertexAIClaude3_5_Sonnet_20241022()
+      lm = vertexai.VertexAIClaude35Sonnet_20241022()
       lm('hi')
 
-    model = vertexai.VertexAIClaude3_5_Sonnet_20241022(project='langfun')
-
+    model = vertexai.VertexAIClaude35Sonnet_20241022(project='langfun')
+    self.assertEqual(model.resource_id, 'vertexai://claude-3-5-sonnet-20241022')
+    # Map a Anthropic model back to VertexAI model.
+    self.assertTrue(
+        vertexai.VertexAIAnthropic(
+            'claude-3-5-sonnet-20241022', project='langfun'
+        ).model,
+        'claude-3-5-sonnet-v2@20241022',
+    )
     # NOTE(daiyip): For OSS users, default credentials are not available unless
     # users have already set up their GCP project. Therefore we ignore the
     # exception here.
@@ -99,6 +108,24 @@ class VertexAIAnthropicTest(unittest.TestCase):
         },
     )
 
+  def test_lm_get(self):
+    self.assertIsInstance(
+        lf.LanguageModel.get('gemini-2.0-flash'),
+        vertexai.VertexAIGemini,
+    )
+    self.assertIsInstance(
+        lf.LanguageModel.get('claude-3-5-sonnet-v2@20241022'),
+        vertexai.VertexAIAnthropic,
+    )
+    self.assertIsInstance(
+        lf.LanguageModel.get('llama-3.1-405b-instruct-maas'),
+        vertexai.VertexAILlama,
+    )
+    self.assertIsInstance(
+        lf.LanguageModel.get('mistral-large-2411'),
+        vertexai.VertexAIMistral,
+    )
+
 
 if __name__ == '__main__':
   unittest.main()

{langfun-0.1.2.dev202502110804.dist-info → langfun-0.1.2.dev202502120804.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: langfun
-Version: 0.1.2.dev202502110804
+Version: 0.1.2.dev202502120804
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.2.dev202502110804.dist-info → langfun-0.1.2.dev202502120804.dist-info}/RECORD

@@ -1,5 +1,5 @@
 langfun/__init__.py,sha256=fhfPXpHN7GoGqixpFfqhQkYxFs_siP_LhbjZhd3lhio,2497
-langfun/core/__init__.py,sha256=HOV4RFypDBmQhXouUA2UEassOuYTol0und3EtSa3D5c,4543
+langfun/core/__init__.py,sha256=S6YS15fLy5M3L3qaSa3XhDgku-klUm1LJxFPB-Wl-tY,4596
 langfun/core/component.py,sha256=g1kQM0bryYYYWVDrSMnHfc74wIBbpfe5_B3s-UIP5GE,3028
 langfun/core/component_test.py,sha256=0CxTgjAud3aj8wBauFhG2FHDqrxCTl4OI4gzQTad-40,9254
 langfun/core/concurrent.py,sha256=zY-pXqlGqss_GI20tM1gXvyW8QepVPUuFNmutcIdhbI,32760
@@ -8,8 +8,8 @@ langfun/core/console.py,sha256=V_mOiFi9oGh8gLsUeR56pdFDkuvYOpvQt7DY1KUTWTA,2535
 langfun/core/console_test.py,sha256=pBOcuNMJdVELywvroptfcRtJMsegMm3wSlHAL2TdxVk,1679
 langfun/core/langfunc.py,sha256=G50YgoVZ0y1GFw2ev41MlOqr6qa8YakbvNC0h_E0PiA,11140
 langfun/core/langfunc_test.py,sha256=fKIAqcSNI_7M6nwoZW77HEam8Oa6vcWhsCNgVJanzb4,8822
-langfun/core/language_model.py,sha256=Hfs-SIxiO2rK7BvPmnOuYTrT4jbmFiV31NMmI-mkPzY,36635
-langfun/core/language_model_test.py,sha256=jic5w8w9Mie43e6zwjXDwLV-8Cqc1HxHP0GKkXzeL8I,33666
+langfun/core/language_model.py,sha256=RWTyPXGxcP0hlTuGgKAcdyrYSDAAAbErZb4U7ypRXqg,45669
+langfun/core/language_model_test.py,sha256=iA5uo7rIj2jAtCYzMzhyNg1fWqE2Onn60bOO58q72C0,36454
 langfun/core/logging.py,sha256=W3mLEMXdo210Q5OX3a1ZTc4nU-xMy73-IfNKnsA-RFo,8051
 langfun/core/logging_test.py,sha256=N7-YvSXC8zvnr2SNwWHOykn1CFmqvIuTLDgn41Ku9JU,6642
 langfun/core/memory.py,sha256=f-asN1F7Vehgdn_fK84v73GrEUOxRtaW934keutTKjk,2416
@@ -73,31 +73,31 @@ langfun/core/eval/v2/reporting.py,sha256=7rL9LLmGYnQ5HIjqRqsOMkUlBl4BmFPEL6Vlofq
 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=qoFPJhueF6tuZSfi5QyiXXDhFSqrHQxid1bjb9dJm3Y,8178
-langfun/core/llms/anthropic.py,sha256=z_DWDpR1VKNzv6wq-9CXLzWdqCDXRKuVFacJNpgBqAs,10826
-langfun/core/llms/anthropic_test.py,sha256=zZ2eSP8hhVv-RDSWxT7wX-NS5DfGfQmCjS9P0pusAHM,6556
+langfun/core/llms/__init__.py,sha256=gph3HngayM-WExoaoKM-S6Ke7_MRtvEojKoaf9LusJ0,7740
+langfun/core/llms/anthropic.py,sha256=nov4_EoqDL0_DAW93Eb5p-fHE-AfvzA8Q7PtYfcQdMg,20548
+langfun/core/llms/anthropic_test.py,sha256=SSK7OTx3gMYE1NMAi_PqQqeNsCkZAcVJvl_OCEOhyzk,7145
 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/deepseek.py,sha256=9EbuNrngd5BwpsvsEfkW2XQdq3K23S3-jlI8RbMNZF4,4141
-langfun/core/llms/deepseek_test.py,sha256=dS72i52bwMpCN4dJDvpJI59AnNChpwxS5eYYFrhGh90,1843
-langfun/core/llms/fake.py,sha256=gCHBYBLvBCsC78HI1hpoqXCS-p1FMTgY1P1qh_sGBPk,3070
+langfun/core/llms/deepseek.py,sha256=jvTxdXPr-vH6HNakn_Ootx1heDg8Fen2FUkUW36bpCs,5247
+langfun/core/llms/deepseek_test.py,sha256=DvROWPlDuow5E1lfoSkhyGt_ELA19JoQoDsTnRgDtTg,1847
+langfun/core/llms/fake.py,sha256=xmgCkk9y0I4x0IT32SZ9_OT27aLadXH8PRiYNo5VTd4,3265
 langfun/core/llms/fake_test.py,sha256=2h13qkwEz_JR0mtUDPxdAhQo7MueXaFSwsD2DIRDW9g,7653
-langfun/core/llms/gemini.py,sha256=7qxdj9tDsG68dstkXAcvnqAMYv0tYkLW8bVLjxLNYqk,19752
-langfun/core/llms/gemini_test.py,sha256=2ERhYWCJwnfDTQbCaZHFuB1TdWJFrOBS7yyCBInIdQk,6129
-langfun/core/llms/google_genai.py,sha256=WbVZAIDhh7mKy15LZMXOEniQ8Obd-hjiy2KviUCfmXQ,4261
-langfun/core/llms/google_genai_test.py,sha256=JZf_cbQ4GGGpwiQCLjFJn7V4jxBBqgZhIx91AzbGKVo,1250
-langfun/core/llms/groq.py,sha256=oGxyyCi5TtMVu2POdO8pXS7pK4Es56FtqjghZIGYopc,7579
-langfun/core/llms/groq_test.py,sha256=9aOD3nO4dEoH57B-5iOr5DQjG0vyv1jzFtAe7HfoiNg,1963
+langfun/core/llms/gemini.py,sha256=95PD90TuR3N8m9L50LHIuCByPTMi2qHaYaCwy7d0qYI,22511
+langfun/core/llms/gemini_test.py,sha256=ZzQmgB5bxJdtD2nVOz0WQ1zjNCea_jM5Yuqs8kchnGw,6621
+langfun/core/llms/google_genai.py,sha256=mzUdWAShRzz2k_foUo0eIdNOyHs_xjEtbewqYyuD0k0,4768
+langfun/core/llms/google_genai_test.py,sha256=NKNtpebArQ9ZR7Qsnhd2prFIpMjleojy6o6VMXkJ1zY,1502
+langfun/core/llms/groq.py,sha256=S9V10kFo3cgX89qPgt_umq-SpRnxEDLTt_hJmpERfbo,12066
+langfun/core/llms/groq_test.py,sha256=P4EgexCqsh4K2x11w0UL_vz-YYNaPdQU0WsDAdnTRQ8,2045
 langfun/core/llms/llama_cpp.py,sha256=Z7P3gc4xeIjc2bX0Ey1y5EUYJVMnMa2Q67PZ9iye9sE,1409
 langfun/core/llms/llama_cpp_test.py,sha256=wfTO7nmUwL65U2kK9P9fcMt92JjNDuVia4G1E7znf_4,1086
-langfun/core/llms/openai.py,sha256=bIJae2UDO_uwoZGZRYP58GM4FdP4w9qjQhBIYJyGH34,17089
-langfun/core/llms/openai_compatible.py,sha256=MF9JCc0uTPkIK95uvQTMIBXk4z_xKQMZqQHFaVmXwcA,5898
-langfun/core/llms/openai_compatible_test.py,sha256=0uFYhCiuHo2Wrlgj16-GRG6rW8P6EaHCUguL3jS0QJ8,16708
-langfun/core/llms/openai_test.py,sha256=m85YjGCvWvV5ZYagjC0FqI0FcqyCEVCbUUs8Wm3iUrc,2475
-langfun/core/llms/rest.py,sha256=46aGJghpQunwc3ju47V0N-1a1X4DydhwRBtRP9LBH2E,3828
+langfun/core/llms/openai.py,sha256=_ptVgolpK2atVn6c_7TgO3wHyOvagWtfDMCv2KXv5LY,38553
+langfun/core/llms/openai_compatible.py,sha256=lhq9_XLVUoGa5AIPXTbqGqWvm2P3dkHFSUQ7XnyLozw,5442
+langfun/core/llms/openai_compatible_test.py,sha256=I5WWL3lRo-WXnSoUKLkIEjXfwjoiHRX9o0dj0j09jsk,17024
+langfun/core/llms/openai_test.py,sha256=gwuO6aoa296iM2welWV9ua4KF8gEVGsEPakgbtkWkFQ,2687
+langfun/core/llms/rest.py,sha256=xdR4ar4y7YkeZTs_BHUyNOdhqoghztMcqyz1f9kTXH8,4054
 langfun/core/llms/rest_test.py,sha256=zWGiI08f9gXsoQPJS9TlX1zD2uQLrJUB-1VpAJXRHfs,3475
-langfun/core/llms/vertexai.py,sha256=CYi_KbsNcMW4O4NR_zvBF0RykvuF3ACO0W4-uYV-1cE,16516
-langfun/core/llms/vertexai_test.py,sha256=9LYv2equht9Vq-kwzdSZerENZjV9LYb4IIplSnfV5Os,3334
+langfun/core/llms/vertexai.py,sha256=UBOt58FHlzwM1EwYBt2ADk7K2AFg7HMdI2aEwQZLAc4,17473
+langfun/core/llms/vertexai_test.py,sha256=dOprP_uLNmXHYxMoX_hMPMsjKR-e_B5nKHjhlMCQoOQ,4252
 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=i58oiQL28RDsq37dwqgVpC2mBETJjIEFS20yHiV5MKU,5185
@@ -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.dev202502110804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202502110804.dist-info/METADATA,sha256=vNl3CaNv4v8tHGItFTnZlRgk0hM9cmZBxmO0A8RRdIg,8172
-langfun-0.1.2.dev202502110804.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-langfun-0.1.2.dev202502110804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202502110804.dist-info/RECORD,,
+langfun-0.1.2.dev202502120804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202502120804.dist-info/METADATA,sha256=qY1L92qXnvb7WBDdQ_xQTWQvv4wZq4-TqqhHp1i9wHs,8172
+langfun-0.1.2.dev202502120804.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+langfun-0.1.2.dev202502120804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202502120804.dist-info/RECORD,,