langfun 0.0.2.dev20240429__py3-none-any.whl → 0.1.2.dev202501150804__py3-none-any.whl

langfun/core/concurrent_test.py

@@ -94,7 +94,7 @@ class RetryErrorTest(unittest.TestCase):
     )
 
 
-class WithRetryTest(unittest.TestCase):
+class RetryTest(unittest.TestCase):
 
   def assert_retry(self, func, expected_attempts, expected_wait_intervals):
     with pg.catch_errors(concurrent.RetryError) as error_context:
@@ -138,10 +138,19 @@ class WithRetryTest(unittest.TestCase):
       raise ValueError('Intentional error.')
 
     foo_with_retry = concurrent.with_retry(
-        foo, ValueError, max_attempts=4, retry_interval=1
+        foo, ValueError, max_attempts=4, retry_interval=1,
     )
     self.assert_retry(foo_with_retry, 4, [1, 2, 4])
 
+  def test_retry_with_max_retry_interval(self):
+    def foo():
+      raise ValueError('Intentional error.')
+
+    foo_with_retry = concurrent.with_retry(
+        foo, ValueError, max_attempts=4, retry_interval=1, max_retry_interval=3,
+    )
+    self.assert_retry(foo_with_retry, 4, [1, 2, 3])
+
   def test_retry_with_uncaught_exception(self):
     def foo():
       raise ValueError('Intentional error.')
@@ -153,6 +162,31 @@ class WithRetryTest(unittest.TestCase):
     with self.assertRaises(ValueError):
       foo_with_retry()
 
+  def test_retry_with_job(self):
+    count = 0
+
+    def foo():
+      nonlocal count
+      count += 1
+      if count < 3:
+        raise ValueError('Foo temporary error.')
+      return 'Success'
+
+    job = concurrent.Job(
+        foo,
+        retry_on_errors=ValueError,
+        retry_interval=1,
+    )
+    job()
+    self.assertEqual(job.result, 'Success')
+    self.assertEqual(
+        [retry_entry.wait_interval for retry_entry in job.retry_entries],
+        [0, 1, 2],
+    )
+    self.assertIsInstance(job.retry_entries[0].error, ValueError)
+    self.assertIsInstance(job.retry_entries[1].error, ValueError)
+    self.assertIsNone(job.retry_entries[2].error)
+
 
 class ConcurrentExecuteTest(unittest.TestCase):
 
@@ -208,8 +242,8 @@ class ProgressTest(unittest.TestCase):
     def fun2(unused_x):
       raise ValueError('Intentional error.')
 
-    job1 = concurrent.Job(fun, 1)
-    job2 = concurrent.Job(fun2, 2)
+    job1 = concurrent.Job(fun, (1,))
+    job2 = concurrent.Job(fun2, (2,))
     job1()
     job2()
 
@@ -233,6 +267,57 @@ class ProgressTest(unittest.TestCase):
     self.assertIs(p.last_error, job2.error)
 
 
+class ProgressControlTest(unittest.TestCase):
+
+  def test_noop(self):
+    concurrent.progress_bar = None
+    ctrl = concurrent._progress_control(100, 'noop', 'blue', None)
+    self.assertIsInstance(ctrl, concurrent._NoopProgressControl)
+    string_io = io.StringIO()
+    with contextlib.redirect_stderr(string_io):
+      ctrl.update(1)
+      ctrl.refresh()
+    self.assertEqual(string_io.getvalue(), '')
+    concurrent.progress_bar = 'tqdm'
+
+  def test_console(self):
+    concurrent.progress_bar = 'console'
+    ctrl = concurrent._progress_control(100, 'foo', 'blue', None)
+    self.assertIsInstance(ctrl, concurrent._ConsoleProgressControl)
+    string_io = io.StringIO()
+    with contextlib.redirect_stderr(string_io):
+      ctrl.set_status('bar')
+      ctrl.update(10)
+      ctrl.refresh()
+    self.assertEqual(
+        string_io.getvalue(),
+        '\x1b[1m\x1b[31mfoo\x1b[0m: \x1b[34m10% (10/100)\x1b[0m : bar\n'
+    )
+    concurrent.progress_bar = 'tqdm'
+
+  def test_tqdm(self):
+    concurrent.progress_bar = 'tqdm'
+    string_io = io.StringIO()
+    with contextlib.redirect_stderr(string_io):
+      ctrl = concurrent._progress_control(100, 'foo', 'blue', None)
+      self.assertIsInstance(ctrl, concurrent._TqdmProgressControl)
+      ctrl.update(10)
+      ctrl.refresh()
+    self.assertIn('10/100', string_io.getvalue())
+
+    tqdm = concurrent.tqdm
+    concurrent.tqdm = None
+    with self.assertRaisesRegex(RuntimeError, 'install package "tqdm"'):
+      _ = concurrent._progress_control(100, 'foo', 'blue', None)
+    concurrent.tqdm = tqdm
+
+  def test_unsupported(self):
+    concurrent.progress_bar = 'unknown'
+    with self.assertRaisesRegex(ValueError, 'Unsupported progress bar type'):
+      _ = concurrent._progress_control(100, 'foo', 'blue', None)
+    concurrent.progress_bar = 'tqdm'
+
+
 class ProgressBarTest(unittest.TestCase):
 
   def test_multithread_support(self):
