langfun 0.1.2.dev202509120804__py3-none-any.whl → 0.1.2.dev202512040805__py3-none-any.whl

langfun/core/agentic/action.py

@@ -14,8 +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
@@ -35,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
 
@@ -148,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)
   ```
 
@@ -187,11 +195,23 @@ class Action(pg.Object):
     self._session = None
     self._invocation: ActionInvocation | None = None
 
+    # NOTE(daiyip): Users could use `self._state` to keep track of the state
+    # during the execution of the action.
+    # Strictly speaking action state should better fit into
+    # ActionInvocation, we make it a property of Action as it's usually
+    # initialized before `__call__` is called.
+    self._state: Any = None
+
   @property
   def session(self) -> Optional['Session']:
     """Returns the session started by this action."""
     return self._session
 
+  @property
+  def state(self) -> Any:
+    """Returns the state of the action."""
+    return self._state
+
   @property
   def result(self) -> Any:
     """Returns the result of the action."""
@@ -295,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,
@@ -315,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
 
@@ -347,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:
@@ -357,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
@@ -428,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."""
@@ -448,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
@@ -460,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
@@ -525,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:
@@ -761,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,
@@ -808,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.
   #
@@ -850,9 +1176,20 @@ 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."""
-  action: 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,
+      'The action being invoked.'
+  ]
 
   result: Annotated[
       Any,
@@ -900,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']:
@@ -931,6 +1269,20 @@ class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
     """Returns True if the action invocation has an error."""
     return self.error is not None
 
+  @property
+  def state(self) -> Any:
+    """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."""
@@ -982,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:
@@ -1140,31 +1494,346 @@ class RootAction(Action):
     raise NotImplementedError('Shall not be called.')
 
 
+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'
+  ) -> None:
+    """Called when a session starts."""
+
+  def on_session_end(
+      self,
+      session: 'Session'
+  ) -> None:
+    """Called when a session ends."""
+
+  def on_action_start(
+      self,
+      session: 'Session',
+      action: ActionInvocation
+  ) -> None:
+    """Called when an action starts."""
+
+  def on_action_end(
+      self,
+      session: 'Session',
+      action: ActionInvocation
+  ) -> None:
+    """Called when an action ends."""
+
+  def on_action_progress(
+      self,
+      session: 'Session',
+      action: ActionInvocation,
+      title: str,
+      **kwargs
+  ) -> None:
+    """Called when an action progress is updated."""
+
+  def on_query_start(
+      self,
+      session: 'Session',
+      action: ActionInvocation,
+      query: lf_structured.QueryInvocation,
+  ) -> None:
+    """Called when a query starts."""
+
+  def on_query_end(
+      self,
+      session: 'Session',
+      action: ActionInvocation,
+      query: lf_structured.QueryInvocation,
+  ) -> None:
+    """Called when a query ends."""
+
+
+@dataclasses.dataclass
+class SessionEventHandlerChain(SessionEventHandler):
+  """A session event handler that chains multiple event handlers."""
+
+  handlers: list[SessionEventHandler]
+
+  def on_session_start(self, session: 'Session') -> None:
+    """Called when a session starts."""
+    for handler in self.handlers:
+      handler.on_session_start(session)
+
+  def on_session_end(self, session: 'Session') -> None:
+    """Called when a session ends."""
+    for handler in self.handlers:
+      handler.on_session_end(session)
+
+  def on_action_start(
+      self,
+      session: 'Session',
+      action: ActionInvocation) -> None:
+    """Called when an action starts."""
+    for handler in self.handlers:
+      handler.on_action_start(session, action)
+
+  def on_action_end(
+      self,
+      session: 'Session',
+      action: ActionInvocation) -> None:
+    """Called when an action ends."""
+    for handler in self.handlers:
+      handler.on_action_end(session, action)
+
+  def on_action_progress(
+      self,
+      session: 'Session',
+      action: ActionInvocation,
+      title: str,
+      **kwargs
+  ) -> None:
+    """Called when an action progress is updated."""
+    for handler in self.handlers:
+      handler.on_action_progress(session, action, title, **kwargs)
+
+  def on_query_start(
+      self,
+      session: 'Session',
+      action: ActionInvocation,
+      query: lf_structured.QueryInvocation,
+  ) -> None:
+    """Called when a query starts."""
+    for handler in self.handlers:
+      handler.on_query_start(session, action, query)
+
+  def on_query_end(
+      self,
+      session: 'Session',
+      action: ActionInvocation,
+      query: lf_structured.QueryInvocation,
+  ) -> None:
+    """Called when a query ends."""
+    for handler in self.handlers:
+      handler.on_query_end(session, action, query)
+
+
+@dataclasses.dataclass
+class SessionLogging(SessionEventHandler):
+  """An event handler that logs Session events."""
+
+  verbose: bool = False
+
+  def on_session_end(self, session: 'Session'):
+    if session.has_error:
+      session.error(
+          f'Trajectory failed in {session.elapse:.2f} seconds.',
+          error=session.final_error,
+          metadata=session.root.metadata,
+          keep=True,
+      )
+    elif self.verbose:
+      session.info(
+          f'Trajectory succeeded in {session.elapse:.2f} seconds.',
+          result=session.final_result,
+          metadata=session.root.metadata,
+          keep=False,
+      )
+
+  def on_action_start(
+      self,
+      session: 'Session',
+      action: ActionInvocation
+  ) -> None:
+    if self.verbose:
+      session.info(
+          'Action execution started.',
+          action=action.action,
+          keep=False,
+      )
+
+  def on_action_end(
+      self,
+      session: 'Session',
+      action: ActionInvocation
+  ) -> None:
+    if action.has_error:
+      session.warning(
+          (
+              f'Action execution failed in '
+              f'{action.execution.elapse:.2f} seconds.'
+          ),
+          action=action.action,
+          error=action.error,
+          keep=True,
+      )
+    elif self.verbose:
+      session.info(
+          (
+              f'Action execution succeeded in '
+              f'{action.execution.elapse:.2f} seconds.'
+          ),
+          action=action.action,
+          result=action.result,
+          keep=False,
+      )
+
+  def on_query_start(
+      self,
+      session: 'Session',
+      action: ActionInvocation,
+      query: lf_structured.QueryInvocation,
+  ) -> None:
+    if self.verbose:
+      session.info(
+          'Querying LLM started.',
+          lm=query.lm.model_id,
+          output_type=(
+              lf_structured.annotation(query.schema.spec)
+              if query.schema is not None else None
+          ),
+          keep=False,
+      )
+
+  def on_query_end(
+      self,
+      session: 'Session',
+      action: ActionInvocation,
+      query: lf_structured.QueryInvocation,
+  ) -> None:
+    if query.has_error:
+      session.warning(
+          (
+              f'Querying LLM failed in '
+              f'{time.time() - query.start_time:.2f} seconds.'
+          ),
+          lm=query.lm.model_id,
+          output_type=(
+              lf_structured.annotation(query.schema.spec)
+              if query.schema is not None else None
+          ),
+          error=query.error,
+          keep=True,
+      )
+    elif self.verbose:
+      session.info(
+          (
+              f'Querying LLM succeeded in '
+              f'{time.time() - query.start_time:.2f} seconds.'
+          ),
+          lm=query.lm.model_id,
+          output_type=(
+              lf_structured.annotation(query.schema.spec)
+              if query.schema is not None else None
+          ),
+          keep=False,
+      )
+
+
 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,
       'The root action invocation of the session.'
-  ] = ActionInvocation(RootAction())
+  ]
 
   id: Annotated[
       str | None,
-      'An optional identifier for the sessin, which will be used for logging.'
-  ] = None
+      'An optional identifier for the session, which will be used for logging.'
+  ]
 
-  verbose: Annotated[
-      bool,
-      (
-          'If True, the session will be logged with verbose action and query '
-          'activities.'
-      )
-  ] = False
+  @pg.explicit_method_override
+  def __init__(
+      self,
+      id: str | None = None,   # pylint: disable=redefined-builtin
+      *,
+      verbose: bool = False,
+      event_handler: SessionEventHandler | None = None,
+      root: ActionInvocation | None = None,
+      **kwargs
+  ):
+    super().__init__(
+        id=id,
+        root=root or ActionInvocation(RootAction()),
+        **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."""
@@ -1250,6 +1919,7 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
   def start(self) -> None:
     """Starts the session."""
     self.root.execution.start()
+    self.event_handler.on_session_start(self)
 
   def end(
       self,
@@ -1258,21 +1928,8 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       metadata: dict[str, Any] | None = None,
   ) -> None:
     """Ends the session."""
-    if error is not None:
-      self.error(
-          f'Trajectory failed in {self.elapse:.2f} seconds.',
-          error=error,
-          metadata=metadata,
-          keep=True,
-      )
-    elif self.verbose:
-      self.info(
-          f'Trajectory succeeded in {self.elapse:.2f} seconds.',
-          result=result,
-          metadata=metadata,
-          keep=False,
-      )
     self.root.end(result, error, metadata)
+    self.event_handler.on_session_end(self)
 
   def check_execution_time(self) -> None:
     """Checks the execution time of the current action."""
@@ -1291,6 +1948,20 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
           'seconds.'
       )
 
+  def update_progress(self, title: str, **kwargs: Any) -> None:
+    """Updates the progress of current action's execution.
+
+    Args:
+      title: The title of the progress update.
+      **kwargs: Additional keyword arguments to pass to the event handler.
+    """
+    self.event_handler.on_action_progress(
+        self,
+        self._current_action,
+        title,
+        **kwargs
+    )
+
   def __enter__(self):
     """Enters the session."""
     self.start()
@@ -1350,34 +2021,10 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       self._current_execution = invocation.execution
       # Start the execution of the current action.
       self._current_action.start()
-      if self.verbose:
-        self.info(
-            'Action execution started.',
-            action=invocation.action,
-            keep=False,
-        )
+      self.event_handler.on_action_start(self, self._current_action)
       yield invocation
     finally:
-      if invocation.has_error:
-        self.warning(
-            (
-                f'Action execution failed in '
-                f'{invocation.execution.elapse:.2f} seconds.'
-            ),
-            action=invocation.action,
-            error=invocation.error,
-            keep=True,
-        )
-      elif self.verbose:
-        self.info(
-            (
-                f'Action execution succeeded in '
-                f'{invocation.execution.elapse:.2f} seconds.'
-            ),
-            action=invocation.action,
-            result=invocation.result,
-            keep=False,
-        )
+      self.event_handler.on_action_end(self, self._current_action)
       self._current_execution = parent_execution
       self._current_action = parent_action
 
@@ -1403,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`
@@ -1425,51 +2079,21 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
           skip_notification=False, raise_on_no_change=False
       )
       execution.append(invocation)
