langfun 0.1.2.dev202508210804__py3-none-any.whl → 0.1.2.dev202508230803__py3-none-any.whl

langfun/core/agentic/__init__.py

@@ -17,6 +17,9 @@
 # pylint: disable=g-importing-member
 # pylint: disable=g-import-not-at-top
 
+from langfun.core.agentic.action import ActionError
+from langfun.core.agentic.action import ActionTimeoutError
+
 from langfun.core.agentic.action import Action
 from langfun.core.agentic.action import ActionInvocation
 from langfun.core.agentic.action import ExecutionTrace

langfun/core/agentic/action.py

@@ -26,6 +26,14 @@ from langfun.core import structured as lf_structured
 import pyglove as pg
 
 
+class ActionError(Exception):   # pylint: disable=g-bad-exception-name
+  """Base class for common action errors."""
+
+
+class ActionTimeoutError(ActionError):
+  """Raised when an action exceeds the max execution time."""
+
+
 class Action(pg.Object):
   """Base class for Langfun's agentic actions.
 
@@ -203,6 +211,7 @@ class Action(pg.Object):
       self,
       session: Optional['Session'] = None,
       *,
+      max_execution_time: float | None = None,
       show_progress: bool = True,
       verbose: bool = False,
       **kwargs
@@ -212,6 +221,7 @@ class Action(pg.Object):
     return await lf.invoke_async(
         self.__call__,
         session,
+        max_execution_time=max_execution_time,
         show_progress=show_progress,
         verbose=verbose,
         **kwargs
@@ -221,11 +231,29 @@ class Action(pg.Object):
       self,
       session: Optional['Session'] = None,
       *,
+      max_execution_time: float | None = None,
       show_progress: bool = True,
       verbose: bool = False,
       **kwargs
   ) -> Any:
-    """Executes the action."""
+    """Executes the action.
+
+    Args:
+      session: The session to use for the action.
+      max_execution_time: The max allowed execution time in seconds for the
+        action. The effective `max_execution_time` is the smaller of the
+        remaining time of the parent invocation and the `max_execution_time`
+        for the current invocation. Since we are running action within the
+        current thread, there is no guarantee that the action will be executed
+        within this time limit, but we will try to stop the action as soon as
+        the next operation on the session is called.
+      show_progress: Whether to show the progress of the action.
+      verbose: Whether to log the action execution.
+      **kwargs: Additional keyword arguments to pass to the action.
+
+    Returns:
+      The result of the action.
+    """
     if session is None:
       session = Session(verbose=verbose)
       session.start()
@@ -238,7 +266,7 @@ class Action(pg.Object):
     else:
       self._session = None
 
-    with session.track_action(self):
+    with session.track_action(self, max_execution_time=max_execution_time):
       try:
         result = self.call(session=session, **kwargs)
         self._invocation.end(result)
@@ -841,6 +869,21 @@ class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
       'Error from the action if failed.'
   ] = None
 
+  max_execution_time: Annotated[
+      float | None,
+      (
+          'The maximum allowed execution time for the action. '
+          'It is set when the action or any of its parent actions has been '
+          'called with a `max_execution_time` argument. '
+          'The value is the remaining time of the parent invocation or '
+          'the `max_execution_time` for the current invocation, '
+          'whichever is smaller. Since we are running action within the '
+          'current thread, there is no guarantee that the action will be '
+          'executed within this time limit, but we will try to stop the action '
+          'as soon as the next operation on the session is called.'
+      )
+  ] = None
+
   execution: Annotated[
       ExecutionTrace,
       'The execution sequence of the action.'
@@ -863,6 +906,13 @@ class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
     """Returns the parent action invocation."""
     return self.sym_ancestor(lambda x: isinstance(x, ActionInvocation))
 
+  @property
+  def max_remaining_execution_time(self) -> float | None:
+    """Returns the remaining execution time for the action."""
+    if self.max_execution_time is None:
+      return None
+    return max(0, self.max_execution_time - self.execution.elapse)
+
   @functools.cached_property
   def id(self) -> str:
     """Returns the id of the action invocation."""
@@ -1224,6 +1274,23 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       )
     self.root.end(result, error, metadata)
 
+  def check_execution_time(self) -> None:
+    """Checks the execution time of the current action."""
+    action = self._current_action
+    if action is not None and 0 == action.max_remaining_execution_time:
+      # We raise error on the top-most action which has a time limit.
+      current_action = action
+      parent_action = current_action.parent_action
+      while parent_action is not None and parent_action.max_execution_time:
+        current_action = parent_action
+        parent_action = current_action.parent_action
+      raise ActionTimeoutError(
+          f'Action {current_action.id} '
+          f'({current_action.action.__class__.__name__}) has exceeded its '
+          f'maximum execution time of {current_action.max_execution_time} '
+          'seconds.'
+      )
+
   def __enter__(self):
     """Enters the session."""
     self.start()
@@ -1254,7 +1321,11 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
   #
 
   @contextlib.contextmanager
-  def track_action(self, action: Action) -> Iterator[ActionInvocation]:
+  def track_action(
+      self,
+      action: Action,
+      max_execution_time: float | None = None
+  ) -> Iterator[ActionInvocation]:
     """Track the execution of an action."""
     if not self.root.execution.has_started:
       raise ValueError(
@@ -1263,7 +1334,13 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
           'signal the start and end of the session.'
       )
 
-    invocation = ActionInvocation(pg.maybe_ref(action))
+    # Early terminate the action if the execution time is exceeded.
+    self.check_execution_time()
+
+    invocation = ActionInvocation(
+        pg.maybe_ref(action),
+        max_execution_time=self._child_max_execution_time(max_execution_time)
+    )
     parent_action = self._current_action
     parent_execution = self._current_execution
     parent_execution.append(invocation)
@@ -1339,6 +1416,9 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       call.
     """
     def _query_start(invocation: lf_structured.QueryInvocation):
