langfun 0.1.2.dev202509120804__py3-none-any.whl → 0.1.2.dev202512150805__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)
@@ -51,6 +52,7 @@ class Foo(action_lib.Action):
     with session.track_phase('prepare'):
       session.info('Begin Foo', x=1)
       time.sleep(self.simulate_execution_time[0])
+      Bar()(session, lm=lm)
       session.query(
           'foo',
           schema=int if self.simulate_query_error else None,
@@ -64,14 +66,21 @@ class Foo(action_lib.Action):
     def _sub_task(i):
       session.add_metadata(**{f'subtask_{i}': i})
       time.sleep(self.simulate_execution_time[2])
+      Bar()(session, lm=lm)
       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]
@@ -81,6 +90,50 @@ class Foo(action_lib.Action):
     lf_structured.query('additional query', lm=lm)
 
 
+class ExecutionUnitPositionTest(unittest.TestCase):
+
+  def test_basics(self):
+    pos1 = action_lib.ExecutionUnit.Position(None, 0)
+    self.assertEqual(repr(pos1), 'Position(0)')
+    self.assertEqual(str(pos1), '')
+    self.assertIsNone(pos1.parent)
+    self.assertEqual(pos1.index, 0)
+    self.assertEqual(pos1.indices(), (0,))
+    self.assertEqual(pos1, (0,))
+    self.assertEqual(pos1, '')
+    self.assertEqual(pos1, action_lib.ExecutionUnit.Position(None, 0))
+    self.assertNotEqual(pos1, 1)
+    self.assertNotEqual(pos1, (1,))
+    self.assertNotEqual(pos1, action_lib.ExecutionUnit.Position(None, 1))
+
+    pos2 = action_lib.ExecutionUnit.Position(pos1, 0)
+    self.assertEqual(repr(pos2), 'Position(0, 0)')
+    self.assertEqual(str(pos2), '1')
+    self.assertEqual(pos2, '1')
+    self.assertEqual(pos2.parent, pos1)
+    self.assertEqual(pos2.index, 0)
+    self.assertEqual(pos2.indices(), (0, 0))
+    self.assertNotEqual(pos1, pos2)
+    self.assertLess(pos1, pos2)
+    self.assertGreater(pos2, pos1)
+    self.assertEqual(
+        hash(pos2),
+        hash(
+            action_lib.ExecutionUnit.Position(
+                action_lib.ExecutionUnit.Position(None, 0), 0
+            )
+        )
+    )
+
+    pos3 = action_lib.ExecutionUnit.Position(pos2, 0)
+    self.assertEqual(str(pos3), '1.1')
+    self.assertEqual(pos3, '1.1')
+    self.assertEqual(pos3.parent, pos2)
+    self.assertEqual(pos3.index, 0)
+    self.assertEqual(pos3.indices(), (0, 0, 0))
+    self.assertEqual(pos3.to_str(separator='>'), '1>1')
+
+
 class ActionInvocationTest(unittest.TestCase):
 
   def test_basics(self):
@@ -101,9 +154,7 @@ class ExecutionTraceTest(unittest.TestCase):
     self.assertEqual(execution.id, '')
 
     root = action_lib.ActionInvocation(action=action_lib.RootAction())
-    action_invocation = action_lib.ActionInvocation(
-        action=Foo(1)
-    )
+    action_invocation = action_lib.ActionInvocation(action=Foo(1))
     root.execution.append(action_invocation)
     self.assertEqual(action_invocation.execution.id, '/a1')
 
@@ -118,10 +169,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 +182,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)
@@ -143,6 +197,7 @@ class SessionTest(unittest.TestCase):
 
     self.assertIsInstance(session.root.action, action_lib.RootAction)
     self.assertIs(session.current_action, session.root)
+    self.assertIs(session.metadata, session.root.metadata)
 
     #
     # Inspecting the root invocation.
@@ -165,20 +220,25 @@ class SessionTest(unittest.TestCase):
     )
 
     # The root space should have one action (foo), no queries, and no logs.
