langfun 0.1.2.dev202411060804__py3-none-any.whl → 0.1.2.dev202411090804__py3-none-any.whl

langfun/core/eval/scoring.py

@@ -85,12 +85,12 @@ class Scoring(base.Evaluation):
     del progress
     return {
         'Average Score': {self.avg_score},
-        'Scored': '%.2f%% (%d/%d)' % (
+        'Scored': '%.3f%% (%d/%d)' % (
             self.score_rate * 100,
             self.num_scored,
             self.num_completed,
         ),
-        'Failed': '%.2f%% (%d/%d)' % (
+        'Failed': '%.3f%% (%d/%d)' % (
             self.failure_rate * 100,
             self.num_failures,
             self.num_completed,
@@ -101,8 +101,8 @@ class Scoring(base.Evaluation):
     assert self.result is not None
     m = self.result.metrics
     return (
-        'COMPLETED(%s): AvgScore=%f Scored=%.2f%% (%d/%d) '
-        'Failures=%.2f%% (%d/%d)'
+        'COMPLETED(%s): AvgScore=%f Scored=%.3f%% (%d/%d) '
+        'Failures=%.3f%% (%d/%d)'
     ) % (
         run_status,
         m.avg_score,
@@ -155,12 +155,12 @@ class Scoring(base.Evaluation):
   def _render_result_row(self, s: io.StringIO):
     super()._render_result_row(s)
     s.write(
-        '<td><span style="color:blue">%.2f</span></td>' % self.avg_score
+        '<td><span style="color:blue">%.3f</span></td>' % self.avg_score
     )
     s.write(
         '<td><span style="color:red">%s</span>%s</td>'
         % (
-            '%.2f%% ' % (self.score_rate * 100),
+            '%.3f%% ' % (self.score_rate * 100),
             '<a href="%s">(%d/%d)</a>'
             % (self.scored_link, self.num_scored, self.num_completed),
         )
@@ -173,7 +173,7 @@ class Scoring(base.Evaluation):
     self._render_link(
         s,
         'Average score (%d/%d)' % (m.num_scored, m.total),
-        '%.2f (%.2f%%)' % (m.avg_score, m.score_rate * 100),
+        '%.3f (%.3f%%)' % (m.avg_score, m.score_rate * 100),
         'color:green',
         lambda: self.scored_link,
     )

langfun/core/llms/__init__.py

@@ -95,6 +95,9 @@ 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_Haiku_20241022
 
 from langfun.core.llms.groq import Groq
 from langfun.core.llms.groq import GroqLlama3_2_3B

langfun/core/llms/anthropic.py

@@ -14,8 +14,9 @@
 """Language models from Anthropic."""
 
 import base64
+import functools
 import os
-from typing import Annotated, Any
+from typing import Annotated, Any, Literal
 
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
@@ -23,6 +24,20 @@ 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 = None  # 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
@@ -30,20 +45,43 @@ SUPPORTED_MODELS_AND_SETTINGS = {
     #     as RPM/TPM of the largest-available model (Claude-3-Opus).
     # Price in US dollars at https://www.anthropic.com/pricing
     # as of 2024-10-10.
+    # Anthropic models hosted on VertexAI.
+    'claude-3-5-sonnet-v2@20241022': pg.Dict(
+        max_tokens=8192,
+        rpm=1000,
+        tpm=100000,
+        cost_per_1k_input_tokens=0.003,
+        cost_per_1k_output_tokens=0.015,
+    ),
+    'claude-3-5-haiku@20241022': pg.Dict(
+        max_tokens=8192,
+        rpm=1000,
+        tpm=100000,
+        cost_per_1k_input_tokens=0.001,
+        cost_per_1k_output_tokens=0.005,
+    ),
+    # Anthropic hosted models.
     'claude-3-5-sonnet-20241022': pg.Dict(
-        max_tokens=4096,
+        max_tokens=8192,
         rpm=4000,
         tpm=400000,
         cost_per_1k_input_tokens=0.003,
         cost_per_1k_output_tokens=0.015,
     ),
     'claude-3-5-sonnet-20240620': pg.Dict(
-        max_tokens=4096,
+        max_tokens=8192,
         rpm=4000,
         tpm=400000,
         cost_per_1k_input_tokens=0.003,
         cost_per_1k_output_tokens=0.015,
     ),
+    'claude-3-5-haiku-20241022': pg.Dict(
+        max_tokens=8192,
+        rpm=4000,
+        tpm=400000,
+        cost_per_1k_input_tokens=0.001,
+        cost_per_1k_output_tokens=0.005,
+    ),
     'claude-3-opus-20240229': pg.Dict(
         max_tokens=4096,
         rpm=4000,
@@ -327,3 +365,100 @@ 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_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_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

@@ -18,6 +18,10 @@ import os
 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
@@ -182,5 +186,50 @@ 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/rest.py

@@ -98,6 +98,7 @@ class REST(lf.LanguageModel):
         500,  # Server side issue (might be bug).
         502,  # Bad gateway (upstream issue, might retry).
         503,  # Servers currently under load, retry after a brief wait.
+        529,  # Overloaded, retry after a brief wait.
     ):
       error_cls = lf.TemporaryLMError
     else:

langfun/core/templates/conversation.py

@@ -38,6 +38,11 @@ class Conversation(Completion):
       '(Optional) Preamble before beginning the conversation.',
   ] = None
 
+  role: Annotated[
+      str | None,
+      '(Optional) User defined role for the AI response in the conversation.',
+  ] = None
+
   conversation_context: Annotated[
       lf.LangFunc | None,
       (
@@ -71,6 +76,10 @@ class Conversation(Completion):
     with lf.context(**kwargs):
       # Call LM based on the prompt generated from `input_message`.
       lm_response = super().__call__()
+      if self.role is not None:
+        lm_response.rebind(
+            sender=self.role, skip_notification=True, raise_on_no_change=False
+        )
 
       # Add current turn to memory.
       self.add(self.input_message, lm_response)

langfun/core/templates/conversation_test.py

@@ -83,6 +83,7 @@ class ConversationTest(unittest.TestCase):
   def test_call(self):
     c = Conversation(
         lm=QuestionCounter(),
+        role='Agent',
         preamble="You are a helpful and joyful AI bot. Now let's chat.",
     )
     # First round.
@@ -102,7 +103,7 @@ class ConversationTest(unittest.TestCase):
         inspect.cleandoc("""
         You are a helpful and joyful AI bot. Now let's chat.
         User: Hello
-        AI: Response 1.
+        Agent: Response 1.
         User: How are you?
         """),
     )
@@ -114,9 +115,9 @@ class ConversationTest(unittest.TestCase):
         inspect.cleandoc("""
         You are a helpful and joyful AI bot. Now let's chat.
         User: Hello
-        AI: Response 1.
+        Agent: Response 1.
         User: How are you?
-        AI: Response 2.
+        Agent: Response 2.
         User: Okay, bye.
         """),
     )

{langfun-0.1.2.dev202411060804.dist-info → langfun-0.1.2.dev202411090804.dist-info}/METADATA

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

{langfun-0.1.2.dev202411060804.dist-info → langfun-0.1.2.dev202411090804.dist-info}/RECORD

@@ -50,11 +50,11 @@ langfun/core/eval/matching.py,sha256=UnjdM_ebPqXKJamY4lvL3AYxrMIz3LqkjRTnHJ5xsYc
 langfun/core/eval/matching_test.py,sha256=QCoYEuf4b_1bkHqUCuRzKMbXHrV3AB2FCOBivo1stC4,5249
 langfun/core/eval/patching.py,sha256=R0s2eAd1m97exQt06dmUL0V_MBG0W2Hxg7fhNB7cXW0,3866
 langfun/core/eval/patching_test.py,sha256=8kCd54Egjju22FMgtJuxEsrXkW8ifs-UUBHtrCG1L6w,4775
-langfun/core/eval/scoring.py,sha256=SUdMzOkP0n2qGaSuUA4VwFiTw36jgMvgCJHPJS4yYDw,6254
+langfun/core/eval/scoring.py,sha256=B69IsIxiPs1xZcOBFIhZF70YmDue2Siik-CPL2bh33s,6254
 langfun/core/eval/scoring_test.py,sha256=O8olHbrUEg60gMxwOkWzKBJZpZoUlmVnBANX5Se2SXM,4546
-langfun/core/llms/__init__.py,sha256=SFa3qKHSq_P_bV_SkIEVasm58J6wvwLVeoh0rhhDy9o,5985
-langfun/core/llms/anthropic.py,sha256=Wwpe7WZT6NhHELVuQxowVYWCwfZJCE-igR10cyJUCLc,9785
-langfun/core/llms/anthropic_test.py,sha256=9l0WuM0As7X78lbbnpuvguMBMVSHX2QiAZuwjbj0HbA,6555
+langfun/core/llms/__init__.py,sha256=ATayGMO5rQzkPBAf7zl3vfgRQ8fth33TFVXv_tvraCQ,6190
+langfun/core/llms/anthropic.py,sha256=XPQxjfe9O4b-CygCgqvQU0MPSfe1rU7uErNbo8zth7Q,13606
+langfun/core/llms/anthropic_test.py,sha256=-2U4kc_pgBM7wqxu8RuxzyHPGww1EAWqKUvN4PW8Btw,8058
 langfun/core/llms/fake.py,sha256=gCHBYBLvBCsC78HI1hpoqXCS-p1FMTgY1P1qh_sGBPk,3070
 langfun/core/llms/fake_test.py,sha256=2h13qkwEz_JR0mtUDPxdAhQo7MueXaFSwsD2DIRDW9g,7653
 langfun/core/llms/google_genai.py,sha256=btUIfWteBoj8Jl0j8d3e8hyI6p3Biq4rldlQYctVQfg,10936
@@ -65,7 +65,7 @@ langfun/core/llms/llama_cpp.py,sha256=9tXQntSCDtjTF3bnyJrAPCr4N6wycy5nXYvp9uduyg
 langfun/core/llms/llama_cpp_test.py,sha256=MWO_qaOeKjRniGjcaWPDScd7HPaIJemqUZoslrt4FPs,1806
 langfun/core/llms/openai.py,sha256=qrAiJxE0tS7ZqjaVzRgVJKtMtoo1Z5TYpvi5ikTwPpw,22716
 langfun/core/llms/openai_test.py,sha256=_8cd3VRNEUfE0-Ko1RiM6MlC5hjalRj7nYTJNhG1p3E,18907
-langfun/core/llms/rest.py,sha256=laopuq-zD8V-3Y6eFDngftHEbE66VlUkCD2-rvvRaLU,3388
+langfun/core/llms/rest.py,sha256=sWbYUV8S3SuOg9giq7xwD-xDRfaF7NP_ig7bI52-Rj4,3442
 langfun/core/llms/rest_test.py,sha256=NZ3Nf0XQVpT9kLP5cBVo_yBHLI7vWTYhWQxYEJVMGs4,3472
 langfun/core/llms/vertexai.py,sha256=eUFU0JjgpTVCAvQVEWiGqZrlE3Dye-EkzZCF1P8nwc4,18866
 langfun/core/llms/vertexai_test.py,sha256=7uBVOF5VF86xQ9HFAbSTh4J-0NjYLnuotBS1YRm-vgw,10529
@@ -113,14 +113,14 @@ langfun/core/structured/tokenization_test.py,sha256=dVW30kGYkX2HNtiRZe1oTmXFP7iI
 langfun/core/templates/__init__.py,sha256=bO0eMsVJbi7sxEB2YlInKRQ2EVP-RyyKUwcD-8msuN4,927
 langfun/core/templates/completion.py,sha256=mUqZHOEV3ag6-A08XghpeEltcrBvCDxXP004eDDfeag,1931
 langfun/core/templates/completion_test.py,sha256=vGnjnM38UHyVDUyaUYtmp20s9KBGOdbPVsX-H-ET11E,1636
-langfun/core/templates/conversation.py,sha256=HVih0atzCFpYbE-2bmO1VewhKGyKYAoCty-IZYCjf00,2589
-langfun/core/templates/conversation_test.py,sha256=RryYyIhfc34dLWOs6GfPQ8HU8mXpKGL87UgVkysfIMo,3911
+langfun/core/templates/conversation.py,sha256=dTy_fs3g92d3vP09bACGJ_zJ3b8-OzA8vsKgUDoB76o,2866
+langfun/core/templates/conversation_test.py,sha256=HCD5f1sgmSGqL3OGe7hEctjRdcFjcmkz3NV1ArexNYI,3942
 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=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
-langfun-0.1.2.dev202411060804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202411060804.dist-info/METADATA,sha256=7pmVZg70qMsD08wH-RC6Fi7ljIKiMOtVEHyUBslyo3o,8890
-langfun-0.1.2.dev202411060804.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-langfun-0.1.2.dev202411060804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202411060804.dist-info/RECORD,,
+langfun-0.1.2.dev202411090804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202411090804.dist-info/METADATA,sha256=Ez5vUJUc5-XDAdDMyR6Wcgj4fiDL1RwJpu1NI3i3lwE,8890
+langfun-0.1.2.dev202411090804.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+langfun-0.1.2.dev202411090804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202411090804.dist-info/RECORD,,