langfun 0.1.2.dev202410100804__py3-none-any.whl → 0.1.2.dev202410120803__py3-none-any.whl

langfun/core/llms/vertexai.py

@@ -40,24 +40,106 @@ except ImportError:
   Credentials = Any
 
 
+# https://cloud.google.com/vertex-ai/generative-ai/pricing
+# describes that the average number of characters per token is about 4.
+AVGERAGE_CHARS_PER_TOEKN = 4
+
+
+# Price in US dollars,
+# from https://cloud.google.com/vertex-ai/generative-ai/pricing
+# as of 2024-10-10.
 SUPPORTED_MODELS_AND_SETTINGS = {
-    'gemini-1.5-pro-001': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-pro-002': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-flash-002': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-flash-001': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-pro': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-flash': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-pro-latest': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-flash-latest': pg.Dict(api='gemini', rpm=500),
-    'gemini-1.5-pro-preview-0514': pg.Dict(api='gemini', rpm=50),
-    'gemini-1.5-pro-preview-0409': pg.Dict(api='gemini', rpm=50),
-    'gemini-1.5-flash-preview-0514': pg.Dict(api='gemini', rpm=200),
-    'gemini-1.0-pro': pg.Dict(api='gemini', rpm=300),
-    'gemini-1.0-pro-vision': pg.Dict(api='gemini', rpm=100),
+    'gemini-1.5-pro-001': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.0003125,
+        cost_per_1k_output_chars=0.00125,
+    ),
+    'gemini-1.5-pro-002': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.0003125,
+        cost_per_1k_output_chars=0.00125,
+    ),
+    'gemini-1.5-flash-002': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.00001875,
+        cost_per_1k_output_chars=0.000075,
+    ),
+    'gemini-1.5-flash-001': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.00001875,
+        cost_per_1k_output_chars=0.000075,
+    ),
+    'gemini-1.5-pro': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.0003125,
+        cost_per_1k_output_chars=0.00125,
+    ),
+    'gemini-1.5-flash': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.00001875,
+        cost_per_1k_output_chars=0.000075,
+    ),
+    'gemini-1.5-pro-latest': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.0003125,
+        cost_per_1k_output_chars=0.00125,
+    ),
+    'gemini-1.5-flash-latest': pg.Dict(
+        api='gemini',
+        rpm=500,
+        cost_per_1k_input_chars=0.00001875,
+        cost_per_1k_output_chars=0.000075,
+    ),
+    'gemini-1.5-pro-preview-0514': pg.Dict(
+        api='gemini',
+        rpm=50,
+        cost_per_1k_input_chars=0.0003125,
+        cost_per_1k_output_chars=0.00125,
+    ),
+    'gemini-1.5-pro-preview-0409': pg.Dict(
+        api='gemini',
+        rpm=50,
+        cost_per_1k_input_chars=0.0003125,
+        cost_per_1k_output_chars=0.00125,
+    ),
+    'gemini-1.5-flash-preview-0514': pg.Dict(
+        api='gemini',
+        rpm=200,
+        cost_per_1k_input_chars=0.00001875,
+        cost_per_1k_output_chars=0.000075,
+    ),
+    'gemini-1.0-pro': pg.Dict(
+        api='gemini',
+        rpm=300,
+        cost_per_1k_input_chars=0.000125,
+        cost_per_1k_output_chars=0.000375,
+    ),
+    'gemini-1.0-pro-vision': pg.Dict(
+        api='gemini',
+        rpm=100,
+        cost_per_1k_input_chars=0.000125,
+        cost_per_1k_output_chars=0.000375,
+    ),
     # PaLM APIs.
-    'text-bison': pg.Dict(api='palm', rpm=1600),
-    'text-bison-32k': pg.Dict(api='palm', rpm=300),
-    'text-unicorn': pg.Dict(api='palm', rpm=100),
+    'text-bison': pg.Dict(
+        api='palm',
+        rpm=1600
+    ),
+    'text-bison-32k': pg.Dict(
+        api='palm',
+        rpm=300
+    ),
+    'text-unicorn': pg.Dict(
+        api='palm',
+        rpm=100
+    ),
     # Endpoint
     # TODO(chengrun): Set a more appropriate rpm for endpoint.
     'custom': pg.Dict(api='endpoint', rpm=20),
@@ -161,6 +243,25 @@ class VertexAI(lf.LanguageModel):
         tokens_per_min=0,
     )
 
