langfun 0.1.2.dev202509300805__py3-none-any.whl → 0.1.2.dev202510010805__py3-none-any.whl

langfun/core/llms/gemini.py

@@ -752,6 +752,11 @@ class Gemini(rest.REST):
         prompt.as_format('gemini', chunk_preprocessor=modality_conversion)
     )
     request['contents'] = contents
+    # Users could use `metadata_gemini_tools` to pass Gemini tools. For example,
+    # for enabling Search Grounding, users could pass:
+    # metadata_gemini_tools=[{'google_search': {}}]
+    if tools := prompt.metadata.get('gemini_tools'):
+      request['tools'] = tools
     return request
 
   def _generation_config(

langfun/core/llms/rest.py

@@ -98,7 +98,9 @@ class REST(lf.LanguageModel):
       raise lf.TemporaryLMError(str(e)) from e
     except (
         requests.exceptions.ConnectionError,
+        requests.exceptions.ChunkedEncodingError,
         ConnectionError,
+        ConnectionResetError,
     ) as e:
       error_message = str(e)
       if 'REJECTED_CLIENT_THROTTLED' in error_message:
@@ -107,6 +109,8 @@ class REST(lf.LanguageModel):
         raise lf.TemporaryLMError(error_message) from e
       if 'UNREACHABLE_ERROR' in error_message:
         raise lf.TemporaryLMError(error_message) from e
+      if 'Connection reset by peer' in error_message:
+        raise lf.TemporaryLMError(error_message) from e
       raise lf.LMError(error_message) from e
 
   def _error(self, status_code: int, content: str) -> lf.LMError:

langfun/env/base_environment.py

@@ -359,6 +359,7 @@ class BaseEnvironment(interface.Environment):
       return
 
     self._set_status(self.Status.SHUTTING_DOWN)
+    self.on_shutting_down()
 
     shutting_down_time = time.time()
     try:
@@ -419,7 +420,6 @@ class BaseEnvironment(interface.Environment):
             )
             # Append is atomic and does not require locking.
             self._sandbox_pool.append(sandbox)
-            self._offline_start_time = None
             return sandbox
           except (
               interface.EnvironmentError, interface.SandboxStateError
@@ -432,18 +432,28 @@ class BaseEnvironment(interface.Environment):
       set_acquired: bool = False,
   ) -> base_sandbox.BaseSandbox:
     """Brings up a new sandbox."""
-    sandbox = self._create_sandbox(
-        sandbox_id=sandbox_id,
-        reusable=self.enable_pooling,
-        proactive_session_setup=self.proactive_session_setup,
-        keepalive_interval=self.sandbox_keepalive_interval,
-    )
-    for handler in self.event_handlers:
-      sandbox.add_event_handler(handler)
-    sandbox.start()
-    if set_acquired:
-      sandbox.set_acquired()
-    return sandbox
+    env_error = None
+    try:
+      sandbox = self._create_sandbox(
+          sandbox_id=sandbox_id,
+          reusable=self.enable_pooling,
+          proactive_session_setup=self.proactive_session_setup,
+          keepalive_interval=self.sandbox_keepalive_interval,
+      )
+      for handler in self.event_handlers:
+        sandbox.add_event_handler(handler)
+      sandbox.start()
+      if set_acquired:
+        sandbox.set_acquired()
+      return sandbox
+    except (interface.EnvironmentError, interface.SandboxStateError) as e:
+      env_error = e
+      raise e
+    finally:
+      if env_error is None:
+        self._offline_start_time = None
+      elif self._offline_start_time is None:
+        self._offline_start_time = time.time()
 
   def _bring_up_sandbox_with_retry(
       self,
@@ -486,8 +496,6 @@ class BaseEnvironment(interface.Environment):
       shutdown_env_upon_outage: bool = True
   ):
     """Raises error if the grace period has passed or wait for retry."""
-    if self._offline_start_time is None:
-      self._offline_start_time = time.time()
     if self.offline_duration > self.outage_grace_period:
       if shutdown_env_upon_outage:
         self.shutdown()
@@ -505,39 +513,47 @@ class BaseEnvironment(interface.Environment):
     """Housekeeping loop for the environment."""
     while self._status not in (self.Status.SHUTTING_DOWN, self.Status.OFFLINE):
       housekeep_start_time = time.time()
+
+      is_online = True
       dead_pool_indices = [
           i for i, s in enumerate(self._sandbox_pool)
           if s.status == interface.Sandbox.Status.OFFLINE
       ]
