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

langfun/core/llms/vertexai.py

@@ -0,0 +1,192 @@
+# Copyright 2023 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.
+"""Vertex AI generative models."""
+
+import functools
+import os
+from typing import Annotated, Any
+
+import langfun.core as lf
+from langfun.core.llms import gemini
+import pyglove as pg
+
+try:
+  # pylint: disable=g-import-not-at-top
+  from google import auth as google_auth
+  from google.auth import credentials as credentials_lib
+  from google.auth.transport import requests as auth_requests
+  # pylint: enable=g-import-not-at-top
+
+  Credentials = credentials_lib.Credentials
+except ImportError:
+  google_auth = None
+  credentials_lib = None
+  auth_requests = None
+  Credentials = Any
+
+
+@lf.use_init_args(['model'])
+@pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
+class VertexAI(gemini.Gemini):
+  """Language model served on VertexAI with REST API."""
+
+  project: Annotated[
+      str | None,
+      (
+          'Vertex AI project ID. Or set from environment variable '
+          'VERTEXAI_PROJECT.'
+      ),
+  ] = None
+
+  location: Annotated[
+      str | None,
+      (
+          'Vertex AI service location. Or set from environment variable '
+          'VERTEXAI_LOCATION.'
+      ),
+  ] = None
+
+  credentials: Annotated[
+      Credentials | None,
+      (
+          'Credentials to use. If None, the default credentials to the '
+          'environment will be used.'
+      ),
+  ] = None
+
+  def _on_bound(self):
+    super()._on_bound()
+    if google_auth is None:
+      raise ValueError(
+          'Please install "langfun[llm-google-vertex]" to use Vertex AI models.'
+      )
+    self._project = None
+    self._credentials = None
+
+  def _initialize(self):
+    project = self.project or os.environ.get('VERTEXAI_PROJECT', None)
+    if not project:
+      raise ValueError(
+          'Please specify `project` during `__init__` or set environment '
+          'variable `VERTEXAI_PROJECT` with your Vertex AI project ID.'
+      )
+
+    location = self.location or os.environ.get('VERTEXAI_LOCATION', None)
+    if not location:
+      raise ValueError(
+          'Please specify `location` during `__init__` or set environment '
+          'variable `VERTEXAI_LOCATION` with your Vertex AI service location.'
+      )
+
+    self._project = project
+    self._location = location
+
+    credentials = self.credentials
+    if credentials is None:
+      # Use default credentials.
+      credentials = google_auth.default(
+          scopes=['https://www.googleapis.com/auth/cloud-platform']
+      )
+    self._credentials = credentials
+
+  @property
+  def model_id(self) -> str:
+    """Returns a string to identify the model."""
+    return f'VertexAI({self.model})'
+
+  @functools.cached_property
+  def _session(self):
+    assert self._api_initialized
+    assert self._credentials is not None
+    assert auth_requests is not None
+    s = auth_requests.AuthorizedSession(self._credentials)
+    s.headers.update(self.headers or {})
+    return s
+
+  @property
+  def api_endpoint(self) -> str:
+    assert self._api_initialized
+    return (
+        f'https://{self._location}-aiplatform.googleapis.com/v1/projects/'
+        f'{self._project}/locations/{self._location}/publishers/google/'
+        f'models/{self.model}:generateContent'
+    )
+
+
+class VertexAIGeminiFlash2_0ThinkingExp_20241219(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini Flash 2.0 Thinking model launched on 12/19/2024."""
+
+  api_version = 'v1alpha'
+  model = 'gemini-2.0-flash-thinking-exp-1219'
+  timeout = None
+
+
+class VertexAIGeminiFlash2_0Exp(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 2.0 Flash model."""
+
+  model = 'gemini-2.0-flash-exp'
+
+
+class VertexAIGeminiExp_20241206(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini Experimental model launched on 12/06/2024."""
+
+  model = 'gemini-exp-1206'
+
+
+class VertexAIGeminiExp_20241114(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini Experimental model launched on 11/14/2024."""
+
+  model = 'gemini-exp-1114'
+
+
+class VertexAIGeminiPro1_5(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.5 Pro model."""
+
+  model = 'gemini-1.5-pro-latest'
+
+
+class VertexAIGeminiPro1_5_002(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.5 Pro model."""
+
+  model = 'gemini-1.5-pro-002'
+
+
+class VertexAIGeminiPro1_5_001(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.5 Pro model."""
+
+  model = 'gemini-1.5-pro-001'
+
+
+class VertexAIGeminiFlash1_5(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.5 Flash model."""
+
+  model = 'gemini-1.5-flash'
+
+
+class VertexAIGeminiFlash1_5_002(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.5 Flash model."""
+
+  model = 'gemini-1.5-flash-002'
+
+
+class VertexAIGeminiFlash1_5_001(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.5 Flash model."""
+
+  model = 'gemini-1.5-flash-001'
+
+
+class VertexAIGeminiPro1(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.0 Pro model."""
+
+  model = 'gemini-1.0-pro'

langfun/core/llms/vertexai_test.py

@@ -0,0 +1,52 @@
+# 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 VertexAI models."""
+
+import os
+import unittest
+from unittest import mock
+
+from langfun.core.llms import vertexai
+
+
+class VertexAITest(unittest.TestCase):
+  """Tests for Vertex model with REST API."""
+
+  @mock.patch.object(vertexai.VertexAI, 'credentials', new=True)
+  def test_project_and_location_check(self):
+    with self.assertRaisesRegex(ValueError, 'Please specify `project`'):
+      _ = vertexai.VertexAIGeminiPro1()._api_initialized
+
+    with self.assertRaisesRegex(ValueError, 'Please specify `location`'):
+      _ = vertexai.VertexAIGeminiPro1(project='abc')._api_initialized
+
+    self.assertTrue(
+        vertexai.VertexAIGeminiPro1(
+            project='abc', location='us-central1'
+        )._api_initialized
+    )
+
+    os.environ['VERTEXAI_PROJECT'] = 'abc'
+    os.environ['VERTEXAI_LOCATION'] = 'us-central1'
+    model = vertexai.VertexAIGeminiPro1()
+    self.assertTrue(model.model_id.startswith('VertexAI('))
+    self.assertIn('us-central1', model.api_endpoint)
+    self.assertTrue(model._api_initialized)
+    self.assertIsNotNone(model._session)
+    del os.environ['VERTEXAI_PROJECT']
+    del os.environ['VERTEXAI_LOCATION']
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/logging.py

@@ -0,0 +1,284 @@
+# 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 contextlib
+import datetime
+import functools
+import typing
+from typing import Any, Iterator, Literal
+
+from langfun.core import component
+from langfun.core import console
+import pyglove as pg
+
+
+LogLevel = Literal['debug', 'info', 'error', 'warning', 'fatal']
+_LOG_LEVELS = list(typing.get_args(LogLevel))
+
+
+@contextlib.contextmanager
+def use_log_level(log_level: LogLevel = 'info') -> Iterator[None]:
+  """Contextmanager to enable logging at a given level."""
+  with component.context(__event_log_level__=log_level):
+    try:
+      yield
+    finally:
+      pass
+
+
+def get_log_level() -> LogLevel:
+  """Gets the current minimum log level."""
+  return component.context_value('__event_log_level__', 'info')
+
+
+class LogEntry(pg.Object, pg.views.HtmlTreeView.Extension):
+  """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 format(self,
+             compact: bool = False,
+             verbose: bool = True,
+             root_indent: int = 0,
+             *,
+             text_format: bool = True,
+             **kwargs):
+    if text_format:
+      s = f"""{self.time.strftime('%H:%M:%S')} {self.level.upper()} - {self.message}"""
+      if self.metadata:
+        s += f' (metadata: {self.metadata!r})'
+      return s
+    return super().format(
+        compact=compact,
+        verbose=verbose,
+        root_indent=root_indent,
+        **kwargs
+    )
+
+  def _html_tree_view_summary(
+      self,
+      view: pg.views.HtmlTreeView,
+      title: str | pg.Html | None = None,
+      max_summary_len_for_str: int = 80,
+      **kwargs
+  ) -> pg.Html | None:
+    if len(self.message) > max_summary_len_for_str:
+      message = self.message[:max_summary_len_for_str] + '...'
+    else:
+      message = self.message
+
+    s = pg.Html(
+        pg.Html.element(
+            'span',
+            [self.time.strftime('%H:%M:%S')],
+            css_classes=['log-time']
+        ),
+        pg.Html.element(
+            'span',
+            [pg.Html.escape(message)],
+            css_classes=['log-summary'],
+        ),
+    )
+    return view.summary(
+        self,
+        title=title or s,
+        max_summary_len_for_str=max_summary_len_for_str,
+        **kwargs,
+    )
+
+  # pytype: disable=annotation-type-mismatch
+  def _html_tree_view_content(
+      self,
+      view: pg.views.HtmlTreeView,
+      root_path: pg.KeyPath | None = None,
+      max_summary_len_for_str: int = 80,
+      collapse_level: int | None = 1,
+      extra_flags: dict[str, Any] | None = None,
+      **kwargs
+  ) -> pg.Html:
+    # pytype: enable=annotation-type-mismatch
+    extra_flags = extra_flags if extra_flags is not None else {}
+    collapse_log_metadata_level: int | None = extra_flags.get(
+        'collapse_log_metadata_level', None
+    )
+    def render_message_text():
+      if len(self.message) < max_summary_len_for_str:
+        return None
+      return pg.Html.element(
+          'span',
+          [pg.Html.escape(self.message)],
+          css_classes=['log-text'],
+      )
+
+    def render_metadata():
+      if not self.metadata:
+        return None
+      return pg.Html.element(
+          'div',
+          [
+              view.render(
+                  self.metadata,
+                  name='metadata',
+                  root_path=pg.KeyPath('metadata', root_path),
+                  parent=self,
+                  collapse_level=view.get_collapse_level(
+                      (collapse_level, -1), collapse_log_metadata_level,
+                  ),
+                  max_summary_len_for_str=max_summary_len_for_str,
+                  extra_flags=extra_flags,
+                  **view.get_passthrough_kwargs(**kwargs),
+              )
+          ],
+          css_classes=['log-metadata'],
+      )
+
+    return pg.Html.element(
+        'div',
+        [
+            render_message_text(),
+            render_metadata(),
+        ],
+        css_classes=['complex_value'],
+    )
+
+  def _html_tree_view_config(self) -> dict[str, Any]:
+    return pg.views.HtmlTreeView.get_kwargs(
+        super()._html_tree_view_config(),
+        dict(
+            css_classes=[f'log-{self.level}'],
+        )
+    )
+
+  @classmethod
+  @functools.cache
+  def _html_tree_view_css_styles(cls) -> list[str]:
+    return super()._html_tree_view_css_styles() + [
+        """
+        /* Langfun LogEntry styles. */
+        .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 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,125 @@
+# 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 datetime
+import inspect
+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):
+    def assert_color(entry, color):
+      self.assertIn(f'background-color: {color}', entry._repr_html_())
+
+    assert_color(logging.debug('hi', indent=0), '#EEEEEE')
+    assert_color(logging.info('hi', indent=1), '#A3E4D7')
+    assert_color(logging.warning('hi', x=1, y=2), '#F8C471')
+    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_format(self):
+    time = datetime.datetime(2024, 10, 10, 12, 30, 45)
+    self.assertEqual(
+        str(
+            logging.LogEntry(
+                level='info', message='hello\nworld',
+                time=time, metadata=dict(x=1),
+            )
+        ),
+        '12:30:45 INFO - hello\nworld (metadata: {x=1})',
+    )
+    self.assertIn(
+        'LogEntry(',
+        logging.LogEntry(
+            level='info', message='hello\nworld',
+            time=time, metadata=dict(x=1),
+        ).format(text_format=False),
+    )
+
+  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 log-info"><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=dict(z=1), y=2)
+        ).to_html(
+            extra_flags=dict(
+                collapse_log_metadata_level=1,
+            ),
+            max_summary_len_for_str=10,
+            enable_summary_tooltip=False,
+        ),
+        """
+        <details open class="pyglove log-entry log-error"><summary><div class="summary-title log-error"><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<span class="tooltip">metadata</span></div><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><details class="pyglove dict"><summary><div class="summary-name">x<span class="tooltip">metadata.x</span></div><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><details open class="pyglove int"><summary><div class="summary-name">z<span class="tooltip">metadata.x.z</span></div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details></div></details><details open class="pyglove int"><summary><div class="summary-name">y<span class="tooltip">metadata.y</span></div><div class="summary-title">int</div></summary><span class="simple-value int">2</span></details></div></details></div></div></details>
+        """
+    )
+    self.assert_html_content(
+        logging.LogEntry(
+            level='error', message='This is a longer message: 5 + 2 > 3',
+            time=time, metadata=dict(x=dict(z=1), y=2)
+        ).to_html(
+            extra_flags=dict(
+                max_summary_len_for_str=10,
+            ),
+            enable_summary_tooltip=False,
+        ),
+        """
+        <details open class="pyglove log-entry log-error"><summary><div class="summary-title log-error"><span class="log-time">12:30:45</span><span class="log-summary">This is a longer message: 5 + 2 &gt; 3</span></div></summary><div class="complex_value"><div class="log-metadata"><details open class="pyglove dict"><summary><div class="summary-name">metadata<span class="tooltip">metadata</span></div><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><details open class="pyglove dict"><summary><div class="summary-name">x<span class="tooltip">metadata.x</span></div><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><details open class="pyglove int"><summary><div class="summary-name">z<span class="tooltip">metadata.x.z</span></div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details></div></details><details open class="pyglove int"><summary><div class="summary-name">y<span class="tooltip">metadata.y</span></div><div class="summary-title">int</div></summary><span class="simple-value int">2</span></details></div></details></div></div></details>
+        """
+    )
+
+if __name__ == '__main__':
+  unittest.main()