langfun 0.1.2.dev202501150804__py3-none-any.whl → 0.1.2.dev202501170804__py3-none-any.whl

langfun/core/__init__.py

@@ -75,7 +75,6 @@ from langfun.core.sampling import random_sample
 from langfun.core.concurrent import RetryEntry
 from langfun.core.concurrent import concurrent_execute
 from langfun.core.concurrent import concurrent_map
-from langfun.core.concurrent import with_context_access
 from langfun.core.concurrent import with_retry
 
 # Interface for natural language formattable.

langfun/core/component.py

@@ -13,10 +13,7 @@
 # limitations under the License.
 """langfun Component."""
 
-import contextlib
-import dataclasses
-import threading
-from typing import Annotated, Any, ContextManager, Iterator, Type
+from typing import ContextManager
 import pyglove as pg
 
 
@@ -24,22 +21,9 @@ import pyglove as pg
 RAISE_IF_HAS_ERROR = (pg.MISSING_VALUE,)
 
 
-class Component(pg.Object):
+class Component(pg.ContextualObject):
   """Base class for langfun components."""
 
-  # Override __repr__ format to use inferred values when available.
-  __repr_format_kwargs__ = dict(
-      compact=True,
-      use_inferred=True,
-  )
-
-  # Override __str__ format to use inferred values when available.
-  __str_format_kwargs__ = dict(
-      compact=False,
-      verbose=False,
-      use_inferred=True,
-  )
-
   # Allow symbolic assignment, which invalidates the object and recomputes
   # states upon update.
   allow_symbolic_assignment = True
@@ -75,122 +59,23 @@ class Component(pg.Object):
     if additional_fields:
       cls.update_schema(additional_fields)
 
-  def _on_bound(self):
-    super()._on_bound()
-    self._tls = threading.local()
-
-  def _sym_inferred(self, key: str, **kwargs):
-    """Override to allow attribute to access scoped value.
-
-    Args:
-      key: attribute name.
-      **kwargs: Optional keyword arguments for value inference.
-
-    Returns:
-      The value of the symbolic attribute. If not available, returns the
-        default value.
-
-    Raises:
-      AttributeError: If the attribute does not exist or contextual attribute
-        is not ready.
-    """
-    if key not in self._sym_attributes:
-      raise AttributeError(key)
-
-    # Step 1: Try use value from `self.override`.
-    # The reason is that `self.override` is short-lived and explicitly specified
-    # by the user in scenarios like `LangFunc.render`, which should not be
-    # affected by `lf.context`.
-    v = _get_scoped_value(self._tls, _CONTEXT_OVERRIDES, key)
-    if v is not None:
-      return v.value
-
-    # Step 2: Try use value from `lf.context` with `override_attrs`.
-    # This gives users a chance to override the bound attributes of components
-    # from the top, allowing change of bindings without modifying the code
-    # that produces the components.
-    override = get_contextual_override(key)
-    if override and override.override_attrs:
-      return override.value
-
-    # Step 3: Try use value from the symbolic tree, starting from self to
-    # the root of the tree.
-    # Step 4: If the value is not present, use the value from `context()` (
-    # override_attrs=False).
-    # Step 5: Otherwise use the default value from `ContextualAttribute`.
-    return super()._sym_inferred(key, context_override=override, **kwargs)
-
-  def override(
-      self, **kwargs) -> ContextManager[dict[str, 'ContextualOverride']]:
-    """Context manager to override the attributes of this component."""
-    vs = {k: ContextualOverride(v) for k, v in kwargs.items()}
-    return _contextual_scope(self._tls, _CONTEXT_OVERRIDES, **vs)
-
-  def __getattribute__(self, name: str) -> Any:
-    """Override __getattribute__ to deal with class attribute override."""
-    if not name.startswith('_') and hasattr(self.__class__, name):
-      tls = self.__dict__.get('_tls', None)
-      if tls is not None:
-        v = _get_scoped_value(tls, _CONTEXT_OVERRIDES, name)
-        if v is not None:
-          return v.value
-    return super().__getattribute__(name)
-
-
-_global_tls = threading.local()
-
-_CONTEXT_OVERRIDES = 'context_overrides'
-
 
-@dataclasses.dataclass(frozen=True)
-class ContextualOverride:
-  """Value marker for contextual override for an attribute."""
+# Aliases from PyGlove for ease of access.
+context = pg.contextual_override
+get_contextual_override = pg.utils.get_contextual_override
+context_value = pg.utils.contextual_value
+all_contextual_values = pg.utils.all_contextual_values
+contextual = pg.contextual_attribute
 