+    self.assertEqual(len(root.execution_units), 1)
     self.assertEqual(len(root.actions), 1)
     self.assertEqual(len(root.queries), 0)
     self.assertEqual(len(root.logs), 0)
-    # 1 query from Bar, 2 from Foo and 3 from parallel executions.
-    self.assertEqual(len(session.all_queries), 6)
-    self.assertEqual(len(root.all_queries), 6)
-    # 2 actions: Foo and Bar.
-    self.assertEqual(len(session.all_actions), 2)
-    self.assertEqual(len(root.all_actions), 2)
-    # 1 log from Bar and 1 from Foo.
-    self.assertEqual(len(session.all_logs), 2)
-    self.assertEqual(len(root.all_logs), 2)
+    # 2 query from Bar, 2 from Foo and 2 * 3 from parallel executions.
+    self.assertEqual(len(session.all_queries), 10)
+    self.assertEqual(len(root.all_queries), 10)
+    # 6 actions: Foo and 2 Bar, and 3 Bar from parallel executions.
+    self.assertEqual(len(session.all_actions), 6)
+    self.assertEqual(
+        [str(a.position) for a in session.all_actions],
+        ['1', '1.1', '1.2.1.1', '1.2.2.1', '1.2.3.1', '1.3']
+    )
+    self.assertEqual(len(root.all_actions), 6)
+    # 1 log from Bar and 1 from Foo and 3 from Bar in parallel executions.
+    self.assertEqual(len(session.all_logs), 6)
+    self.assertEqual(len(root.all_logs), 6)
     self.assertIs(session.usage_summary, root.usage_summary)
-    self.assertEqual(root.usage_summary.total.num_requests, 6)
+    self.assertEqual(root.usage_summary.total.num_requests, 10)
 
     # Inspecting the top-level action (Foo)
     foo_invocation = root.execution[0]
@@ -190,15 +250,19 @@ class SessionTest(unittest.TestCase):
 
     # Prepare phase.
     prepare_phase = foo_invocation.execution[0]
+    self.assertIsNone(prepare_phase.position)
     self.assertIsInstance(prepare_phase, action_lib.ExecutionTrace)
     self.assertEqual(prepare_phase.id, 'agent@1:/a1/prepare')
-    self.assertEqual(len(prepare_phase.items), 2)
+    self.assertEqual(len(prepare_phase.items), 3)
     self.assertTrue(prepare_phase.has_started)
     self.assertTrue(prepare_phase.has_stopped)
-    self.assertEqual(prepare_phase.usage_summary.total.num_requests, 1)
+    self.assertEqual(prepare_phase.usage_summary.total.num_requests, 2)
     self.assertIsInstance(prepare_phase.items[0], lf.logging.LogEntry)
-    self.assertIsInstance(prepare_phase.items[1], lf_structured.QueryInvocation)
-    self.assertEqual(prepare_phase.items[1].id, 'agent@1:/a1/prepare/q1')
+    self.assertIsInstance(prepare_phase.items[1], action_lib.ActionInvocation)
+    self.assertIs(prepare_phase.items[1].parent_execution_unit, foo_invocation)
+    self.assertEqual(prepare_phase.items[1].id, 'agent@1:/a1/prepare/a1')
+    self.assertIsInstance(prepare_phase.items[2], lf_structured.QueryInvocation)
+    self.assertEqual(prepare_phase.items[2].id, 'agent@1:/a1/prepare/q1')
 
     # Tracked queries.
     query_invocation = foo_invocation.execution[1]
@@ -220,20 +284,44 @@ class SessionTest(unittest.TestCase):
 
     # Tracked parallel executions.
     parallel_executions = foo_invocation.execution[2]
+    # root (0) > foo (0) > parallel executions (1)
+    self.assertEqual(parallel_executions.position, (0, 0, 1))
     self.assertEqual(parallel_executions.id, 'agent@1:/a1/p1')
     self.assertIsInstance(parallel_executions, action_lib.ParallelExecutions)
