langfun 0.0.2.dev20240605__py3-none-any.whl → 0.0.2.dev20240608__py3-none-any.whl

langfun/core/__init__.py

@@ -101,6 +101,7 @@ from langfun.core.language_model import LanguageModel
 from langfun.core.language_model import LMSample
 from langfun.core.language_model import LMSamplingOptions
 from langfun.core.language_model import LMSamplingUsage
+from langfun.core.language_model import UsageNotAvailable
 from langfun.core.language_model import LMSamplingResult
 from langfun.core.language_model import LMScoringResult
 from langfun.core.language_model import LMCache
@@ -111,12 +112,18 @@ from langfun.core.language_model import RetryableLMError
 from langfun.core.language_model import RateLimitError
 from langfun.core.language_model import TemporaryLMError
 
+# Context manager for tracking usages.
+from langfun.core.language_model import track_usages
+
 # Components for building agents.
 from langfun.core.memory import Memory
 
 # Utility for console output.
 from langfun.core import console
 
+# Utility for event logging.
+from langfun.core import logging
+
 # Import internal modules.
 
 # pylint: enable=g-import-not-at-top

langfun/core/component.py

@@ -210,6 +210,16 @@ def get_contextual_override(var_name: str) -> ContextualOverride | None:
   return _get_scoped_value(_global_tls, _CONTEXT_OVERRIDES, var_name)
 
 
+def context_value(var_name: str, default: Any = RAISE_IF_HAS_ERROR) -> Any:
+  """Returns the value of a variable defined in `lf.context`."""
+  override = get_contextual_override(var_name)
+  if override is None:
+    if default == RAISE_IF_HAS_ERROR:
+      raise KeyError(f'{var_name!r} does not exist in current context.')
+    return default
+  return override.value
+
+
 def all_contextual_values() -> dict[str, Any]:
   """Returns all contextual values provided from `lf.context` in scope."""
   overrides = getattr(_global_tls, _CONTEXT_OVERRIDES, {})

langfun/core/component_test.py

@@ -84,6 +84,11 @@ class ComponentContextTest(unittest.TestCase):
           lf.get_contextual_override('y'),
           lf.ContextualOverride(3, cascade=False, override_attrs=False),
       )
+      self.assertEqual(lf.context_value('x'), 3)
+      self.assertIsNone(lf.context_value('f', None))
+      with self.assertRaisesRegex(KeyError, '.* does not exist'):
+        lf.context_value('f')
+
       self.assertEqual(lf.all_contextual_values(), dict(x=3, y=3, z=3))
 
       # Member attributes take precedence over `lf.context`.

langfun/core/eval/base.py

@@ -549,7 +549,7 @@ class Evaluable(lf.Component):
         )
         s.write(html.escape(pg.format(m.result)))
         s.write('</div>')
-      if 'usage' in m.metadata and m.usage is not None:
+      if m.metadata.get('usage', None):
         s.write(
             '<div style="background-color: #EEEEEE; color: black; '
             'white-space: pre-wrap; padding: 10px; border: 0px solid; '
@@ -1321,7 +1321,7 @@ class Evaluation(Evaluable):
       self._render_summary_metrics(s)
 
       # Summarize average usage.
-      if self.result.usage is not None:
+      if self.result.usage:
         self._render_summary_usage(s)
 
     s.write('</td></tr></table></div>')
@@ -1441,9 +1441,10 @@ class Evaluation(Evaluable):
   def audit_usage(self, message: lf.Message, dryrun: bool = False) -> None:
     del dryrun
     for m in message.trace():
-      if m.metadata.get('usage', None) is not None:
-        self._total_prompt_tokens += m.usage.prompt_tokens
-        self._total_completion_tokens += m.usage.completion_tokens
+      usage = m.metadata.get('usage', None)
+      if usage:
+        self._total_prompt_tokens += usage.prompt_tokens
+        self._total_completion_tokens += usage.completion_tokens
         self._num_usages += 1
 
   def audit_processed(
@@ -1504,7 +1505,7 @@ class Evaluation(Evaluable):
         '<td>Schema</td>'
         '<td>Additional Args</td>'
     )
-    if self.result.usage is not None:
+    if self.result.usage:
       s.write('<td>Usage</td>')
     s.write('<td>OOP Failures</td>')
     s.write('<td>Non-OOP Failures</td>')
@@ -1533,7 +1534,7 @@ class Evaluation(Evaluable):
         f'{_html_repr(self.additional_args, compact=False)}</td>'
     )
     # Usage.