-  # Overridden value.
-  value: Any
-
-  # If True, this override will apply to both current scope and nested scope,
-  # meaning current `lf.context` will take precedence over all nested
-  # `lf.context` on this attribute.
-  cascade: bool = False
-
-  # If True, this override will apply to attributes that already have values.
-  override_attrs: bool = False
-
-
-def context(
-    *,
-    cascade: bool = False,
-    override_attrs: bool = False,
-    **variables,
-) -> ContextManager[dict[str, ContextualOverride]]:
-  """Context manager to provide overriden values for contextual attributes.
-
-  Args:
-    cascade: If True, this override will apply to both current scope and nested
-      scope, meaning that this `lf.context` will take precedence over all
-      nested `lf.context` on the overriden variables.
-    override_attrs: If True, this override will apply to attributes that already
-      have values. Otherwise overridden variables will only be used for
-      contextual attributes whose values are not present.
-    **variables: Key/values as override for contextual attributes.
-
-  Returns:
-    A dict of attribute names to their contextual overrides.
-  """
-  vs = {}
-  for k, v in variables.items():
-    if not isinstance(v, ContextualOverride):
-      v = ContextualOverride(v, cascade, override_attrs)
-    vs[k] = v
-  return _contextual_scope(_global_tls, _CONTEXT_OVERRIDES, **vs)
+# Decorator for setting the positional arguments for Component.
+use_init_args = pg.use_init_args
 
 
 def use_settings(
     *,
     cascade: bool = False,
     **settings,
-) -> ContextManager[dict[str, ContextualOverride]]:
+) -> ContextManager[dict[str, pg.utils.ContextualOverride]]:
   """Shortcut method for overriding component attributes.
 
   Args:
@@ -203,150 +88,3 @@ def use_settings(
     A dict of attribute names to their contextual overrides.
   """
   return context(cascade=cascade, override_attrs=True, **settings)
-
-
-def get_contextual_override(var_name: str) -> ContextualOverride | None:
-  """Returns the overriden contextual value in current scope."""
-  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, {})
-  return {k: v.value for k, v in overrides.items()}
-
-
-@contextlib.contextmanager
-def _contextual_scope(
-    tls: threading.local, tls_key, **variables
-) -> Iterator[dict[str, ContextualOverride]]:
-  """Context manager to set variables within a scope."""
-  previous_values = getattr(tls, tls_key, {})
-  current_values = dict(previous_values)
-  for k, v in variables.items():
-    old_v = current_values.get(k, None)
-    if old_v and old_v.cascade:
-      v = old_v
-    current_values[k] = v
-  try:
-    setattr(tls, tls_key, current_values)
-    yield current_values
-  finally:
-    setattr(tls, tls_key, previous_values)
-
-
-def _get_scoped_value(
-    tls: threading.local, tls_key: str, var_name: str, default: Any = None
-) -> ContextualOverride:
-  """Gets the value for requested variable from current scope."""
-  scoped_values = getattr(tls, tls_key, {})
-  return scoped_values.get(var_name, default)
-
-
-class ContextualAttribute(
-    pg.symbolic.ValueFromParentChain, pg.views.HtmlTreeView.Extension
-):
-  """Attributes whose values are inferred from the context of the component.
-
-  Please see go/langfun-component#attribute-value-retrieval for details.
-  """
-
-  NO_DEFAULT = (pg.MISSING_VALUE,)
-
-  type: Annotated[Type[Any] | None, 'An optional type constraint.'] = None
-
-  default: Any = NO_DEFAULT
-
-  def value_from(
-      self,
-      parent,
-      *,
-      context_override: ContextualOverride | None = None,
-      **kwargs,
-  ):
-    if parent not in (None, self.sym_parent) and isinstance(parent, Component):
-      # Apply original search logic along the component containing chain.
-      return super().value_from(parent, **kwargs)
-    elif parent is None:
-      # When there is no value inferred from the symbolic tree.
-      # Search context override, and then attribute-level default.
-      if context_override:
-        return context_override.value
-      if self.default == ContextualAttribute.NO_DEFAULT:
-        return pg.MISSING_VALUE
-      return self.default
-    else:
-      return pg.MISSING_VALUE
-
-  def _html_tree_view_content(
-      self,
-      *,
-      view: pg.views.HtmlTreeView,
-      parent: Any = None,
-      root_path: pg.KeyPath | None = None,
-      **kwargs,
-  ) -> pg.Html:
-    inferred_value = pg.MISSING_VALUE
-    if isinstance(parent, pg.Symbolic) and root_path:
-      inferred_value = parent.sym_inferred(root_path.key, pg.MISSING_VALUE)
-
-    if inferred_value is not pg.MISSING_VALUE:
-      kwargs.pop('name', None)
-      return view.render(
-          inferred_value, parent=self,
-          root_path=pg.KeyPath('<inferred>', root_path),
-          **view.get_passthrough_kwargs(**kwargs)
-      )
-    return pg.Html.element(
-        'div',
-        [
-            '(not available)',
-        ],
-        css_classes=['unavailable-contextual'],
-    )
-
-  def _html_tree_view_config(self) -> dict[str, Any]:
-    return pg.views.HtmlTreeView.get_kwargs(
-        super()._html_tree_view_config(),
-        dict(
-            collapse_level=1,
-        )
-    )
-
-  @classmethod
-  def _html_tree_view_css_styles(cls) -> list[str]:
-    return super()._html_tree_view_css_styles() + [
-        """
-        .contextual-attribute {
-          color: purple;
-        }
-        .unavailable-contextual {
-          color: gray;
-          font-style: italic;
-        }
-        """
-    ]
-
-
-# NOTE(daiyip): Returning Any instead of `lf.ContextualAttribute` to avoid
-# pytype check error as `contextual()` can be assigned to any type.
-def contextual(
-    type: Type[Any] | None = None,  # pylint: disable=redefined-builtin
-    default: Any = ContextualAttribute.NO_DEFAULT,
-) -> Any:
-  """Value marker for a contextual attribute."""
-  return ContextualAttribute(type=type, default=default, allow_partial=True)
-
-
-# Decorator for setting the positional arguments for Component.
-use_init_args = pg.use_init_args

langfun/core/component_test.py

@@ -73,26 +73,7 @@ class ComponentContextTest(unittest.TestCase):
     self.assertEqual(a1.y, 2)
     self.assertEqual(a1.z, -1)
 
-    with lf.context(x=3, y=3, z=3) as parent_override:
-      self.assertEqual(
-          parent_override,
-          dict(
-              x=lf.ContextualOverride(3, cascade=False, override_attrs=False),
-              y=lf.ContextualOverride(3, cascade=False, override_attrs=False),
-              z=lf.ContextualOverride(3, cascade=False, override_attrs=False),
-          ),
-      )
-      self.assertEqual(
-          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))
-
+    with lf.context(x=3, y=3, z=3):
       # Member attributes take precedence over `lf.context`.
       self.assertEqual(a1.x, 1)
       self.assertEqual(a1.y, 2)
@@ -109,15 +90,7 @@ class ComponentContextTest(unittest.TestCase):
       self.assertEqual(a1.z, 3)
 
       # Test nested contextual override with override_attrs=True (default).
-      with lf.context(y=4, z=4, override_attrs=True) as nested_override:
-        self.assertEqual(
-            nested_override,
-            dict(
-                x=lf.ContextualOverride(3, cascade=False, override_attrs=False),
-                y=lf.ContextualOverride(4, cascade=False, override_attrs=True),
-                z=lf.ContextualOverride(4, cascade=False, override_attrs=True),
-            ),
-        )
+      with lf.context(y=4, z=4, override_attrs=True):
 
         # Member attribute is not overriden as current scope does not override
         # `x``.

langfun/core/concurrent.py

@@ -25,7 +25,6 @@ import threading
 import time
 from typing import Annotated, Any, Callable, Iterable, Iterator, Literal, Sequence, Tuple, Type, Union
 
-from langfun.core import component
 import pyglove as pg
 
 
@@ -39,18 +38,6 @@ except ImportError:
   tqdm = None
 
 
-def with_context_access(func: Callable[..., Any]) -> Callable[..., Any]:
-  """Derives a user function with the access to the current context."""
-  with component.context() as current_context:
-    pass
-
-  def _func(*args, **kwargs) -> Any:
-    with component.context(**current_context):
-      return func(*args, **kwargs)
-
-  return _func
-
-
 class RetryError(RuntimeError):
   """Retry error."""
 
@@ -249,7 +236,8 @@ def concurrent_execute(
   try:
     executed_jobs = list(
         executor.map(
-            lambda job: job(), [with_context_access(job) for job in jobs]
+            lambda job: job(),
+            [pg.with_contextual_override(job) for job in jobs]
         )
     )
     for job in executed_jobs:
@@ -736,9 +724,7 @@ def concurrent_map(
         retry_interval=retry_interval,
         exponential_backoff=exponential_backoff,
     )
-    future = executor.submit(
-        with_context_access(job),
-    )
+    future = executor.submit(pg.with_contextual_override(job))
     pending_futures.append(future)
     future_to_job[future] = job
     total += 1

langfun/core/concurrent_test.py

@@ -29,22 +29,6 @@ class A(component.Component):
   y: int = component.contextual()
 
 
-class WithContextAccessTest(unittest.TestCase):
-
-  def test_context_access(self):
-    inputs = [A(1), A(2)]
-    with futures.ThreadPoolExecutor() as executor:
-      with component.context(y=3):
-        self.assertEqual(
-            list(
-                executor.map(
-                    concurrent.with_context_access(lambda x: x.y), inputs
-                )
-            ),
-            [3, 3],
-        )
-
-
 class RetryErrorTest(unittest.TestCase):
 
   def test_basics(self):

langfun/core/llms/__init__.py

@@ -30,6 +30,9 @@ from langfun.core.llms.compositional import RandomChoice
 # REST-based models.
 from langfun.core.llms.rest import REST
 
+# VertexAI-based models.
+from langfun.core.llms.vertexai import VertexAI
+
 # Gemini models.
 from langfun.core.llms.google_genai import GenAI
 from langfun.core.llms.google_genai import GeminiFlash2_0ThinkingExp_20241219
@@ -44,7 +47,7 @@ from langfun.core.llms.google_genai import GeminiFlash1_5_002
 from langfun.core.llms.google_genai import GeminiFlash1_5_001
 from langfun.core.llms.google_genai import GeminiPro1
 
-from langfun.core.llms.vertexai import VertexAI
+from langfun.core.llms.vertexai import VertexAIGemini
 from langfun.core.llms.vertexai import VertexAIGeminiFlash2_0ThinkingExp_20241219
 from langfun.core.llms.vertexai import VertexAIGeminiFlash2_0Exp
 from langfun.core.llms.vertexai import VertexAIGeminiExp_20241206
@@ -114,6 +117,8 @@ from langfun.core.llms.openai import Gpt3Curie
 from langfun.core.llms.openai import Gpt3Babbage
 from langfun.core.llms.openai import Gpt3Ada
 
+# Anthropic models.
+
 from langfun.core.llms.anthropic import Anthropic
 from langfun.core.llms.anthropic import Claude35Sonnet
 from langfun.core.llms.anthropic import Claude35Sonnet20241022
@@ -121,11 +126,14 @@ from langfun.core.llms.anthropic import Claude35Sonnet20240620
 from langfun.core.llms.anthropic import Claude3Opus
 from langfun.core.llms.anthropic import Claude3Sonnet
 from langfun.core.llms.anthropic import Claude3Haiku
-from langfun.core.llms.anthropic import VertexAIAnthropic
-from langfun.core.llms.anthropic import VertexAIClaude3_5_Sonnet_20241022
-from langfun.core.llms.anthropic import VertexAIClaude3_5_Sonnet_20240620
-from langfun.core.llms.anthropic import VertexAIClaude3_5_Haiku_20241022
-from langfun.core.llms.anthropic import VertexAIClaude3_Opus_20240229
+
+from langfun.core.llms.vertexai import VertexAIAnthropic
+from langfun.core.llms.vertexai import VertexAIClaude3_5_Sonnet_20241022
+from langfun.core.llms.vertexai import VertexAIClaude3_5_Sonnet_20240620
+from langfun.core.llms.vertexai import VertexAIClaude3_5_Haiku_20241022
+from langfun.core.llms.vertexai import VertexAIClaude3_Opus_20240229
+
+# Misc open source models.
 
 from langfun.core.llms.groq import Groq
 from langfun.core.llms.groq import GroqLlama3_2_3B

langfun/core/llms/anthropic.py

@@ -14,9 +14,8 @@
 """Language models from Anthropic."""
 
 import base64
-import functools
 import os
-from typing import Annotated, Any, Literal
+from typing import Annotated, Any
 
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
@@ -24,20 +23,6 @@ from langfun.core.llms import rest
 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
-  Credentials = credentials_lib.Credentials
-  # pylint: enable=g-import-not-at-top
-except ImportError:
-  google_auth = None
-  auth_requests = None
-  credentials_lib = None
-  Credentials = Any  # pylint: disable=invalid-name
-
-
 SUPPORTED_MODELS_AND_SETTINGS = {
     # See https://docs.anthropic.com/claude/docs/models-overview
     # Rate limits from https://docs.anthropic.com/claude/reference/rate-limits
@@ -379,110 +364,3 @@ class Claude21(Anthropic):
 class ClaudeInstant(Anthropic):
   """Cheapest small and fast model, 100K context window."""
   model = 'claude-instant-1.2'
-
-
-#
-# Authropic models on VertexAI.
-#
-
-
-class VertexAIAnthropic(Anthropic):
-  """Anthropic models on VertexAI."""
-
-  project: Annotated[
-      str | None,
-      'Google Cloud project ID.',
-  ] = None
-
-  location: Annotated[
-      Literal['us-east5', 'europe-west1'],
-      'GCP location with Anthropic models hosted.'
-  ] = 'us-east5'
-
-  credentials: Annotated[
-      Credentials | None,    # pytype: disable=invalid-annotation
-      (
-          'Credentials to use. If None, the default credentials '
-          'to the environment will be used.'
-      ),
-  ] = None
-
-  api_version = 'vertex-2023-10-16'
-
-  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.'
-      )
-    self._project = project
-    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
-
-  @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 headers(self):
-    return {
-        'Content-Type': 'application/json; charset=utf-8',
-    }
-
-  @property
-  def api_endpoint(self) -> str:
-    return (
-        f'https://{self.location}-aiplatform.googleapis.com/v1/projects/'
-        f'{self._project}/locations/{self.location}/publishers/anthropic/'
-        f'models/{self.model}:streamRawPredict'
-    )
-
-  def request(
-      self,
-      prompt: lf.Message,
-      sampling_options: lf.LMSamplingOptions
-  ):
-    request = super().request(prompt, sampling_options)
-    request['anthropic_version'] = self.api_version
-    del request['model']
-    return request
-
-
-class VertexAIClaude3_Opus_20240229(VertexAIAnthropic):  # pylint: disable=invalid-name
-  """Anthropic's Claude 3 Opus model on VertexAI."""
-  model = 'claude-3-opus@20240229'
-
-
-class VertexAIClaude3_5_Sonnet_20241022(VertexAIAnthropic):  # pylint: disable=invalid-name
-  """Anthropic's Claude 3.5 Sonnet model on VertexAI."""
-  model = 'claude-3-5-sonnet-v2@20241022'
-
-
-class VertexAIClaude3_5_Sonnet_20240620(VertexAIAnthropic):  # pylint: disable=invalid-name
-  """Anthropic's Claude 3.5 Sonnet model on VertexAI."""
-  model = 'claude-3-5-sonnet@20240620'
-
-
-class VertexAIClaude3_5_Haiku_20241022(VertexAIAnthropic):  # pylint: disable=invalid-name
-  """Anthropic's Claude 3.5 Haiku model on VertexAI."""
-  model = 'claude-3-5-haiku@20241022'

langfun/core/llms/anthropic_test.py

@@ -19,9 +19,6 @@ from typing import Any
 import unittest
 from unittest import mock
 
-from google.auth import exceptions
-from langfun.core import language_model
-from langfun.core import message as lf_message
 from langfun.core import modalities as lf_modalities
 from langfun.core.llms import anthropic
 import pyglove as pg
@@ -186,50 +183,5 @@ class AnthropicTest(unittest.TestCase):
           lm('hello', max_attempts=1)
 
 
-class VertexAIAnthropicTest(unittest.TestCase):
-  """Tests for VertexAI Anthropic models."""
-
-  def test_basics(self):
-    with self.assertRaisesRegex(ValueError, 'Please specify `project`'):
-      lm = anthropic.VertexAIClaude3_5_Sonnet_20241022()
-      lm('hi')
-
-    model = anthropic.VertexAIClaude3_5_Sonnet_20241022(project='langfun')
-
-    # NOTE(daiyip): For OSS users, default credentials are not available unless
-    # users have already set up their GCP project. Therefore we ignore the
-    # exception here.
-    try:
-      model._initialize()
-    except exceptions.DefaultCredentialsError:
-      pass
-
-    self.assertEqual(
-        model.api_endpoint,
-        (
-            'https://us-east5-aiplatform.googleapis.com/v1/projects/'
-            'langfun/locations/us-east5/publishers/anthropic/'
-            'models/claude-3-5-sonnet-v2@20241022:streamRawPredict'
-        )
-    )
-    request = model.request(
-        lf_message.UserMessage('hi'),
-        language_model.LMSamplingOptions(temperature=0.0),
-    )
-    self.assertEqual(
-        request,
-        {
-            'anthropic_version': 'vertex-2023-10-16',
-            'max_tokens': 8192,
-            'messages': [
-                {'content': [{'text': 'hi', 'type': 'text'}], 'role': 'user'}
-            ],
-            'stream': False,
-            'temperature': 0.0,
-            'top_k': 40,
-        },
-    )
-
-
 if __name__ == '__main__':
   unittest.main()

langfun/core/llms/vertexai.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The Langfun Authors
+# Copyright 2025 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.
@@ -15,10 +15,12 @@
 
 import functools
 import os
-from typing import Annotated, Any
+from typing import Annotated, Any, Literal
 
 import langfun.core as lf
+from langfun.core.llms import anthropic
 from langfun.core.llms import gemini
+from langfun.core.llms import rest
 import pyglove as pg
 
 try:
@@ -36,10 +38,21 @@ except ImportError:
   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."""
+@pg.use_init_args(['api_endpoint'])
+class VertexAI(rest.REST):
+  """Base class for VertexAI models.
+
+  This class handles the authentication of vertex AI models. Subclasses
+  should implement `request` and `result` methods, as well as the `api_endpoint`
+  property. Or let users to provide them as __init__ arguments.
+
+  Please check out VertexAIGemini in `gemini.py` as an example.
+  """
+
+  model: Annotated[
+      str | None,
+      'Model ID.'
+  ] = None
 
   project: Annotated[
       str | None,
@@ -114,6 +127,17 @@ class VertexAI(gemini.Gemini):
     s.headers.update(self.headers or {})
     return s
 
+
+#
+# Gemini models served by Vertex AI.
+#
+
+
+@pg.use_init_args(['model'])
+@pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
+class VertexAIGemini(VertexAI, gemini.Gemini):
+  """Gemini models served by Vertex AI.."""
+
   @property
   def api_endpoint(self) -> str:
     assert self._api_initialized
@@ -124,7 +148,7 @@ class VertexAI(gemini.Gemini):
     )
 
 
-class VertexAIGeminiFlash2_0ThinkingExp_20241219(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash2_0ThinkingExp_20241219(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini Flash 2.0 Thinking model launched on 12/19/2024."""
 
   api_version = 'v1alpha'
@@ -132,61 +156,128 @@ class VertexAIGeminiFlash2_0ThinkingExp_20241219(VertexAI):  # pylint: disable=i
   timeout = None
 
 
-class VertexAIGeminiFlash2_0Exp(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash2_0Exp(VertexAIGemini):  # 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
+class VertexAIGeminiExp_20241206(VertexAIGemini):  # 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
+class VertexAIGeminiExp_20241114(VertexAIGemini):  # 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
+class VertexAIGeminiPro1_5(VertexAIGemini):  # 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
+class VertexAIGeminiPro1_5_002(VertexAIGemini):  # 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
+class VertexAIGeminiPro1_5_001(VertexAIGemini):  # 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
+class VertexAIGeminiFlash1_5(VertexAIGemini):  # 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
+class VertexAIGeminiFlash1_5_002(VertexAIGemini):  # 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
+class VertexAIGeminiFlash1_5_001(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-001'
 
 
-class VertexAIGeminiPro1(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.0 Pro model."""
 
   model = 'gemini-1.0-pro'
+
+
+#
+# Anthropic models on Vertex AI.
+#
+
+
+@pg.use_init_args(['model'])
+@pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
+class VertexAIAnthropic(VertexAI, anthropic.Anthropic):
+  """Anthropic models on VertexAI."""
+
+  location: Annotated[
+      Literal['us-east5', 'europe-west1'],
+      'GCP location with Anthropic models hosted.'
+  ] = 'us-east5'
+
+  api_version = 'vertex-2023-10-16'
+
+  @property
+  def headers(self):
+    return {
+        'Content-Type': 'application/json; charset=utf-8',
+    }
+
+  @property
+  def api_endpoint(self) -> str:
+    return (
+        f'https://{self.location}-aiplatform.googleapis.com/v1/projects/'
+        f'{self._project}/locations/{self.location}/publishers/anthropic/'
+        f'models/{self.model}:streamRawPredict'
+    )
+
+  def request(
+      self,
+      prompt: lf.Message,
+      sampling_options: lf.LMSamplingOptions
+  ):
+    request = super().request(prompt, sampling_options)
+    request['anthropic_version'] = self.api_version
+    del request['model']
+    return request
+
+
+# pylint: disable=invalid-name
+
+
+class VertexAIClaude3_Opus_20240229(VertexAIAnthropic):
+  """Anthropic's Claude 3 Opus model on VertexAI."""
+  model = 'claude-3-opus@20240229'
+
+
+class VertexAIClaude3_5_Sonnet_20241022(VertexAIAnthropic):
+  """Anthropic's Claude 3.5 Sonnet model on VertexAI."""
+  model = 'claude-3-5-sonnet-v2@20241022'
+
+
+class VertexAIClaude3_5_Sonnet_20240620(VertexAIAnthropic):
+  """Anthropic's Claude 3.5 Sonnet model on VertexAI."""
+  model = 'claude-3-5-sonnet@20240620'
+
+
+class VertexAIClaude3_5_Haiku_20241022(VertexAIAnthropic):
+  """Anthropic's Claude 3.5 Haiku model on VertexAI."""
+  model = 'claude-3-5-haiku@20241022'
+
+# pylint: enable=invalid-name

langfun/core/llms/vertexai_test.py

@@ -17,6 +17,8 @@ import os
 import unittest
 from unittest import mock
 
+from google.auth import exceptions
+import langfun.core as lf
 from langfun.core.llms import vertexai
 
 
@@ -48,5 +50,55 @@ class VertexAITest(unittest.TestCase):
     del os.environ['VERTEXAI_LOCATION']
 
 
+class VertexAIAnthropicTest(unittest.TestCase):
+  """Tests for VertexAI Anthropic models."""
+
+  def test_basics(self):
+    with self.assertRaisesRegex(ValueError, 'Please specify `project`'):
+      lm = vertexai.VertexAIClaude3_5_Sonnet_20241022()
+      lm('hi')
+
+    model = vertexai.VertexAIClaude3_5_Sonnet_20241022(project='langfun')
+
+    # NOTE(daiyip): For OSS users, default credentials are not available unless
+    # users have already set up their GCP project. Therefore we ignore the
+    # exception here.
+    try:
+      model._initialize()
+    except exceptions.DefaultCredentialsError:
+      pass
+
+    self.assertEqual(
+        model.api_endpoint,
+        (
+            'https://us-east5-aiplatform.googleapis.com/v1/projects/'
+            'langfun/locations/us-east5/publishers/anthropic/'
+            'models/claude-3-5-sonnet-v2@20241022:streamRawPredict'
+        )
+    )
+    self.assertEqual(
+        model.headers,
+        {
+            'Content-Type': 'application/json; charset=utf-8',
+        },
+    )
+    request = model.request(
+        lf.UserMessage('hi'), lf.LMSamplingOptions(temperature=0.0),
+    )
+    self.assertEqual(
+        request,
+        {
+            'anthropic_version': 'vertex-2023-10-16',
+            'max_tokens': 8192,
+            'messages': [
+                {'content': [{'text': 'hi', 'type': 'text'}], 'role': 'user'}
+            ],
+            'stream': False,
+            'temperature': 0.0,
+            'top_k': 40,
+        },
+    )
+
+
 if __name__ == '__main__':
   unittest.main()

{langfun-0.1.2.dev202501150804.dist-info → langfun-0.1.2.dev202501170804.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: langfun
-Version: 0.1.2.dev202501150804
+Version: 0.1.2.dev202501170804
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.2.dev202501150804.dist-info → langfun-0.1.2.dev202501170804.dist-info}/RECORD

@@ -1,9 +1,9 @@
 langfun/__init__.py,sha256=fhfPXpHN7GoGqixpFfqhQkYxFs_siP_LhbjZhd3lhio,2497
-langfun/core/__init__.py,sha256=oo81OUA4Xy2q7icxcxAMicVzbla5gMDf9Nzw_iIAkis,4627
-langfun/core/component.py,sha256=HVrEoTL1Y01iqOHC3FYdbAOnffqfHHtGJXoK1vkdEwo,11583
-langfun/core/component_test.py,sha256=sG-T2wpvBfHqWGZE7sc4NayJj2aj5QFBzSwFiwrGEIc,10376
-langfun/core/concurrent.py,sha256=4CqOlH9YrtrK-Jijp09hLX7AYxRYL8slpy-tL2aqnIM,33121
-langfun/core/concurrent_test.py,sha256=0ucEYxPgTErOOK1TIN-TSqZphHS1Vl1VssyQP49cpOI,17997
+langfun/core/__init__.py,sha256=Cy3BU8R9VHnJmbxiN9QXXBb1K09ZoESYIJGf_uwJrDs,4571
+langfun/core/component.py,sha256=g1kQM0bryYYYWVDrSMnHfc74wIBbpfe5_B3s-UIP5GE,3028
+langfun/core/component_test.py,sha256=0CxTgjAud3aj8wBauFhG2FHDqrxCTl4OI4gzQTad-40,9254
+langfun/core/concurrent.py,sha256=zY-pXqlGqss_GI20tM1gXvyW8QepVPUuFNmutcIdhbI,32760
+langfun/core/concurrent_test.py,sha256=rc5T-2giWgtbwNuN6gmei7Uwo66HsJeeRtXZCpya_QU,17590
 langfun/core/console.py,sha256=V_mOiFi9oGh8gLsUeR56pdFDkuvYOpvQt7DY1KUTWTA,2535
 langfun/core/console_test.py,sha256=pBOcuNMJdVELywvroptfcRtJMsegMm3wSlHAL2TdxVk,1679
 langfun/core/langfunc.py,sha256=G50YgoVZ0y1GFw2ev41MlOqr6qa8YakbvNC0h_E0PiA,11140
@@ -71,9 +71,9 @@ langfun/core/eval/v2/reporting.py,sha256=QOp5jX761Esvi5w_UIRLDqPY_XRO6ru02-DOrdq
 langfun/core/eval/v2/reporting_test.py,sha256=UmYSAQvD3AIXsSyWQ-WD2uLtEISYpmBeoKY5u5Qwc8E,5696
 langfun/core/eval/v2/runners.py,sha256=DKEmSlGXjOXKWFdBhTpLy7tMsBHZHd1Brl3hWIngsSQ,15931
 langfun/core/eval/v2/runners_test.py,sha256=A37fKK2MvAVTiShsg_laluJzJ9AuAQn52k7HPbfD0Ks,11666
-langfun/core/llms/__init__.py,sha256=Ntr0kvHc17VEZ5EV9fCoYY1kzRvQxCoZrtDRYNiMWCs,6742
-langfun/core/llms/anthropic.py,sha256=a5MmnFsBA0CbfvwzXT1v_0fqLRMrhUNdh1tx6469PQ4,14357
-langfun/core/llms/anthropic_test.py,sha256=-2U4kc_pgBM7wqxu8RuxzyHPGww1EAWqKUvN4PW8Btw,8058
+langfun/core/llms/__init__.py,sha256=wA6t_E3peTYTjsW6uOHnOs9wjQ_Tj1WYlhVVLk2Sjcg,6867
+langfun/core/llms/anthropic.py,sha256=z_DWDpR1VKNzv6wq-9CXLzWdqCDXRKuVFacJNpgBqAs,10826
+langfun/core/llms/anthropic_test.py,sha256=zZ2eSP8hhVv-RDSWxT7wX-NS5DfGfQmCjS9P0pusAHM,6556
 langfun/core/llms/compositional.py,sha256=csW_FLlgL-tpeyCOTVvfUQkMa_zCN5Y2I-YbSNuK27U,2872
 langfun/core/llms/compositional_test.py,sha256=4eTnOer-DncRKGaIJW2ZQQMLnt5r2R0UIx_DYOvGAQo,2027
 langfun/core/llms/deepseek.py,sha256=Y7DlLUWrukbPVyBMesppd-m75Q-PxD0b3KnMKaoY_8I,3744
@@ -94,8 +94,8 @@ langfun/core/llms/openai_compatible_test.py,sha256=0uFYhCiuHo2Wrlgj16-GRG6rW8P6E
 langfun/core/llms/openai_test.py,sha256=m85YjGCvWvV5ZYagjC0FqI0FcqyCEVCbUUs8Wm3iUrc,2475
 langfun/core/llms/rest.py,sha256=sWbYUV8S3SuOg9giq7xwD-xDRfaF7NP_ig7bI52-Rj4,3442
 langfun/core/llms/rest_test.py,sha256=zWGiI08f9gXsoQPJS9TlX1zD2uQLrJUB-1VpAJXRHfs,3475
-langfun/core/llms/vertexai.py,sha256=MuwLPTJ6-9x2uRDCSM1_biPK6M76FFlL1ezf5OmobDA,5504
-langfun/core/llms/vertexai_test.py,sha256=iXjmQs7TNiwcueoaRGpdp4KnASkDJaTP__Z9QroN8zQ,1787
+langfun/core/llms/vertexai.py,sha256=JV9iHsCM3Ee-4nE1ENNkTXIYGxjCHxrEeir175YpCM8,7869
+langfun/core/llms/vertexai_test.py,sha256=6eLQOyeL5iGZOIWb39sFcf1TgYD_6TBGYdMO4UIvhf4,3333
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=rt3zwmyw0y9jsSGW-ZbV1vAfLxQ7_3AVk0l2EySlse4,3918
 langfun/core/llms/cache/in_memory.py,sha256=i58oiQL28RDsq37dwqgVpC2mBETJjIEFS20yHiV5MKU,5185
@@ -146,8 +146,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=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
-langfun-0.1.2.dev202501150804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202501150804.dist-info/METADATA,sha256=HAqG3RU6S4kwv8Knx7vDVp01x8Kuc7It1QnYg0oeol8,8172
-langfun-0.1.2.dev202501150804.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-langfun-0.1.2.dev202501150804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202501150804.dist-info/RECORD,,
+langfun-0.1.2.dev202501170804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202501170804.dist-info/METADATA,sha256=X3MDNl6D6StuwltvUclYhE20uKNQ2x8lY3CkPggJyI4,8172
+langfun-0.1.2.dev202501170804.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+langfun-0.1.2.dev202501170804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202501170804.dist-info/RECORD,,