langfun 0.1.2.dev202501160804__py3-none-any.whl → 0.1.2.dev202501180803__py3-none-any.whl

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
@@ -57,9 +64,6 @@ from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_002
 from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_001
 from langfun.core.llms.vertexai import VertexAIGeminiPro1
 
-# Base for OpenAI-compatible models.
-from langfun.core.llms.openai_compatible import OpenAICompatible
-
 # OpenAI models.
 from langfun.core.llms.openai import OpenAI
 
@@ -114,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
@@ -135,18 +153,28 @@ 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
 
-# DeepSeek models.
-from langfun.core.llms.deepseek import DeepSeek
-from langfun.core.llms.deepseek import DeepSeekChat
 
 # Placeholder for Google-internal imports.
 

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/gemini.py

@@ -380,7 +380,7 @@ class Gemini(rest.REST):
     return (
         cost_per_1m_input_tokens * num_input_tokens
         + cost_per_1m_output_tokens * num_output_tokens
-    ) / 1000_1000
+    ) / 1000_000
 
   @property
   def model_id(self) -> str:

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,13 @@
 
 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 openai_compatible
+from langfun.core.llms import rest
 import pyglove as pg
 
 try:
@@ -36,10 +39,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,
@@ -95,7 +109,7 @@ class VertexAI(gemini.Gemini):
     credentials = self.credentials
     if credentials is None:
       # Use default credentials.