-    if self.result.usage is not None:
+    if self.result.usage:
       s.write('<td>')
       self._render_summary_usage(s)
       s.write('</td>')

langfun/core/langfunc_test.py

@@ -87,7 +87,11 @@ class LangFuncCallTest(unittest.TestCase):
 
     r = l()
     self.assertEqual(
-        r, message.AIMessage('Hello!!!', score=0.0, logprobs=None, usage=None)
+        r,
+        message.AIMessage(
+            'Hello!!!', score=0.0, logprobs=None,
+            usage=language_model.UsageNotAvailable()
+        )
     )
     self.assertEqual(r.tags, ['lm-response', 'lm-output'])
     self.assertEqual(r.source, message.UserMessage('Hello'))
@@ -112,7 +116,11 @@ class LangFuncCallTest(unittest.TestCase):
       self.assertEqual(l.render(), 'Hello')
       r = l()
       self.assertEqual(
-          r, message.AIMessage('Hello!!!', score=0.0, logprobs=None, usage=None)
+          r,
+          message.AIMessage(
+              'Hello!!!', score=0.0, logprobs=None,
+              usage=language_model.UsageNotAvailable()
+          )
       )
       self.assertEqual(r.tags, ['lm-response', 'lm-output'])
 

langfun/core/language_model.py

@@ -14,10 +14,11 @@
 """Interface for language model."""
 
 import abc
+import contextlib
 import dataclasses
 import enum
 import time
-from typing import Annotated, Any, Callable, Sequence, Tuple, Type, Union
+from typing import Annotated, Any, Callable, Iterator, Sequence, Tuple, Type, Union
 from langfun.core import component
 from langfun.core import concurrent
 from langfun.core import console
@@ -84,6 +85,23 @@ class LMSamplingUsage(pg.Object):
   completion_tokens: int
   total_tokens: int
 
+  def __add__(self, other: 'LMSamplingUsage') -> 'LMSamplingUsage':
+    return LMSamplingUsage(
+        prompt_tokens=self.prompt_tokens + other.prompt_tokens,
+        completion_tokens=self.completion_tokens + other.completion_tokens,
+        total_tokens=self.total_tokens + other.total_tokens,
+    )
+
+
+class UsageNotAvailable(LMSamplingUsage):
+  """Usage information not available."""
+  prompt_tokens: pg.typing.Int(0).freeze()        # pytype: disable=invalid-annotation
+  completion_tokens: pg.typing.Int(0).freeze()   # pytype: disable=invalid-annotation
+  total_tokens: pg.typing.Int(0).freeze()        # pytype: disable=invalid-annotation
+
+  def __bool__(self) -> bool:
+    return False
+
 
 class LMSamplingResult(pg.Object):
   """Language model response."""
@@ -97,9 +115,9 @@ class LMSamplingResult(pg.Object):
   ] = []
 
   usage: Annotated[
-      LMSamplingUsage | None,
+      LMSamplingUsage,
       'Usage information. Currently only OpenAI models are supported.',
-  ] = None
+  ] = UsageNotAvailable()
 
 
 class LMSamplingOptions(component.Component):
@@ -398,7 +416,7 @@ class LanguageModel(component.Component):
           # which is accurate when n=1. For n > 1, we average the usage across
           # multiple samples.
           usage = result.usage
-          if len(result.samples) == 1 or usage is None:
+          if len(result.samples) == 1 or not usage:
             response.metadata.usage = usage
           else:
             n = len(result.samples)
@@ -408,6 +426,13 @@ class LanguageModel(component.Component):
                 total_tokens=usage.total_tokens // n,
             )
 
+          # Track usage.
+          trackers = component.context_value('__usage_trackers__', [])
+          if trackers:
+            model_id = self.model_id
+            for tracker in trackers:
+              tracker.track(model_id, usage)
+
           # Track the prompt for corresponding response.
           response.source = prompt
 
@@ -511,7 +536,7 @@ class LanguageModel(component.Component):
       prompt: message_lib.Message,
       response: message_lib.Message,
       call_counter: int,
-      usage: LMSamplingUsage | None,
+      usage: LMSamplingUsage,
       elapse: float,
   ) -> None:
     """Outputs debugging information."""