+  def estimate_cost(
+      self,
+      num_input_tokens: int,
+      num_output_tokens: int
+  ) -> float | None:
+    """Estimate the cost based on usage."""
+    cost_per_1k_input_chars = SUPPORTED_MODELS_AND_SETTINGS[self.model].get(
+        'cost_per_1k_input_chars', None
+    )
+    cost_per_1k_output_chars = SUPPORTED_MODELS_AND_SETTINGS[self.model].get(
+        'cost_per_1k_output_chars', None
+    )
+    if cost_per_1k_output_chars is None or cost_per_1k_input_chars is None:
+      return None
+    return (
+        cost_per_1k_input_chars * num_input_tokens
+        + cost_per_1k_output_chars * num_output_tokens
+    ) * AVGERAGE_CHARS_PER_TOEKN / 1000
+
   def _generation_config(
       self, prompt: lf.Message, options: lf.LMSamplingOptions
   ) -> Any:  # generative_models.GenerationConfig
@@ -285,6 +386,10 @@ class VertexAI(lf.LanguageModel):
         prompt_tokens=usage_metadata.prompt_token_count,
         completion_tokens=usage_metadata.candidates_token_count,
         total_tokens=usage_metadata.total_token_count,
+        estimated_cost=self.estimate_cost(
+            num_input_tokens=usage_metadata.prompt_token_count,
+            num_output_tokens=usage_metadata.candidates_token_count,
+        ),
     )
     return lf.LMSamplingResult(
         [

langfun/core/logging.py

@@ -13,16 +13,13 @@
 # limitations under the License.
 """Langfun event logging."""
 
-from collections.abc import Iterator
 import contextlib
 import datetime
-import io
 import typing
-from typing import Any, Literal
+from typing import Any, Iterator, Literal, Sequence
 
 from langfun.core import component
 from langfun.core import console
-from langfun.core import repr_utils
 import pyglove as pg
 
 
@@ -56,49 +53,159 @@ class LogEntry(pg.Object):
   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(repr_utils.html_repr(self.metadata))
-    s.write('</div>')
-    return s.getvalue()
-
-  @property
-  def _message_text_bgcolor(self) -> str:
-    match self.level:
-      case 'debug':
-        return '#EEEEEE'
-      case 'info':
-        return '#A3E4D7'
-      case 'warning':
-        return '#F8C471'
-      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>'
+  def _html_tree_view_summary(
+      self,
+      view: pg.views.HtmlTreeView,
+      title: str | pg.Html | None = None,
+      max_str_len_for_summary: int = pg.View.PresetArgValue(80),  # pytype: disable=annotation-type-mismatch
+      **kwargs
+      ) -> str:
+    if len(self.message) > max_str_len_for_summary:
+      message = self.message[:max_str_len_for_summary] + '...'
+    else:
+      message = self.message
+
+    s = pg.Html(
+        pg.Html.element(
+            'span',
+            [self.time.strftime('%H:%M:%S')],
+            css_class=['log-time']
+        ),
+        pg.Html.element(
+            'span',
+            [pg.Html.escape(message)],
+            css_class=['log-summary'],
+        ),
+    )
+    return view.summary(
+        self,
+        title=title or s,
+        max_str_len_for_summary=max_str_len_for_summary,
+        **kwargs,
     )
 
-  @property
-  def _message_display(self) -> str:
-    return repr_utils.html_round_text(
-        self._time_display + '&nbsp;' + self.message,
-        background_color=self._message_text_bgcolor,
+  # pytype: disable=annotation-type-mismatch
+  def _html_tree_view_content(
+      self,
+      view: pg.views.HtmlTreeView,
+      root_path: pg.KeyPath,
+      collapse_log_metadata_level: int | None = pg.View.PresetArgValue(0),
+      max_str_len_for_summary: int = pg.View.PresetArgValue(80),
+      collapse_level: int | None = pg.View.PresetArgValue(1),
+      **kwargs
+  ) -> pg.Html:
+    # pytype: enable=annotation-type-mismatch
+    def render_message_text():
+      if len(self.message) < max_str_len_for_summary:
+        return None
+      return pg.Html.element(
+          'span',
+          [pg.Html.escape(self.message)],
+          css_class=['log-text'],
+      )
+
+    def render_metadata():
+      if not self.metadata:
+        return None
+      child_path = root_path + 'metadata'
+      return pg.Html.element(
+          'div',
+          [
+              view.render(
+                  self.metadata,
+                  name='metadata',
+                  root_path=child_path,
+                  parent=self,
+                  collapse_level=(
+                      view.max_collapse_level(
+                          collapse_level,
+                          collapse_log_metadata_level,
+                          child_path
+                      )
+                  )
+              )
+          ],
+          css_class=['log-metadata'],
+      )
+
+    return pg.Html.element(
+        'div',
+        [
+            render_message_text(),
+            render_metadata(),
+        ],
+        css_class=['complex_value'],
     )
 
+  def _html_style(self) -> list[str]:
+    return super()._html_style() + [
+        """
+        .log-time {
+          color: #222;
+          font-size: 12px;
+          padding-right: 10px;
+        }
+        .log-summary {
+          font-weight: normal;
+          font-style: italic;
+          padding: 4px;
+        }
+        .log-debug > summary > .summary_title::before {
+          content: '🛠️ '
+        }
+        .log-info > summary > .summary_title::before {
+          content: '💡 '
+        }
+        .log-warning > summary > .summary_title::before {
+          content: '❗ '
+        }
+        .log-error > summary > .summary_title::before {
+          content: '❌ '
+        }
+        .log-fatal > summary > .summary_title::before {
+          content: '💀 '
+        }
+        .log-text {
+          display: block;
+          color: black;
+          font-style: italic;
+          padding: 20px;
+          border-radius: 5px;
+          background: rgba(255, 255, 255, 0.5);
+          white-space: pre-wrap;
+        }
+        details.log-entry {
+          margin: 0px 0px 10px;
+          border: 0px;
+        }
+        div.log-metadata {
+          margin: 10px 0px 0px 0px; 
+        }
+        .log-metadata > details {
+          background-color: rgba(255, 255, 255, 0.5);
+          border: 1px solid transparent;
+        }
+        .log-debug {
+          background-color: #EEEEEE
+        }
+        .log-warning {
+          background-color: #F8C471
+        }
+        .log-info {
+          background-color: #A3E4D7
+        }
+        .log-error {
+          background-color: #F5C6CB
+        }
+        .log-fatal {
+          background-color: #F19CBB
+        }
+        """
+    ]
+
+  def _html_element_class(self) -> Sequence[str] | None:
+    return super()._html_element_class() + [f'log-{self.level}']
+
 
 def log(level: LogLevel,
         message: str,

langfun/core/logging_test.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 """Tests for langfun.core.logging."""
 
+import datetime
+import inspect
 import unittest
 
 from langfun.core import logging
@@ -52,6 +54,37 @@ class LoggingTest(unittest.TestCase):
     assert_color(logging.error('hi', indent=2, x=1, y=2), '#F5C6CB')
     assert_color(logging.fatal('hi', indent=2, x=1, y=2), '#F19CBB')
 
+  def assert_html_content(self, html, expected):
+    expected = inspect.cleandoc(expected).strip()
+    actual = html.content.strip()
+    if actual != expected:
+      print(actual)
+    self.assertEqual(actual, expected)
+
+  def test_html(self):
+    time = datetime.datetime(2024, 10, 10, 12, 30, 45)
+    self.assert_html_content(
+        logging.LogEntry(
+            level='info', message='5 + 2 > 3',
+            time=time, metadata={}
+        ).to_html(enable_summary_tooltip=False),
+        """
+        <details open class="pyglove log-entry log-info"><summary><div class="summary_title"><span class="log-time">12:30:45</span><span class="log-summary">5 + 2 &gt; 3</span></div></summary><div class="complex_value"></div></details>
+        """
+    )
+    self.assert_html_content(
+        logging.LogEntry(
+            level='error', message='This is a longer message: 5 + 2 > 3',
+            time=time, metadata=dict(x=1, y=2)
+        ).to_html(
+            max_str_len_for_summary=10,
+            enable_summary_tooltip=False,
+            collapse_log_metadata_level=1
+        ),
+        """
+        <details open class="pyglove log-entry log-error"><summary><div class="summary_title"><span class="log-time">12:30:45</span><span class="log-summary">This is a ...</span></div></summary><div class="complex_value"><span class="log-text">This is a longer message: 5 + 2 &gt; 3</span><div class="log-metadata"><details open class="pyglove dict"><summary><div class="summary_name">metadata</div><div class="summary_title">Dict(...)</div></summary><div class="complex_value dict"><table><tr><td><span class="object_key str">x</span><span class="tooltip key-path">metadata.x</span></td><td><div><span class="simple_value int">1</span></div></td></tr><tr><td><span class="object_key str">y</span><span class="tooltip key-path">metadata.y</span></td><td><div><span class="simple_value int">2</span></div></td></tr></table></div></details></div></div></details>
+        """
+    )
 
 if __name__ == '__main__':
   unittest.main()