langfun 0.1.1.dev20240826__py3-none-any.whl → 0.1.1.dev20240828__py3-none-any.whl

langfun/core/concurrent.py

@@ -13,17 +13,30 @@
 # limitations under the License.
 """Utility library for handling concurrency in langfun."""
 
+import abc
 import collections
 import concurrent.futures
 import dataclasses
+import io
 import random
+import sys
 import threading
 import time
 from typing import Any, Callable, Iterable, Iterator, Literal, Sequence, Tuple, Type, Union
 
 from langfun.core import component
+from langfun.core import text_formatting
 import pyglove as pg
-from tqdm import auto as tqdm
+
+
+progress_bar: Literal['tqdm', 'console', None] = None
+
+try:
+  from tqdm import auto as tqdm   # pylint: disable=g-import-not-at-top
+  progress_bar = 'tqdm'
+except ImportError as e:
+  progress_bar = 'console'
+  tqdm = None
 
 
 def with_context_access(func: Callable[..., Any]) -> Callable[..., Any]:
@@ -142,7 +155,6 @@ def with_retry(
         attempt = 1
       return base_interval() * (2 ** (attempt - 1))
 
-    wait_interval = None
     wait_intervals = []
     errors = []
     while True:
@@ -356,17 +368,17 @@ class ProgressBar:
     label: str | None
     total: int
     color: str | None = None
-    postfix: dict[str, str] | None = None
+    status: dict[str, Any] | None = None
 
   @dataclasses.dataclass
   class Update:
     """Progress bar update."""
     bar_id: int
     delta: int
-    postfix: Union[dict[str, str], str, None] = None
+    status: Union[dict[str, Any], str, None] = None
     color: str | None = None
 
-  _progress_bars: dict[int, tqdm.tqdm] = {}
+  _progress_bars: dict[int, '_ProgressControl'] = {}
   _install_requests: list[tuple[int, Settings]] = []
   _updates: collections.deque[Update] = collections.deque()
   _uninstall_requests: list[int] = []
@@ -378,11 +390,11 @@ class ProgressBar:
       label: str | None,
       total: int,
       color: str | None = None,
-      postfix: dict[str, str] | None = None,
+      status: dict[str, Any] | None = None,
       ) -> int:
     """Installs a progress bar and returns a reference id."""
     with cls._lock:
-      settings = ProgressBar.Settings(label, total, color, postfix)
+      settings = ProgressBar.Settings(label, total, color, status)
       bar_id = id(settings)
       cls._install_requests.append((bar_id, settings))
       return bar_id
@@ -392,15 +404,17 @@ class ProgressBar:
       cls,
       bar_id: int,
       delta: int = 0,
-      postfix: Union[dict[str, str], str, None] = None,
+      status: Union[dict[str, Any], str, None] = None,
       color: str | None = None,
       refresh: bool = True,
       ) -> None:
     """Report the progress for a label."""
+    if status is not None and not isinstance(status, (str, dict)):
+      raise ValueError(f'Unsupported status: {status}')
     with cls._lock:
       cls._updates.append(
           ProgressBar.Update(
-              bar_id=bar_id, delta=delta, postfix=postfix, color=color,
+              bar_id=bar_id, delta=delta, status=status, color=color,
           )
       )
     if refresh:
@@ -422,11 +436,11 @@ class ProgressBar:
       # Process install requests.
       if cls._install_requests:
         for bar_id, settings in cls._install_requests:
-          cls._progress_bars[bar_id] = tqdm.tqdm(
+          cls._progress_bars[bar_id] = _progress_control(
               total=settings.total,
-              desc=settings.label,
-              colour=settings.color,
-              postfix=settings.postfix)
+              label=settings.label,
+              color=settings.color,
+              status=settings.status)
         cls._install_requests.clear()
 
       # Process updates.
@@ -441,15 +455,11 @@ class ProgressBar:
         if update.delta > 0:
           bar.update(update.delta)
 
-        if isinstance(update.postfix, str):
-          bar.set_postfix_str(update.postfix, refresh=False)
-        elif isinstance(update.postfix, dict):
-          bar.set_postfix(update.postfix, refresh=False)
-        elif update.postfix is not None:
-          raise ValueError(f'Unsupported postfix: {update.postfix}')
+        if update.status is not None:
+          bar.set_status(update.status)
 
         if update.color is not None:
-          bar.colour = update.color
+          bar.set_color(update.color)
         updated_bars.add(bar)
 
       # Refresh each updated bar just once.
@@ -603,7 +613,7 @@ def concurrent_map(
         if len(error_text) >= 64:
           error_text = error_text[:64] + '...'
         status['LastError'] = error_text
-      ProgressBar.update(bar_id, delta=1, postfix=status)
+      ProgressBar.update(bar_id, delta=1, status=status)
 
   try:
     if ordered:
@@ -729,5 +739,122 @@ class ExecutorPool:
       raise ValueError(f'Unsupported value: {maybe_executor}.')
 
 
+class _ProgressControl(pg.Object):
+  """Abstract progress control."""
+  # Disable symbolic comparison so the hash is based on object address.
+  use_symbolic_comparison = False
+
+  total: int
+  label: str | None
+  color: str | None
+  status: str | dict[str, Any] | None
+
+  def set_color(self, color: str | None):
+    with pg.notify_on_change(False):
+      self.rebind(color=color)
+
+  def set_status(self, status: str | dict[str, Any] | None):
+    with pg.notify_on_change(False):
+      self.rebind(status=status)
+
+  @abc.abstractmethod
+  def update(self, delta):
+    """Update progress."""
+
+  @abc.abstractmethod
+  def refresh(self) -> None:
+    """Refresh progress bar."""
+
+
+class _TqdmProgressControl(_ProgressControl):
+  """Tqdm-based progress control."""
+
+  def _on_bound(self):
+    super()._on_bound()
+    assert tqdm is not None
+    self._tqdm = tqdm.tqdm(
+        total=self.total,
+        desc=self.label,
+        colour=self.color,
+        postfix=self.status,
+    )
+
+  def update(self, delta: int) -> None:
+    self._tqdm.update(delta)
+
+  def refresh(self):
+    self._tqdm.set_description(self.label, refresh=False)
+    if isinstance(self.status, str):
+      self._tqdm.set_postfix_str(self.status, refresh=False)
+    else:
+      self._tqdm.set_postfix(self.status, refresh=False)
+    self._tqdm.colour = self.color
+    self._tqdm.refresh()
+
+
+class _ConsoleProgressControl(_ProgressControl):
+  """Simple progress control by printing the status to the console."""
+
+  def _on_bound(self):
+    super()._on_bound()
+    self._progress = 0
+
+  def update(self, delta: int) -> None:
+    self._progress += delta
+
+  def refresh(self):
+    s = io.StringIO()
+    if self.label is not None:
+      s.write(text_formatting.colored(self.label, 'red', styles=['bold']))
+      s.write(': ')
+    s.write(
+        text_formatting.colored(
+            '%d%% (%d/%d)' %
+            (
+                self._progress * 100 // self.total,
+                self._progress,
+                self.total,
+            ),
+            color=self.color or 'green'
+        )
+    )
+    if self.status is not None:
+      status = repr(self.status) if isinstance(
+          self.status, dict) else self.status
+      s.write(f' : {status}')
+    sys.stderr.write(s.getvalue() + '\n')
+
+
+class _NoopProgressControl(_ProgressControl):
+  """No-op progress control."""
+
+  def update(self, delta: int) -> None:
+    pass
+
+  def refresh(self) -> None:
+    pass
+
+
+def _progress_control(
+    total: int,
+    label: str | None,
+    color: str | None,
+    status: str | dict[str, Any] | None,
+) -> _ProgressControl:
+  """Creates a process control."""
+  if progress_bar == 'tqdm':
+    if not tqdm:
+      raise RuntimeError(
+          'Please install package "tqdm" to use `tqdm` progress bar.'
+      )
+    return _TqdmProgressControl(total, label, color, status)
+  elif progress_bar == 'console':
+    return _ConsoleProgressControl(total, label, color, status)
+  elif progress_bar is None:
+    return _NoopProgressControl(total, label, color, status)
+  else:
+    raise ValueError(f'Unsupported progress bar type: {progress_bar}')
+
+
 # The global executor pool based on resource IDs.
 _executor_pool = ExecutorPool()

langfun/core/concurrent_test.py

@@ -233,6 +233,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 +292,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,12 +305,12 @@ 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)
-      concurrent.ProgressBar.update(bar_id, 1, postfix='hello')
-      concurrent.ProgressBar.update(bar_id, color='lightgreen')
-      concurrent.ProgressBar.update(bar_id, 2, postfix=dict(x=1))
-      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())

langfun/core/eval/base.py

@@ -242,7 +242,7 @@ class Evaluable(lf.Component):
       ):
         if show_progress:
           lf.concurrent.ProgressBar.update(
-              progress_bar, postfix='LOADING SAVED RESULTS...', color='yellow'
+              progress_bar, status='LOADING SAVED RESULTS...', color='yellow'
           )
         if self.try_load_result():
           run_status = 'CACHED'
@@ -271,7 +271,7 @@ class Evaluable(lf.Component):
         if should_save:
           if show_progress:
             lf.concurrent.ProgressBar.update(
-                progress_bar, postfix='SAVING RESULTS...', color='yellow'
+                progress_bar, status='SAVING RESULTS...', color='yellow'
             )
 
           # Save evaluation results.
@@ -284,7 +284,7 @@ class Evaluable(lf.Component):
       if show_progress:
         lf.concurrent.ProgressBar.update(
             progress_bar,
-            postfix=self._completion_status(run_status),
+            status=self._completion_status(run_status),
             color='green',
         )
     else:
@@ -340,7 +340,7 @@ class Evaluable(lf.Component):
               f'[#{leaf.index} - {leaf.node.id}]',
               total=leaf.node.num_examples if leaf.enabled else 0,
               color='cyan' if leaf.enabled else 'yellow',
-              postfix=None if leaf.enabled else 'SKIPPED.')
+              status=None if leaf.enabled else 'SKIPPED.')
 
         # Run leaf groups in parallel.
         try:
@@ -354,7 +354,7 @@ class Evaluable(lf.Component):
           # Save results for non-leaf nodes.
           lf.concurrent.ProgressBar.update(
               overview_bar,
-              postfix='SAVING RESULTS...',
+              status='SAVING RESULTS...',
               color='yellow')
 
           for node in self.nonleaf_nodes:
@@ -364,7 +364,7 @@ class Evaluable(lf.Component):
 
           if should_save and summary:
             lf.concurrent.ProgressBar.update(
-                overview_bar, postfix='FINALIZING SUMMARY...'
+                overview_bar, status='FINALIZING SUMMARY...'
             )
 
             summary.save(os.path.join(self.root_dir, Evaluable.SUMMARY_HTML))
@@ -378,7 +378,7 @@ class Evaluable(lf.Component):
           # Signal all task completed by making the bar green.
           lf.concurrent.ProgressBar.update(
               overview_bar,
-              postfix='COMPLETED',
+              status='COMPLETED',
               color='green')
 
         finally:
@@ -1261,7 +1261,7 @@ class Evaluation(Evaluable):
 
   def finalize(self) -> pg.Dict:
     """Finalizes the evaluation result."""
-    if self.cache:
+    if self.cache is not None:
       cache_stats = dict(
           use_cache=True,
           num_queries=self.cache.stats.num_queries,

langfun/core/eval/base_test.py

@@ -564,7 +564,7 @@ class SuiteTest(unittest.TestCase):
                 schema_fn='answer_schema()',
             ),
             cache_stats=dict(
-                use_cache=True, num_queries=4, num_hits=1, num_updates=3
+                use_cache=True, num_queries=4, num_hits=0, num_updates=4
             ),
             metrics=dict(
                 total=2,

langfun/core/langfunc.py

@@ -269,6 +269,9 @@ class LangFunc(
         # Send rendered text to LM.
         lm_output = self.lm(lm_input, cache_seed=cache_seed)
 
+        # Attach cache seed.
+        lm_input.metadata.cache_seed = cache_seed
+
         # Transform the output message.
         lm_output = self.transform_output(lm_output)
         lm_output.tag(message_lib.Message.TAG_LM_OUTPUT)

langfun/core/langfunc_test.py

@@ -94,11 +94,13 @@ class LangFuncCallTest(unittest.TestCase):
         )
     )
     self.assertEqual(r.tags, ['lm-response', 'lm-output'])
-    self.assertEqual(r.source, message.UserMessage('Hello'))
+    self.assertEqual(
+        r.source,
+        message.UserMessage('Hello', metadata=dict(cache_seed=0))
+    )
     self.assertEqual(r.source.tags, ['rendered', 'lm-input'])
 
     self.assertEqual(str(l), 'Hello')
-    print(repr(l))
     self.assertEqual(
         repr(l),
         "LangFunc(template_str='Hello', clean=True,"
@@ -114,7 +116,7 @@ class LangFuncCallTest(unittest.TestCase):
       self.assertEqual(l, 'Hello')
       self.assertEqual(l.natural_language_format(), 'Hello')
       self.assertEqual(l.render(), 'Hello')
-      r = l()
+      r = l(cache_seed=1)
       self.assertEqual(
           r,
           message.AIMessage(
@@ -123,6 +125,7 @@ class LangFuncCallTest(unittest.TestCase):
           )
       )
       self.assertEqual(r.tags, ['lm-response', 'lm-output'])
+      self.assertEqual(r.source.metadata.cache_seed, 1)
 
     self.assertEqual(str(l), 'Hello')
 

langfun/core/language_model.py

@@ -234,6 +234,7 @@ class LMCache(pg.Object):
     num_hit_expires: int = 0
     num_misses: int = 0
     num_updates: int = 0
+    num_deletes: int = 0
 
   @abc.abstractmethod
   def get(
@@ -251,6 +252,15 @@ class LMCache(pg.Object):
   ) -> None:
     """Puts the result of a prompt generated by a language model in cache."""
 
+  @abc.abstractmethod
+  def delete(
+      self,
+      lm: 'LanguageModel',
+      prompt: message_lib.Message,
+      seed: int,
+  ) -> bool:
+    """Deletes the result of a prompt generated by a language model in cache."""
+
   @property
   @abc.abstractmethod
   def stats(self) -> Stats:

langfun/core/llms/cache/base.py

@@ -60,13 +60,16 @@ class LMCacheBase(lf.LMCache):
       self, lm: lf.LanguageModel, prompt: lf.Message, seed: int
   ) -> lf.LMSamplingResult | None:
     """Gets the cached result of a prompt generated by a language model."""
-    entry = self._get(lm.model_id, self._key(lm, prompt, seed))
+    key = self._key(lm, prompt, seed)
+    entry = self._get(lm.model_id, key)
     self._stats.num_queries += 1
     if entry is None:
       self._stats.num_misses += 1
       return None
     if entry.expire is not None and entry.expire < datetime.datetime.now():
       self._stats.num_hit_expires += 1
+      self._stats.num_deletes += 1
+      assert self._delete(lm.model_id, key)
       return None
     self._stats.num_hits += 1
     return entry.result
@@ -86,6 +89,18 @@ class LMCacheBase(lf.LMCache):
     self._put(lm.model_id, self._key(lm, prompt, seed), entry)
     self._stats.num_updates += 1
 