-      credentials = google_auth.default(
+      credentials, _ = google_auth.default(
           scopes=['https://www.googleapis.com/auth/cloud-platform']
       )
     self._credentials = credentials
@@ -114,6 +128,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 +149,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 +157,405 @@ 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
+
+#
+# Llama models on Vertex AI.
+# pylint: disable=line-too-long
+# Pricing: https://cloud.google.com/vertex-ai/generative-ai/pricing?_gl=1*ukuk6u*_ga*MjEzMjc4NjM2My4xNzMzODg4OTg3*_ga_WH2QY8WWF5*MTczNzEzNDU1Mi4xMjQuMS4xNzM3MTM0NzczLjU5LjAuMA..#meta-models
+# pylint: enable=line-too-long
+
+LLAMA_MODELS = {
+    'llama-3.2-90b-vision-instruct-maas': pg.Dict(
+        latest_update='2024-09-25',
+        in_service=True,
+        rpm=0,
+        tpm=0,
+        # Free during preview.
+        cost_per_1m_input_tokens=None,
+        cost_per_1m_output_tokens=None,
+    ),
+    'llama-3.1-405b-instruct-maas': pg.Dict(
+        latest_update='2024-09-25',
+        in_service=True,
+        rpm=0,
+        tpm=0,
+        # GA.
+        cost_per_1m_input_tokens=5,
+        cost_per_1m_output_tokens=16,
+    ),
+    'llama-3.1-70b-instruct-maas': pg.Dict(
+        latest_update='2024-09-25',
+        in_service=True,
+        rpm=0,
+        tpm=0,
+        # Free during preview.
+        cost_per_1m_input_tokens=None,
+        cost_per_1m_output_tokens=None,
+    ),
+    'llama-3.1-8b-instruct-maas': pg.Dict(
+        latest_update='2024-09-25',
+        in_service=True,
+        rpm=0,
+        tpm=0,
+        # Free during preview.
+        cost_per_1m_input_tokens=None,
+        cost_per_1m_output_tokens=None,
+    )
+}
+
+
+@pg.use_init_args(['model'])
+@pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
+class VertexAILlama(VertexAI, openai_compatible.OpenAICompatible):
+  """Llama models on VertexAI."""
+
+  model: pg.typing.Annotated[
+      pg.typing.Enum(pg.MISSING_VALUE, list(LLAMA_MODELS.keys())),
+      'Llama model ID.',
+  ]
+
+  locations: Annotated[
+      Literal['us-central1'],
+      (
+          'GCP locations with Llama models hosted. '
+          'See https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/llama#regions-quotas'
+      )
+  ] = 'us-central1'
+
+  @property
+  def api_endpoint(self) -> str:
+    assert self._api_initialized
+    return (
+        f'https://{self._location}-aiplatform.googleapis.com/v1beta1/projects/'
+        f'{self._project}/locations/{self._location}/endpoints/'
+        f'openapi/chat/completions'
+    )
+
+  def request(
+      self,
+      prompt: lf.Message,
+      sampling_options: lf.LMSamplingOptions
+  ):
+    request = super().request(prompt, sampling_options)
+    request['model'] = f'meta/{self.model}'
+    return request
+
+  @property
+  def max_concurrency(self) -> int:
+    rpm = LLAMA_MODELS[self.model].get('rpm', 0)
+    tpm = LLAMA_MODELS[self.model].get('tpm', 0)
+    return self.rate_to_max_concurrency(
+        requests_per_min=rpm, tokens_per_min=tpm
+    )
+
+  def estimate_cost(
+      self,
+      num_input_tokens: int,
+      num_output_tokens: int
+  ) -> float | None:
+    """Estimate the cost based on usage."""
+    cost_per_1m_input_tokens = LLAMA_MODELS[self.model].get(
+        'cost_per_1m_input_tokens', None
+    )
+    cost_per_1m_output_tokens = LLAMA_MODELS[self.model].get(
+        'cost_per_1m_output_tokens', None
+    )
+    if cost_per_1m_output_tokens is None or cost_per_1m_input_tokens is None:
+      return None
+    return (
+        cost_per_1m_input_tokens * num_input_tokens
+        + cost_per_1m_output_tokens * num_output_tokens
+    ) / 1000_000
+
+
+# pylint: disable=invalid-name
+class VertexAILlama3_2_90B(VertexAILlama):
+  """Llama 3.2 90B vision instruct model on VertexAI."""
+
+  model = 'llama-3.2-90b-vision-instruct-maas'
+
+
+class VertexAILlama3_1_405B(VertexAILlama):
+  """Llama 3.1 405B vision instruct model on VertexAI."""
+
+  model = 'llama-3.1-405b-instruct-maas'
+
+
+class VertexAILlama3_1_70B(VertexAILlama):
+  """Llama 3.1 70B vision instruct model on VertexAI."""
+
+  model = 'llama-3.1-70b-instruct-maas'
+
+
+class VertexAILlama3_1_8B(VertexAILlama):
+  """Llama 3.1 8B vision instruct model on VertexAI."""
+
+  model = 'llama-3.1-8b-instruct-maas'
+# pylint: enable=invalid-name
+
+#
+# Mistral models on Vertex AI.
+# pylint: disable=line-too-long
+# Pricing: https://cloud.google.com/vertex-ai/generative-ai/pricing?_gl=1*ukuk6u*_ga*MjEzMjc4NjM2My4xNzMzODg4OTg3*_ga_WH2QY8WWF5*MTczNzEzNDU1Mi4xMjQuMS4xNzM3MTM0NzczLjU5LjAuMA..#mistral-models
+# pylint: enable=line-too-long
+
+
+MISTRAL_MODELS = {
+    'mistral-large-2411': pg.Dict(
+        latest_update='2024-11-21',
+        in_service=True,
+        rpm=0,
+        tpm=0,
+        # GA.
+        cost_per_1m_input_tokens=2,
+        cost_per_1m_output_tokens=6,
+    ),
+    'mistral-large@2407': pg.Dict(
+        latest_update='2024-07-24',
+        in_service=True,
+        rpm=0,
+        tpm=0,
+        # GA.
+        cost_per_1m_input_tokens=2,
+        cost_per_1m_output_tokens=6,
+    ),
+    'mistral-nemo@2407': pg.Dict(
+        latest_update='2024-07-24',
+        in_service=True,
+        rpm=0,
+        tpm=0,
+        # GA.
+        cost_per_1m_input_tokens=0.15,
+        cost_per_1m_output_tokens=0.15,
+    ),
+    'codestral-2501': pg.Dict(
+        latest_update='2025-01-13',
+        in_service=True,
+        rpm=0,
+        tpm=0,
+        # GA.
+        cost_per_1m_input_tokens=0.3,
+        cost_per_1m_output_tokens=0.9,
+    ),
+    'codestral@2405': pg.Dict(
+        latest_update='2024-05-29',
+        in_service=True,
+        rpm=0,
+        tpm=0,
+        # GA.
+        cost_per_1m_input_tokens=0.2,
+        cost_per_1m_output_tokens=0.6,
+    ),
+}
+
+
+@pg.use_init_args(['model'])
+@pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
+class VertexAIMistral(VertexAI, openai_compatible.OpenAICompatible):
+  """Mistral AI models on VertexAI."""
+
+  model: pg.typing.Annotated[
+      pg.typing.Enum(pg.MISSING_VALUE, list(MISTRAL_MODELS.keys())),
+      'Mistral model ID.',
+  ]
+
+  locations: Annotated[
+      Literal['us-central1', 'europe-west4'],
+      (
+          'GCP locations with Mistral models hosted. '
+          'See https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/mistral#regions-quotas'
+      )
+  ] = 'us-central1'
+
+  @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/mistralai/'
+        f'models/{self.model}:rawPredict'
+    )
+
+  @property
+  def max_concurrency(self) -> int:
+    rpm = MISTRAL_MODELS[self.model].get('rpm', 0)
+    tpm = MISTRAL_MODELS[self.model].get('tpm', 0)
+    return self.rate_to_max_concurrency(
+        requests_per_min=rpm, tokens_per_min=tpm
+    )
+
+  def estimate_cost(
+      self,
+      num_input_tokens: int,
+      num_output_tokens: int
+  ) -> float | None:
+    """Estimate the cost based on usage."""
+    cost_per_1m_input_tokens = MISTRAL_MODELS[self.model].get(
+        'cost_per_1m_input_tokens', None
+    )
+    cost_per_1m_output_tokens = MISTRAL_MODELS[self.model].get(
+        'cost_per_1m_output_tokens', None
+    )
+    if cost_per_1m_output_tokens is None or cost_per_1m_input_tokens is None:
+      return None
+    return (
+        cost_per_1m_input_tokens * num_input_tokens
+        + cost_per_1m_output_tokens * num_output_tokens
+    ) / 1000_000
+
+
+# pylint: disable=invalid-name
+class VertexAIMistralLarge_20241121(VertexAIMistral):
+  """Mistral Large model on VertexAI released on 2024/11/21."""
+
+  model = 'mistral-large-2411'
+
+
+class VertexAIMistralLarge_20240724(VertexAIMistral):
+  """Mistral Large model on VertexAI released on 2024/07/24."""
+
+  model = 'mistral-large@2407'
+
+
+class VertexAIMistralNemo_20240724(VertexAIMistral):
+  """Mistral Nemo model on VertexAI released on 2024/07/24."""
+
+  model = 'mistral-nemo@2407'
+
+
+class VertexAICodestral_20250113(VertexAIMistral):
+  """Mistral Nemo model on VertexAI released on 2024/07/24."""
+
+  model = 'codestral-2501'
+
+
+class VertexAICodestral_20240529(VertexAIMistral):
+  """Mistral Nemo model on VertexAI released on 2024/05/29."""
+
+  model = 'codestral@2405'
+# 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.dev202501160804.dist-info → langfun-0.1.2.dev202501180803.dist-info}/METADATA

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

