langfun 0.1.2.dev202511090804__tar.gz → 0.1.2.dev202512180805__tar.gz

{langfun-0.1.2.dev202511090804 → langfun-0.1.2.dev202512180805}/PKG-INFO

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

{langfun-0.1.2.dev202511090804 → langfun-0.1.2.dev202512180805}/langfun/core/__init__.py

@@ -121,6 +121,7 @@ from langfun.core.language_model import LMDebugMode
 from langfun.core.language_model import LMError
 from langfun.core.language_model import LMInputError
 from langfun.core.language_model import ContextLimitError
+from langfun.core.language_model import EmptyGenerationError
 from langfun.core.language_model import RetryableLMError
 from langfun.core.language_model import RateLimitError
 from langfun.core.language_model import TemporaryLMError

{langfun-0.1.2.dev202511090804 → langfun-0.1.2.dev202512180805}/langfun/core/agentic/__init__.py

@@ -21,9 +21,12 @@ 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 ExecutionUnit
 from langfun.core.agentic.action import ActionInvocation
-from langfun.core.agentic.action import ExecutionTrace
 from langfun.core.agentic.action import ParallelExecutions
+
+from langfun.core.agentic.action import ExecutionTrace
 from langfun.core.agentic.action import Session
 
 from langfun.core.agentic.action import SessionEventHandler

{langfun-0.1.2.dev202511090804 → langfun-0.1.2.dev202512180805}/langfun/core/agentic/action.py

@@ -14,9 +14,11 @@
 """Base classes for agentic actions."""
 
 import abc
+import collections
 import contextlib
 import dataclasses
 import functools
+import itertools
 import threading
 import time
 import typing
@@ -286,6 +288,8 @@ class Action(pg.Object):
 
     with session.track_action(self, max_execution_time=max_execution_time):
       try:
+        # Early terminate the action if the execution time is exceeded.
+        session.check_execution_time()
         result = self.call(session=session, **kwargs)
         self._invocation.end(result)
       except BaseException as e:
@@ -313,12 +317,176 @@ class Action(pg.Object):
     """
 
 
+#
+# Execution tracking.
+#
+
+
+class ExecutionUnit(pg.Object):
+  """Base class for execution units in an agentic trajectory.
+
+  An `ExecutionUnit` represents a logical step or container in the agent's
+  execution flow. It serves as the common interface for top-level executable
+  items.
+
+  The concrete subclasses of `ExecutionUnit` are typically:
+  * **`ActionInvocation`**: Represents a single, specific action executed by
+    the agent.
+  * **`ParallelExecutions`**: Represents a container for a group of
+    `ExecutionUnits` that were executed concurrently.
+
+  Users can retrieve the immediate child execution units from an `ExecutionUnit`
+  object. To access the leaf nodes of the execution tree, use `all_actions`,
+  `all_queries`, and `all_logs` instead.
+
+  Each unit exposes a **`position`** property to reveal its specific location
+  within the execution hierarchy (e.g., '1.2.3').
+
+  Users could use **`parent_execution_unit`** to get the parent execution unit
+  of the current execution unit.
+  """
+
+  @dataclasses.dataclass
+  class Position:
+    """The position of an executed unit under current session."""
+
+    parent: Optional['ExecutionUnit.Position'] = None
+    index: int = 0
+
+    def indices(self) -> tuple[int, ...]:
+      """Returns the indices from root to current execution unit."""
+      # A deque is efficient for adding items to the front.
+      path = collections.deque()
+      current_pos = self
+
+      # Traverse up from the current position to the root.
+      while current_pos.parent is not None:
+        path.appendleft(current_pos.index)
+        current_pos = current_pos.parent
+
+      path.appendleft(current_pos.index)
+      return tuple(path)
+
+    def to_str(
+        self,
+        *,
+        index_base: int = 1,
+        separator: str = '.',
+        **kwargs
+    ) -> str:
+      """Returns a string description of the position."""
+      # For root action, we return empty string as it's position descriptor.
+      if self.parent is None:
+        return ''
+      parent_descriptor = self.parent.to_str(
+          index_base=index_base,
+          separator=separator,
+          **kwargs
+      )
+      if not parent_descriptor:
+        return f'{self.index + index_base}'
+      return parent_descriptor + separator + f'{self.index + index_base}'
+
+    def __repr__(self) -> str:
+      return f'Position({", ".join(str(x) for x in self.indices())})'
+
+    def __str__(self) -> str:
+      return self.to_str()
+
+    def __eq__(self, other: 'ExecutionUnit.Position') -> bool:
+      if isinstance(other, ExecutionUnit.Position):
+        return self.indices() == other.indices()
+      if isinstance(other, tuple):
+        return self.indices() == other
+      if isinstance(other, str):
+        return str(self) == other
+      return False
+
+    def __ne__(self, other: 'ExecutionUnit.Position') -> bool:
+      return not self == other
+
+    def __hash__(self) -> int:
+      return hash(self.indices())
+
+    def __lt__(self, other: 'ExecutionUnit.Position') -> bool:
+      return self.indices() < other.indices()
+
+    def __gt__(self, other: 'ExecutionUnit.Position') -> bool:
+      return self.indices() > other.indices()
+
+  def _on_parent_change(self, *args, **kwargs):
+    super()._on_parent_change(*args, **kwargs)
+    self.__dict__.pop('parent_execution_unit', None)
+    self.__dict__.pop('position', None)
+
+  @functools.cached_property
+  def parent_execution_unit(self) -> Optional['ExecutionUnit']:
+    """Returns the parent execution unit of the current execution unit."""
+    parent_trace = self.sym_ancestor(lambda x: isinstance(x, ExecutionTrace))
+    assert isinstance(parent_trace, ExecutionTrace), (
+        'Execution unit is not associated with any `ExecutionTrace`: '
+        f'{self}'
+    )
+    return parent_trace.parent_execution_unit
+
+  @functools.cached_property
+  def position(self) -> Position:
+    """Returns the execution position of the action."""
+    parent_trace = self.sym_ancestor(lambda x: isinstance(x, ExecutionTrace))
+    while parent_trace is not None:
+      parent_position = parent_trace.position
+      if parent_position is not None:
+        return ExecutionUnit.Position(
+            parent_position, parent_trace.indexof(self, ExecutionUnit)
+        )
+      parent_trace = parent_trace.sym_ancestor(
+          lambda x: isinstance(x, ExecutionTrace)
+      )
+    return ExecutionUnit.Position(None, 0)
+
+  @property
+  @abc.abstractmethod
+  def execution_units(
+      self,
+  ) -> list['ExecutionUnit']:
+    """Returns immediate child execution items."""
+
+  @property
+  @abc.abstractmethod
+  def queries(self) -> list[lf_structured.QueryInvocation]:
+    """Returns queries issued by the execution item."""
+
+  @property
+  @abc.abstractmethod
+  def actions(self) -> list['ActionInvocation']:
+    """Returns immediate child action invocations."""
+
+  @property
+  @abc.abstractmethod
+  def logs(self) -> list[lf.logging.LogEntry]:
+    """Returns immediate logs under current execution item."""
+
+  @property
+  @abc.abstractmethod
+  def all_queries(self) -> list[lf_structured.QueryInvocation]:
+    """Returns all queries from the subtree."""
+
+  @property
+  @abc.abstractmethod
+  def all_actions(self) -> list['ActionInvocation']:
+    """Returns all action invocations from the subtree."""
+
+  @property
+  @abc.abstractmethod
+  def all_logs(self) -> list[lf.logging.LogEntry]:
+    """Returns all logs from the subtree."""
+
+
 # Type definition for traced item during execution.
 TracedItem = Union[
+    ExecutionUnit,
     lf_structured.QueryInvocation,
-    'ActionInvocation',
     'ExecutionTrace',
-    'ParallelExecutions',
     # NOTE(daiyip): Consider remove log entry once we migrate existing agents.
     lf.logging.LogEntry,
 ]
@@ -372,8 +540,14 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
   def _on_parent_change(self, *args, **kwargs):
     super()._on_parent_change(*args, **kwargs)
     self.__dict__.pop('id', None)
+    self.__dict__.pop('parent_execution_unit', None)
+    self.__dict__.pop('position', None)
 
-  def indexof(self, item: TracedItem, count_item_cls: Type[Any]) -> int:
+  def indexof(
+      self,
+      item: TracedItem,
+      count_item_cls: Type[Any] | tuple[Type[Any], ...]
+  ) -> int:
     """Returns the index of the child item of given type."""
     pos = 0
     for x in self._iter_children(count_item_cls):
@@ -382,6 +556,52 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
       pos += 1
     return -1
 
+  @functools.cached_property
+  def parent_execution_unit(self) -> ExecutionUnit:
+    """Returns the parent execution unit of the current execution trace."""
+    parent = self.sym_parent
+    if isinstance(parent, ActionInvocation):
+      # Current execution trace is the body of an action.
+      return parent
+    elif isinstance(parent, pg.List):
+      container = parent.sym_parent
+      if isinstance(container, ParallelExecutions):
+        return container
+      elif isinstance(container, ExecutionTrace):
+        return container.parent_execution_unit
+    assert False, (
+        'Execution trace is not associated with any `ActionInvocation` or '
+        f'`ParallelExecutions`: {self}'
+    )
+
+  @functools.cached_property
+  def position(self) -> ExecutionUnit.Position | None:
+    """Returns the execution position of the execution trace.
+
+    Returns:
+      The execution position of the execution trace, or None if the execution
+      trace is either not associated with an execution unit or the execution
+      trace is a phase under another execution trace.
+    """
+    parent = self.sym_parent
+    if isinstance(parent, ActionInvocation):
+      # Current execution trace is the body of an action.
+      return parent.position
+    elif isinstance(parent, pg.List):
+      container = parent.sym_parent
+      if isinstance(container, ParallelExecutions):
+        return ExecutionUnit.Position(
+            container.position, self.sym_path.key
+        )
+      elif isinstance(container, ExecutionTrace):
+        # When execution trace is a phase under another execution trace,
+        # we return None as the position.
+        return None
+    assert False, (
+        'Execution trace is not associated with any `ActionInvocation` or '
+        f'`ParallelExecutions`: {self}'
+    )
+
   @functools.cached_property
   def id(self) -> str:
     parent = self.sym_parent
@@ -453,6 +673,11 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
     """Returns action invocations from the sequence."""
     return list(self._iter_children(ActionInvocation))
 
+  @property
+  def execution_units(self) -> list[ExecutionUnit]:
+    """Returns parallel executions from the sequence."""
+    return list(self._iter_children(ExecutionUnit))
+
   @property
   def logs(self) -> list[lf.logging.LogEntry]:
     """Returns logs from the sequence."""
@@ -473,7 +698,9 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
     """Returns all logs from current trace and its child execution items."""
     return list(self._iter_subtree(lf.logging.LogEntry))
 
-  def _iter_children(self, item_cls: Type[Any]) -> Iterator[TracedItem]:
+  def _iter_children(
+      self, item_cls: Type[Any] | tuple[Type[Any], ...]
+  ) -> Iterator[TracedItem]:
     for item in self.items:
       if isinstance(item, item_cls):
         yield item
@@ -485,7 +712,10 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
           for x in branch._iter_children(item_cls):  # pylint: disable=protected-access
             yield x
 
-  def _iter_subtree(self, item_cls: Type[Any]) -> Iterator[TracedItem]:
+  def _iter_subtree(
+      self,
+      item_cls: Type[Any] | tuple[Type[Any], ...]
+  ) -> Iterator[TracedItem]:
     for item in self.items:
       if isinstance(item, item_cls):
         yield item
@@ -798,7 +1028,7 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
     ]
 
 
-class ParallelExecutions(pg.Object, pg.views.html.HtmlTreeView.Extension):
+class ParallelExecutions(ExecutionUnit, pg.views.html.HtmlTreeView.Extension):
   """A container for multiple parallel execution traces.
 
   When `session.concurrent_map` is used, it creates a `ParallelExecutions`
@@ -850,8 +1080,64 @@ class ParallelExecutions(pg.Object, pg.views.html.HtmlTreeView.Extension):
       self.branches.append(branch)
       if self._tab_control is not None:
         self._tab_control.append(self._branch_tab(branch))
+
+      # Invalidate cached properties.
+      self.__dict__.pop('all_queries', None)
+      self.__dict__.pop('all_actions', None)
+      self.__dict__.pop('all_logs', None)
       return branch
 
+  #
+  # ExecutionUnit interface.
+  #
+
+  @property
+  def execution_units(self) -> list[ExecutionUnit]:
+    """Returns immediate child execution items from execution sequence."""
+    return []
+
+  @property
+  def queries(self) -> list[lf_structured.QueryInvocation]:
+    """Returns immediate queries made by the parallel execution."""
+    return []
+
+  @property
+  def actions(self) -> list['ActionInvocation']:
+    """Returns immediate child action invocations."""
+    return []
+
+  @property
+  def logs(self) -> list[lf.logging.LogEntry]:
+    """Returns immediate child logs from execution sequence."""
+    return []
+
+  @functools.cached_property
+  def all_queries(self) -> list[lf_structured.QueryInvocation]:
+    """Returns all queries made by the action and its child execution items."""
+    return list(
+        itertools.chain.from_iterable(
+            branch.all_queries for branch in self.branches
+        )
+    )
+
+  @functools.cached_property
+  def all_actions(self) -> list['ActionInvocation']:
+    """Returns all actions made by the action and its child execution items."""
+    return list(
+        itertools.chain.from_iterable(
+            branch.all_actions for branch in self.branches
+        )
+    )
+
+  @property
+  def all_logs(self) -> list[lf.logging.LogEntry]:
+    """Returns all logs made by the action and its child execution items."""
+    return list(
+        itertools.chain.from_iterable(
+            branch.all_logs for branch in self.branches
+        )
+    )
+
   #
   # HTML views.
   #
