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

langfun/core/agentic/action_eval.py

@@ -24,7 +24,14 @@ import pyglove as pg
 
 
 class ActionEval(lf.eval.v2.Evaluation):
-  """Agent evaluation."""
+  """Evaluation for agentic actions.
+
+  `ActionEval` is a specialized evaluation class for executing and evaluating
+  agentic actions based on provided inputs. Each input example is expected to
+  contain an `action` attribute. The `process` method executes the action
+  within a dedicated `Session`, captures the final result, and returns it
+  along with the session details in the metadata.
+  """
 
   action_args: Annotated[
       dict[str, Any],
@@ -68,7 +75,7 @@ class ExampleView(pg.Object):
 class ActionEvalV1(lf_eval.Matching):
   """Base class for action evaluations.
 
-  The input function should returns a list of pg.Dict, with `action` and
+  The input function should return a list of pg.Dict, with `action` and
   `groundtruth` fields.
   """
   # We override the schema and prompt to dummy values since they are not used.

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)
@@ -66,12 +67,18 @@ class Foo(action_lib.Action):
       time.sleep(self.simulate_execution_time[2])
       return lf_structured.query(f'subtask_{i}', lm=lm)
 
+    self._state = []
     for i, output, error in session.concurrent_map(
-        _sub_task, range(3), max_workers=2, silence_on_errors=None,
+        _sub_task,
+        range(3),
+        max_workers=2,
+        ordered=True,
+        silence_on_errors=None,
     ):
       assert isinstance(i, int), i
       assert isinstance(output, str), output
       assert error is None, error
+      self._state.append(i)
     return self.x + Bar(
         simulate_action_error=self.simulate_action_error,
         simulate_execution_time=self.simulate_execution_time[3]
@@ -118,10 +125,11 @@ class SessionTest(unittest.TestCase):
     foo = Foo(1)
     self.assertIsNone(foo.session)
     self.assertIsNone(foo.invocation)
+    self.assertIsNone(foo.state)
     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)
@@ -130,12 +138,14 @@ 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)
     self.assertEqual(result, 3)
     self.assertIsNone(foo.session)
+    self.assertEqual(foo.state, [0, 1, 2])
+    self.assertIs(foo.invocation.state, foo.state)
     self.assertEqual(foo.result, 3)
     self.assertEqual(
         foo.metadata, dict(note='foo', subtask_0=0, subtask_1=1, subtask_2=2)
@@ -366,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)
@@ -386,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)
@@ -399,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')
@@ -480,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)
@@ -493,6 +530,31 @@ class SessionTest(unittest.TestCase):
     self.assertIn('agent@', session.id)
     self.assertIsInstance(session.as_message(), lf.AIMessage)
 
+  def test_query_with_track_if(self):
+    lm = fake.StaticResponse('lm response')
+    session = action_lib.Session()
+
+    # Render session to trigger javascript updates to the HTML when
+    # operating on the session.
+    _ = session.to_html()
+    with session:
+      # This query will succeed.
+      session.query(
+          'prompt1',
+          schema=None,
+          lm=lm,
+          track_if=lambda q: not q.has_error,
+          default=None)
+      # This query will fail during parsing.
+      session.query(
+          'prompt2',
+          schema=int,
+          lm=lm,
+          track_if=lambda q: not q.has_error,
+          default=None)
+    self.assertEqual(len(session.root.queries), 1)
+    self.assertIsNone(session.root.queries[0].error)
+
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/async_support.py

@@ -11,18 +11,117 @@
 # 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."""
+"""Utilities for asynchronous programming in Langfun."""
 
 import asyncio
-from typing import Any, Callable
+import contextlib
+from typing import Any, Awaitable, Callable, Iterator
+import anyio
 import pyglove as pg
 
 
 async def invoke_async(
-    callable_object: Callable[..., Any], *args, **kwargs
+    sync_callable: Callable[..., Any], *args, **kwargs
 ) -> Any:
-  """Invokes a callable asynchronously with `lf.context` manager enabled."""
+  """Invokes a sync callable asynchronously in a separate thread.
+
+  This is useful for wrapping a sync function into an async function,
+  allowing multiple calls of the sync function to run concurrently.
+  `lf.context` will be propagated to the thread that runs the sync callable.
+
+  Args:
+    sync_callable: The sync callable to invoke.
+    *args: Positional arguments to pass to the callable.
+    **kwargs: Keyword arguments to pass to the callable.
+
+  Returns:
+    An awaitable that resolves to the return value of the sync_callable.
+  """
   return await asyncio.to_thread(
       # Enable `lf.context` manager for async calls.
-      pg.with_contextual_override(callable_object), *args, **kwargs
+      pg.with_contextual_override(sync_callable), *args, **kwargs
   )
