langfun 0.1.2.dev202510150805__tar.gz → 0.1.2.dev202510170805__tar.gz

{langfun-0.1.2.dev202510150805 → langfun-0.1.2.dev202510170805}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langfun
-Version: 0.1.2.dev202510150805
+Version: 0.1.2.dev202510170805
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors
@@ -21,11 +21,13 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: Software Development :: Libraries
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: anyio>=4.7.0
 Requires-Dist: jinja2>=3.1.2
 Requires-Dist: puremagic>=1.20
 Requires-Dist: pyglove>=0.4.5.dev202507140812
 Requires-Dist: requests>=2.31.0
 Provides-Extra: all
+Requires-Dist: anyio>=4.7.0; extra == "all"
 Requires-Dist: jinja2>=3.1.2; extra == "all"
 Requires-Dist: puremagic>=1.20; extra == "all"
 Requires-Dist: pyglove>=0.4.5.dev202507140812; extra == "all"

{langfun-0.1.2.dev202510150805 → langfun-0.1.2.dev202510170805}/langfun/core/__init__.py

@@ -40,9 +40,13 @@ from langfun.core.component import context
 as_context = context
 use_context = context
 
-# Invoke a callable object asynchronously.
+# Support for async IO.
 from langfun.core.async_support import invoke_async
 
+# Adaptors for async function/context manager to sync versions.
+from langfun.core.async_support import invoke_sync
+from langfun.core.async_support import sync_context_manager
+
 # Shortcut function for overriding components attributes, usually for
 # override settings.
 from langfun.core.component import use_settings

{langfun-0.1.2.dev202510150805 → langfun-0.1.2.dev202510170805}/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.dev202510150805 → langfun-0.1.2.dev202510170805}/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.dev202510150805 → langfun-0.1.2.dev202510170805}/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.dev202510170805/langfun/core/async_support.py

@@ -0,0 +1,98 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Utility for async IO in Langfun."""
+
+import asyncio
+import contextlib
+from typing import Any, Awaitable, Callable, Iterator
+import anyio
+import pyglove as pg
+
+
+async def invoke_async(
+    sync_callable: Callable[..., Any], *args, **kwargs
+) -> Any:
+  """Invokes a callable asynchronously with `lf.context` manager enabled."""
+  return await asyncio.to_thread(
+      # Enable `lf.context` manager for async calls.
+      pg.with_contextual_override(sync_callable), *args, **kwargs
+  )
+
+
+def invoke_sync(
+    async_callable: Callable[..., Awaitable[Any]],
+    *args,
+    **kwargs
+) -> Any:
+  """Invokes a async callable synchronously."""
+  async def _invoke():
+    return await async_callable(*args, **kwargs)
+  invoke_fn = pg.with_contextual_override(_invoke)
+  blocking_portal = pg.utils.thread_local_get('__blocking_portal__', None)
+  if blocking_portal is None:
+    return anyio.run(invoke_fn)
+  return blocking_portal.call(invoke_fn)
+
+
+@contextlib.contextmanager
+def sync_context_manager(
+    async_context_manager: contextlib.AbstractAsyncContextManager[Any]
+) -> Iterator[Any]:
+  """Adapts an async context manager to a sync context manager.
+
+  sync_context_manager installs a blocking portal in current thread to run the
+  async context manager in a blocking way. It's useful for running async code in
+  sync context managers, e.g. `sync_context_manager` can be nested and share the
+  same event loop.
+
+  Example:
+
+    ```python
+    @contextlib.asynccontextmanager
+    async def foo(x):
+      try:
+        yield x
+      finally:
+        pass
+
+    with lf.sync_context_manager(foo(x)) as x
+      with lf.sync_context_manager(foo(y)) as y:
+        ...
+    ```
+
+  Args:
+    async_context_manager: The async context manager to adapt.
+
+  Yields:
+    The value yielded by the async context manager.
+  """
+  blocking_portal = pg.utils.thread_local_get('__blocking_portal__', None)
+  portal_exit_stack = None
+
+  try:
+    if blocking_portal is None:
+      portal_exit_stack = contextlib.ExitStack()
+      blocking_portal = portal_exit_stack.enter_context(
+          anyio.from_thread.start_blocking_portal()
+      )
+      pg.utils.thread_local_set('__blocking_portal__', blocking_portal)
+    context_manager = blocking_portal.wrap_async_context_manager(
+        async_context_manager
+    )
+    with context_manager as value:
+      yield value
+  finally:
+    if portal_exit_stack is not None:
+      portal_exit_stack.close()
+      pg.utils.thread_local_del('__blocking_portal__')

{langfun-0.1.2.dev202510150805 → langfun-0.1.2.dev202510170805}/langfun/core/async_support_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import contextlib
 import time
 import unittest
 
@@ -34,6 +35,28 @@ class AsyncSupportTest(unittest.TestCase):
     with pg.contextual_override(z=3):
       self.assertEqual(asyncio.run(r), 6)
 
+  def test_invoke_sync(self):
+    @contextlib.asynccontextmanager
+    async def bar(x):
+      try:
+        yield x
+      finally:
+        pass
+
+    async def foo(x, *, y):
+      time.sleep(2)
+      return x + y + pg.contextual_value('z', 0)
+
+    with pg.contextual_override(z=3):
+      with async_support.sync_context_manager(bar(1)) as x:
+        self.assertEqual(x, 1)
+        with async_support.sync_context_manager(bar(2)) as y:
+          self.assertEqual(y, 2)
+          self.assertEqual(async_support.invoke_sync(foo, 1, y=2), 6)
+
+    with pg.contextual_override(z=2):
+      self.assertEqual(async_support.invoke_sync(foo, 1, y=2), 5)
+
 
 if __name__ == '__main__':
   unittest.main()