langfun 0.1.2.dev202510230805__py3-none-any.whl → 0.1.2.dev202511270805__py3-none-any.whl

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
@@ -36,7 +38,12 @@ class ActionTimeoutError(ActionError):
 
 
 class Action(pg.Object):
-  """Base class for Langfun's agentic actions.
+  """Base class for agentic actions.
+
+  An `Action` represents a single, executable step or task that an agent can
+  perform, such as calling a tool, querying a language model, or returning a
+  final answer. Actions are designed to be composable and trackable within a
+  `Session`.
 
   # Developing Actions
 
@@ -149,7 +156,7 @@ class Action(pg.Object):
 
   # Explicitly create and pass a session.
   with lf.Session(id='my_agent_session') as session:
-    result = calc(session=session) # Pass the session explicitly
+    result = calc(session=session)  # Pass the session explicitly
     print(result)
   ```
 
@@ -308,19 +315,190 @@ 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,
 ]
 
 
 class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
-  """Trace of the execution of an action."""
+  """Trace of an execution, containing queries, logs, and sub-actions.
+
+  `ExecutionTrace` records the sequence of operations performed during an
+  action's execution or within a specific phase of execution (demarcated by
+  `session.track_phase`). It captures `lf.query` calls, log entries, and
+  nested `ActionInvocation` objects in the order they occurred. It also
+  aggregates LLM usage summaries from its child items.
+  """
 
   name: Annotated[
       str | None,
@@ -328,7 +506,7 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
           'The name of the execution trace. If None, the trace is unnamed, '
           'which is the case for the top-level trace of an action. An '
           'execution trace could have sub-traces, called phases, which are '
-          'created and named by `session.phase()` context manager.'
+          'created and named by `session.track_phase()` context manager.'
       )
   ] = None
 
@@ -360,9 +538,15 @@ 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:
-    """Returns the index of the child items of given type."""
+  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):
       if x is item:
@@ -370,6 +554,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
@@ -441,6 +671,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."""
@@ -461,7 +696,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
@@ -473,7 +710,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
@@ -538,6 +778,18 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
           remove_class=['not-started'],
       )
 
+  def remove(self, item: TracedItem) -> None:
+    """Removes an item from the sequence."""
+    index = self.items.index(item)
+    if index == -1:
+      raise ValueError(f'Item not found in execution trace: {item!r}')
+
+    with pg.notify_on_change(False):
+      self.items.pop(index)
+
+    if self._tab_control is not None:
+      self._tab_control.remove(index)
+
   def extend(self, items: Iterable[TracedItem]) -> None:
     """Extends the sequence with a list of items."""
     for item in items:
@@ -774,8 +1026,13 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
     ]
 
 
-class ParallelExecutions(pg.Object, pg.views.html.HtmlTreeView.Extension):
-  """A class for encapsulating parallel execution traces."""
+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`
+  object to hold an `ExecutionTrace` for each parallel branch of execution,
+  allowing inspection of parallel workflows.
+  """
 
   name: Annotated[
       str | None,
@@ -821,8 +1078,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.
   #
@@ -863,8 +1176,15 @@ class ParallelExecutions(pg.Object, pg.views.html.HtmlTreeView.Extension):
     )
 
 
-class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
-  """A class for capturing the invocation of an action."""
+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
+  the `Action` object itself, its result or error, associated metadata,
+  and an `ExecutionTrace` detailing the steps taken during its execution
+  (queries, logs, sub-actions). Invocations form a tree structure within a
+  `Session`, reflecting the hierarchy of agentic operations.
+  """
 
   action: Annotated[
       Action,
@@ -917,6 +1237,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']:
@@ -953,6 +1274,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."""
@@ -1004,10 +1334,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:
@@ -1165,6 +1497,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'
@@ -1394,7 +1739,50 @@ class SessionLogging(SessionEventHandler):
 
 
 class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
-  """Session for performing an agentic task."""
+  """Manages the execution trajectory of agentic actions.
+
+  A `Session` tracks the execution of a root `Action` and all its
+  sub-actions, including LLM queries (`lf.query`), logging messages,
+  and nested actions. It provides a complete, hierarchical trace of an
+  agent's workflow, which is important for debugging, analysis, and
+  visualization.
+
+  Sessions can be created implicitly when an action is called without an
+  active session, or explicitly for more control.
+
+  **1. Implicit Session:**
+  When an action is called without a session, Langfun creates one automatically.
+
+  ```python
+  action = MyAction()
+  action()
+  session = action.session  # Access the implicit session
+  ```
+
+  **2. Explicit Session:**
+  Use a `with` statement to manage a session explicitly. This is useful for
+  setting session IDs or capturing the trajectory of multiple top-level actions.
+
+  ```python
+  with lf.Session(id='my-session') as session:
+    action1()
+    action2()
+  ```
+
+  **3. Accessing Trajectory:**
+  The `session.root` attribute provides access to the `ActionInvocation` tree.
+
+  ```python
+  with lf.Session() as session:
+    my_action()
+
+  # Get all queries in the session
+  print(session.all_queries)
+
+  # Get all top-level action calls in the session
+  print(session.root.actions)
+  ```
+  """
 
   root: Annotated[
       ActionInvocation,
@@ -1406,11 +1794,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,
@@ -1424,14 +1807,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."""
@@ -1547,7 +1949,7 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       )
 
   def update_progress(self, title: str, **kwargs: Any) -> None:
-    """Update the progress of current action's execution.
+    """Updates the progress of current action's execution.
 
     Args:
       title: The title of the progress update.
@@ -1648,13 +2050,20 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
   @contextlib.contextmanager
   def track_queries(
       self,
-      phase: str | None = None
+      phase: str | None = None,
+      track_if: Callable[
+          [lf_structured.QueryInvocation],
+          bool
+      ] | None = None,
   ) -> Iterator[list[lf_structured.QueryInvocation]]:
     """Tracks `lf.query` made within the context.
 
     Args:
       phase: The name of a new phase to track the queries in. If not provided,
         the queries will be tracked in the parent phase.
+      track_if: A function that takes a `lf_structured.QueryInvocation` and
+        returns True if the query should be included in the result. If None,
+        all queries (including failed queries) will be included.
 
     Yields:
       A list of `lf.QueryInvocation` objects, each for a single `lf.query`
@@ -1673,6 +2082,11 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       self.event_handler.on_query_start(self, self._current_action, invocation)
 
     def _query_end(invocation: lf_structured.QueryInvocation):
+      if track_if is not None and not track_if(invocation):
+        self._current_execution.remove(invocation)
+      # Even if the query is not included in the execution trace, we still
+      # count the usage summary to the current execution and trigger the
+      # event handler to log the query.
       self._current_execution.merge_usage_summary(invocation.usage_summary)
       self.event_handler.on_query_end(self, self._current_action, invocation)
 
@@ -1705,8 +2119,9 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       *,
       lm: lf.LanguageModel,
       examples: list[lf_structured.MappingExample] | None = None,
+      track_if: Callable[[lf_structured.QueryInvocation], bool] | None = None,
       **kwargs
-      ) -> Any:
+  ) -> Any:
     """Calls `lf.query` and associates it with the current invocation.
 
     The following code are equivalent:
@@ -1731,12 +2146,15 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       default: The default value to return if the query fails.
       lm: The language model to use for the query.
       examples: The examples to use for the query.
+      track_if: A function that takes a `lf_structured.QueryInvocation`
+        and returns True if the query should be tracked.
+        If None, all queries (including failed queries) will be tracked.
       **kwargs: Additional keyword arguments to pass to `lf.query`.
 
     Returns:
       The result of the query.
     """
-    with self.track_queries():
+    with self.track_queries(track_if=track_if):
       return lf_structured.query(
           prompt,
           schema=schema,