@@ -529,12 +554,13 @@ class LanguageModel(component.Component):
       self._debug_response(response, call_counter, usage, elapse)
 
   def _debug_model_info(
-      self, call_counter: int, usage: LMSamplingUsage | None) -> None:
+      self, call_counter: int, usage: LMSamplingUsage) -> None:
     """Outputs debugging information about the model."""
     title_suffix = ''
-    if usage and usage.total_tokens != 0:
+    if usage.total_tokens != 0:
       title_suffix = console.colored(
-          f' (total {usage.total_tokens} tokens)', 'red')
+          f' (total {usage.total_tokens} tokens)', 'red'
+      )
 
     console.write(
         self.format(compact=True, use_inferred=True),
@@ -546,11 +572,11 @@ class LanguageModel(component.Component):
       self,
       prompt: message_lib.Message,
       call_counter: int,
-      usage: LMSamplingUsage | None,
+      usage: LMSamplingUsage,
   ) -> None:
     """Outputs debugging information about the prompt."""
     title_suffix = ''
-    if usage and usage.prompt_tokens != 0:
+    if usage.prompt_tokens != 0:
       title_suffix = console.colored(f' ({usage.prompt_tokens} tokens)', 'red')
 
     console.write(
@@ -574,12 +600,12 @@ class LanguageModel(component.Component):
       self,
       response: message_lib.Message,
       call_counter: int,
-      usage: LMSamplingUsage | None,
+      usage: LMSamplingUsage,
       elapse: float
   ) -> None:
     """Outputs debugging information about the response."""
     title_suffix = ' ('
-    if usage and usage.completion_tokens != 0:
+    if usage.completion_tokens != 0:
       title_suffix += f'{usage.completion_tokens} tokens '
     title_suffix += f'in {elapse:.2f} seconds)'
     title_suffix = console.colored(title_suffix, 'red')
@@ -641,7 +667,7 @@ class LanguageModel(component.Component):
       debug = LMDebugMode.ALL if debug else LMDebugMode.NONE
 
     if debug & LMDebugMode.INFO:
-      self._debug_model_info(call_counter, None)
+      self._debug_model_info(call_counter, UsageNotAvailable())
 
     if debug & LMDebugMode.PROMPT:
       console.write(
@@ -692,3 +718,60 @@ class LanguageModel(component.Component):
       return max(int(requests_per_min / 60), 1)  # Max concurrency can't be zero
     else:
       return DEFAULT_MAX_CONCURRENCY  # Default of 1
+
+
+class _UsageTracker:
+  """Usage tracker."""
+
+  def __init__(self, model_ids: set[str] | None):
+    self.model_ids = model_ids
+    self.usages = {
+        m: LMSamplingUsage(0, 0, 0) for m in model_ids
+    } if model_ids else {}
+
+  def track(self, model_id: str, usage: LMSamplingUsage):
+    if self.model_ids is not None and model_id not in self.model_ids:
+      return
+    if not isinstance(usage, UsageNotAvailable) and model_id in self.usages:
+      self.usages[model_id] += usage
+    else:
+      self.usages[model_id] = usage
+
+
+@contextlib.contextmanager
+def track_usages(
+    *lm: Union[str, LanguageModel]
+) -> Iterator[dict[str, LMSamplingUsage]]:
+  """Context manager to track the usages of all language models in scope.
+
+  `lf.track_usages` works with threads spawned by `lf.concurrent_map` and
+  `lf.concurrent_execute`.
+
+  Example:
+    ```
+    lm = lf.llms.GeminiPro1()
+    with lf.track_usages() as usages:
+      # invoke any code that will call LLMs.
+
+    print(usages[lm.model_id])
+    ```
+
+  Args:
+    *lm: The language model(s) to track. If None, track all models in scope.
+
+  Yields:
+    A dictionary of model ID to usage. If a model does not supports usage
+    counting, the dict entry will be None.
+  """
+  if not lm:
+    model_ids = None
+  else:
+    model_ids = [m.model_id if isinstance(m, LanguageModel) else m for m in lm]
+
+  trackers = component.context_value('__usage_trackers__', [])
+  tracker = _UsageTracker(set(model_ids) if model_ids else None)
+  with component.context(__usage_trackers__=trackers + [tracker]):
+    try:
+      yield tracker.usages
+    finally:
+      pass

langfun/core/language_model_test.py

@@ -27,8 +27,8 @@ import pyglove as pg
 @pg.use_init_args(['failures_before_attempt'])
 class MockModel(lm_lib.LanguageModel):
   """A mock model that echo back user prompts."""
-
   failures_before_attempt: int = 0
+  name: str = 'MockModel'
 
   def _sample(self,
               prompts: list[message_lib.Message]
@@ -63,6 +63,10 @@ class MockModel(lm_lib.LanguageModel):
         retry_interval=1,
     )(prompts)
 
+  @property
+  def model_id(self) -> str:
+    return self.name
+
 
 class MockScoringModel(MockModel):
 
@@ -580,6 +584,34 @@ class LanguageModelTest(unittest.TestCase):
     self.assertEqual(lm.rate_to_max_concurrency(requests_per_min=1), 1)
     self.assertEqual(lm.rate_to_max_concurrency(tokens_per_min=1), 1)
 
+  def test_track_usages(self):
+    lm = MockModel(name='model1')
+    lm2 = MockModel(name='model2')
+    with lm_lib.track_usages() as usages1:
+      _ = lm('hi')
+      with lm_lib.track_usages(lm2) as usages2:
+        with lm_lib.track_usages('model1') as usages3:
+          with lm_lib.track_usages('model1', lm2) as usages4:
+            def call_lm(prompt):
+              _ = lm.sample([prompt] * 2)
+            lm2('hi')
+            list(concurrent.concurrent_map(call_lm, ['hi', 'hello']))
+
+    self.assertEqual(usages2, {
+        'model2': lm_lib.LMSamplingUsage(100, 100, 200),
+    })
+    self.assertEqual(usages3, {
+        'model1': lm_lib.LMSamplingUsage(100 * 4, 100 * 4, 200 * 4),
+    })
+    self.assertEqual(usages4, {
+        'model1': lm_lib.LMSamplingUsage(100 * 4, 100 * 4, 200 * 4),
+        'model2': lm_lib.LMSamplingUsage(100, 100, 200),
+    })
+    self.assertEqual(usages1, {
+        'model1': lm_lib.LMSamplingUsage(100 * 5, 100 * 5, 200 * 5),
+        'model2': lm_lib.LMSamplingUsage(100, 100, 200),
+    })
+
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/llms/openai.py

@@ -202,15 +202,14 @@ class OpenAI(lf.LanguageModel):
             lf.LMSample(choice.text.strip(), score=choice.logprobs or 0.0)
         )
 
+      n = len(samples_by_index)
       usage = lf.LMSamplingUsage(
-          prompt_tokens=response.usage.prompt_tokens,
-          completion_tokens=response.usage.completion_tokens,
-          total_tokens=response.usage.total_tokens,
+          prompt_tokens=response.usage.prompt_tokens // n,
+          completion_tokens=response.usage.completion_tokens // n,
+          total_tokens=response.usage.total_tokens // n,
       )
       return [
-          lf.LMSamplingResult(
-              samples_by_index[index], usage=usage if index == 0 else None
-          )
+          lf.LMSamplingResult(samples_by_index[index], usage=usage)
           for index in sorted(samples_by_index.keys())
       ]
 

langfun/core/llms/openai_test.py

@@ -191,9 +191,9 @@ class OpenAITest(unittest.TestCase):
                         score=0.0,
                         logprobs=None,
                         usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66
+                            prompt_tokens=16,
+                            completion_tokens=16,
+                            total_tokens=33
                         ),
                         tags=[lf.Message.TAG_LM_RESPONSE],
                     ),
@@ -206,9 +206,9 @@ class OpenAITest(unittest.TestCase):
                         score=0.1,
                         logprobs=None,
                         usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66
+                            prompt_tokens=16,
+                            completion_tokens=16,
+                            total_tokens=33
                         ),
                         tags=[lf.Message.TAG_LM_RESPONSE],
                     ),
@@ -221,9 +221,9 @@ class OpenAITest(unittest.TestCase):
                         score=0.2,
                         logprobs=None,
                         usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66
+                            prompt_tokens=16,
+                            completion_tokens=16,
+                            total_tokens=33
                         ),
                         tags=[lf.Message.TAG_LM_RESPONSE],
                     ),
@@ -232,7 +232,7 @@ class OpenAITest(unittest.TestCase):
                 ),
             ],
             usage=lf.LMSamplingUsage(
-                prompt_tokens=100, completion_tokens=100, total_tokens=200
+                prompt_tokens=50, completion_tokens=50, total_tokens=100
             ),
         ),
     )
@@ -245,7 +245,11 @@ class OpenAITest(unittest.TestCase):
                         'Sample 0 for prompt 1.',
                         score=0.0,
                         logprobs=None,
-                        usage=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=16,
+                            completion_tokens=16,
+                            total_tokens=33
+                        ),
                         tags=[lf.Message.TAG_LM_RESPONSE],
                     ),
                     score=0.0,
@@ -256,7 +260,11 @@ class OpenAITest(unittest.TestCase):
                         'Sample 1 for prompt 1.',
                         score=0.1,
                         logprobs=None,
-                        usage=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=16,
+                            completion_tokens=16,
+                            total_tokens=33
+                        ),
                         tags=[lf.Message.TAG_LM_RESPONSE],
                     ),
                     score=0.1,
@@ -267,13 +275,20 @@ class OpenAITest(unittest.TestCase):
                         'Sample 2 for prompt 1.',
                         score=0.2,
                         logprobs=None,
-                        usage=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=16,
+                            completion_tokens=16,
+                            total_tokens=33
+                        ),
                         tags=[lf.Message.TAG_LM_RESPONSE],
                     ),
                     score=0.2,
                     logprobs=None,
                 ),
             ],
+            usage=lf.LMSamplingUsage(
+                prompt_tokens=50, completion_tokens=50, total_tokens=100
+            ),
         ),
     )
 

langfun/core/llms/rest_test.py

@@ -89,7 +89,7 @@ class RestTest(unittest.TestCase):
               "max_tokens=4096, stop=['\\n']."
           ),
       )