+    self.assertIs(
+        parallel_executions.all_actions[0].parent_execution_unit,
+        parallel_executions
+    )
+    self.assertIs(
+        parallel_executions.all_actions[0].parent_action,
+        foo_invocation
+    )
     self.assertEqual(len(parallel_executions), 3)
     self.assertEqual(parallel_executions[0].id, 'agent@1:/a1/p1/b1')
     self.assertEqual(parallel_executions[1].id, 'agent@1:/a1/p1/b2')
     self.assertEqual(parallel_executions[2].id, 'agent@1:/a1/p1/b3')
+    self.assertEqual(len(parallel_executions[0].execution_units), 1)
+    self.assertEqual(len(parallel_executions[1].execution_units), 1)
+    self.assertEqual(len(parallel_executions[2].execution_units), 1)
     self.assertEqual(len(parallel_executions[0].queries), 1)
+    self.assertEqual(len(parallel_executions[0].all_queries), 2)
     self.assertEqual(len(parallel_executions[1].queries), 1)
+    self.assertEqual(len(parallel_executions[1].all_queries), 2)
     self.assertEqual(len(parallel_executions[2].queries), 1)
+    self.assertEqual(len(parallel_executions[2].all_queries), 2)
+    self.assertEqual(len(parallel_executions.execution_units), 0)
+    self.assertEqual(len(parallel_executions.actions), 0)
+    self.assertEqual(len(parallel_executions.queries), 0)
+    self.assertEqual(len(parallel_executions.logs), 0)
+    self.assertEqual(len(parallel_executions.all_actions), 3)
+    self.assertEqual(len(parallel_executions.all_queries), 6)
+    self.assertEqual(len(parallel_executions.all_logs), 3)
 
     # Invocation to Bar.
     bar_invocation = foo_invocation.execution[3]
     self.assertIs(bar_invocation.parent_action, foo_invocation)
-    self.assertEqual(bar_invocation.id, 'agent@1:/a1/a1')
+    self.assertIs(bar_invocation.parent_execution_unit, foo_invocation)
+    self.assertEqual(bar_invocation.id, 'agent@1:/a1/a5')
     self.assertIsInstance(bar_invocation, action_lib.ActionInvocation)
     self.assertIsInstance(bar_invocation.action, Bar)
     self.assertEqual(bar_invocation.result, 2)
@@ -366,7 +454,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 +474,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 +487,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 +568,58 @@ 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_session_start(self, session):
+        session.add_metadata(progresses=pg.Ref(self.progresses))
+
+      def on_action_progress(self, session, action, title, **kwargs):
+        self.progresses.append((action.id, title))
+
+    handler = MyActionHandler()
+    self.assertIs(handler.get(MyActionHandler), handler)
+    self.assertIsNone(handler.get(action_lib.SessionLogging))
+
+    handler_chain = action_lib.SessionEventHandlerChain(
+        handlers=[handler, action_lib.SessionLogging()]
+    )
+    self.assertIs(handler_chain.get(MyActionHandler), handler)
+    self.assertIs(
+        handler_chain.get(action_lib.SessionLogging),
+        handler_chain.handlers[1]
+    )
+
+    session = action_lib.Session(
+        id='agent@1',
+        event_handler=handler_chain
+    )
+    bar = Bar()
+    with session:
+      bar(session, lm=fake.StaticResponse('lm response'))
+      session.update_progress('Trajectory completed')
+
+    self.assertIs(session.metadata['progresses'], handler.progresses)
+    self.assertEqual(handler.progresses, [
+        ('agent@1:/a1', 'Query completed'),
+        ('agent@1:', 'Trajectory completed'),
+    ])
+
+  def test_clone(self):
+    event_handler = action_lib.SessionLogging()
+    session = action_lib.Session(event_handler=event_handler)
+    other = session.clone()
+    self.assertIsNot(session, other)
+    self.assertIs(other.event_handler, event_handler)
+
+    other = session.clone(deep=True)
+    self.assertIsNot(session, other)
+    self.assertIsNot(other.event_handler, session.event_handler)
+
   def test_log(self):
     session = action_lib.Session()
     session.debug('hi', x=1, y=2)
@@ -493,6 +633,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: