langfun 0.0.2.dev20240606__py3-none-any.whl → 0.0.2.dev20240607__py3-none-any.whl

langfun/core/__init__.py

@@ -111,12 +111,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/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,13 @@ 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 LMSamplingResult(pg.Object):
   """Language model response."""
@@ -408,6 +416,16 @@ class LanguageModel(component.Component):
                 total_tokens=usage.total_tokens // n,
             )
 
+          # Track usage.
+          if usage:
+            tracked_usages = component.context_value('__tracked_usages__', [])
+            model_id = self.model_id
+            for usage_dict in tracked_usages:
+              if model_id in usage_dict:
+                usage_dict[model_id] += usage
+              else:
+                usage_dict[model_id] = usage
+
           # Track the prompt for corresponding response.
           response.source = prompt
 
@@ -692,3 +710,15 @@ 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
+
+
+@contextlib.contextmanager
+def track_usages() -> Iterator[dict[str, LMSamplingUsage]]:
+  """Context manager to track the usage of language models by model ID."""
+  tracked_usages = component.context_value('__tracked_usages__', [])
+  current_usage = dict()
+  with component.context(__tracked_usages__=tracked_usages + [current_usage]):
+    try:
+      yield current_usage
+    finally:
+      pass

langfun/core/language_model_test.py

@@ -580,6 +580,18 @@ 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):
+    with lm_lib.track_usages() as usages1:
+      lm = MockModel()
+      _ = lm('hi')
+      with lm_lib.track_usages() as usages2:
+        _ = lm('hi')
+    self.assertEqual(usages2, {
+        'MockModel': lm_lib.LMSamplingUsage(100, 100, 200),
+    })
+    self.assertEqual(usages1, {
+        'MockModel': lm_lib.LMSamplingUsage(200, 200, 400),
+    })
 
 if __name__ == '__main__':
   unittest.main()

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-0.0.2.dev20240606.dist-info → langfun-0.0.2.dev20240607.dist-info}/METADATA

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

{langfun-0.0.2.dev20240606.dist-info → langfun-0.0.2.dev20240607.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=F3WGAww--u0CdQkg4ENBudsd0ZeGdecscF0R3YXSWmE,4670
+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/language_model.py,sha256=geHX9bHLjWGlubgGvSLTGluFyB7x_NIk09p5KPkp1Wg,22328
+langfun/core/language_model_test.py,sha256=UfAg7ExCWSVqFpFGFOvQWSwO343FcvzVImb_F3hTfa0,20517
+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
@@ -113,8 +115,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=DYVrkk7uNKCqJGEHH31HssU2BPuMItU1vJLzfcXIlYg,2156
-langfun-0.0.2.dev20240606.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.0.2.dev20240606.dist-info/METADATA,sha256=3uuzrvWTcGuq7AMpsCkUFvOlk6U_R_mFLOpU4ZFS9Wk,3550
-langfun-0.0.2.dev20240606.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-langfun-0.0.2.dev20240606.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.0.2.dev20240606.dist-info/RECORD,,
+langfun-0.0.2.dev20240607.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.0.2.dev20240607.dist-info/METADATA,sha256=rFmOTbbcV4Jlik6ma8c87fUqXmYpscGfb3_FcMN2Nf0,3550
+langfun-0.0.2.dev20240607.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+langfun-0.0.2.dev20240607.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.0.2.dev20240607.dist-info/RECORD,,