@@ -241,13 +326,12 @@ class ProgressBarTest(unittest.TestCase):
       bar_id = concurrent.ProgressBar.install(None, 5)
       def fun(x):
         del x
-        concurrent.ProgressBar.update(bar_id, 1, postfix=None)
+        concurrent.ProgressBar.update(bar_id, 1, status=None)
 
       for _ in concurrent.concurrent_execute(fun, range(5)):
         concurrent.ProgressBar.refresh()
       concurrent.ProgressBar.uninstall(bar_id)
     output_str = string_io.getvalue()
-    print(output_str)
     self.assertIn('100%', output_str)
     self.assertIn('5/5', output_str)
 
@@ -255,19 +339,19 @@ class ProgressBarTest(unittest.TestCase):
     string_io = io.StringIO()
     with contextlib.redirect_stderr(string_io):
       bar_id = concurrent.ProgressBar.install(None, 4)
-      concurrent.ProgressBar.update(bar_id, 1, postfix=None)
-      self.assertIn('1/4', string_io.getvalue())
-      concurrent.ProgressBar.update(bar_id, 1, postfix='hello')
-      self.assertIn('2/4', string_io.getvalue())
-      self.assertIn('hello', string_io.getvalue())
-      concurrent.ProgressBar.update(bar_id, color='lightgreen')
-      self.assertIn('2/4', string_io.getvalue())
-      concurrent.ProgressBar.update(bar_id, 2, postfix=dict(x=1))
-      self.assertIn('4/4', string_io.getvalue())
-      self.assertIn('x=1', string_io.getvalue())
-      with self.assertRaisesRegex(ValueError, 'Unsupported postfix'):
-        concurrent.ProgressBar.update(bar_id, 0, postfix=1)
+      concurrent.ProgressBar.update(bar_id, 1, status=None)
+      concurrent.ProgressBar.update(bar_id, 1, status='hello')
+      concurrent.ProgressBar.update(bar_id, color='green')
+      concurrent.ProgressBar.update(bar_id, 2, status=dict(x=1))
+      with self.assertRaisesRegex(ValueError, 'Unsupported status'):
+        concurrent.ProgressBar.update(bar_id, 0, status=1)
       concurrent.ProgressBar.uninstall(bar_id)
+    self.assertIn('1/4', string_io.getvalue())
+    self.assertIn('2/4', string_io.getvalue())
+    self.assertIn('hello', string_io.getvalue())
+    self.assertNotIn('3/4', string_io.getvalue())
+    self.assertIn('4/4', string_io.getvalue())
+    self.assertIn('x=1', string_io.getvalue())
 
 
 class ConcurrentMapTest(unittest.TestCase):
@@ -479,10 +563,11 @@ class ConcurrentMapTest(unittest.TestCase):
 
   def test_concurrent_map_with_showing_progress(self):
     def fun(x):
-      if x == 2:
-        raise ValueError('Intentional error.')
-      time.sleep(x)
-      return x
+      with pg.timeit('foo'):
+        if x == 2:
+          raise ValueError('Intentional error.')
+        time.sleep(x)
+        return x
 
     string_io = io.StringIO()
     with contextlib.redirect_stderr(string_io):
@@ -499,7 +584,9 @@ class ConcurrentMapTest(unittest.TestCase):
             (3, pg.MISSING_VALUE),
         ],
     )
