langfun 0.1.2.dev202511010804__py3-none-any.whl → 0.1.2.dev202511030805__py3-none-any.whl

langfun/env/event_handlers/chain.py

@@ -84,43 +84,75 @@ class EventHandlerChain(interface.EventHandler):
 
   def on_sandbox_start(
       self,
-      environment: interface.Environment,
       sandbox: interface.Sandbox,
       duration: float,
       error: BaseException | None
   ) -> None:
     for handler in self._handlers:
-      handler.on_sandbox_start(environment, sandbox, duration, error)
+      handler.on_sandbox_start(sandbox, duration, error)
 
   def on_sandbox_status_change(
       self,
-      environment: interface.Environment,
       sandbox: interface.Sandbox,
       old_status: interface.Sandbox.Status,
       new_status: interface.Sandbox.Status,
       span: float
   ) -> None:
     for handler in self._handlers:
-      handler.on_sandbox_status_change(
-          environment, sandbox, old_status, new_status, span
-      )
+      handler.on_sandbox_status_change(sandbox, old_status, new_status, span)
 
   def on_sandbox_shutdown(
       self,
-      environment: interface.Environment,
       sandbox: interface.Sandbox,
       duration: float,
       lifetime: float,
       error: BaseException | None
   ) -> None:
     for handler in self._handlers:
-      handler.on_sandbox_shutdown(
-          environment, sandbox, duration, lifetime, error
+      handler.on_sandbox_shutdown(sandbox, duration, lifetime, error)
+
+  def on_sandbox_session_start(
+      self,
+      sandbox: interface.Sandbox,
+      session_id: str,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session starts."""
+    for handler in self._handlers:
+      handler.on_sandbox_session_start(sandbox, session_id, duration, error)
+
+  def on_sandbox_session_end(
+      self,
+      sandbox: interface.Sandbox,
+      session_id: str,
+      duration: float,
+      lifetime: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session ends."""
+    for handler in self._handlers:
+      handler.on_sandbox_session_end(
+          sandbox, session_id, duration, lifetime, error
+      )
+
+  def on_sandbox_activity(
+      self,
+      name: str,
+      sandbox: interface.Sandbox,
+      session_id: str | None,
+      duration: float,
+      error: BaseException | None,
+      **kwargs
+  ) -> None:
+    """Called when a sandbox activity is performed."""
+    for handler in self._handlers:
+      handler.on_sandbox_activity(
+          name, sandbox, session_id, duration, error, **kwargs
       )
 
   def on_sandbox_housekeep(
       self,
-      environment: interface.Environment,
       sandbox: interface.Sandbox,
       counter: int,
       duration: float,
@@ -129,42 +161,30 @@ class EventHandlerChain(interface.EventHandler):
   ) -> None:
     """Called when a sandbox feature is housekeeping."""
     for handler in self._handlers:
-      handler.on_sandbox_housekeep(
-          environment, sandbox, counter, duration, error, **kwargs
-      )
+      handler.on_sandbox_housekeep(sandbox, counter, duration, error, **kwargs)
 
   def on_feature_setup(
       self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
       feature: interface.Feature,
       duration: float,
       error: BaseException | None
   ) -> None:
     """Called when a sandbox feature is setup."""
     for handler in self._handlers:
-      handler.on_feature_setup(
-          environment, sandbox, feature, duration, error
-      )
+      handler.on_feature_setup(feature, duration, error)
 
   def on_feature_teardown(
       self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
       feature: interface.Feature,
       duration: float,
       error: BaseException | None
   ) -> None:
     """Called when a sandbox feature is teardown."""
     for handler in self._handlers:
-      handler.on_feature_teardown(
-          environment, sandbox, feature, duration, error
-      )
+      handler.on_feature_teardown(feature, duration, error)
 
   def on_feature_setup_session(
       self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
       feature: interface.Feature,
       session_id: str | None,
       duration: float,
@@ -172,14 +192,10 @@ class EventHandlerChain(interface.EventHandler):
   ) -> None:
     """Called when a sandbox feature is setup."""
     for handler in self._handlers:
-      handler.on_feature_setup_session(
-          environment, sandbox, feature, session_id, duration, error
-      )
+      handler.on_feature_setup_session(feature, session_id, duration, error)
 
   def on_feature_teardown_session(
       self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
       feature: interface.Feature,
       session_id: str,
       duration: float,
@@ -187,69 +203,31 @@ class EventHandlerChain(interface.EventHandler):
   ) -> None:
     """Called when a sandbox feature is teardown."""
     for handler in self._handlers:
-      handler.on_feature_teardown_session(
-          environment, sandbox, feature, session_id, duration, error
-      )
+      handler.on_feature_teardown_session(feature, session_id, duration, error)
 
-  def on_feature_housekeep(
+  def on_feature_activity(
       self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
+      name: str,
       feature: interface.Feature,
-      counter: int,
+      session_id: str | None,
       duration: float,
       error: BaseException | None,
       **kwargs
   ) -> None:
-    """Called when a sandbox feature is housekeeping."""
-    for handler in self._handlers:
-      handler.on_feature_housekeep(
-          environment, sandbox, feature, counter, duration, error, **kwargs
-      )
-
-  def on_session_start(
-      self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
-      session_id: str,
-      duration: float,
-      error: BaseException | None
-  ) -> None:
-    """Called when a sandbox session starts."""
-    for handler in self._handlers:
-      handler.on_session_start(
-          environment, sandbox, session_id, duration, error
-      )
-
-  def on_session_end(
-      self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
-      session_id: str,
-      duration: float,
-      lifetime: float,
-      error: BaseException | None
-  ) -> None:
-    """Called when a sandbox session ends."""
+    """Called when a feature activity is performed."""
     for handler in self._handlers:
-      handler.on_session_end(
-          environment, sandbox, session_id, duration, lifetime, error
+      handler.on_feature_activity(
+          name, feature, session_id, duration, error, **kwargs
       )
 
-  def on_sandbox_activity(
+  def on_feature_housekeep(
       self,
-      name: str,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
-      feature: interface.Feature | None,
-      session_id: str | None,
+      feature: interface.Feature,
+      counter: int,
       duration: float,
       error: BaseException | None,
       **kwargs
   ) -> None:
-    """Called when a sandbox activity is performed."""
+    """Called when a sandbox feature is housekeeping."""
     for handler in self._handlers:
-      handler.on_sandbox_activity(
-          name, environment, sandbox, feature, session_id, duration, error,
-          **kwargs
-      )
+      handler.on_feature_housekeep(feature, counter, duration, error, **kwargs)

langfun/env/event_handlers/chain_test.py

@@ -49,15 +49,14 @@ class MockEventHandler(interface.EventHandler):
         'on_environment_shutdown', environment, duration, lifetime, error
     )
 
-  def on_sandbox_start(self, environment, sandbox, duration, error):
-    self._record_call('on_sandbox_start', environment, sandbox, duration, error)
+  def on_sandbox_start(self, sandbox, duration, error):
+    self._record_call('on_sandbox_start', sandbox, duration, error)
 
   def on_sandbox_status_change(
-      self, environment, sandbox, old_status, new_status, span
+      self, sandbox, old_status, new_status, span
   ):
     self._record_call(
         'on_sandbox_status_change',
-        environment,
         sandbox,
         old_status,
         new_status,
@@ -65,113 +64,100 @@ class MockEventHandler(interface.EventHandler):
     )
 
   def on_sandbox_shutdown(
-      self, environment, sandbox, duration, lifetime, error
+      self, sandbox, duration, lifetime, error
   ):
+    self._record_call('on_sandbox_shutdown', sandbox, duration, lifetime, error)
+
+  def on_sandbox_session_start(self, sandbox, session_id, duration, error):
     self._record_call(
-        'on_sandbox_shutdown', environment, sandbox, duration, lifetime, error
+        'on_sandbox_session_start', sandbox, session_id, duration, error
     )
 
-  def on_sandbox_housekeep(
-      self, environment, sandbox, counter, duration, error, **kwargs
+  def on_sandbox_session_end(
+      self, sandbox, session_id, duration, lifetime, error
   ):
     self._record_call(
-        'on_sandbox_housekeep',
-        environment,
+        'on_sandbox_session_end',
         sandbox,
-        counter,
+        session_id,
         duration,
+        lifetime,
         error,
-        **kwargs,
-    )
-
-  def on_feature_setup(self, environment, sandbox, feature, duration, error):
-    self._record_call(
-        'on_feature_setup', environment, sandbox, feature, duration, error
-    )
-
-  def on_feature_teardown(self, environment, sandbox, feature, duration, error):
-    self._record_call(
-        'on_feature_teardown', environment, sandbox, feature, duration, error
     )
 
-  def on_feature_setup_session(
-      self, environment, sandbox, feature, session_id, duration, error
+  def on_sandbox_activity(
+      self,
+      name,
+      sandbox,
+      session_id,
+      duration,
+      error,
+      **kwargs,
   ):
     self._record_call(
-        'on_feature_setup_session',
-        environment,
+        'on_sandbox_activity',
+        name,
         sandbox,
-        feature,
         session_id,
         duration,
         error,
+        **kwargs,
     )
 
-  def on_feature_teardown_session(
-      self, environment, sandbox, feature, session_id, duration, error
-  ):
+  def on_sandbox_housekeep(self, sandbox, counter, duration, error, **kwargs):
     self._record_call(
-        'on_feature_teardown_session',
-        environment,
+        'on_sandbox_housekeep',
         sandbox,
-        feature,
-        session_id,
+        counter,
         duration,
         error,
+        **kwargs,
     )
 
-  def on_feature_housekeep(
-      self, environment, sandbox, feature, counter, duration, error, **kwargs
-  ):
+  def on_feature_setup(self, feature, duration, error):
+    self._record_call('on_feature_setup', feature, duration, error)
+
+  def on_feature_teardown(self, feature, duration, error):
+    self._record_call('on_feature_teardown', feature, duration, error)
+
+  def on_feature_setup_session(self, feature, session_id, duration, error):
     self._record_call(
-        'on_feature_housekeep',
-        environment,
-        sandbox,
+        'on_feature_setup_session',
         feature,
-        counter,
+        session_id,
         duration,
         error,
-        **kwargs,
     )
 
-  def on_session_start(
-      self, environment, sandbox, session_id, duration, error
-  ):
+  def on_feature_teardown_session(self, feature, session_id, duration, error):
     self._record_call(
-        'on_session_start', environment, sandbox, session_id, duration, error
+        'on_feature_teardown_session',
+        feature,
+        session_id,
+        duration,
+        error,
     )
 
-  def on_session_end(
-      self, environment, sandbox, session_id, duration, lifetime, error
+  def on_feature_activity(
+      self, name, feature, session_id, duration, error, **kwargs,
   ):
     self._record_call(
-        'on_session_end',
-        environment,
-        sandbox,
+        'on_feature_activity',
+        name,
+        feature,
         session_id,
         duration,
-        lifetime,
         error,
+        **kwargs,
     )
 
-  def on_sandbox_activity(
-      self,
-      name,
-      environment,
-      sandbox,
-      feature,
-      session_id,
-      duration,
-      error,
-      **kwargs,
+  def on_feature_housekeep(
+      self, feature, counter, duration, error, **kwargs
   ):
     self._record_call(
-        'on_sandbox_activity',
-        name,
-        environment,
-        sandbox,
+        'on_feature_housekeep',
         feature,
-        session_id,
+        counter,
         duration,
         error,
         **kwargs,
@@ -194,26 +180,19 @@ class EventHandlerChainTest(unittest.TestCase):
     chain_handler.on_environment_start(env, 2.0, None)
     chain_handler.on_environment_housekeep(env, 1, 3.0, None, a=1)
     chain_handler.on_environment_shutdown(env, 4.0, 5.0, None)
-    chain_handler.on_sandbox_start(env, sandbox, 6.0, None)
-    chain_handler.on_sandbox_status_change(env, sandbox, 'old', 'new', 7.0)
-    chain_handler.on_sandbox_shutdown(env, sandbox, 8.0, 9.0, None)
-    chain_handler.on_sandbox_housekeep(env, sandbox, 2, 10.0, None, b=2)
-    chain_handler.on_feature_setup(env, sandbox, feature, 11.0, None)
-    chain_handler.on_feature_teardown(env, sandbox, feature, 12.0, None)
-    chain_handler.on_feature_setup_session(
-        env, sandbox, feature, 's1', 13.0, None
-    )
-    chain_handler.on_feature_teardown_session(
-        env, sandbox, feature, 's1', 14.0, None
-    )
-    chain_handler.on_feature_housekeep(
-        env, sandbox, feature, 3, 15.0, None, c=3
-    )
-    chain_handler.on_session_start(env, sandbox, 's2', 16.0, None)
-    chain_handler.on_session_end(env, sandbox, 's2', 17.0, 18.0, None)
-    chain_handler.on_sandbox_activity(
-        'act', env, sandbox, feature, 's2', 19.0, None, d=4
-    )
+    chain_handler.on_sandbox_start(sandbox, 6.0, None)
+    chain_handler.on_sandbox_status_change(sandbox, 'old', 'new', 7.0)
+    chain_handler.on_sandbox_shutdown(sandbox, 8.0, 9.0, None)
+    chain_handler.on_sandbox_session_start(sandbox, 's2', 16.0, None)
+    chain_handler.on_sandbox_session_end(sandbox, 's2', 17.0, 18.0, None)
+    chain_handler.on_sandbox_activity('act', sandbox, 's2', 19.0, None, d=4)
+    chain_handler.on_sandbox_housekeep(sandbox, 2, 10.0, None, b=2)
+    chain_handler.on_feature_setup(feature, 11.0, None)
+    chain_handler.on_feature_teardown(feature, 12.0, None)
+    chain_handler.on_feature_setup_session(feature, 's1', 13.0, None)
+    chain_handler.on_feature_teardown_session(feature, 's1', 14.0, None)
+    chain_handler.on_feature_activity('act', feature, 's1', 16.0, None, d=5)
+    chain_handler.on_feature_housekeep(feature, 3, 15.0, None, c=3)
 
     self.assertEqual(handler1.calls, handler2.calls)
     self.assertEqual(
@@ -224,34 +203,27 @@ class EventHandlerChainTest(unittest.TestCase):
             ('on_environment_start', (env, 2.0, None), {}),
             ('on_environment_housekeep', (env, 1, 3.0, None), {'a': 1}),
             ('on_environment_shutdown', (env, 4.0, 5.0, None), {}),
-            ('on_sandbox_start', (env, sandbox, 6.0, None), {}),
-            ('on_sandbox_status_change', (env, sandbox, 'old', 'new', 7.0), {}),
-            ('on_sandbox_shutdown', (env, sandbox, 8.0, 9.0, None), {}),
-            ('on_sandbox_housekeep', (env, sandbox, 2, 10.0, None), {'b': 2}),
-            ('on_feature_setup', (env, sandbox, feature, 11.0, None), {}),
-            ('on_feature_teardown', (env, sandbox, feature, 12.0, None), {}),
-            (
-                'on_feature_setup_session',
-                (env, sandbox, feature, 's1', 13.0, None),
-                {},
-            ),
-            (
-                'on_feature_teardown_session',
-                (env, sandbox, feature, 's1', 14.0, None),
-                {},
-            ),
-            (
-                'on_feature_housekeep',
-                (env, sandbox, feature, 3, 15.0, None),
-                {'c': 3},
-            ),
-            ('on_session_start', (env, sandbox, 's2', 16.0, None), {}),
-            ('on_session_end', (env, sandbox, 's2', 17.0, 18.0, None), {}),
+            ('on_sandbox_start', (sandbox, 6.0, None), {}),
+            ('on_sandbox_status_change', (sandbox, 'old', 'new', 7.0), {}),
+            ('on_sandbox_shutdown', (sandbox, 8.0, 9.0, None), {}),
+            ('on_sandbox_session_start', (sandbox, 's2', 16.0, None), {}),
+            ('on_sandbox_session_end', (sandbox, 's2', 17.0, 18.0, None), {}),
             (
                 'on_sandbox_activity',
-                ('act', env, sandbox, feature, 's2', 19.0, None),
+                ('act', sandbox, 's2', 19.0, None),
                 {'d': 4},
             ),
+            ('on_sandbox_housekeep', (sandbox, 2, 10.0, None), {'b': 2}),
+            ('on_feature_setup', (feature, 11.0, None), {}),
+            ('on_feature_teardown', (feature, 12.0, None), {}),
+            ('on_feature_setup_session', (feature, 's1', 13.0, None), {}),
+            ('on_feature_teardown_session', (feature, 's1', 14.0, None), {}),
+            (
+                'on_feature_activity',
+                ('act', feature, 's1', 16.0, None),
+                {'d': 5}
+            ),
+            ('on_feature_housekeep', (feature, 3, 15.0, None), {'c': 3}),
         ],
     )