langfun 0.1.2.dev202509290805__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

@@ -221,7 +221,7 @@ class BaseEnvironment(interface.Environment):
       self._sandbox_pool = [None] * self.min_pool_size
       for i, sandbox, _ in lf.concurrent_map(
           lambda i: self._bring_up_sandbox_with_retry(
-              sandbox_id=str(i), shutdown_env_upon_outage=False
+              sandbox_id=f'{i}:0', shutdown_env_upon_outage=False
           ),
           range(self.min_pool_size),
           silence_on_errors=None,
@@ -335,6 +335,7 @@ class BaseEnvironment(interface.Environment):
         f'it is in {self._status.value!r} status.'
     )
 
+    self.on_starting()
     starting_time = time.time()
     try:
       self._start()
@@ -358,12 +359,14 @@ class BaseEnvironment(interface.Environment):
       return
 
     self._set_status(self.Status.SHUTTING_DOWN)
+    self.on_shutting_down()
 
+    shutting_down_time = time.time()
     try:
       self._shutdown()
-      self.on_shutdown()
+      self.on_shutdown(duration=time.time() - shutting_down_time)
     except BaseException as e:  # pylint: disable=broad-except
-      self.on_shutdown(error=e)
+      self.on_shutdown(duration=time.time() - shutting_down_time, error=e)
       raise e
 
   #
@@ -412,11 +415,11 @@ class BaseEnvironment(interface.Environment):
         else:
           try:
             sandbox = self._bring_up_sandbox(
-                sandbox_id=str(self._increment_sandbox_id()), set_acquired=True,
+                sandbox_id=f'{self._increment_sandbox_id()}:0',
+                set_acquired=True,
             )
             # 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
@@ -429,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,
@@ -483,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()
@@ -502,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: '
@@ -544,26 +563,42 @@ class BaseEnvironment(interface.Environment):
         len(dead_pool_indices),
     )
     def _replace(i: int):
-      self._sandbox_pool[i] = self._bring_up_sandbox_with_retry(
-          str(i), shutdown_env_upon_outage=False
-      )
+      generation = int(self._sandbox_pool[i].id.sandbox_id.split(':')[1])
+      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.
   #
 
+  def on_starting(self) -> None:
+    """Called when the environment is getting started."""
+    for handler in self.event_handlers:
+      handler.on_environment_starting(self)
+
   def on_start(
       self,
       duration: float, error: BaseException | None = None
@@ -575,15 +610,26 @@ 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_shutdown(self, error: BaseException | None = None) -> None:
+  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,
+      duration: float,
+      error: BaseException | None = None) -> None:
     """Called when the environment is shutdown."""
     lifetime = (time.time() - self.start_time) if self.start_time else 0.0
     for handler in self.event_handlers:
-      handler.on_environment_shutdown(self, lifetime, error)
+      handler.on_environment_shutdown(self, duration, lifetime, error)

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

@@ -418,6 +418,7 @@ class BaseSandbox(interface.Sandbox):
       self.end_session(shutdown_sandbox=True)
       return
 
+    shutting_down_time = time.time()
     self._set_status(interface.Sandbox.Status.SHUTTING_DOWN)
 
     if (self._housekeep_thread is not None
@@ -429,7 +430,10 @@ class BaseSandbox(interface.Sandbox):
     try:
       self._shutdown()
       self._set_status(interface.Sandbox.Status.OFFLINE)
-      self.on_shutdown(teardown_error)
+      self.on_shutdown(
+          duration=time.time() - shutting_down_time,
+          error=teardown_error
+      )
       shutdown_error = None
     except BaseException as e:  # pylint: disable=broad-except
       shutdown_error = e
@@ -439,7 +443,10 @@ class BaseSandbox(interface.Sandbox):
           '[%s]: Sandbox shutdown with error: %s',
           self.id, e
       )
-      self.on_shutdown(teardown_error or shutdown_error)
+      self.on_shutdown(
+          duration=time.time() - shutting_down_time,
+          error=teardown_error or shutdown_error
+      )
 
     # We raise non-state errors to the user following timely order, so the user
     # code could be surfaced and handled properly.
@@ -591,6 +598,7 @@ class BaseSandbox(interface.Sandbox):
     # Set sandbox status to EXITING_SESSION to avoid re-entry.
     self._set_status(self.Status.EXITING_SESSION)
     shutdown_sandbox = shutdown_sandbox or not self.reusable
+    ending_time = time.time()
 
     # Teardown features for the current session.
     end_session_error = self._end_session()
@@ -617,7 +625,9 @@ class BaseSandbox(interface.Sandbox):
             self.shutdown()
 
         # End session before setting up the next session.
-        self.on_session_end(previous_session_id)
+        self.on_session_end(
+            previous_session_id, duration=time.time() - ending_time
+        )
 
         # Mark the sandbox as setting up to prevent it from being acquired by
         # other threads.
@@ -628,7 +638,9 @@ class BaseSandbox(interface.Sandbox):
         threading.Thread(target=_setup_next_session).start()
       else:
         # End session before reporting sandbox status change.
-        self.on_session_end(previous_session_id)
+        self.on_session_end(
+            previous_session_id, duration=time.time() - ending_time
+        )
 
         # If shutdown is requested, mark the sandbox as acquired to prevent it
         # from being acquired by other threads.
@@ -639,7 +651,11 @@ class BaseSandbox(interface.Sandbox):
 
     # Otherwise, shutdown the sandbox.
     else:
-      self.on_session_end(previous_session_id, self.state_errors[0])
+      self.on_session_end(
+          previous_session_id,
+          duration=time.time() - ending_time,
+          error=self.state_errors[0]
+      )
       self._set_status(interface.Sandbox.Status.ACQUIRED)
       shutdown_sandbox = True
 
@@ -690,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()
@@ -762,25 +777,32 @@ class BaseSandbox(interface.Sandbox):
           time.time() - self._status_start_time
       )
 
-  def on_shutdown(self, error: BaseException | None = None) -> None:
+  def on_shutdown(
+      self,
+      duration: float,
+      error: BaseException | None = None
+  ) -> None:
     """Called when the sandbox is shutdown."""
     if self._start_time is None:
       lifetime = 0.0
     else:
       lifetime = time.time() - self._start_time
     for handler in self._event_handlers:
-      handler.on_sandbox_shutdown(self.environment, self, lifetime, error)
+      handler.on_sandbox_shutdown(
+          self.environment, self, duration, lifetime, error
+      )
 
   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(
@@ -880,13 +902,14 @@ class BaseSandbox(interface.Sandbox):
   def on_session_end(
       self,
       session_id: str,
+      duration: float,
       error: BaseException | None = None
   ) -> None:
     """Called when the user session ends."""
     lifetime = time.time() - self._session_start_time
     for handler in self._event_handlers:
       handler.on_session_end(
-          self.environment, self, session_id, lifetime, error
+          self.environment, self, session_id, duration, lifetime, error
       )