@@ -892,7 +1178,7 @@ class ParallelExecutions(pg.Object, pg.views.html.HtmlTreeView.Extension):
     )
 
 
-class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
+class ActionInvocation(ExecutionUnit, pg.views.html.HtmlTreeView.Extension):
   """An invocation of an action, capturing its execution and result.
 
   `ActionInvocation` represents a single call to an `Action`. It contains
@@ -953,6 +1239,7 @@ class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
   def _on_parent_change(self, *args, **kwargs):
     super()._on_parent_change(*args, **kwargs)
     self.__dict__.pop('id', None)
+    self.__dict__.pop('position', None)
 
   @property
   def parent_action(self) -> Optional['ActionInvocation']:
@@ -989,6 +1276,15 @@ class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
     """Returns the state of the action."""
     return self.action.state
 
+  #
+  # Implement `ExecutionUnit` interface.
+  #
+
+  @property
+  def execution_units(self) -> list[ExecutionUnit]:
+    """Returns immediate child execution items from execution sequence."""
+    return self.execution.execution_units
+
   @property
   def logs(self) -> list[lf.logging.LogEntry]:
     """Returns immediate child logs from execution sequence."""
@@ -1040,10 +1336,12 @@ class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
       metadata: dict[str, Any] | None = None,
   ) -> None:
     """Ends the execution of the action with result and metadata."""
-    rebind_dict = dict(result=result, error=error)
-    if metadata is not None:
-      rebind_dict['metadata'] = metadata
-    self.rebind(**rebind_dict, skip_notification=True, raise_on_no_change=False)
+    with pg.notify_on_change(False):
+      self.result = result
+      self.error = error
+      if metadata:
+        self.metadata.update(metadata)
+
     self.execution.stop()
     if self._tab_control is not None:
       if self.metadata:
@@ -1201,6 +1499,19 @@ class RootAction(Action):
 class SessionEventHandler:
   """Interface for handling session events."""
 
+  def get(
+      self,
+      session_cls: type['SessionEventHandler']
+  ) -> Optional['SessionEventHandler']:
+    """Returns this or a child event handler for the given session class."""
+    if isinstance(self, session_cls):
+      return self
+    elif isinstance(self, SessionEventHandlerChain):
+      for handler in self.handlers:
+        if v := handler.get(session_cls):
+          return v
+    return None
+
   def on_session_start(
       self,
       session: 'Session'
@@ -1485,11 +1796,6 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       'An optional identifier for the session, which will be used for logging.'
   ]
 
-  event_handler: Annotated[
-      SessionEventHandler,
-      'Event handler for the session.'
-  ]
-
   @pg.explicit_method_override
   def __init__(
       self,
@@ -1503,14 +1809,33 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
     super().__init__(
         id=id,
         root=root or ActionInvocation(RootAction()),
-        event_handler=event_handler or SessionLogging(verbose=verbose),
         **kwargs
     )
+    self._event_handler = event_handler or SessionLogging(verbose=verbose)
+
+  @property
+  def event_handler(self) -> SessionEventHandler:
+    """Returns the event handler for the session."""
+    return self._event_handler
+
+  def _sym_clone(self, deep: bool, memo: Any = None) -> 'Session':
+    other = super()._sym_clone(deep=deep, memo=memo)
+    if deep:
+      event_handler = pg.clone(self.event_handler, deep=deep, memo=memo)
+    else:
+      event_handler = self.event_handler
+    other._event_handler = event_handler  # pylint: disable=protected-access
+    return other
 
   #
   # Shortcut methods for accessing the root action invocation.
   #
 
+  @property
+  def metadata(self) -> dict[str, Any]:
+    """Returns metadata associated with the root of the session."""
+    return self.root.metadata
+
   @property
   def all_queries(self) -> list[lf_structured.QueryInvocation]:
     """Returns all queries made by the session."""
@@ -1682,9 +2007,6 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
           'signal the start and end of the session.'
       )
 
-    # 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)