langfun 0.1.2.dev202510140804__tar.gz → 0.1.2.dev202510160805__tar.gz

{langfun-0.1.2.dev202510140804 → langfun-0.1.2.dev202510160805}/PKG-INFO

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

{langfun-0.1.2.dev202510140804 → langfun-0.1.2.dev202510160805}/langfun/core/agentic/__init__.py

@@ -26,6 +26,10 @@ from langfun.core.agentic.action import ExecutionTrace
 from langfun.core.agentic.action import ParallelExecutions
 from langfun.core.agentic.action import Session
 
+from langfun.core.agentic.action import SessionEventHandler
+from langfun.core.agentic.action import SessionEventHandlerChain
+from langfun.core.agentic.action import SessionLogging
+
 from langfun.core.agentic.action_eval import ActionEval
 from langfun.core.agentic.action_eval import ActionEvalV1
 

{langfun-0.1.2.dev202510140804 → langfun-0.1.2.dev202510160805}/langfun/core/agentic/action.py

@@ -15,6 +15,7 @@
 
 import abc
 import contextlib
+import dataclasses
 import functools
 import threading
 import time
@@ -1161,26 +1162,271 @@ 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."""
 
   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.
@@ -1271,6 +1517,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,
@@ -1279,21 +1526,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."""
@@ -1312,6 +1546,20 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
           'seconds.'
       )
 
+  def update_progress(self, title: str, **kwargs: Any) -> None:
+    """Update 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()
@@ -1371,34 +1619,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
 
@@ -1446,51 +1670,16 @@ 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):
       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

{langfun-0.1.2.dev202510140804 → langfun-0.1.2.dev202510160805}/langfun/core/agentic/action_test.py

@@ -34,6 +34,7 @@ class Bar(action_lib.Action):
     time.sleep(self.simulate_execution_time)
     session.query('bar', lm=lm)
     session.add_metadata(note='bar')
+    session.update_progress('Query completed')
     if self.simulate_action_error:
       raise ValueError('Bar error')
     return 2 + pg.contextual_value('baz', 0)
@@ -128,7 +129,7 @@ class SessionTest(unittest.TestCase):
     self.assertIsNone(foo.result)
     self.assertIsNone(foo.metadata)
 
-    session = action_lib.Session(id='agent@1')
+    session = action_lib.Session(id='agent@1', verbose=True)
     self.assertEqual(session.id, 'agent@1')
     self.assertFalse(session.has_started)
     self.assertFalse(session.has_stopped)
@@ -137,7 +138,7 @@ class SessionTest(unittest.TestCase):
     _ = session.to_html()
 
     with session:
-      result = foo(session, lm=lm, verbose=True)
+      result = foo(session, lm=lm)
 
     self.assertTrue(session.has_started)
     self.assertTrue(session.has_stopped)
@@ -375,7 +376,7 @@ class SessionTest(unittest.TestCase):
     self.assertFalse(session.has_stopped)
 
     session.start()
-    result = foo(session, lm=lm, verbose=True)
+    result = foo(session, lm=lm)
     session.end(result)
 
     self.assertTrue(session.has_started)
@@ -395,7 +396,7 @@ class SessionTest(unittest.TestCase):
     session = action_lib.Session(id='agent@1')
     with self.assertRaisesRegex(ValueError, 'Bar error'):
       with session:
-        foo(session, lm=lm, verbose=True)
+        foo(session, lm=lm)
     self.assertTrue(session.has_started)
     self.assertTrue(session.has_stopped)
     self.assertTrue(session.has_error)
@@ -408,7 +409,7 @@ class SessionTest(unittest.TestCase):
     foo = Foo(1, simulate_action_error=True)
     session = action_lib.Session(id='agent@1')
     with self.assertRaisesRegex(ValueError, 'Please call `Session.start'):
-      foo(session, lm=lm, verbose=True)
+      foo(session, lm=lm)
 
   def test_succeed_with_multiple_actions(self):
     lm = fake.StaticResponse('lm response')
@@ -489,6 +490,33 @@ class SessionTest(unittest.TestCase):
     ):
       foo(lm=lm, max_execution_time=1.0)
 
+  def test_event_handler(self):
+
+    class MyActionHandler(pg.Object, action_lib.SessionEventHandler):
+      def _on_bound(self):
+        super()._on_bound()
+        self.progresses = []
+
+      def on_action_progress(self, session, action, title, **kwargs):
+        self.progresses.append((action.id, title))
+
+    handler = MyActionHandler()
+    session = action_lib.Session(
+        id='agent@1',
+        event_handler=action_lib.SessionEventHandlerChain(
+            handlers=[handler, action_lib.SessionLogging()]
+        )
+    )
+    bar = Bar()
+    with session:
+      bar(session, lm=fake.StaticResponse('lm response'))
+      session.update_progress('Trajectory completed')
+
+    self.assertEqual(handler.progresses, [
+        ('agent@1:/a1', 'Query completed'),
+        ('agent@1:', 'Trajectory completed'),
+    ])
+
   def test_log(self):
     session = action_lib.Session()
     session.debug('hi', x=1, y=2)

{langfun-0.1.2.dev202510140804 → langfun-0.1.2.dev202510160805}/langfun/core/eval/v2/progress_tracking_test.py

@@ -14,6 +14,7 @@
 import contextlib
 import io
 import os
+import sys
 import tempfile
 import unittest
 
@@ -49,6 +50,7 @@ class TqdmProgressTrackerTest(unittest.TestCase):
     string_io = io.StringIO()
     with contextlib.redirect_stderr(string_io):
       _ = experiment.run(root_dir, 'new', plugins=[])
+    sys.stderr.flush()
     self.assertIn('All: 100%', string_io.getvalue())
 
   def test_with_example_ids(self):
@@ -59,6 +61,7 @@ class TqdmProgressTrackerTest(unittest.TestCase):
     string_io = io.StringIO()
     with contextlib.redirect_stderr(string_io):
       _ = experiment.run(root_dir, 'new', example_ids=[1], plugins=[])
+    sys.stderr.flush()
     self.assertIn('All: 100%', string_io.getvalue())
 
 

{langfun-0.1.2.dev202510140804 → langfun-0.1.2.dev202510160805}/langfun/core/structured/querying.py

@@ -1209,14 +1209,14 @@ class _QueryTracker:
       )
   ] = True
 
-  start_callabck: Annotated[
+  start_callback: Annotated[
       Callable[[QueryInvocation], None] | None,
       (
           'A callback function to be called when a query is started.'
       )
   ] = None
 
-  end_callabck: Annotated[
+  end_callback: Annotated[
       Callable[[QueryInvocation], None] | None,
       (
           'A callback function to be called when a query is completed.'
@@ -1232,21 +1232,21 @@ class _QueryTracker:
 
   def track(self, invocation: QueryInvocation) -> None:
     self.tracked_queries.append(invocation)
-    if self.start_callabck is not None:
-      self.start_callabck(invocation)
+    if self.start_callback is not None:
+      self.start_callback(invocation)
 
   def mark_completed(self, invocation: QueryInvocation) -> None:
     assert invocation in self.tracked_queries, invocation
-    if self.end_callabck is not None:
-      self.end_callabck(invocation)
+    if self.end_callback is not None:
+      self.end_callback(invocation)
 
 
 @contextlib.contextmanager
 def track_queries(
     include_child_scopes: bool = True,
     *,
-    start_callabck: Callable[[QueryInvocation], None] | None = None,
-    end_callabck: Callable[[QueryInvocation], None] | None = None,
+    start_callback: Callable[[QueryInvocation], None] | None = None,
+    end_callback: Callable[[QueryInvocation], None] | None = None,
 ) -> Iterator[list[QueryInvocation]]:
   """Track all queries made during the context.
 