+      # Early terminate the action if the execution time is exceeded.
+      self.check_execution_time()
+
       execution = self._current_execution
       invocation.rebind(
           id=f'{execution.id}/q{len(execution.queries) + 1}',
@@ -1469,6 +1549,7 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       ] = Exception
   ) -> Iterator[Any]:
     """Starts and tracks parallel execution with `lf.concurrent_map`."""
+    self.check_execution_time()
     parallel_inputs = list(parallel_inputs)
     parallel_execution = ParallelExecutions(name=phase)
     self._current_execution.append(parallel_execution)
@@ -1492,90 +1573,23 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
         _map_single,
         parallel_inputs,
         max_workers=max_workers,
-        timeout=timeout,
+        timeout=self._child_max_execution_time(timeout),
         silence_on_errors=silence_on_errors
     ):
       yield input_value, result, error
 
-  # NOTE(daiyip): Clean up `query_prompt` and `query_output` once TS
-  # code migration is done.
-  def query_prompt(
-      self,
-      prompt: Union[str, lf.Template, Any],
-      schema: Union[
-          lf_structured.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
-      ] = None,
-      **kwargs,
-  ) -> Any:
-    """Calls `lf.query_prompt` and associates it with the current invocation.
-
-    The following code are equivalent:
-
-      Code 1:
-      ```
-      session.query_prompt(...)
-      ```
-
-      Code 2:
-      ```
-      with session.track_queries() as queries:
-        output = lf.query_prompt(...)
-      ```
-    The former is preferred when `lf.query_prompt` is directly called by the
-    action.
-    If `lf.query_prompt` is called by a function that does not have access to
-    the
-    session, the latter should be used.
-
-    Args:
-      prompt: The prompt to query.
-      schema: The schema to use for the query.
-      **kwargs: Additional keyword arguments to pass to `lf.query_prompt`.
-
-    Returns:
-      The result of the query.
-    """
-    with self.track_queries():
-      return lf_structured.query_prompt(prompt, schema=schema, **kwargs)
-
-  def query_output(
-      self,
-      response: Union[str, lf.Template, Any],
-      schema: Union[
-          lf_structured.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
-      ] = None,
-      **kwargs,
-  ) -> Any:
-    """Calls `lf.query_output` and associates it with the current invocation.
-
-    The following code are equivalent:
-
-      Code 1:
-      ```
-      session.query_output(...)
-      ```
-
-      Code 2:
-      ```
-      with session.track_queries() as queries:
-        output = lf.query_output(...)
-      ```
-    The former is preferred when `lf.query_output` is directly called by the
-    action.
-    If `lf.query_output` is called by a function that does not have access to
-    the
-    session, the latter should be used.
-
-    Args:
-      response: The response to query.
-      schema: The schema to use for the query.
-      **kwargs: Additional keyword arguments to pass to `lf.query_prompt`.
-
-    Returns:
-      The result of the query.
-    """
-    with self.track_queries():
-      return lf_structured.query_output(response, schema=schema, **kwargs)
+  def _child_max_execution_time(
+      self, max_execution_time: float | None
+  ) -> float | None:
+    """Returns the max execution time for the child action."""
+    max_remaining = None
+    if self._current_action is not None:
+      max_remaining = self._current_action.max_remaining_execution_time
+    if max_remaining is None:
+      return max_execution_time
+    if max_execution_time is None:
+      return max_remaining
+    return min(max_remaining, max_execution_time)
 
   def _log(
       self,
@@ -1597,6 +1611,9 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       **kwargs: Additional keyword arguments to pass to `lf.logging.log` as
         metadata to show.
     """
+    # Early terminate the action if the execution time is exceeded.
+    self.check_execution_time()
+
     execution = self._current_execution
     if for_action is None:
       for_action = self._current_action

langfun/core/agentic/action_test.py

@@ -14,6 +14,7 @@
 """Tests for base action."""
 
 import asyncio
+import time
 import unittest
 
 import langfun.core as lf
@@ -25,10 +26,12 @@ import pyglove as pg
 
 class Bar(action_lib.Action):
   simulate_action_error: bool = False
+  simulate_execution_time: float = 0
 
   def call(self, session, *, lm, **kwargs):
     assert session.current_action.action is self
     session.info('Begin Bar')
+    time.sleep(self.simulate_execution_time)
     session.query('bar', lm=lm)
     session.add_metadata(note='bar')
     if self.simulate_action_error:
@@ -40,22 +43,27 @@ class Foo(action_lib.Action):
   x: int
   simulate_action_error: bool = False
   simulate_query_error: bool = False
+  simulate_execution_time: list[float] = [0, 0, 0, 0]
+  max_bar_execution_time: float | None = None
 
   def call(self, session, *, lm, **kwargs):
     assert session.current_action.action is self
     with session.track_phase('prepare'):
       session.info('Begin Foo', x=1)
+      time.sleep(self.simulate_execution_time[0])
       session.query(
           'foo',
           schema=int if self.simulate_query_error else None,
           lm=lm
       )
     with session.track_queries():
+      time.sleep(self.simulate_execution_time[1])
       self.make_additional_query(lm)
     session.add_metadata(note='foo')
 
     def _sub_task(i):
       session.add_metadata(**{f'subtask_{i}': i})
+      time.sleep(self.simulate_execution_time[2])
       return lf_structured.query(f'subtask_{i}', lm=lm)
 
     for i, output, error in session.concurrent_map(
@@ -65,8 +73,9 @@ class Foo(action_lib.Action):
       assert isinstance(output, str), output
       assert error is None, error
     return self.x + Bar(
-        simulate_action_error=self.simulate_action_error
-    )(session, lm=lm)
+        simulate_action_error=self.simulate_action_error,
+        simulate_execution_time=self.simulate_execution_time[3]
+    )(session, lm=lm, max_execution_time=self.max_bar_execution_time)
 
   def make_additional_query(self, lm):
     lf_structured.query('additional query', lm=lm)
@@ -426,6 +435,51 @@ class SessionTest(unittest.TestCase):
     self.assertFalse(session.root.execution[0].has_error)
     self.assertTrue(session.root.execution[1].has_error)
 
+  def test_max_execution_time(self):
+    lm = fake.StaticResponse('lm response')
+    bar = Bar(simulate_execution_time=1)
+    with self.assertRaisesRegex(
+        action_lib.ActionTimeoutError,
+        'Action .*Bar.*has exceeded .* 0.5 seconds'
+    ):
+      bar(lm=lm, max_execution_time=0.5)
+
+    foo = Foo(1, simulate_execution_time=[0, 0, 0, 1])
+    with self.assertRaisesRegex(
+        action_lib.ActionTimeoutError,
+        'Action .*Foo.* has exceeded .* 0.5 seconds'
+    ):
+      foo(lm=lm, max_execution_time=0.5)
+
+    # Timeout within concurrent_map.
+    foo = Foo(1, simulate_execution_time=[0, 0, 1, 0])
+    with self.assertRaisesRegex(
+        action_lib.ActionTimeoutError,
+        'Action .*Foo.* has exceeded .* 0.5 seconds'
+    ):
+      foo(lm=lm, max_execution_time=0.5)
+
+    # Timeout within bar.
+    foo = Foo(
+        1, simulate_execution_time=[0, 0, 0, 1], max_bar_execution_time=0.5
+    )
+    with self.assertRaisesRegex(
+        action_lib.ActionTimeoutError,
+        'Action .*Bar.* has exceeded .* 0.5 seconds'
+    ):
+      foo(lm=lm)
+
+    # Timeout within bar, however the effective max_execution_time of bar is the
+    # remaining time of the parent action as it's smaller (0.5 < 1).
+    foo = Foo(
+        1, simulate_execution_time=[0, 0.5, 0, 1.0], max_bar_execution_time=1.0
+    )
+    with self.assertRaisesRegex(
+        action_lib.ActionTimeoutError,
+        'Action .*Foo.* has exceeded .*1.0 seconds'
+    ):
+      foo(lm=lm, max_execution_time=1.0)
+
   def test_log(self):
     session = action_lib.Session()
     session.debug('hi', x=1, y=2)

{langfun-0.1.2.dev202508210804.dist-info → langfun-0.1.2.dev202508230803.dist-info}/METADATA

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

{langfun-0.1.2.dev202508210804.dist-info → langfun-0.1.2.dev202508230803.dist-info}/RECORD

@@ -34,11 +34,11 @@ langfun/core/subscription.py,sha256=euawEuSZP-BHydaT-AQpfYFL0m5pWPGcW0upFhrojqc,
 langfun/core/subscription_test.py,sha256=Y4ZdbZEwm83YNZBxHff0QR4QUa4rdaNXA3_jfIcArBo,8717
 langfun/core/template.py,sha256=GSOZ3OcmRtw-q05GdHrE-Y4-2MGDYsTRKmWGvVPbdhE,24962
 langfun/core/template_test.py,sha256=AQv_m9qE93WxhEhSlm1xaBgB4hu0UVtA53dljngkUW0,17090
-langfun/core/agentic/__init__.py,sha256=qR3jlfUO4rhIoYdRDLz-d22YZf3FvU4FW88vsjiGDQQ,1224
-langfun/core/agentic/action.py,sha256=HvCaYb3x9tpa0xhDuFtEGyZ9du_edDpIdaQYPPo87ts,51097
+langfun/core/agentic/__init__.py,sha256=s9zRiAOtiTvp_sWeyGcykjlo6rez8asvLK7tiOELWEU,1336
+langfun/core/agentic/action.py,sha256=CgOmp9ht4J6N1bg0DSqkP1G2YYP3HPi1Z7Es5qqhevs,52824
 langfun/core/agentic/action_eval.py,sha256=YTilyUEkJl_8FVMgdfO17PurWWaEJ6oA15CuefJJRLk,4887
 langfun/core/agentic/action_eval_test.py,sha256=7AkOwNbUX-ZgR1R0a7bvUZ5abNTUV7blf_8Mnrwb-II,2811
-langfun/core/agentic/action_test.py,sha256=5ZUCbyuJDwNCUE_hl4qt5epzQQTRdZEYN5W8J2iKRQo,16084
+langfun/core/agentic/action_test.py,sha256=rorIo58S1o27CZrtw05b49mLZOaiitl_xGanGZa5HXo,18028
 langfun/core/coding/__init__.py,sha256=5utju_fwEsImaiftx4oXKl9FAM8p281k8-Esdh_-m1w,835
 langfun/core/coding/python/__init__.py,sha256=yTXm92oLpQb4A-fZ2qy-bzfhPYND7B-oidtbv1PNaX0,1678
 langfun/core/coding/python/correction.py,sha256=4PD76Xfv36Xrm8Ji3-GgGDNImtcDqWfMw3z6ircJMlM,7285
@@ -165,8 +165,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.dev202508210804.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202508210804.dist-info/METADATA,sha256=C-d5bGlbkHHdz0Hd0-qdvG7fnmNMmJ8pU3Hg5zDKhqU,7380
-langfun-0.1.2.dev202508210804.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-langfun-0.1.2.dev202508210804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202508210804.dist-info/RECORD,,
+langfun-0.1.2.dev202508230803.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202508230803.dist-info/METADATA,sha256=GjLNmc4jIoMKc8vc79MpHhBbQnaVFt-N5myH4qdLnw8,7380
+langfun-0.1.2.dev202508230803.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+langfun-0.1.2.dev202508230803.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202508230803.dist-info/RECORD,,