+
+
+def invoke_sync(
+    async_callable: Callable[..., Awaitable[Any]],
+    *args,
+    **kwargs
+) -> Any:
+  """Invokes an async callable synchronously.
+
+  This is useful for calling an async function from a sync context.
+  If there is an existing async event loop in current thread managed by
+  `lf.sync_context_manager`, it will be used for running the async callable.
+  Otherwise, `anyio.run` will be used to run the async callable in a new
+  event loop.
+  `lf.context` will be propagated to the async callable.
+
+  Args:
+    async_callable: The async callable to invoke.
+    *args: Positional arguments to pass to the callable.
+    **kwargs: Keyword arguments to pass to the callable.
+
+  Returns:
+    The return value of the async_callable.
+  """
+  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/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()

langfun/core/coding/python/correction.py

@@ -19,13 +19,23 @@ import pyglove as pg
 
 
 class CodeWithError(pg.Object):
-  """Python code with error."""
+  """A structure representing Python code along with an execution error.
+
+  This is used as input to a language model for error correction, providing
+  the model with the code that failed and the error message it produced.
+  """
 
   code: str
   error: str
 
 
 class CorrectedCode(pg.Object):
+  """A structure containing corrected Python code.
+
+  This is used as the output schema when asking a language model to correct
+  code, expecting the model to return the fixed code in the `corrected_code`
+  field.
+  """
   corrected_code: str
 
 