+  def delete(
+      self,
+      lm: lf.LanguageModel,
+      prompt: lf.Message,
+      seed: int,
+  ) -> bool:
+    """Deletes the result of a prompt generated by a language model in cache."""
+    deleted = self._delete(lm.model_id, self._key(lm, prompt, seed))
+    if deleted:
+      self._stats.num_deletes += 1
+    return deleted
+
   @abc.abstractmethod
   def _get(self, model_id: str, key: str) -> LMCacheEntry | None:
     """Returns a LM cache entry associated with the key."""
@@ -94,6 +109,10 @@ class LMCacheBase(lf.LMCache):
   def _put(self, model_id: str, key: str, entry: LMCacheEntry) -> None:
     """Puts a LM cache entry associated with the key."""
 
+  @abc.abstractmethod
+  def _delete(self, model_id: str, key: str) -> bool:
+    """Deletes a LM cache entry associated with the key."""
+
   def _sym_clone(self, deep: bool, memo: Any = None) -> 'LMCacheBase':
     v = super()._sym_clone(deep, memo)
     v._stats = self._stats  # pylint: disable=protected-access
@@ -102,4 +121,4 @@ class LMCacheBase(lf.LMCache):
 
 def default_key(lm: lf.LanguageModel, prompt: lf.Message, seed: int) -> Any:
   """Default key for LM cache."""
-  return (prompt.text, lm.sampling_options.cache_key(), seed)
+  return (prompt.text_with_modality_hash, lm.sampling_options.cache_key(), seed)

langfun/core/llms/cache/in_memory.py

@@ -99,6 +99,13 @@ class InMemory(base.LMCacheBase):
     """Puts a LM cache entry associated with the key."""
     self._cache[model_id][key] = entry
 
+  def _delete(self, model_id: str, key: str) -> bool:
+    """Deletes a LM cache entry associated with the key."""
+    model_cache = self._cache.get(model_id, None)
+    if model_cache is None:
+      return False
+    return model_cache.pop(key, None) is not None
+
   def reset(self, model_id: str | None = None) -> None:
     """Resets the cache."""
     if model_id is not None:

langfun/core/llms/cache/in_memory_test.py

@@ -148,6 +148,50 @@ class InMemoryLMCacheTest(unittest.TestCase):
     self.assertIs(copy.deepcopy(cache)._cache, cache._cache)
     self.assertIs(copy.deepcopy(cache)._stats, cache._stats)
 
+    self.assertFalse(
+        cache.delete(fake.StaticResponse('hi'), lf.UserMessage('c'), seed=0)
+    )
+    self.assertFalse(cache.delete(lm, lf.UserMessage('c'), seed=1))
+    self.assertFalse(cache.delete(lm, lf.UserMessage('d'), seed=0))
+    self.assertTrue(cache.delete(lm, lf.UserMessage('c'), seed=0))
+    self.assertEqual(
+        list(cache.keys('StaticSequence')),
+        [
+            ('a', (None, None, 1, 40, None, None), 0),
+            ('a', (None, None, 1, 40, None, None), 1),
+            ('b', (None, None, 1, 40, None, None), 0),
+        ],
+    )
+    self.assertEqual(cache.stats.num_deletes, 1)
+
+  def test_cache_with_modalities(self):
+
+    class CustomModality(lf.Modality):
+      content: str
+
+      def to_bytes(self):
+        return self.content.encode()
+
+    cache = in_memory.InMemory()
+    lm = fake.StaticSequence(['1', '2', '3', '4', '5', '6'], cache=cache)
+    lm(lf.UserMessage('hi <<[[image]]>>', image=CustomModality('foo')))
+    lm(lf.UserMessage('hi <<[[image]]>>', image=CustomModality('bar')))
+    self.assertEqual(
+        list(cache.keys()),
+        [
+            (
+                'hi <<[[image]]>><image>acbd18db</image>',
+                (None, None, 1, 40, None, None),
+                0,
+            ),
+            (
+                'hi <<[[image]]>><image>37b51d19</image>',
+                (None, None, 1, 40, None, None),
+                0,
+            ),
+        ],
+    )
+
   def test_ttl(self):
     cache = in_memory.InMemory(ttl=1)
     lm = fake.StaticSequence(['1', '2', '3'], cache=cache)
@@ -160,6 +204,7 @@ class InMemoryLMCacheTest(unittest.TestCase):
     self.assertEqual(cache.stats.num_hits, 1)
     self.assertEqual(cache.stats.num_hit_expires, 1)
     self.assertEqual(cache.stats.num_misses, 1)
+    self.assertEqual(cache.stats.num_deletes, 1)
 
   def test_different_sampling_options(self):
     cache = in_memory.InMemory()

langfun/core/message.py

@@ -278,6 +278,16 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
   # API for supporting modalities.
   #
 
+  @property
+  def text_with_modality_hash(self) -> str:
+    """Returns text with modality object placeheld by their 8-byte MD5 hash."""
+    parts = [self.text]
+    for name, modality_obj in self.referred_modalities().items():
+      parts.append(
+          f'<{name}>{modality_obj.hash}</{name}>'
+      )
+    return ''.join(parts)
+
   def get_modality(
       self, var_name: str, default: Any = None, from_message_chain: bool = True
   ) -> modality.Modality | None:

