langfun 0.1.2.dev202505110803__tar.gz → 0.1.2.dev202505130804__tar.gz

{langfun-0.1.2.dev202505110803 → langfun-0.1.2.dev202505130804}/PKG-INFO

@@ -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 → langfun-0.1.2.dev202505130804}/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-0.1.2.dev202505110803 → langfun-0.1.2.dev202505130804}/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 → langfun-0.1.2.dev202505130804}/langfun.egg-info/PKG-INFO

@@ -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