langfun 0.1.2.dev202508250805__py3-none-any.whl → 0.1.2.dev202511110805__py3-none-any.whl

langfun/core/agentic/action.py

@@ -15,6 +15,7 @@
 
 import abc
 import contextlib
+import dataclasses
 import functools
 import threading
 import time
@@ -35,7 +36,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 +154,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 +193,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."""
@@ -307,7 +325,14 @@ TracedItem = Union[
 
 
 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 +340,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
 
@@ -349,7 +374,7 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
     self.__dict__.pop('id', None)
 
   def indexof(self, item: TracedItem, count_item_cls: Type[Any]) -> int:
-    """Returns the index of the child items of given type."""
+    """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:
@@ -525,6 +550,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:
@@ -762,7 +799,12 @@ 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."""
+  """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,
@@ -851,8 +893,19 @@ 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
+  """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,
@@ -931,6 +984,11 @@ 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
+
   @property
   def logs(self) -> list[lf.logging.LogEntry]:
     """Returns immediate child logs from execution sequence."""
@@ -1140,26 +1198,314 @@ class RootAction(Action):
     raise NotImplementedError('Shall not be called.')
 
 
+class SessionEventHandler:
+  """Interface for handling session events."""
+
+  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
+  event_handler: Annotated[
+      SessionEventHandler,
+      'Event handler for the session.'
+  ]
+
+  @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()),
+        event_handler=event_handler or SessionLogging(verbose=verbose),
+        **kwargs
+    )
 
   #
   # Shortcut methods for accessing the root action invocation.
@@ -1250,6 +1596,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 +1605,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 +1625,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 +1698,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 +1727,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 +1756,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 +1796,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 +1823,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,
@@ -1546,7 +1851,9 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       timeout: int | None = None,
       silence_on_errors: Union[
           Type[BaseException], tuple[Type[BaseException], ...], None
-      ] = Exception
+      ] = Exception,
+      ordered: bool = False,
+      show_progress: bool | int = False,
   ) -> Iterator[Any]:
     """Starts and tracks parallel execution with `lf.concurrent_map`."""
     self.check_execution_time()
@@ -1574,7 +1881,9 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
         parallel_inputs,
         max_workers=max_workers,
         timeout=self._child_max_execution_time(timeout),
-        silence_on_errors=silence_on_errors
+        silence_on_errors=silence_on_errors,
+        ordered=ordered,
+        show_progress=show_progress,
     ):
       yield input_value, result, error