langfun/core/message_test.py

@@ -282,6 +282,20 @@ class MessageTest(unittest.TestCase):
         },
     )
 
+  def test_text_with_modality_hash(self):
+    m = message.UserMessage(
+        'hi, this is a <<[[img1]]>> and <<[[x.img2]]>>',
+        img1=CustomModality('foo'),
+        x=dict(img2=CustomModality('bar')),
+    )
+    self.assertEqual(
+        m.text_with_modality_hash,
+        (
+            'hi, this is a <<[[img1]]>> and <<[[x.img2]]>>'
+            '<img1>acbd18db</img1><x.img2>37b51d19</x.img2>'
+        )
+    )
+
   def test_chunking(self):
     m = message.UserMessage(
         inspect.cleandoc("""

langfun/core/modality.py

@@ -14,6 +14,8 @@
 """Interface for modality (e.g. Image, Video, etc.)."""
 
 import abc
+import functools
+import hashlib
 from typing import Any, ContextManager
 from langfun.core import component
 import pyglove as pg
@@ -35,6 +37,11 @@ class Modality(component.Component):
   REF_START = '<<[['
   REF_END = ']]>>'
 
+  def _on_bound(self):
+    super()._on_bound()
+    # Invalidate cached hash if modality member is changed.
+    self.__dict__.pop('hash', None)
+
   def format(self, *args, **kwargs) -> str:
     if self.referred_name is None or not pg.object_utils.thread_local_get(
         _TLS_MODALITY_AS_REF, False
@@ -46,6 +53,11 @@ class Modality(component.Component):
   def to_bytes(self) -> bytes:
     """Returns content in bytes."""
 
+  @functools.cached_property
+  def hash(self) -> str:
+    """Returns a 8-byte MD5 hash as the identifier for this modality object."""
+    return hashlib.md5(self.to_bytes()).hexdigest()[:8]
+
   @classmethod
   def text_marker(cls, var_name: str) -> str:
     """Returns a marker in the text for this object."""

langfun/core/modality_test.py

@@ -32,6 +32,7 @@ class ModalityTest(unittest.TestCase):
     v = CustomModality('a')
     self.assertIsNone(v.referred_name)
     self.assertEqual(str(v), "CustomModality(\n  content = 'a'\n)")
+    self.assertEqual(v.hash, '0cc175b9')
 
     _ = pg.Dict(metadata=pg.Dict(x=pg.Dict(metadata=pg.Dict(y=v))))
     self.assertEqual(v.referred_name, 'x.metadata.y')

langfun/core/structured/mapping.py

@@ -352,6 +352,12 @@ class Mapping(lf.LangFunc):
       lm_output = self.postprocess_response(lm_output)
       lm_output.result = self.postprocess_result(self.parse_result(lm_output))
     except Exception as e:  # pylint: disable=broad-exception-caught
+      if (self.lm.cache is not None
+          and lm_output.lm_input.cache_seed is not None):
+        success = self.lm.cache.delete(
+            self.lm, lm_output.lm_input, lm_output.lm_input.cache_seed
+        )
+        assert success
       if self.default == lf.RAISE_IF_HAS_ERROR:
         raise MappingError(lm_output, e) from e
       lm_output.result = self.default

langfun/core/structured/prompting_test.py

@@ -19,6 +19,7 @@ import unittest
 import langfun.core as lf
 from langfun.core import modalities
 from langfun.core.llms import fake
+from langfun.core.llms.cache import in_memory
 from langfun.core.structured import mapping
 from langfun.core.structured import prompting
 import pyglove as pg
@@ -799,15 +800,18 @@ class QueryStructureJsonTest(unittest.TestCase):
       self.assertIsNone(r.result[0].hotel)
 
   def test_bad_transform(self):
-    with lf.context(
-        lm=fake.StaticSequence(['3']),
-        override_attrs=True,
-    ):
-      with self.assertRaisesRegex(
-          mapping.MappingError,
-          'No JSON dict in the output',
+    with in_memory.lm_cache() as cache:
+      with lf.context(
+          lm=fake.StaticSequence(['3']),
+          override_attrs=True,
       ):
-        prompting.query('Compute 1 + 2', int, protocol='json')
+        with self.assertRaisesRegex(
+            mapping.MappingError,
+            'No JSON dict in the output',
+        ):
+          prompting.query('Compute 1 + 2', int, protocol='json', cache_seed=1)
+      # Make sure bad mapping does not impact cache.
+      self.assertEqual(len(cache), 0)
 
   def test_query(self):
     lm = fake.StaticSequence(['{"result": 1}'])

langfun/core/text_formatting.py

@@ -16,7 +16,11 @@
 import io
 import re
 from typing import Any
-import termcolor
+
+try:
+  import termcolor      # pylint: disable=g-import-not-at-top
+except ImportError:
+  termcolor = None
 
 
 # Regular expression for ANSI color characters.
@@ -49,6 +53,8 @@ def colored(
   Returns:
     A string with ANSI color characters embracing the entire text.
   """
+  if not termcolor:
+    return text
   return termcolor.colored(
       text,
       color=color,

langfun/core/text_formatting_test.py

@@ -42,6 +42,24 @@ class TextFormattingTest(unittest.TestCase):
     )
     self.assertEqual(text_formatting.decolored(colored_text), original_text)
 
+  def test_colored_without_termcolor(self):
+    termcolor = text_formatting.termcolor
+    text_formatting.termcolor = None
+    original_text = inspect.cleandoc("""
+        Hi {{ foo }}
+        {# print x if x is present #}
+        {% if x %}
+        {{ x }}
+        {% endif %}
+        """)
+
+    colored_text = text_formatting.colored_template(
+        text_formatting.colored(original_text, color='blue')
+    )
+    self.assertEqual(colored_text, original_text)
+    self.assertEqual(text_formatting.decolored(colored_text), original_text)
+    text_formatting.termcolor = termcolor
+
 
 if __name__ == '__main__':
   unittest.main()

{langfun-0.1.1.dev20240826.dist-info → langfun-0.1.1.dev20240828.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.1.1.dev20240826
+Version: 0.1.1.dev20240828
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.1.dev20240826.dist-info → langfun-0.1.1.dev20240828.dist-info}/RECORD

@@ -2,21 +2,21 @@ langfun/__init__.py,sha256=ZqERg4fvvtFwr5L41Lfev9FSuCBm8PbQrZmcKvvloxE,2274
 langfun/core/__init__.py,sha256=r86kuy-BiJIveqnXx5OklUUXtMG3q79nWRBum6zFOCQ,4835
 langfun/core/component.py,sha256=Icyoj9ICoJoK2r2PHbrFXbxnseOr9QZZOvKWklLWNo8,10276
 langfun/core/component_test.py,sha256=q15Xn51cVTu2RKxZ9U5VQgT3bm6RQ4638bKhWBtvW5o,8220
-langfun/core/concurrent.py,sha256=TRc49pJ3HQro2kb5FtcWkHjhBm8UcgE8RJybU5cU3-0,24537
-langfun/core/concurrent_test.py,sha256=9HWnS4ja5Ji1ALfJs2zIt7w4w0oNm9dOZM7bgzoX5ag,15176
+langfun/core/concurrent.py,sha256=qZteOkA7jzLt7a5AVmNqDc5ZEvj_hwiLECqGMVYdzSU,27756
+langfun/core/concurrent_test.py,sha256=ULVrE6JGzsClRBLHAlqLF1A2mfxMO2X0245qYTYQ76o,16943
 langfun/core/console.py,sha256=bk5rNPNm9rMGW5YT2HixxU04p2umnoabn5SDz6Dqe88,2317
 langfun/core/console_test.py,sha256=5SYJdxpJGLgdSSQqqMPoA1X6jpsLD8rgcyk-EgI65oE,1077
-langfun/core/langfunc.py,sha256=RvIcRjIq0jWYRu1xim-FYe4HSrt97r3GMBO_PuagUmw,11060
-langfun/core/langfunc_test.py,sha256=lyt-UzkD8972cxZwzCkps0_RMLeSsOBrcUFIW-fB6us,8653
-langfun/core/language_model.py,sha256=oGni82fhYB3kUsL0okzvIXkKgXEMHVE-c0jR5LRmsIc,26039
+langfun/core/langfunc.py,sha256=G50YgoVZ0y1GFw2ev41MlOqr6qa8YakbvNC0h_E0PiA,11140
+langfun/core/langfunc_test.py,sha256=ZLlj6ysNIWUlm0jcq6PbNUwdJ2XstW5QQT0L2s5GlGY,8753
+langfun/core/language_model.py,sha256=o1MPDdiZ7eRUhcCjxLMT6IYYXge6nM7elvCWQ-NANaU,26286
 langfun/core/language_model_test.py,sha256=ebJ1vnaxKSKvlwi6v07yHjn91xMiDw2bQ9DBnyVorYw,23303
 langfun/core/logging.py,sha256=oDSeqGIQogZJ6xuPTcr9mkmLC2YnLP67UHtTdWbbiVY,4250
 langfun/core/logging_test.py,sha256=poSsNGKi6G9LWOcWnTY0BQjj0BtaQknH-NK6FcQrVT4,2152
 langfun/core/memory.py,sha256=f-asN1F7Vehgdn_fK84v73GrEUOxRtaW934keutTKjk,2416
-langfun/core/message.py,sha256=XmVIuj4g9t2P3pdpisDfvxQnmVReBaIj4iU4MBeWskI,18255
-langfun/core/message_test.py,sha256=7l3gSVqzVl4bUGwPU-L0gMPbjFODW0MVjWHzW4JlaaQ,12085
-langfun/core/modality.py,sha256=Tla4t86DUYHpbZ2G7dy1r19fTj_Ga5XOvlYp6lbWa-Q,3512
-langfun/core/modality_test.py,sha256=HyZ5xONKQ0Fw18SzoWAq-Ob9njOXIIjBo1hNtw-rudw,2400
+langfun/core/message.py,sha256=7UOxNMA1Le0ZGkleryqwXcWojZ-l_hku5Sc58BsIOGw,18586
+langfun/core/message_test.py,sha256=2o4WvE-WL67OsVY-O5g__67OIJ73vvg6MuojSc2uVRs,12504
+langfun/core/modality.py,sha256=g9wGx347oVofAJlMu_CpzTMIyTT9DJW8NfO4E-d-oNM,3879
+langfun/core/modality_test.py,sha256=7SwhixFME2Q1sIXRgJx97EZFiIyC31A9NVr6_nDtFv4,2441
 langfun/core/natural_language.py,sha256=3ynSnaYQnjE60LIPK5fyMgdIjubnPYZwzGq4rWPeloE,1177
 langfun/core/natural_language_test.py,sha256=LHGU_1ytbkGuSZQFIFP7vP3dBlcY4-A12fT6dbjUA0E,1424
 langfun/core/repr_utils.py,sha256=Y6ccoQUMpRxDv_jUy2QtnP9cdz3QBjJtTIgxGIU-kfM,5537
@@ -27,8 +27,8 @@ langfun/core/subscription.py,sha256=euawEuSZP-BHydaT-AQpfYFL0m5pWPGcW0upFhrojqc,
 langfun/core/subscription_test.py,sha256=Y4ZdbZEwm83YNZBxHff0QR4QUa4rdaNXA3_jfIcArBo,8717
 langfun/core/template.py,sha256=UhNNGUDJ4StUhPBKzHmjym36khxHOGWGr9MDxBwgxQA,22284
 langfun/core/template_test.py,sha256=cb005qM_3dvJikaPDB3rirfIrMAIIVRpoiXGREGJ43o,15468
-langfun/core/text_formatting.py,sha256=ytjj7opnRJ6w-pkglL2CZUyfYDXLpNf65E42LBb31gc,5158
-langfun/core/text_formatting_test.py,sha256=nyKC6tn2L4hPJiqQHgxcbQsJJi4A4Nbj8FiO8iT6B80,1514
+langfun/core/text_formatting.py,sha256=d7t9vaY6aCn1dkfkikpNYnBy5E_i93vHbfyDWFclGZU,5284
+langfun/core/text_formatting_test.py,sha256=ck0Xzdd4YF4CtCUj7VE0GybfbAyKQ8p3xkM1FBGrqIk,2096
 langfun/core/coding/__init__.py,sha256=5utju_fwEsImaiftx4oXKl9FAM8p281k8-Esdh_-m1w,835
 langfun/core/coding/python/__init__.py,sha256=MJ-vubliz-ebrZH3OBRKBwMi0S9-FrhGCp8YQLR6_I4,1776
 langfun/core/coding/python/correction.py,sha256=a2aFUt9ocbXTCR6Z6OGNjQZDI1LfU0PBkSe7hJB8dEM,6589
@@ -44,8 +44,8 @@ langfun/core/coding/python/parsing_test.py,sha256=9vAWF484kWIm6JZq8NFiMgKUDhXV-d
 langfun/core/coding/python/permissions.py,sha256=1QWGHvzL8MM0Ok_auQ9tURqZHtdOfJaDpBzZ29GUE-c,2544
 langfun/core/coding/python/permissions_test.py,sha256=w5EDb8QxpxgJyZkojyzVWQvDfg366zn99-g__6TbPQ0,2699
 langfun/core/eval/__init__.py,sha256=Ogdr9OtTywhhLPHi3AZzOD2mXX2oyaHWflrSTMm96uA,1899
-langfun/core/eval/base.py,sha256=BiWColibVo9-4P27Z0hIWXe8_UPocJTSTUdKeOPVwxI,74746
-langfun/core/eval/base_test.py,sha256=p1EfqviHMz_ppQY8FU67h5OCgL0tzhLvXzGIsq0sVyI,26930
+langfun/core/eval/base.py,sha256=YU_R9pKITKWyWa5-dSY8Q5O9pK8xioWX1ZdqYVOZcK8,74751
+langfun/core/eval/base_test.py,sha256=VEraWaRybSxOCOcZrZouNkiroDEPR6uyFBJoAz-1pQg,26930
 langfun/core/eval/matching.py,sha256=9GX8HfO9jKxgNLAivgy5K88Xhoh6Z7Pptq65pe7vht8,9762
 langfun/core/eval/matching_test.py,sha256=f7iVyXH5KGJBWt4Wp14Bt9J3X59A6Ayfog9MbuFvPew,5532
 langfun/core/eval/patching.py,sha256=R0s2eAd1m97exQt06dmUL0V_MBG0W2Hxg7fhNB7cXW0,3866
@@ -70,9 +70,9 @@ langfun/core/llms/rest_test.py,sha256=NZ3Nf0XQVpT9kLP5cBVo_yBHLI7vWTYhWQxYEJVMGs
 langfun/core/llms/vertexai.py,sha256=tXAnP357XhcsETTnk6M-hH4xyFi7tk6fsaf3tjzsY6E,14501
 langfun/core/llms/vertexai_test.py,sha256=EPR-mB2hNUpvpf7E8m_k5bh04epdQTVUuYU6hPgZyu8,10321
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
-langfun/core/llms/cache/base.py,sha256=cFfYvOIUae842pncqCAsRvqXCk2AnAsRYVx0mcIoAeY,3338
-langfun/core/llms/cache/in_memory.py,sha256=YfFyJEhLs73cUiB0ZfhMxYpdE8Iuxxw-dvMFwGHTSHw,4742
-langfun/core/llms/cache/in_memory_test.py,sha256=D-n26h__rVXQO51WRFhRfq5sw1oifRLx2SvCQWuNEm8,8747
+langfun/core/llms/cache/base.py,sha256=rt3zwmyw0y9jsSGW-ZbV1vAfLxQ7_3AVk0l2EySlse4,3918
+langfun/core/llms/cache/in_memory.py,sha256=l6b-iU9OTfTRo9Zmg4VrQIuArs4cCJDOpXiEpvNocjo,5004
+langfun/core/llms/cache/in_memory_test.py,sha256=ic6Z_crpuYdkUVPBbB9DomOddRxQdf57a2L_Sv5nHtM,10217
 langfun/core/memories/__init__.py,sha256=HpghfZ-w1NQqzJXBx8Lz0daRhB2rcy2r9Xm491SBhC4,773
 langfun/core/memories/conversation_history.py,sha256=c9amD8hCxGFiZuVAzkP0dOMWSp8L90uvwkOejjuBqO0,1835
 langfun/core/memories/conversation_history_test.py,sha256=AaW8aNoFjxNusanwJDV0r3384Mg0eAweGmPx5DIkM0Y,2052
@@ -96,12 +96,12 @@ langfun/core/structured/description.py,sha256=SXW4MJvshFjbR-0gw6rE21o6WXq12UlRXa
 langfun/core/structured/description_test.py,sha256=UtZGjSFUaQ6130t1E5tcL7ODu0xIefkapb53TbnqsK8,7362
 langfun/core/structured/function_generation.py,sha256=pFgS3vcRAWiuFBol2x5Eeip3XqoudONsOpeJpWyjT3s,7479
 langfun/core/structured/function_generation_test.py,sha256=ZJI-aaGgWWszn92u7h5IZ9Pl70N2DgAGGJrIxPzsvwg,10065
-langfun/core/structured/mapping.py,sha256=CsflMwm5cKJYZ2ag-neroA4CQlhu2wjFRSxKpd_qQDQ,11778
+langfun/core/structured/mapping.py,sha256=dKOCvIA_kCQ88KoCnP5k0iOe9xRt8WLC2kbT6qPqrd8,12016
 langfun/core/structured/mapping_test.py,sha256=zQoVx3kAD5oSm_OJAQA6q41NXLLyn8qs6CIVJgAoP_w,4489
 langfun/core/structured/parsing.py,sha256=keoVqEfzAbdULh6GawWFsTQzU91MzJXYFZjXGXLaD8g,11492
 langfun/core/structured/parsing_test.py,sha256=34wDrXaQ-EYhJLfDL8mX9K53oQMSzh5pVYdKjnESmK8,20895
 langfun/core/structured/prompting.py,sha256=_U6Z65AwXvVvfaQFCY9GawB_QV9S3u7P7BOU2URABmw,8873
-langfun/core/structured/prompting_test.py,sha256=Yc5GevbTyxtW6FAORXiqzakcuoJF-IPjoZ8kvLJW5is,23057
+langfun/core/structured/prompting_test.py,sha256=SyyBMmw-gwnrnr1MrwkjZiQaMq85kc9MHoYwzCITa10,23266
 langfun/core/structured/schema.py,sha256=oiT4P4Q9pG-QOnFzxETN2EQZqNln8nG4zAJHxcmeX9U,27729
 langfun/core/structured/schema_generation.py,sha256=U3nRQsqmMZg_qIVDh2fiY3K4JLfsAL1LcKzIFP1iXFg,5316
 langfun/core/structured/schema_generation_test.py,sha256=RM9s71kMNg2jTePwInkiW9fK1ACN37eyPeF8OII-0zw,2950
@@ -119,8 +119,8 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=rBW2Qr8yi-aWYwoTwRR-n1peKyMX9QXPZXURjLgoiRs,2264
-langfun-0.1.1.dev20240826.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.1.dev20240826.dist-info/METADATA,sha256=khcCyGAm6QiMy38AI0NPoJUE979BHRBpe-SzVHk8I7A,5234
-langfun-0.1.1.dev20240826.dist-info/WHEEL,sha256=Mdi9PDNwEZptOjTlUcAth7XJDFtKrHYaQMPulZeBCiQ,91
-langfun-0.1.1.dev20240826.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.1.dev20240826.dist-info/RECORD,,
+langfun-0.1.1.dev20240828.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.1.dev20240828.dist-info/METADATA,sha256=zPpmQkxNgUcb7ilrk9JmSBVnq_ISkkujk3bPQ2nPKDQ,5234
+langfun-0.1.1.dev20240828.dist-info/WHEEL,sha256=UvcQYKBHoFqaQd6LKyqHw9fxEolWLQnlzP0h_LgJAfI,91
+langfun-0.1.1.dev20240828.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.1.dev20240828.dist-info/RECORD,,

{langfun-0.1.1.dev20240826.dist-info → langfun-0.1.1.dev20240828.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (73.0.1)
+Generator: setuptools (74.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any