langfun 0.1.2.dev202502060804__py3-none-any.whl → 0.1.2.dev202502080803__py3-none-any.whl

langfun/core/llms/gemini.py

@@ -90,7 +90,7 @@ SUPPORTED_MODELS_AND_SETTINGS = {
         experimental=True,
         in_service=True,
         supported_modalities=ALL_MODALITIES,
-        rpm_free=2000,
+        rpm_free=10_000,
         tpm_free=40_000_000,
         rpm_paid=0,
         tpm_paid=0,
@@ -106,32 +106,32 @@ SUPPORTED_MODELS_AND_SETTINGS = {
         experimental=True,
         in_service=True,
         supported_modalities=ALL_MODALITIES,
-        rpm_free=2000,
+        rpm_free=10_000,
         tpm_free=40_000_000,
         rpm_paid=0,
         tpm_paid=0,
-        cost_per_1m_input_tokens_up_to_128k=0,
-        cost_per_1m_output_tokens_up_to_128k=0,
-        cost_per_1m_cached_tokens_up_to_128k=0,
-        cost_per_1m_input_tokens_longer_than_128k=0,
-        cost_per_1m_output_tokens_longer_than_128k=0,
-        cost_per_1m_cached_tokens_longer_than_128k=0,
+        cost_per_1m_input_tokens_up_to_128k=0.1,
+        cost_per_1m_output_tokens_up_to_128k=0.4,
+        cost_per_1m_cached_tokens_up_to_128k=0.025,
+        cost_per_1m_input_tokens_longer_than_128k=0.1,
+        cost_per_1m_output_tokens_longer_than_128k=0.4,
+        cost_per_1m_cached_tokens_longer_than_128k=0.025,
     ),
     'gemini-2.0-pro-exp-02-05': pg.Dict(
         latest_update='2025-02-05',
         experimental=True,
         in_service=True,
         supported_modalities=ALL_MODALITIES,
-        rpm_free=1000,
-        tpm_free=40_000_000,
+        rpm_free=10,
+        tpm_free=4_000_000,
         rpm_paid=0,
         tpm_paid=0,
-        cost_per_1m_input_tokens_up_to_128k=0,
-        cost_per_1m_output_tokens_up_to_128k=0,
-        cost_per_1m_cached_tokens_up_to_128k=0,
-        cost_per_1m_input_tokens_longer_than_128k=0,
-        cost_per_1m_output_tokens_longer_than_128k=0,
-        cost_per_1m_cached_tokens_longer_than_128k=0,
+        cost_per_1m_input_tokens_up_to_128k=0.1,
+        cost_per_1m_output_tokens_up_to_128k=0.4,
+        cost_per_1m_cached_tokens_up_to_128k=0.025,
+        cost_per_1m_input_tokens_longer_than_128k=0.1,
+        cost_per_1m_output_tokens_longer_than_128k=0.4,
+        cost_per_1m_cached_tokens_longer_than_128k=0.025,
     ),
     'gemini-2.0-flash-thinking-exp-01-21': pg.Dict(
         latest_update='2024-01-21',

langfun/core/llms/rest.py

@@ -63,15 +63,14 @@ class REST(lf.LanguageModel):
   def _initialize(self) -> None:
     """Initializes the API. Subclasses can override."""
 
-  @functools.cached_property
-  def _session(self) -> requests.Session:
+  def session(self) -> requests.Session:
     assert self._api_initialized
-    s = self._create_session()
+    s = self._session()
     # Placeholder for Google-internal session adapter.
     s.headers.update(self.headers or {})
     return s
 
-  def _create_session(self) -> requests.Session:
+  def _session(self) -> requests.Session:
     """Creates a new session."""
     return requests.Session()
 
@@ -88,12 +87,14 @@ class REST(lf.LanguageModel):
 
   def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
     try:
-      response = self._session.post(
-          self.api_endpoint,
-          json=self.request(prompt, self.sampling_options),
-          timeout=self.timeout,
-      )
-      return self._parse_response(response)
+      with self.session() as session:
+        return self._parse_response(
+            session.post(
+                self.api_endpoint,
+                json=self.request(prompt, self.sampling_options),
+                timeout=self.timeout,
+            )
+        )
     except (requests.exceptions.ReadTimeout,
             requests.exceptions.ConnectTimeout) as e:
       raise lf.TemporaryLMError(str(e)) from e

langfun/core/llms/vertexai.py

@@ -118,7 +118,7 @@ class VertexAI(rest.REST):
     """Returns a string to identify the model."""
     return f'VertexAI({self.model})'
 
-  def _create_session(self):
+  def _session(self):
     assert self._credentials is not None
     assert auth_requests is not None
     return auth_requests.AuthorizedSession(self._credentials)
@@ -198,37 +198,43 @@ class VertexAIGeminiExp_20241114(VertexAIGemini):  # pylint: disable=invalid-nam
 class VertexAIGeminiPro1_5(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
-  model = 'gemini-1.5-pro-latest'
+  model = 'gemini-1.5-pro-002'
+  location = 'us-central1'
 
 
 class VertexAIGeminiPro1_5_002(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-002'
+  location = 'us-central1'
 
 
 class VertexAIGeminiPro1_5_001(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-001'
+  location = 'us-central1'
 
 
 class VertexAIGeminiFlash1_5(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash'
+  location = 'us-central1'
 
 
 class VertexAIGeminiFlash1_5_002(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-002'
+  location = 'us-central1'
 
 
 class VertexAIGeminiFlash1_5_001(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-001'
+  location = 'us-central1'
 
 
 class VertexAIGeminiPro1(VertexAIGemini):  # pylint: disable=invalid-name

langfun/core/llms/vertexai_test.py

@@ -45,7 +45,7 @@ class VertexAITest(unittest.TestCase):
     self.assertTrue(model.model_id.startswith('VertexAI('))
     self.assertIn('us-central1', model.api_endpoint)
     self.assertTrue(model._api_initialized)
-    self.assertIsNotNone(model._session)
+    self.assertIsNotNone(model.session())
     del os.environ['VERTEXAI_PROJECT']
     del os.environ['VERTEXAI_LOCATION']
 

{langfun-0.1.2.dev202502060804.dist-info → langfun-0.1.2.dev202502080803.dist-info}/METADATA

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

{langfun-0.1.2.dev202502060804.dist-info → langfun-0.1.2.dev202502080803.dist-info}/RECORD

@@ -82,7 +82,7 @@ langfun/core/llms/deepseek.py,sha256=9EbuNrngd5BwpsvsEfkW2XQdq3K23S3-jlI8RbMNZF4
 langfun/core/llms/deepseek_test.py,sha256=dS72i52bwMpCN4dJDvpJI59AnNChpwxS5eYYFrhGh90,1843
 langfun/core/llms/fake.py,sha256=gCHBYBLvBCsC78HI1hpoqXCS-p1FMTgY1P1qh_sGBPk,3070
 langfun/core/llms/fake_test.py,sha256=2h13qkwEz_JR0mtUDPxdAhQo7MueXaFSwsD2DIRDW9g,7653
-langfun/core/llms/gemini.py,sha256=oYcpQ3I7-c1MeSC7mj9iKtAZTdbWKI9Fzp84drvR5QY,19703
+langfun/core/llms/gemini.py,sha256=jSwj9TQMDHkByr09hftq3CQozifjqwPwHfn-A64MIBg,19736
 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
@@ -94,10 +94,10 @@ langfun/core/llms/openai.py,sha256=bIJae2UDO_uwoZGZRYP58GM4FdP4w9qjQhBIYJyGH34,1
 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=2N-zg9wuB-XeQvISc-xWOBCfBVvjMc4p3LO_rD23z7c,3818
+langfun/core/llms/rest.py,sha256=46aGJghpQunwc3ju47V0N-1a1X4DydhwRBtRP9LBH2E,3828
 langfun/core/llms/rest_test.py,sha256=zWGiI08f9gXsoQPJS9TlX1zD2uQLrJUB-1VpAJXRHfs,3475
-langfun/core/llms/vertexai.py,sha256=zLiHHrm5gpE6DgC5k0cxugD2qui9MF8bhg-4aJDXB24,16364
-langfun/core/llms/vertexai_test.py,sha256=6eLQOyeL5iGZOIWb39sFcf1TgYD_6TBGYdMO4UIvhf4,3333
+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/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.dev202502060804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202502060804.dist-info/METADATA,sha256=tw5p6p7Pp2I6CbNWAJx1EcB-umB-X10YeiJnw8LS65o,8172
-langfun-0.1.2.dev202502060804.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-langfun-0.1.2.dev202502060804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202502060804.dist-info/RECORD,,
+langfun-0.1.2.dev202502080803.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202502080803.dist-info/METADATA,sha256=HwxvNwOrL6oliIMpZ9vm2dAU8E7BY5FGU6-8JpOmbkg,8172
+langfun-0.1.2.dev202502080803.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+langfun-0.1.2.dev202502080803.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202502080803.dist-info/RECORD,,