-    self.assertIn('100%', string_io.getvalue())
+    output = string_io.getvalue()
+    self.assertIn('100%', output)
+    self.assertIn('TimeIt=foo (', output)
 
   def test_concurrent_map_with_showing_progress_and_status_fn(self):
     def fun(x):
@@ -526,8 +613,8 @@ class ConcurrentMapTest(unittest.TestCase):
             (3, pg.MISSING_VALUE),
         ],
     )
-    self.assertIn('100%', string_io.getvalue())
     concurrent.ProgressBar.uninstall(bar_id)
+    self.assertIn('100%', string_io.getvalue())
 
 
 class ExecutorPoolTest(unittest.TestCase):

langfun/core/console.py

@@ -15,7 +15,7 @@
 
 import sys
 from typing import Any
-from langfun.core.text_formatting import colored
+import pyglove as pg
 
 
 def write(
@@ -42,10 +42,14 @@ def write(
   """
   # Print title if present.
   if title is not None:
-    print(colored(title, styles=['bold']))
+    print(pg.colored(title, styles=['bold']))
 
   # Print body.
-  print(colored(str(value), color=color, background=background, styles=styles))
+  print(
+      pg.colored(
+          str(value), color=color, background=background, styles=styles
+      )
+  )
 
 
 try:
@@ -59,12 +63,20 @@ def under_notebook() -> bool:
   return bool(_notebook)
 
 
-def display(value: Any, clear: bool = False) -> None:  # pylint: disable=redefined-outer-name
+def display(value: Any, clear: bool = False) -> Any:  # pylint: disable=redefined-outer-name
   """Displays object in current notebook cell."""
   if _notebook is not None:
     if clear:
       _notebook.clear_output()
-    _notebook.display(value)
+    return _notebook.display(value)
+  return None
+
+
+def run_script(javascript: str) -> Any:
+  """Runs JavaScript in current notebook cell."""
+  if _notebook is not None:
+    return _notebook.display(_notebook.Javascript(javascript))
+  return
 
 
 def clear() -> None:

langfun/core/console_test.py

@@ -18,6 +18,7 @@ import io
 import unittest
 
 from langfun.core import console
+import pyglove as pg
 
 
 class ConsoleTest(unittest.TestCase):
@@ -32,6 +33,22 @@ class ConsoleTest(unittest.TestCase):
 
   def test_under_notebook(self):
     self.assertFalse(console.under_notebook())
+    console._notebook = True
+    self.assertTrue(console.under_notebook())
+    console._notebook = None
+
+  def test_notebook_interaction(self):
+    console._notebook = pg.Dict(
+        display=lambda x: x, Javascript=lambda x: x, clear_output=lambda: None)
+    self.assertEqual(console.display('hi', clear=True), 'hi')
+    self.assertEqual(
+        console.run_script('console.log("hi")'),
+        'console.log("hi")'
+    )
+    console.clear()
+    console._notebook = None
+    self.assertIsNone(console.display('hi'))
+    self.assertIsNone(console.run_script('console.log("hi")'))
 
 
 if __name__ == '__main__':

langfun/core/eval/__init__.py

@@ -16,7 +16,13 @@
 # pylint: disable=g-importing-member
 # pylint: disable=g-bad-import-order
 
-from langfun.core.eval.base import app_run
+from langfun.core.eval import v2
+
+from langfun.core.eval.base import register
+from langfun.core.eval.base import registered_names
+from langfun.core.eval.base import get_evaluations
+from langfun.core.eval.base import get
+from langfun.core.eval.base import run
 
 from langfun.core.eval.base import Evaluable
 from langfun.core.eval.base import Evaluation
@@ -34,6 +40,15 @@ from langfun.core.eval.base import as_inputs
 from langfun.core.eval.matching import Matching
 from langfun.core.eval.scoring import Scoring
 
+# Experiment patching.
+from langfun.core.eval.patching import patch_member
+
+from langfun.core.eval.patching import patch_lm
+from langfun.core.eval.patching import patch_parsing_lm
+from langfun.core.eval.patching import patch_inputs
+from langfun.core.eval.patching import patch_prompt
+from langfun.core.eval.patching import patch_schema_fn
 
+# Placeholder for Google-internal imports.
 # pylint: enable=g-bad-import-order
 # pylint: enable=g-importing-member