-      if dead_pool_indices and not self._replace_dead_sandboxes(
-          dead_pool_indices
-      ):
+      replaced_indices = []
+
+      if dead_pool_indices:
+        replaced_indices = self._replace_dead_sandboxes(dead_pool_indices)
+        if not replaced_indices:
+          is_online = self.offline_duration < self.outage_grace_period
+
+      self._housekeep_counter += 1
+      duration = time.time() - housekeep_start_time
+      kwargs = dict(
+          dead_pool_indices=dead_pool_indices,
+          replaced_indices=replaced_indices,
+          offline_duration=self.offline_duration,
+      )
+      if is_online:
+        self.on_housekeep(duration, **kwargs)
+        time.sleep(self.housekeep_interval)
+      else:
         self.shutdown()
-        self._housekeep_counter += 1
         self.on_housekeep(
-            time.time() - housekeep_start_time,
+            duration,
             interface.EnvironmentOutageError(
-                environment=self,
-                offline_duration=self.offline_duration
-            )
+                environment=self, offline_duration=self.offline_duration
+            ),
+            **kwargs
         )
-        break
-      self._housekeep_counter += 1
-      self.on_housekeep(time.time() - housekeep_start_time)
-      time.sleep(self.housekeep_interval)
 
-  def _replace_dead_sandboxes(self, dead_pool_indices: list[int]) -> bool:
+  def _replace_dead_sandboxes(self, dead_pool_indices: list[int]) -> list[int]:
     """Replaces a dead sandbox with a new one.
 
     Args:
       dead_pool_indices: The indices of the dead sandboxes to replace.
 
     Returns:
-      Whether all of the dead sandboxes are replaced successfully.
-
-    Raises:
-      interface.EnvironmentOutageError: If the XBox sandboxes cannot be created
-        within the wait time specified by `xbox_outage_grace_period`.
+      Successfully replaced indices.
     """
     pg.logging.warning(
         '[%s]: %s maintenance: '
@@ -548,21 +564,31 @@ class BaseEnvironment(interface.Environment):
     )
     def _replace(i: int):
       generation = int(self._sandbox_pool[i].id.sandbox_id.split(':')[1])
-      self._sandbox_pool[i] = self._bring_up_sandbox_with_retry(
-          f'{i}:{generation + 1}', shutdown_env_upon_outage=False
-      )
+      self._sandbox_pool[i] = self._bring_up_sandbox(f'{i}:{generation + 1}')
 
     # TODO(daiyip): Consider to loose the condition to allow some dead
     # sandboxes to be replaced successfully.
-    return not any([
-        error for _, _, error in lf.concurrent_map(
-            _replace, dead_pool_indices,
-            max_workers=min(
-                self.pool_operation_max_parallelism,
-                len(dead_pool_indices)
-            ),
-        )
-    ])
+    replaced_indices = []
+    for index, _, error in lf.concurrent_map(
+        _replace, dead_pool_indices,
+        max_workers=min(
+            self.pool_operation_max_parallelism,
+            len(dead_pool_indices)
+        ),
+    ):
+      if error is None:
+        replaced_indices.append(index)
+
+    pg.logging.warning(
+        '[%s]: %s maintenance: '
+        '%d/%d dead sandbox(es) have been replaced with new ones. (slots=%s)',
+        self.id,
+        self.__class__.__name__,
+        len(replaced_indices),
+        len(dead_pool_indices),
+        replaced_indices
+    )
+    return replaced_indices
 
   #
   # Event handlers subclasses can override.
@@ -584,12 +610,20 @@ class BaseEnvironment(interface.Environment):
   def on_housekeep(
       self,
       duration: float,
-      error: BaseException | None = None
+      error: BaseException | None = None,
+      **kwargs
   ) -> None:
     """Called when the environment finishes a round of housekeeping."""
     housekeep_counter = self.housekeep_counter
     for handler in self.event_handlers:
-      handler.on_environment_housekeep(self, housekeep_counter, duration, error)
+      handler.on_environment_housekeep(
+          self, housekeep_counter, duration, error, **kwargs
+      )
+
+  def on_shutting_down(self) -> None:
+    """Called when the environment is shutting down."""
+    for handler in self.event_handlers:
+      handler.on_environment_shutting_down(self, self.offline_duration)
 
   def on_shutdown(
       self,

langfun/env/base_feature.py

@@ -186,11 +186,12 @@ class BaseFeature(interface.Feature):
   def on_housekeep(
       self,
       duration: float,
-      error: BaseException | None = None
+      error: BaseException | None = None,
+      **kwargs
   ) -> None:
     """Called when the feature has done housekeeping."""
     self.sandbox.on_feature_housekeep(
-        self, self._housekeep_counter, duration, error
+        self, self._housekeep_counter, duration, error, **kwargs
     )
 
   def on_setup_session(

langfun/env/base_sandbox.py

@@ -706,7 +706,6 @@ class BaseSandbox(interface.Sandbox):
     while self._status not in (self.Status.SHUTTING_DOWN, self.Status.OFFLINE):
       housekeep_start = time.time()
       if self.keepalive_interval is not None:
-
         if time.time() - last_ping > self.keepalive_interval:
           try:
             self.ping()
@@ -796,13 +795,14 @@ class BaseSandbox(interface.Sandbox):
   def on_housekeep(
       self,
       duration: float,
-      error: BaseException | None = None
+      error: BaseException | None = None,
+      **kwargs
   ) -> None:
     """Called when the sandbox finishes a round of housekeeping."""
     counter = self._housekeep_counter
     for handler in self._event_handlers:
       handler.on_sandbox_housekeep(
-          self.environment, self, counter, duration, error
+          self.environment, self, counter, duration, error, **kwargs
       )
 
   def on_feature_setup(

langfun/env/base_test.py

@@ -280,7 +280,11 @@ class EnvironmentTests(unittest.TestCase):
         with self.assertRaises(interface.SandboxStateError):
           sb.shell('bad command', raise_error=interface.SandboxStateError)
       self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
-      env.wait_for_housekeeping()
+      sb_offline_time = time.time()
+      while time.time() - sb_offline_time < 10:
+        if not env.is_online:
+          break
+        time.sleep(0.5)
       self.assertFalse(env.is_online)
 
 
@@ -1177,6 +1181,7 @@ class SandboxActivityTests(unittest.TestCase):
         sb.shell('test_feature')
         sb.remove_event_handler(event_handler)
       events = list(event_handler.logs)
+      sb.wait_until_not(interface.Sandbox.Status.SETTING_UP)
       self.assertGreater(len(events), 0)
       with env.sandbox('session2') as sb:
         sb.shell('test_feature')

langfun/env/event_handlers/base.py

@@ -76,7 +76,16 @@ class _FeatureEventHandler:
       duration: float,
       error: BaseException | None
   ) -> None:
-    """Called when a sandbox feature is setup."""
+    """Called when a sandbox feature is setup.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      feature: The feature.
+      duration: The feature setup duration in seconds.
+      error: The error happened during the feature setup. If None,
+        the feature setup performed normally.
+    """
 
   def on_feature_teardown(
       self,
@@ -86,7 +95,16 @@ class _FeatureEventHandler:
       duration: float,
       error: BaseException | None
   ) -> None:
-    """Called when a sandbox feature is teardown."""
+    """Called when a sandbox feature is teardown.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      feature: The feature.
+      duration: The feature teardown duration in seconds.
+      error: The error happened during the feature teardown. If None,
+        the feature teardown performed normally.
+    """
 
   def on_feature_teardown_session(
       self,
@@ -97,7 +115,17 @@ class _FeatureEventHandler:
       duration: float,
       error: BaseException | None
   ) -> None:
-    """Called when a feature is teardown with a session."""
+    """Called when a feature is teardown with a session.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      feature: The feature.
+      session_id: The session ID.
+      duration: The feature teardown session duration in seconds.
+      error: The error happened during the feature teardown session. If
+        None, the feature teardown session performed normally.
+    """
 
   def on_feature_setup_session(
       self,
@@ -106,9 +134,19 @@ class _FeatureEventHandler:
       feature: Feature,
       session_id: str | None,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
   ) -> None:
-    """Called when a feature is setup with a session."""
+    """Called when a feature is setup with a session.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      feature: The feature.
+      session_id: The session ID.
+      duration: The feature setup session duration in seconds.
+      error: The error happened during the feature setup session. If
+        None, the feature setup session performed normally.
+    """
 
   def on_feature_housekeep(
       self,
@@ -117,9 +155,21 @@ class _FeatureEventHandler:
       feature: Feature,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs,
   ) -> None:
-    """Called when a sandbox feature is housekeeping."""
+    """Called when a sandbox feature is housekeeping.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      feature: The feature.
+      counter: Zero-based counter of the housekeeping round.
+      duration: The feature housekeeping duration in seconds.
+      error: The error happened during the feature housekeeping. If None, the
+        feature housekeeping normally.
+      **kwargs: Feature-specific properties computed during housekeeping.
+    """
 
 
 class _SandboxEventHandler(_FeatureEventHandler, _SessionEventHandler):
@@ -210,7 +260,8 @@ class _SandboxEventHandler(_FeatureEventHandler, _SessionEventHandler):
       sandbox: Sandbox,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs
   ) -> None:
     """Called when a sandbox finishes a round of housekeeping.
 
@@ -221,6 +272,7 @@ class _SandboxEventHandler(_FeatureEventHandler, _SessionEventHandler):
       duration: The sandbox housekeeping duration in seconds.
       error: The error that caused the sandbox to housekeeping. If None, the
         sandbox housekeeping normally.
+      **kwargs: Sandbox-specific properties computed during housekeeping.
     """
 
 
@@ -254,7 +306,8 @@ class EventHandler(_SandboxEventHandler):
       environment: Environment,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs
   ) -> None:
     """Called when the environment finishes a round of housekeeping.
 
@@ -264,6 +317,19 @@ class EventHandler(_SandboxEventHandler):
       duration: The environment start duration in seconds.
       error: The error that failed the housekeeping. If None, the
         housekeeping succeeded.
+      **kwargs: Environment-specific properties computed during housekeeping.
+    """
+
+  def on_environment_shutting_down(
+      self,
+      environment: Environment,
+      offline_duration: float,
+  ) -> None:
+    """Called when the environment is shutting down.
+
+    Args:
+      environment: The environment.
+      offline_duration: The environment offline duration in seconds.
     """
 
   def on_environment_shutdown(

langfun/env/event_handlers/event_logger.py

@@ -105,7 +105,9 @@ class EventLogger(pg.Object, base.EventHandler):
       error: BaseException | None,
   ) -> str:
     if error is not None:
-      message = f'{message} with error: {error}'
+      message = (
+          f'{message} with error: {pg.utils.ErrorInfo.from_exception(error)}'
+      )
     return message
 
   def _keep(
@@ -133,6 +135,20 @@ class EventLogger(pg.Object, base.EventHandler):
         styles=['bold'],
     )
 
+  def on_environment_shutting_down(
+      self,
+      environment: base.Environment,
+      offline_duration: float,
+  ) -> None:
+    """Called when the environment is shutting down."""
+    self._print(
+        f'[{environment.id}] environment shutting down '
+        f'(offline_duration={offline_duration:.2f} seconds)',
+        error=None,
+        color='green',
+        styles=['bold'],
+    )
+
   def on_environment_start(
       self,
       environment: base.Environment,
@@ -153,13 +169,15 @@ class EventLogger(pg.Object, base.EventHandler):
       environment: base.Environment,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs
   ) -> None:
     """Called when the environment is housekeeping."""
     if self.housekeep_status:
       self._print(
-          f'[{environment.id}] environment housekeeping complete'
-          f'(counter={counter}, duration={duration:.2f} seconds)',
+          f'[{environment.id}] environment housekeeping complete '
+          f'(counter={counter}, duration={duration:.2f} seconds, '
+          f'housekeep_info={kwargs})',
           error=error,
           color='green',
       )
@@ -246,13 +264,15 @@ class EventLogger(pg.Object, base.EventHandler):
       sandbox: base.Sandbox,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs
   ) -> None:
     """Called when a sandbox feature is housekeeping."""
     if self.sandbox_status and self.housekeep_status:
       self._print(
           f'[{sandbox.id}] sandbox housekeeping complete '
-          f'(counter={counter}, duration={duration:.2f} seconds)',
+          f'(counter={counter}, duration={duration:.2f} seconds, '
+          f'housekeep_info={kwargs})',
           error=error,
           color='white',
       )
@@ -334,13 +354,15 @@ class EventLogger(pg.Object, base.EventHandler):
       feature: base.Feature,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs
   ) -> None:
     """Called when a sandbox feature is housekeeping."""
     if self.feature_status and self.housekeep_status:
       self._print(
           f'[{sandbox.id}/<idle>/{feature.name}] feature housekeeping complete '
-          f'(counter={counter}, (duration={duration:.2f} seconds)',
+          f'(counter={counter}, (duration={duration:.2f} seconds, '
+          f'housekeep_info={kwargs})',
           error=error,
           color='white',
       )

langfun/env/event_handlers/event_logger_test.py

@@ -125,6 +125,7 @@ class EventLoggerTest(unittest.TestCase):
         unexpected_substrings=[
             'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment housekeeping',
             'sandbox started',
@@ -143,6 +144,7 @@ class EventLoggerTest(unittest.TestCase):
         expected_substrings=[
             'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment housekeeping',
         ],
@@ -166,6 +168,7 @@ class EventLoggerTest(unittest.TestCase):
         expected_substrings=[
             'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment housekeeping',
             'feature setup complete',
@@ -191,6 +194,7 @@ class EventLoggerTest(unittest.TestCase):
         expected_substrings=[
             'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment housekeeping',
             'environment stats',
@@ -217,6 +221,7 @@ class EventLoggerTest(unittest.TestCase):
         expected_substrings=[
             'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment stats',
             'environment housekeeping',
@@ -243,6 +248,7 @@ class EventLoggerTest(unittest.TestCase):
         expected_substrings=[
             'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment stats',
             'sandbox started',
@@ -269,6 +275,7 @@ class EventLoggerTest(unittest.TestCase):
         expected_substrings=[
             'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment housekeeping',
         ],