-      self.assertIsNone(response.usage)
+      self.assertEqual(response.usage, lf.UsageNotAvailable())
 
   def test_call_errors(self):
     for status_code, error_type, error_message in [

langfun/core/logging.py

@@ -0,0 +1,168 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Langfun event logging."""
+
+import datetime
+import io
+import typing
+from typing import Any, Literal, ContextManager
+
+from langfun.core import console
+import pyglove as pg
+
+
+LogLevel = Literal['debug', 'info', 'error', 'warning', 'fatal']
+_LOG_LEVELS = list(typing.get_args(LogLevel))
+_TLS_KEY_MIN_LOG_LEVEL = '_event_log_level'
+
+
+def use_log_level(log_level: LogLevel | None = 'info') -> ContextManager[None]:
+  """Contextmanager to enable logging at a given level."""
+  return pg.object_utils.thread_local_value_scope(
+      _TLS_KEY_MIN_LOG_LEVEL, log_level, 'info')
+
+
+def get_log_level() -> LogLevel | None:
+  """Gets the current minimum log level."""
+  return pg.object_utils.thread_local_get(_TLS_KEY_MIN_LOG_LEVEL, 'info')
+
+
+class LogEntry(pg.Object):
+  """Event log entry."""
+  time: datetime.datetime
+  level: LogLevel
+  message: str
+  metadata: dict[str, Any] = pg.Dict()
+  indent: int = 0
+
+  def should_output(self, min_log_level: LogLevel) -> bool:
+    return _LOG_LEVELS.index(self.level) >= _LOG_LEVELS.index(min_log_level)
+
+  def _repr_html_(self) -> str:
+    s = io.StringIO()
+    padding_left = 50 * self.indent
+    s.write(f'<div style="padding-left: {padding_left}px;">')
+    s.write(self._message_display)
+    if self.metadata:
+      s.write('<div style="padding-left: 20px; margin-top: 10px">')
+      s.write('<table style="border-top: 1px solid #EEEEEE;">')
+      for k, v in self.metadata.items():
+        if hasattr(v, '_repr_html_'):
+          cs = v._repr_html_()  # pylint: disable=protected-access
+        else:
+          cs = f'<span style="white-space: pre-wrap">{str(v)}</span>'
+        key_span = self._round_text(k, color='#F1C40F', margin_bottom='0px')
+        s.write(
+            '<tr>'
+            '<td style="padding: 5px; vertical-align: top; '
+            f'border-bottom: 1px solid #EEEEEE">{key_span}</td>'
+            '<td style="padding: 5px; vertical-align: top; '
+            f'border-bottom: 1px solid #EEEEEE">{cs}</td></tr>'
+        )
+      s.write('</table></div>')
+    return s.getvalue()
+
+  @property
+  def _message_text_color(self) -> str:
+    match self.level:
+      case 'debug':
+        return '#EEEEEE'
+      case 'info':
+        return '#A3E4D7'
+      case 'error':
+        return '#F5C6CB'
+      case 'fatal':
+        return '#F19CBB'
+      case _:
+        raise ValueError(f'Unknown log level: {self.level}')
+
+  @property
+  def _time_display(self) -> str:
+    display_text = self.time.strftime('%H:%M:%S')
+    alt_text = self.time.strftime('%Y-%m-%d %H:%M:%S.%f')
+    return (
+        '<span style="background-color: #BBBBBB; color: white; '
+        'border-radius:5px; padding:0px 5px 0px 5px;" '
+        f'title="{alt_text}">{display_text}</span>'
+    )
+
+  @property
+  def _message_display(self) -> str:
+    return self._round_text(
+        self._time_display + '&nbsp;' + self.message,
+        color=self._message_text_color,
+    )
+
+  def _round_text(
+      self,
+      text: str,
+      *,
+      color: str = '#EEEEEE',
+      display: str = 'inline-block',
+      margin_top: str = '5px',
+      margin_bottom: str = '5px',
+      whitespace: str = 'pre-wrap') -> str:
+    return (
+        f'<span style="background:{color}; display:{display};'
+        f'border-radius:10px; padding:5px; '
+        f'margin-top: {margin_top}; margin-bottom: {margin_bottom}; '
+        f'white-space: {whitespace}">{text}</span>'
+    )
+
+
+def log(level: LogLevel,
+        message: str,
+        *,
+        indent: int = 0,
+        **kwargs) -> LogEntry:
+  """Logs a message."""
+  entry = LogEntry(
+      indent=indent,
+      level=level,
+      time=datetime.datetime.now(),
+      message=message,
+      metadata=kwargs,
+  )
+  if entry.should_output(get_log_level()):
+    if console.under_notebook():
+      console.display(entry)
+    else:
+      # TODO(daiyip): Improve the console output formatting.
+      console.write(entry)
+  return entry
+
+
+def debug(message: str, *, indent: int = 0, **kwargs) -> LogEntry:
+  """Logs a debug message to the session."""
+  return log('debug', message, indent=indent, **kwargs)
+
+
+def info(message: str, *, indent: int = 0, **kwargs) -> LogEntry:
+  """Logs an info message to the session."""
+  return log('info', message, indent=indent, **kwargs)
+
+
+def warning(message: str, *, indent: int = 0, **kwargs) -> LogEntry:
+  """Logs an info message to the session."""
+  return log('warning', message, indent=indent, **kwargs)
+
+
+def error(message: str, *, indent: int = 0, **kwargs) -> LogEntry:
+  """Logs an error message to the session."""
+  return log('error', message, indent=indent, **kwargs)
+
+
+def fatal(message: str, *, indent: int = 0, **kwargs) -> LogEntry:
+  """Logs a fatal message to the session."""
+  return log('fatal', message, indent=indent, **kwargs)

langfun/core/logging_test.py

@@ -0,0 +1,51 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for langfun.core.logging."""
+
+import unittest
+
+from langfun.core import logging
+
+
+class LoggingTest(unittest.TestCase):
+
+  def test_use_log_level(self):
+    self.assertEqual(logging.get_log_level(), 'info')
+    with logging.use_log_level('debug'):
+      self.assertEqual(logging.get_log_level(), 'debug')
+      with logging.use_log_level(None):
+        self.assertIsNone(logging.get_log_level(), None)
+      self.assertEqual(logging.get_log_level(), 'debug')
+    self.assertEqual(logging.get_log_level(), 'info')
+
+  def test_log(self):
+    entry = logging.log('info', 'hi', indent=1, x=1, y=2)
+    self.assertEqual(entry.level, 'info')
+    self.assertEqual(entry.message, 'hi')
+    self.assertEqual(entry.indent, 1)
+    self.assertEqual(entry.metadata, {'x': 1, 'y': 2})
+
+    self.assertEqual(logging.debug('hi').level, 'debug')
+    self.assertEqual(logging.info('hi').level, 'info')
+    self.assertEqual(logging.warning('hi').level, 'warning')
+    self.assertEqual(logging.error('hi').level, 'error')
+    self.assertEqual(logging.fatal('hi').level, 'fatal')
+
+  def test_repr_html(self):
+    entry = logging.log('info', 'hi', indent=1, x=1, y=2)
+    self.assertIn('<div', entry._repr_html_())
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/templates/selfplay_test.py

@@ -57,7 +57,10 @@ class SelfPlayTest(unittest.TestCase):
 
     with lf.context(lm=NumberGuesser(guesses=[50, 20, 5, 10])):
       self.assertEqual(
-          g(), lf.AIMessage('10', score=0.0, logprobs=None, usage=None)
+          g(),
+          lf.AIMessage(
+              '10', score=0.0, logprobs=None, usage=lf.UsageNotAvailable()
+          )
       )
 
     self.assertEqual(g.num_turns, 4)
@@ -67,7 +70,10 @@ class SelfPlayTest(unittest.TestCase):
 
     with lf.context(lm=NumberGuesser(guesses=[50, 20, 5, 2, 5, 4])):
       self.assertEqual(
-          g(), lf.AIMessage('2', score=0.0, logprobs=None, usage=None)
+          g(),
+          lf.AIMessage(
+              '2', score=0.0, logprobs=None, usage=lf.UsageNotAvailable()
+          )
       )
 
     self.assertEqual(g.num_turns, 10)

{langfun-0.0.2.dev20240605.dist-info → langfun-0.0.2.dev20240608.dist-info}/METADATA

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

{langfun-0.0.2.dev20240605.dist-info → langfun-0.0.2.dev20240608.dist-info}/RECORD

@@ -1,15 +1,17 @@
 langfun/__init__.py,sha256=P62MnqA6-f0h8iYfQ3MT6Yg7a4qRnQeb4GrIn6dcSnY,2274
-langfun/core/__init__.py,sha256=ZheiCpop_GAZbVpnSS-uPBJaEEM15Td5xFGGizSGqko,4514
-langfun/core/component.py,sha256=oxesbC0BoE_TbtxwW5x-BAZWxZyyJbuPiX5S38RqCv0,9909
-langfun/core/component_test.py,sha256=uR-_Sz_42Jxc5qzLIB-f5_pXmNwnC01Xlbv5NOQSeSU,8021
+langfun/core/__init__.py,sha256=nFQgTyhNLA83b_V_nPNByyO2JlF576AdS61AfA0SaN8,4728
+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=mwFMZhDUdppnDr7vDSTwcbMHwrdsIoKJwRYNtl4ZWL4,15185
 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=_mfARnakX3oji5HDigFSLMd6yQ2wma-2Mgbztwqn73g,8501
-langfun/core/language_model.py,sha256=PocBg1t3uB0a_bJntLW5aagHhNbZsVdp2iduSBEW6ro,21240
-langfun/core/language_model_test.py,sha256=NZaSUls6cZdtxiqkqumWbtkx9zgNiJlsviYZOWkuHig,20137
+langfun/core/langfunc_test.py,sha256=lyt-UzkD8972cxZwzCkps0_RMLeSsOBrcUFIW-fB6us,8653
+langfun/core/language_model.py,sha256=q41EGhbgoCe68eYWWTAzgt3r6SbdCpINRdXnNGstukQ,23775
+langfun/core/language_model_test.py,sha256=HmA0GACK4-6tCH32TFkfYj9w4CxbrynKqXdKBgiqgwo,21255
+langfun/core/logging.py,sha256=FyZRxUy2TTF6tWLhQCRpCvfH55WGUdNgQjUTK_SQLnY,5320
+langfun/core/logging_test.py,sha256=qvm3RObYP3knO2PnXR9evBRl4gH621GnjnwywbGbRfg,1833
 langfun/core/memory.py,sha256=f-asN1F7Vehgdn_fK84v73GrEUOxRtaW934keutTKjk,2416
 langfun/core/message.py,sha256=Rw3yC9HyGRjMhfDgyNjGlSCALEyDDbJ0_o6qTXeeDiQ,15738
 langfun/core/message_test.py,sha256=b6DDRoQ5j3uK-dc0QPSLelNTKaXX10MxJrRiI61iGX4,9574
@@ -40,7 +42,7 @@ 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=Evt-E4FEhZF2tXL6-byh_AyA7Cc_ZoGmvnN7vkAZedk,1898
-langfun/core/eval/base.py,sha256=zcMPBKmcll5O08waEEnvmkEoXgcINhOat9rRJk8X8b4,74268
+langfun/core/eval/base.py,sha256=GM98Zo4gxZui2ORX6Q7Zr94PfiEViQC5X_qz-uj6b2k,74220
 langfun/core/eval/base_test.py,sha256=cHOTIWVW4Dp8gKKIKcZrAcJ-w84j2GIozTzJoiAX7p4,26743
 langfun/core/eval/matching.py,sha256=Y4vFoNTQEOwko6IA8l9OZ52-vt52e3VGmcTtvLA67wM,9782
 langfun/core/eval/matching_test.py,sha256=f7iVyXH5KGJBWt4Wp14Bt9J3X59A6Ayfog9MbuFvPew,5532
@@ -59,10 +61,10 @@ langfun/core/llms/groq.py,sha256=pqtyOZ_1_OJMOg8xATWT_B_SVbuT9nMRf4VkH9GzW8g,630
 langfun/core/llms/groq_test.py,sha256=GYF_Qtq5S1H1TrKH38t6_lkdroqT7v-joYLDKnmS9e0,5274
 langfun/core/llms/llama_cpp.py,sha256=9tXQntSCDtjTF3bnyJrAPCr4N6wycy5nXYvp9uduygE,2843
 langfun/core/llms/llama_cpp_test.py,sha256=MWO_qaOeKjRniGjcaWPDScd7HPaIJemqUZoslrt4FPs,1806
-langfun/core/llms/openai.py,sha256=IN46gIqfY6aEEfxCPNmyH1hrep3oWBhJDwVFilfqNkM,13657
-langfun/core/llms/openai_test.py,sha256=QWDzTgi8F2Z9u9ip6alK4rDEp_YraVTxWlDX5XOsKJk,14858
+langfun/core/llms/openai.py,sha256=0z9qIH9FlWj9VWUnhOX321T6JHO-vjY2IozT7OVI4GY,13654
+langfun/core/llms/openai_test.py,sha256=3muDTnW7UBOSHq694Fi2bofqhe8Pkj0Tl8IShoLCTOM,15525
 langfun/core/llms/rest.py,sha256=laopuq-zD8V-3Y6eFDngftHEbE66VlUkCD2-rvvRaLU,3388
-langfun/core/llms/rest_test.py,sha256=Zw08Xbl_O2OQuUglLXHsPsY5KW2VEcPGl1gR8PRHuFY,3449
+langfun/core/llms/rest_test.py,sha256=NZ3Nf0XQVpT9kLP5cBVo_yBHLI7vWTYhWQxYEJVMGs4,3472
 langfun/core/llms/vertexai.py,sha256=wIpckH-rMHUBA1vhauQk4LVrSsPQEsVntz7kLDKwm9g,11359
 langfun/core/llms/vertexai_test.py,sha256=G18BG36h5KvmX2zutDTLjtYCRjTuP_nWIFm4FMnLnyY,7651
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
@@ -112,9 +114,9 @@ langfun/core/templates/conversation_test.py,sha256=RryYyIhfc34dLWOs6GfPQ8HU8mXpK
 langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fikKhwhzwhpKI,1460
 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=DYVrkk7uNKCqJGEHH31HssU2BPuMItU1vJLzfcXIlYg,2156
-langfun-0.0.2.dev20240605.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.0.2.dev20240605.dist-info/METADATA,sha256=NMWv4oYMcRuZXl22coMTShqGC8hj_Y2PGfTk7n-Alt0,3550
-langfun-0.0.2.dev20240605.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-langfun-0.0.2.dev20240605.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.0.2.dev20240605.dist-info/RECORD,,
+langfun/core/templates/selfplay_test.py,sha256=rBW2Qr8yi-aWYwoTwRR-n1peKyMX9QXPZXURjLgoiRs,2264
+langfun-0.0.2.dev20240608.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.0.2.dev20240608.dist-info/METADATA,sha256=HhaObQS5WL8Y2uRuwNX389KbnRiRPsV3Vl-XzVGTN0A,3550
+langfun-0.0.2.dev20240608.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+langfun-0.0.2.dev20240608.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.0.2.dev20240608.dist-info/RECORD,,