langfun 0.1.2.dev202501080804__py3-none-any.whl → 0.1.2.dev202501240804__py3-none-any.whl

langfun/core/language_model.py

@@ -81,6 +81,61 @@ class LMSample(pg.Object):
   ] = None
 
 
+class RetryStats(pg.Object):
+  """Retry stats, which is aggregated across multiple retry entries."""
+
+  num_occurences: Annotated[
+      int,
+      'Total number of retry attempts on LLM (excluding the first attempt).',
+  ] = 0
+  total_wait_interval: Annotated[
+      float, 'Total wait interval in seconds due to retry.'
+  ] = 0
+  total_call_interval: Annotated[
+      float, 'Total LLM call interval in seconds.'
+  ] = 0
+  errors: Annotated[
+      dict[str, int],
+      'A Counter of error types encountered during the retry attempts.',
+  ] = {}
+
+  @classmethod
+  def from_retry_entries(
+      cls, retry_entries: Sequence[concurrent.RetryEntry]
+  ) -> 'RetryStats':
+    """Creates a RetryStats from a sequence of RetryEntry."""
+    if not retry_entries:
+      return RetryStats()
+    errors = {}
+    for retry in retry_entries:
+      if retry.error is not None:
+        errors[retry.error.__class__.__name__] = (
+            errors.get(retry.error.__class__.__name__, 0) + 1
+        )
+    return RetryStats(
+        num_occurences=len(retry_entries) - 1,
+        total_wait_interval=sum(e.wait_interval for e in retry_entries),
+        total_call_interval=sum(e.call_interval for e in retry_entries),
+        errors=errors,
+    )
+
+  def __add__(self, other: 'RetryStats') -> 'RetryStats':
+    errors = self.errors.copy()
+    for error, count in other.errors.items():
+      errors[error] = errors.get(error, 0) + count
+    return RetryStats(
+        num_occurences=self.num_occurences + other.num_occurences,
+        total_wait_interval=self.total_wait_interval
+        + other.total_wait_interval,
+        total_call_interval=self.total_call_interval
+        + other.total_call_interval,
+        errors=errors,
+    )
+
+  def __radd__(self, other: 'RetryStats') -> 'RetryStats':
+    return self + other
+
+
 class LMSamplingUsage(pg.Object):
   """Usage information per completion."""
 
@@ -93,8 +148,9 @@ class LMSamplingUsage(pg.Object):
       (
           'Estimated cost in US dollars. If None, cost estimating is not '
           'suppported on the model being queried.'
-      )
+      ),
   ] = None
+  retry_stats: RetryStats = RetryStats()
 
   def __bool__(self) -> bool:
     return self.num_requests > 0
@@ -136,6 +192,7 @@ class LMSamplingUsage(pg.Object):
         total_tokens=self.total_tokens + other.total_tokens,
         num_requests=self.num_requests + other.num_requests,
         estimated_cost=estimated_cost,
+        retry_stats=self.retry_stats + other.retry_stats,
     )
 
   def __radd__(self, other: Optional['LMSamplingUsage']) -> 'LMSamplingUsage':
@@ -511,7 +568,18 @@ class LanguageModel(component.Component):
                 total_tokens=usage.total_tokens // n,
                 estimated_cost=(
                     usage.estimated_cost / n if usage.estimated_cost else None
-                )
+                ),
+                retry_stats=RetryStats(
+                    num_occurences=usage.retry_stats.num_occurences // n,
+                    total_wait_interval=usage.retry_stats.total_wait_interval
+                    / n,
+                    total_call_interval=usage.retry_stats.total_call_interval
+                    / n,
+                    errors={
+                        error: count // n
+                        for error, count in usage.retry_stats.errors.items()
+                    },
+                ),
             )
 
           # Track usage.
@@ -584,16 +652,16 @@ class LanguageModel(component.Component):
 
   def _parallel_execute_with_currency_control(
       self,
-      action: Callable[..., Any],
+      action: Callable[..., LMSamplingResult],
       inputs: Sequence[Any],
       retry_on_errors: Union[
           None,
           Union[Type[BaseException], Tuple[Type[BaseException], str]],
           Sequence[Union[Type[BaseException], Tuple[Type[BaseException], str]]],
       ] = RetryableLMError,
-  ) -> Any:
+  ) -> list[Any]:
     """Helper method for subclasses for implementing _sample."""