-      if self.verbose:
-        self.info(
-            'Querying LLM started.',
-            lm=invocation.lm.model_id,
-            output_type=(
-                lf_structured.annotation(invocation.schema.spec)
-                if invocation.schema is not None else None
-            ),
-            keep=False,
-        )
+      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)
-      if invocation.has_error:
-        self.warning(
-            (
-                f'Querying LLM failed in '
-                f'{time.time() - invocation.start_time:.2f} seconds.'
-            ),
-            lm=invocation.lm.model_id,
-            output_type=(
-                lf_structured.annotation(invocation.schema.spec)
-                if invocation.schema is not None else None
-            ),
-            error=invocation.error,
-            keep=True,
-        )
-      elif self.verbose:
-        self.info(
-            (
-                f'Querying LLM succeeded in '
-                f'{time.time() - invocation.start_time:.2f} seconds.'
-            ),
-            lm=invocation.lm.model_id,
-            output_type=(
-                lf_structured.annotation(invocation.schema.spec)
-                if invocation.schema is not None else None
-            ),
-            keep=False,
-        )
+      self.event_handler.on_query_end(self, self._current_action, invocation)
 
     with self.track_phase(phase), lf_structured.track_queries(
         include_child_scopes=False,
-        start_callabck=_query_start,
-        end_callabck=_query_end,
+        start_callback=_query_start,
+        end_callback=_query_end,
     ) as queries:
       try:
         yield queries
@@ -1495,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:
@@ -1521,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,