{langfun-0.1.2.dev202501160804.dist-info → langfun-0.1.2.dev202501180803.dist-info}/RECORD

@@ -71,16 +71,16 @@ 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=50mJagAgkIhMwhOyHxGq_O5st4HhpnE-okeYzc7GU6c,7667
+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
 langfun/core/llms/deepseek_test.py,sha256=dS72i52bwMpCN4dJDvpJI59AnNChpwxS5eYYFrhGh90,1843
 langfun/core/llms/fake.py,sha256=gCHBYBLvBCsC78HI1hpoqXCS-p1FMTgY1P1qh_sGBPk,3070
 langfun/core/llms/fake_test.py,sha256=2h13qkwEz_JR0mtUDPxdAhQo7MueXaFSwsD2DIRDW9g,7653
-langfun/core/llms/gemini.py,sha256=tfM4vrt0WnvnrxRhWXZWh7Gp8dYYfMnSbi9uOstkSak,17399
+langfun/core/llms/gemini.py,sha256=itwTCmQHRjwSjt7_UzFfaat23gyRL-El4qmJrg-OGVA,17398
 langfun/core/llms/gemini_test.py,sha256=2ERhYWCJwnfDTQbCaZHFuB1TdWJFrOBS7yyCBInIdQk,6129
 langfun/core/llms/google_genai.py,sha256=85Vmx5QmsziON03PRsFQINSu5NF6pAAuFFhUdDteWGc,3662
 langfun/core/llms/google_genai_test.py,sha256=JZf_cbQ4GGGpwiQCLjFJn7V4jxBBqgZhIx91AzbGKVo,1250
@@ -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=SVvLTqQZ6Ha8wZh3azkh4g3O838CpNkuP3XlgIrLMKo,15751
+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.dev202501160804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202501160804.dist-info/METADATA,sha256=_XM3ancZIb8-33gpRxLKmdJOBZsMfd1_2-4otzha19Q,8172
-langfun-0.1.2.dev202501160804.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-langfun-0.1.2.dev202501160804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202501160804.dist-info/RECORD,,
+langfun-0.1.2.dev202501180803.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202501180803.dist-info/METADATA,sha256=W9jkpCCOZx-Tl8sNz3y1IdVZNG48qcjs21airG2TTI0,8172
+langfun-0.1.2.dev202501180803.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+langfun-0.1.2.dev202501180803.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202501180803.dist-info/RECORD,,