-    return concurrent.concurrent_execute(
+    executed_jobs = concurrent.concurrent_execute(
         action,
         inputs,
         executor=self.resource_id if self.max_concurrency else None,
@@ -603,7 +671,15 @@ class LanguageModel(component.Component):
         retry_interval=self.retry_interval,
         exponential_backoff=self.exponential_backoff,
         max_retry_interval=self.max_retry_interval,
+        return_jobs=True,
     )
+    for job in executed_jobs:
+      if isinstance(job.result, LMSamplingResult):
+        job.result.usage.rebind(
+            retry_stats=RetryStats.from_retry_entries(job.retry_entries),
+            skip_notification=True,
+        )
+    return [job.result for job in executed_jobs]
 
   def __call__(
       self, prompt: message_lib.Message, *, cache_seed: int = 0, **kwargs
@@ -653,7 +729,7 @@ class LanguageModel(component.Component):
     """Outputs debugging information about the model."""
     title_suffix = ''
     if usage.total_tokens != 0:
-      title_suffix = console.colored(
+      title_suffix = pg.colored(
           f' (total {usage.total_tokens} tokens)', 'red'
       )
 
@@ -672,7 +748,7 @@ class LanguageModel(component.Component):
     """Outputs debugging information about the prompt."""
     title_suffix = ''
     if usage.prompt_tokens != 0:
-      title_suffix = console.colored(f' ({usage.prompt_tokens} tokens)', 'red')
+      title_suffix = pg.colored(f' ({usage.prompt_tokens} tokens)', 'red')
 
     console.write(
         # We use metadata 'formatted_text' for scenarios where the prompt text
@@ -703,7 +779,7 @@ class LanguageModel(component.Component):
     if usage.completion_tokens != 0:
       title_suffix += f'{usage.completion_tokens} tokens '
     title_suffix += f'in {elapse:.2f} seconds)'
-    title_suffix = console.colored(title_suffix, 'red')
+    title_suffix = pg.colored(title_suffix, 'red')
 
     console.write(
         str(response) + '\n',

langfun/core/language_model_test.py

@@ -35,34 +35,34 @@ class MockModel(lm_lib.LanguageModel):
               ) -> list[lm_lib.LMSamplingResult]:
     context = pg.Dict(attempt=0)
 
-    def fake_sample(prompts):
+    def fake_sample(prompt):
       if context.attempt >= self.failures_before_attempt:
-        return [
-            lm_lib.LMSamplingResult(
-                [
-                    lm_lib.LMSample(  # pylint: disable=g-complex-comprehension
-                        response=prompt.text * self.sampling_options.top_k,
-                        score=self.sampling_options.temperature or -1.0,
-                    )
-                ],
-                usage=lm_lib.LMSamplingUsage(
-                    prompt_tokens=100,
-                    completion_tokens=100,
-                    total_tokens=200,
-                    estimated_cost=1.0,
-                ),
-            )
-            for prompt in prompts
-        ]
-      context.attempt += 1
+        return lm_lib.LMSamplingResult(
+            [
+                lm_lib.LMSample(  # pylint: disable=g-complex-comprehension
+                    response=prompt.text * self.sampling_options.top_k,
+                    score=self.sampling_options.temperature or -1.0,
+                )
+            ],
+            usage=lm_lib.LMSamplingUsage(
+                prompt_tokens=100,
+                completion_tokens=100,
+                total_tokens=200,
+                estimated_cost=1.0,
+            ),
+        )
+      else:
+        context.attempt += 1
       raise ValueError('Failed to sample prompts.')
 
-    return concurrent.with_retry(
-        fake_sample,
-        retry_on_errors=ValueError,
-        max_attempts=self.max_attempts,
-        retry_interval=1,
-    )(prompts)
+    results = self._parallel_execute_with_currency_control(
+        fake_sample, prompts, retry_on_errors=ValueError
+    )
+    for result in results:
+      result.usage.retry_stats.rebind(
+          total_call_interval=0, skip_notification=True
+      )
+    return results
 
   @property
   def model_id(self) -> str:
@@ -448,13 +448,50 @@ class LanguageModelTest(unittest.TestCase):
 
   def test_retry(self):
     lm = MockModel(
-        failures_before_attempt=1, top_k=1,
+        failures_before_attempt=1, top_k=1, max_attempts=2, retry_interval=1
     )
     with self.assertRaisesRegex(
         concurrent.RetryError, 'Calling .* failed after 1 attempts'
     ):
       lm('foo', max_attempts=1)
-    self.assertEqual(lm('foo', max_attempts=2), 'foo')
+
+    usage = lm_lib.LMSamplingUsage(
+        prompt_tokens=100,
+        completion_tokens=100,
+        total_tokens=200,
+        num_requests=1,
+        estimated_cost=1.0,
+        retry_stats=lm_lib.RetryStats(
+            num_occurences=1,
+            total_wait_interval=1,
+            errors={'ValueError': 1},
+        ),
+    )
+    out = lm.sample(['foo'])
+    self.assertEqual(
+        # lm.sample(['foo'], max_attempts=2),
+        out,
+        [
+            lm_lib.LMSamplingResult(
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage(
+                            'foo',
+                            score=-1.0,
+                            logprobs=None,
+                            is_cached=False,
+                            usage=usage,
+                            tags=['lm-response'],
+                        ),
+                        score=-1.0,
+                        logprobs=None,
+                    )
+                ],
+                usage=usage,
+                is_cached=False,
+            )
+        ],
+    )
 
   def test_debug(self):
     class Image(modality.Modality):
@@ -755,16 +792,34 @@ class LMSamplingUsageTest(unittest.TestCase):
 
   def test_add(self):
     usage1 = lm_lib.LMSamplingUsage(100, 200, 300, 4, 5.0)
+    usage1.rebind(retry_stats=lm_lib.RetryStats(1, 3, 4, {'e1': 1}))
     usage2 = lm_lib.LMSamplingUsage(100, 200, 300, 4, 5.0)
     self.assertEqual(usage1 + usage2, usage1 + usage2)
     self.assertIs(usage1 + None, usage1)
     self.assertIs(None + usage1, usage1)
     usage3 = lm_lib.LMSamplingUsage(100, 200, 300, 4, None)
+    usage3.rebind(retry_stats=lm_lib.RetryStats(2, 4, 5, {'e1': 2, 'e2': 3}))
     self.assertEqual(
-        usage1 + usage3, lm_lib.LMSamplingUsage(200, 400, 600, 8, 5.0)
+        usage1 + usage3,
+        lm_lib.LMSamplingUsage(
+            200,
+            400,
+            600,
+            8,
+            5.0,
+            retry_stats=lm_lib.RetryStats(3, 7, 9, {'e1': 3, 'e2': 3}),
+        ),
     )
     self.assertEqual(
-        usage3 + usage1, lm_lib.LMSamplingUsage(200, 400, 600, 8, 5.0)
+        usage3 + usage1,
+        lm_lib.LMSamplingUsage(
+            200,
+            400,
+            600,
+            8,
+            5.0,
+            retry_stats=lm_lib.RetryStats(3, 7, 9, {'e1': 3, 'e2': 3}),
+        ),
     )
 
   def test_usage_not_available(self):

langfun/core/llms/__init__.py

@@ -27,8 +27,15 @@ from langfun.core.llms.fake import StaticSequence
 # Compositional models.
 from langfun.core.llms.compositional import RandomChoice
 
-# REST-based models.
+# Base models by request/response protocol.
 from langfun.core.llms.rest import REST
+from langfun.core.llms.openai_compatible import OpenAICompatible
+from langfun.core.llms.gemini import Gemini
+from langfun.core.llms.anthropic import Anthropic
+
+# Base models by serving platforms.
+from langfun.core.llms.vertexai import VertexAI
+from langfun.core.llms.groq import Groq
 
 # Gemini models.
 from langfun.core.llms.google_genai import GenAI
@@ -44,7 +51,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
@@ -111,20 +118,34 @@ from langfun.core.llms.openai import Gpt3Curie
 from langfun.core.llms.openai import Gpt3Babbage
 from langfun.core.llms.openai import Gpt3Ada
 
-from langfun.core.llms.anthropic import Anthropic
+# Anthropic models.
+
 from langfun.core.llms.anthropic import Claude35Sonnet
 from langfun.core.llms.anthropic import Claude35Sonnet20241022
 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.groq import Groq
+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.
+
+# Gemma models.
+from langfun.core.llms.groq import GroqGemma2_9B_IT
+from langfun.core.llms.groq import GroqGemma_7B_IT
+
+# Llama models.
+from langfun.core.llms.vertexai import VertexAILlama
+from langfun.core.llms.vertexai import VertexAILlama3_2_90B
+from langfun.core.llms.vertexai import VertexAILlama3_1_405B
+from langfun.core.llms.vertexai import VertexAILlama3_1_70B
+from langfun.core.llms.vertexai import VertexAILlama3_1_8B
+
 from langfun.core.llms.groq import GroqLlama3_2_3B
 from langfun.core.llms.groq import GroqLlama3_2_1B
 from langfun.core.llms.groq import GroqLlama3_1_70B
@@ -132,15 +153,29 @@ from langfun.core.llms.groq import GroqLlama3_1_8B
 from langfun.core.llms.groq import GroqLlama3_70B
 from langfun.core.llms.groq import GroqLlama3_8B
 from langfun.core.llms.groq import GroqLlama2_70B
+
+# Mistral models.
+from langfun.core.llms.vertexai import VertexAIMistral
+from langfun.core.llms.vertexai import VertexAIMistralLarge_20241121
+from langfun.core.llms.vertexai import VertexAIMistralLarge_20240724
+from langfun.core.llms.vertexai import VertexAIMistralNemo_20240724
+from langfun.core.llms.vertexai import VertexAICodestral_20250113
+from langfun.core.llms.vertexai import VertexAICodestral_20240529
+
 from langfun.core.llms.groq import GroqMistral_8x7B
-from langfun.core.llms.groq import GroqGemma2_9B_IT
-from langfun.core.llms.groq import GroqGemma_7B_IT
+
+# DeepSeek models.
+from langfun.core.llms.deepseek import DeepSeek
+from langfun.core.llms.deepseek import DeepSeekChat
+
+# Whisper models.
 from langfun.core.llms.groq import GroqWhisper_Large_v3
 from langfun.core.llms.groq import GroqWhisper_Large_v3Turbo
 
 # LLaMA C++ models.
 from langfun.core.llms.llama_cpp import LlamaCppRemote
 
+
 # Placeholder for Google-internal imports.
 
 # Include cache as sub-module.

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()