@@ -1264,8 +1264,8 @@ def track_queries(
     include_child_scopes: If True, the queries made in child scopes will be
       included in the returned list. Otherwise, only the queries made in the
       current scope will be included.
-    start_callabck: A callback function to be called when a query is started.
-    end_callabck: A callback function to be called when a query is completed.
+    start_callback: A callback function to be called when a query is started.
+    end_callback: A callback function to be called when a query is completed.
 
   Yields:
     A list of `QueryInvocation` objects representing the queries made during
@@ -1274,8 +1274,8 @@ def track_queries(
   trackers = lf.context_value('__query_trackers__', [])
   tracker = _QueryTracker(
       include_child_scopes=include_child_scopes,
-      start_callabck=start_callabck,
-      end_callabck=end_callabck
+      start_callback=start_callback,
+      end_callback=end_callback
   )
 
   with lf.context(

{langfun-0.1.2.dev202510140804 → langfun-0.1.2.dev202510160805}/langfun/core/structured/querying_test.py

@@ -1590,7 +1590,7 @@ class TrackQueriesTest(unittest.TestCase):
         'bar',
     ])
     state = {}
-    def start_callabck(query):
+    def start_callback(query):
       self.assertFalse(query.is_completed)
       self.assertIsNone(query.end_time)
       elapse1 = query.elapse
@@ -1615,7 +1615,7 @@ class TrackQueriesTest(unittest.TestCase):
       state['end'] = query
 
     with querying.track_queries(
-        start_callabck=start_callabck, end_callabck=end_callback
+        start_callback=start_callback, end_callback=end_callback
     ) as queries:
       querying.query('foo', lm=lm)
     self.assertIs(state['start'], queries[0])
@@ -1626,7 +1626,7 @@ class TrackQueriesTest(unittest.TestCase):
         'bar',
     ])
     state = {}
-    def start_callabck(query):
+    def start_callback(query):
       self.assertFalse(query.is_completed)
       self.assertIsNone(query.end_time)
       self.assertIsNotNone(query.usage_summary)
@@ -1648,7 +1648,7 @@ class TrackQueriesTest(unittest.TestCase):
 
     with self.assertRaises(mapping.MappingError):
       with querying.track_queries(
-          start_callabck=start_callabck, end_callabck=end_callback
+          start_callback=start_callback, end_callback=end_callback
       ) as queries:
         querying.query('foo', int, lm=lm)
     self.assertIs(state['start'], queries[0])

{langfun-0.1.2.dev202510140804 → langfun-0.1.2.dev202510160805}/langfun.egg-info/PKG-INFO

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