@@ -49,7 +59,7 @@ def run_with_correction(
     code: The source code that may or may not be problematic.
     error: An optional initial error for `code` when it's problematic, usually
       caught from elsewhere when it ran. If None, code will be executed once to
-      verify if its good and obtain a feedback error message.
+      verify if it's good and obtain a feedback error message.
     global_vars: A dict of str to value as the global variables that could be
       accessed within the corrected code.
     lm: Language model to be used. If not specified, it will try to use the `lm`
@@ -57,15 +67,15 @@ def run_with_correction(
     max_attempts: Max number of attempts for the correction.
     sandbox: If True, run code in sandbox; If False, run code in current
       process. If None, run in sandbox first, if the output could not be
-      serialized and pass to current process, run the code again in current
+      serialized and passed to current process, run the code again in current
       process.
     permission: The permission to run the code.
     timeout: The timeout for running the corrected code. If None, there is no
       timeout. Applicable only when sandbox is set to True.
     returns_code: If True, the return value is a tuple of (result, final code).
       Otherwise the return value is the result only.
-    returns_stdout: If True, the stdout (a str) will be returned.
-    outputs_intermediate: If True, intermediate output will be outputted as a
+    returns_stdout: If True, the stdout (a string) will be returned.
+    outputs_intermediate: If True, intermediate output will be output as a
       dict, with the last line's value accessible by key '__result__'. Otherwise
       the value of the last line will be returned.
 
@@ -161,7 +171,7 @@ def correct(
     code: The source code that may or may not be problematic.
     error: An optional initial error for `code` when it's problematic, usually
       caught from elsewhere when it ran. If None, code will be executed once to
-      verify if its good and obtain a feedback error message.
+      verify if it's good and obtain a feedback error message.
     global_vars: A dict of str to value as the global variables that could be
       accessed within the corrected code.
     lm: Language model to be used. If not specified, it will try to use the `lm`
@@ -169,7 +179,7 @@ def correct(
     max_attempts: Max number of attempts for the correction.
     sandbox: If True, run code in sandbox; If False, run code in current
       process. If None, run in sandbox first, if the output could not be
-      serialized and pass to current process, run the code again in current
+      serialized and passed to current process, run the code again in current
       process.
     timeout: The timeout for running the corrected code. If None, there is no
       timeout. Applicable only when sandbox is set to True.
@@ -193,7 +203,7 @@ def correct(
 
 
 def _error_feedback_str(error: Exception) -> str:
-  """Returns the error str for feedback."""
+  """Returns the error string for feedback."""
   if isinstance(error, pg.coding.CodeError):
     return pg.decolor(error.format(include_complete_code=False))
   else:
@@ -201,7 +211,7 @@ def _error_feedback_str(error: Exception) -> str:
 
 
 def _maybe_custom_validate(result: Any) -> Any:
-  """Apply custom validation through __validate_generation__ method."""
+  """Applies custom validation through __validate__ method."""
   if isinstance(result, dict) and "__result__" in result:
     r = result["__result__"]
   else:

langfun/core/coding/python/execution.py

@@ -45,17 +45,17 @@ def evaluate(
     global_vars: An optional dict as the globals that could be referenced by the
       code.
     permission: Permission for the Python code to run.
-    returns_stdout: If True, the stdout (a str) will be returned.
+    returns_stdout: If True, the stdout (a string) will be returned.
     outputs_intermediate: Applicable when returns_stdout is False. If True,
-      intermediate output will be outputted as a dict, with the last line's
-      value accessible by key '__result__' and the std output accessible by
+      intermediate output will be output as a dict, with the last line's
+      value accessible by key '__result__' and the stdout accessible by
       key '__stdout__'. Otherwise the value of the last line will be returned.
 
   Returns:
     The value of the last line of the code block. Or a dict of variable
     names of all locals to their evaluated values as the output of the code to
     run. The value for the last line can be accessed by key '__result__'. Or the
-    stdout as a str.
+    stdout as a string.
   """
   return pg.coding.evaluate(
       parsing.clean(code),
@@ -85,28 +85,30 @@ def run(
 
   Args:
     code: Python code to run.
-    global_vars: An optional dict of
+    global_vars: An optional dict as the globals that could be referenced by the
+      code.
     permission: Permission for the Python code to run.
-    returns_stdout: If True, the stdout (a str) will be returned.
+    returns_stdout: If True, the stdout (a string) will be returned.
     outputs_intermediate: Applicable when returns_stdout is False. If True,
-      intermediate output will be outputted as a dict, with the last line's
-      value accessible by key '__result__' and the std output accessible by
+      intermediate output will be output as a dict, with the last line's
+      value accessible by key '__result__' and the stdout accessible by
       key '__stdout__'. Otherwise the value of the last line will be returned.
     sandbox: If True, run code in sandbox; If False, run code in current
       process. If None, run in sandbox first, if the output could not be
-      serialized and pass to current process, run the code again in current
+      serialized and passed to current process, run the code again in current
       process.
-    timeout: Execution timeout in seconds. If None, wait the code the complete.
+    timeout: Execution timeout in seconds. If None, wait for the code to
+      complete.
 
   Returns:
     The value of the last line of the code block. Or a dict of variable
     names of all locals to their evaluated values as the output of the code to
     run. The value for the last line can be accessed by key '__result__'. Or the
-    stdout as a str.
+    stdout as a string.
 
   Raises:
     TimeoutError: If the execution time exceeds the timeout.
-    Exception: Exception  that are raised from the code.
+    Exception: Exceptions that are raised from the code.
   """
   return pg.coding.run(
       parsing.clean(code),

langfun/core/coding/python/generation.py

@@ -22,9 +22,13 @@ import pyglove as pg
 
 
 class PythonCode(pg.Object):
-  """Symbolic class for Python code.
+  """Represents a piece of Python code that can be executed.
 
-  The value of the last expression of the source will be the returned value.
+  When `PythonCode` is instantiated within a `PythonCode.auto_run()` context,
+  it automatically executes the code and returns the result of the last
+  expression. Otherwise, it acts as a container for the source code, which
+  can be executed by calling the instance. The class also supports automatic
+  error correction via `lf.coding.run_with_correction` when called.
   """
 
   source: Annotated[
@@ -56,7 +60,7 @@ class PythonCode(pg.Object):
         Otherwise, auto call will be disabled.
       sandbox: If True, run code in sandbox; If False, run code in current
         process. If None, run in sandbox first, if the output could not be
-        serialized and pass to current process, run the code again in current
+        serialized and passed to current process, run the code again in current
         process. Applicable when `enabled` is set to True.
       timeout: Timeout in seconds. Applicable when both `enabled` and `sandbox`
         are set to True.
@@ -98,17 +102,17 @@ class PythonCode(pg.Object):
     Args:
       sandbox: If True, run code in sandbox; If False, run code in current
         process. If None, run in sandbox first, if the output could not be
-        serialized and pass to current process, run the code again in current
+        serialized and passed to current process, run the code again in current
         process.
       timeout: Timeout in seconds. If None, there is no timeout. Applicable when
         sandbox is set to True.
       global_vars: Global variables that could be accessed from the source code.
-      returns_stdout: If True, the stdout (a str) will be returned.
+      returns_stdout: If True, the stdout (a string) will be returned.
       outputs_intermediate: Applicable when returns_stdout is False. If True,
-        intermediate output will be outputted as a dict, with the last line's
-        value accessible by key '__result__' and the std output accessible by 
+        intermediate output will be output as a dict, with the last line's
+        value accessible by key '__result__' and the stdout accessible by
         key '__stdout__'. Otherwise the value of the last line will be returned.
-      autofix: Number of attempts to auto fix the generated code. If 0, autofix
+      autofix: Number of attempts to autofix the generated code. If 0, autofix
         is disabled.
       autofix_lm: Language model to be used. If not specified, it will try to
         use the `lm` under `lf.context`.
@@ -117,8 +121,8 @@ class PythonCode(pg.Object):
       The value of the last expression in the source code. Or a dict of local
       variable names defined in the source code to their values if
       `outputs_intermediate` is set to True. The value for the last line can be
-      accessed by key '__result__'. Or the stdout as a str if `returns_stdout`
-      is set to True.
+      accessed by key '__result__'. Or the stdout as a string if
+      `returns_stdout` is set to True.
 
     Raises:
       TimeoutError: If `sandbox` is True and timeout has reached.
@@ -152,12 +156,12 @@ class PythonCode(pg.Object):
     Args:
       sandbox: If True, run code in sandbox; If False, run code in current
         process. If None, run in sandbox first, if the output could not be
-        serialized and pass to current process, run the code again in current
+        serialized and passed to current process, run the code again in current
         process.
       timeout: Timeout in seconds. If None, there is no timeout. Applicable when
         sandbox is set to True.
       global_vars: Global variables that could be accessed from the source code.
-      autofix: Number of attempts to auto fix the generated code. If 0, autofix
+      autofix: Number of attempts to autofix the generated code. If 0, autofix
         is disabled. Auto-fix is not supported for 'json' protocol.
       autofix_lm: Language model to be used. If not specified, it will try to
         use the `lm` under `lf.context`.
@@ -182,10 +186,11 @@ class PythonCode(pg.Object):
 
 
 class PythonFunction(pg.Object):
-  """Generated Python function via source code.
+  """Represents a Python function defined by source code.
 
-  The source code will be directly passed into eval() for execution and the
-  output of the function will be returned.
+  This class takes Python source code that defines a function and makes it
+  callable. The source code is evaluated to create a function object, which
+  can then be invoked like a regular Python function.
   """
 
   name: str
@@ -214,7 +219,7 @@ class PythonFunction(pg.Object):
       *args: Positional arguments that will be passed to the implementation.
       sandbox: If True, run code in sandbox; If False, run code in current
         process. If None, run in sandbox first, if the output could not be
-        serialized and pass to current process, run the code again in current
+        serialized and passed to current process, run the code again in current
         process.
       timeout: Timeout in seconds. If None, there is no timeout. Applicable when
         sandbox is set to True.

langfun/core/coding/python/sandboxing.py

@@ -23,7 +23,14 @@ import pyglove as pg
 
 
 class SandboxOutput(pg.Object):
-  """Sandbox output."""
+  """A structure containing the output from a sandbox execution.
+
+  Attributes:
+    stdout: The standard output captured during execution.
+    stderr: The standard error captured during execution.
+    output_files: A dictionary of file names to their byte content for files
+      generated during execution.
+  """
 
   stdout: Annotated[
       str,
@@ -42,7 +49,14 @@ class SandboxOutput(pg.Object):
 
 
 class BaseSandbox(pg.Object):
-  """Interface and partial implementation for Python sandbox."""
+  """Base class for Python code sandboxing.
+
+  A sandbox provides an isolated environment for executing Python code,
+  typically with restrictions on file system access, network calls, or other
+  potentially harmful operations. This base class defines the interface for
+  sandboxes, including methods for running code (`run`), uploading files
+  (`upload`), and managing the sandbox lifecycle (`setup`, `cleanup`).
+  """
 
   def _on_bound(self):
     super()._on_bound()
@@ -111,7 +125,13 @@ class BaseSandbox(pg.Object):
 
 
 class MultiProcessingSandbox(BaseSandbox):
-  """Sandbox using multiprocessing."""
+  """A sandbox implementation using Python's `multiprocessing`.
+
+  This sandbox executes code in a separate process, providing isolation from
+  the main process. It uses a temporary directory for file operations,
+  which is cleaned up when the sandbox is closed. It relies on
+  `pg.coding.run` with `sandbox=True` for execution.
+  """
 
   def _on_bound(self):
     super()._on_bound()