langfun 0.1.2.dev202505110803__py3-none-any.whl → 0.1.2.dev202505130804__py3-none-any.whl

langfun/core/language_model.py

@@ -1019,11 +1019,21 @@ class LanguageModel(component.Component):
       ] = RetryableLMError,
   ) -> list[Any]:
     """Helper method for subclasses for implementing _sample."""
+    if self.max_concurrency is None:
+      execute = action
+      executor = None
+      max_workers = len(inputs)
+    else:
+      execute = lambda x: _ConcurrencyControl.get(
+          self.resource_id, self.max_concurrency)(action, x)
+      executor = self.resource_id if len(inputs) > 1 else None
+      max_workers = self.max_concurrency
+
     executed_jobs = concurrent.concurrent_execute(
-        action,
+        execute,
         inputs,
-        executor=self.resource_id if self.max_concurrency else None,
-        max_workers=self.max_concurrency or len(inputs),
+        executor=executor,
+        max_workers=max_workers,
         retry_on_errors=retry_on_errors,
         max_attempts=self.max_attempts,
         retry_interval=self.retry_interval,
@@ -1323,6 +1333,46 @@ class LanguageModel(component.Component):
     return None
 
 
+class _ConcurrencyControl:
+  """Controls the max concurrent LLM calls for a given model."""
+
+  _MODEL_CONCURRENCY: ClassVar[dict[str, '_ConcurrencyControl']] = {}
+
+  def __init__(self, max_concurrency: int):
+    self.max_concurrency = max_concurrency
+    self._concurrency = 0
+
+  @property
+  def concurrency(self) -> int:
+    """Returns the current concurrency."""
+    return self._concurrency
+
+  def __call__(self, fn: Callable[..., Any], *args, **kwargs):
+    """Calls the function with concurrency control."""
+    while self._concurrency >= self.max_concurrency:
+      time.sleep(0.01)
+
+    try:
+      # Increment/decrement is atomic in Python, so we don't need to protect it
+      # with a lock.
+      self._concurrency += 1
+      return fn(*args, **kwargs)
+    finally:
+      self._concurrency -= 1
+
+  @classmethod
+  def get(
+      cls, model_id: str, max_concurrency: int | None = None
+  ) -> '_ConcurrencyControl':
+    """Returns the concurrency control for the given model ID."""
+    control = cls._MODEL_CONCURRENCY.get(model_id, None)
+    if control is None:
+      assert max_concurrency is not None
+      control = cls(max_concurrency)
+      cls._MODEL_CONCURRENCY[model_id] = control
+    return control
+
+
 class UsageSummary(pg.Object, pg.views.HtmlTreeView.Extension):
   """Usage sumary."""
 

langfun/core/llms/gemini.py

@@ -678,9 +678,15 @@ class Gemini(rest.REST):
     return config
 
   def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
+    candidates = json.get('candidates')
+    if candidates is None:
+      raise lf.TemporaryLMError(
+          'No candidates found in response. This is a Gemini API issue that '
+          'happens occasionally, and retrying should fix it. '
+      )
     messages = [
         lf.Message.from_value(candidate['content'], format='gemini')
-        for candidate in json.get('candidates', [])
+        for candidate in candidates
     ]
     usage = json['usageMetadata']
     input_tokens = usage['promptTokenCount']

{langfun-0.1.2.dev202505110803.dist-info → langfun-0.1.2.dev202505130804.dist-info}/METADATA

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

{langfun-0.1.2.dev202505110803.dist-info → langfun-0.1.2.dev202505130804.dist-info}/RECORD

@@ -8,7 +8,7 @@ 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=CDn-gJCa5EnjN7cotAVCfSCbuzddq2o-HzEt7kV8HbY,8882
-langfun/core/language_model.py,sha256=c6Vds9qXuC0BiHeZstxCTTsoU8ZL9K4w4co_hXA3-RI,46336
+langfun/core/language_model.py,sha256=dw81qgdl-zn-qOIbRMSurtAmmpiwr2sSYlSlW19nUw0,47829
 langfun/core/language_model_test.py,sha256=iA5uo7rIj2jAtCYzMzhyNg1fWqE2Onn60bOO58q72C0,36454
 langfun/core/logging.py,sha256=ypjSwGzQpIhadewEKkIH5VWS55rfU4e6mlni41PMHlw,9116
 langfun/core/logging_test.py,sha256=vbVGOQxwMmVSiFfbt2897gUt-8nqDpV64jCAeUG_q5U,6924
@@ -92,7 +92,7 @@ langfun/core/llms/deepseek.py,sha256=jvTxdXPr-vH6HNakn_Ootx1heDg8Fen2FUkUW36bpCs
 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=UGI89WjBTZHBAmWbGh4FUCryw6OW3yuiN7htY-Xv_zM,25358
+langfun/core/llms/gemini.py,sha256=ij3NJHhAcs8-H4H2P4ny78q5L9iAmjQC4cu8nUyoPDA,25591
 langfun/core/llms/gemini_test.py,sha256=Ve9X2Wvwu9wVFHpKZDP-qoM1_hzB4kgt6_HR9wxtNkg,7592
 langfun/core/llms/google_genai.py,sha256=YjvyasibmXYPFr_sxdrIdHDPHfaNzMCiJjWJs0iQyms,5188
 langfun/core/llms/google_genai_test.py,sha256=NKNtpebArQ9ZR7Qsnhd2prFIpMjleojy6o6VMXkJ1zY,1502
@@ -156,8 +156,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.dev202505110803.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202505110803.dist-info/METADATA,sha256=j97AkwZGIIKZJQPvRAgH1zf2fbdLkWQ6DwN-ZFXUtbI,8178
-langfun-0.1.2.dev202505110803.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
-langfun-0.1.2.dev202505110803.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202505110803.dist-info/RECORD,,
+langfun-0.1.2.dev202505130804.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202505130804.dist-info/METADATA,sha256=i-rMyyw0gOG8YgoQmvMEXZUGF9OPziH59QA8SJpLS4g,8178
+langfun-0.1.2.dev202505130804.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+langfun-0.1.2.dev202505130804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202505130804.dist-info/RECORD,,