langfun 0.1.2.dev202412010804__py3-none-any.whl → 0.1.2.dev202412030804__py3-none-any.whl

langfun/core/llms/vertexai.py

@@ -13,30 +13,28 @@
 # limitations under the License.
 """Vertex AI generative models."""
 
+import base64
 import functools
 import os
 from typing import Annotated, Any
 
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
+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
-  import vertexai
-  from google.cloud.aiplatform import models as aiplatform_models
-  from vertexai import generative_models
-  from vertexai import language_models
+  from google.auth.transport import requests as auth_requests
   # pylint: enable=g-import-not-at-top
 
   Credentials = credentials_lib.Credentials
 except ImportError:
-  credentials_lib = None   # pylint: disable=invalid-name
-  vertexai = None
-  generative_models = None
-  language_models = None
-  aiplatform_models = None
+  google_auth = None
+  credentials_lib = None
+  auth_requests = None
   Credentials = Any
 
 
@@ -50,122 +48,86 @@ AVGERAGE_CHARS_PER_TOEKN = 4
 # as of 2024-10-10.
 SUPPORTED_MODELS_AND_SETTINGS = {
     'gemini-1.5-pro-001': pg.Dict(
-        api='gemini',
         rpm=100,
         cost_per_1k_input_chars=0.0003125,
         cost_per_1k_output_chars=0.00125,
     ),
     'gemini-1.5-pro-002': pg.Dict(
-        api='gemini',
         rpm=100,
         cost_per_1k_input_chars=0.0003125,
         cost_per_1k_output_chars=0.00125,
     ),
     'gemini-1.5-flash-002': pg.Dict(
-        api='gemini',
         rpm=500,
         cost_per_1k_input_chars=0.00001875,
         cost_per_1k_output_chars=0.000075,
     ),
     'gemini-1.5-flash-001': pg.Dict(
-        api='gemini',
         rpm=500,
         cost_per_1k_input_chars=0.00001875,
         cost_per_1k_output_chars=0.000075,
     ),
     'gemini-1.5-pro': pg.Dict(
-        api='gemini',
         rpm=100,
         cost_per_1k_input_chars=0.0003125,
         cost_per_1k_output_chars=0.00125,
     ),
     'gemini-1.5-flash': pg.Dict(
-        api='gemini',
-        rpm=500,
-        cost_per_1k_input_chars=0.00001875,
-        cost_per_1k_output_chars=0.000075,
-    ),
-    'gemini-1.5-pro-latest': pg.Dict(
-        api='gemini',
-        rpm=100,
-        cost_per_1k_input_chars=0.0003125,
-        cost_per_1k_output_chars=0.00125,
-    ),
-    'gemini-1.5-flash-latest': pg.Dict(
-        api='gemini',
         rpm=500,
         cost_per_1k_input_chars=0.00001875,
         cost_per_1k_output_chars=0.000075,
     ),
     'gemini-1.5-pro-preview-0514': pg.Dict(
-        api='gemini',
         rpm=50,
         cost_per_1k_input_chars=0.0003125,
         cost_per_1k_output_chars=0.00125,
     ),
     'gemini-1.5-pro-preview-0409': pg.Dict(
-        api='gemini',
         rpm=50,
         cost_per_1k_input_chars=0.0003125,
         cost_per_1k_output_chars=0.00125,
     ),
     'gemini-1.5-flash-preview-0514': pg.Dict(
-        api='gemini',
         rpm=200,
         cost_per_1k_input_chars=0.00001875,
         cost_per_1k_output_chars=0.000075,
     ),
     'gemini-1.0-pro': pg.Dict(
-        api='gemini',
         rpm=300,
         cost_per_1k_input_chars=0.000125,
         cost_per_1k_output_chars=0.000375,
     ),
     'gemini-1.0-pro-vision': pg.Dict(
-        api='gemini',
         rpm=100,
         cost_per_1k_input_chars=0.000125,
         cost_per_1k_output_chars=0.000375,
     ),
-    # PaLM APIs.
-    'text-bison': pg.Dict(
-        api='palm',
-        rpm=1600
-    ),
-    'text-bison-32k': pg.Dict(
-        api='palm',
-        rpm=300
-    ),
-    'text-unicorn': pg.Dict(
-        api='palm',
-        rpm=100
-    ),
-    # Endpoint
     # TODO(chengrun): Set a more appropriate rpm for endpoint.
-    'custom': pg.Dict(api='endpoint', rpm=20),
+    'vertexai-endpoint': pg.Dict(
+        rpm=20,
+        cost_per_1k_input_chars=0.0000125,
+        cost_per_1k_output_chars=0.0000375,
+    ),
 }
 
 
 @lf.use_init_args(['model'])
-class VertexAI(lf.LanguageModel):
-  """Language model served on VertexAI."""
+@pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
+class VertexAI(rest.REST):
+  """Language model served on VertexAI with REST API."""
 
   model: pg.typing.Annotated[
       pg.typing.Enum(
           pg.MISSING_VALUE, list(SUPPORTED_MODELS_AND_SETTINGS.keys())
       ),
       (
-          'Vertex AI model name. See '
-          'https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models '
-          'for details.'
+          'Vertex AI model name with REST API support. See '
+          'https://cloud.google.com/vertex-ai/generative-ai/docs/'
+          'model-reference/inference#supported-models'
+          ' for details.'
       ),
   ]
 
-  endpoint_name: pg.typing.Annotated[
-      str | None,
-      'Vertex Endpoint name or ID.',
-  ]
-
   project: Annotated[
       str | None,
       (
@@ -197,14 +159,14 @@ class VertexAI(lf.LanguageModel):
 
   def _on_bound(self):
     super()._on_bound()
-    self.__dict__.pop('_api_initialized', None)
-    if generative_models is None:
-      raise RuntimeError(
+    if google_auth is None:
+      raise ValueError(
           'Please install "langfun[llm-google-vertex]" to use Vertex AI models.'
       )
+    self._project = None
+    self._credentials = None
 
-  @functools.cached_property
-  def _api_initialized(self):
+  def _initialize(self):
     project = self.project or os.environ.get('VERTEXAI_PROJECT', None)
     if not project:
       raise ValueError(
@@ -219,21 +181,14 @@ class VertexAI(lf.LanguageModel):
           'variable `VERTEXAI_LOCATION` with your Vertex AI service location.'
       )
 
+    self._project = project
     credentials = self.credentials
-    # Placeholder for Google-internal credentials.
-    assert vertexai is not None
-    vertexai.init(project=project, location=location, credentials=credentials)
-    return True
-
-  @property
-  def model_id(self) -> str:
-    """Returns a string to identify the model."""
-    return f'VertexAI({self.model})'
-
-  @property
-  def resource_id(self) -> str:
-    """Returns a string to identify the resource for rate control."""
-    return self.model_id
+    if credentials is None:
+      # Use default credentials.
+      credentials = google_auth.default(
+          scopes=['https://www.googleapis.com/auth/cloud-platform']
+      )
+    self._credentials = credentials
 
   @property
   def max_concurrency(self) -> int:
@@ -262,47 +217,75 @@ class VertexAI(lf.LanguageModel):
         + cost_per_1k_output_chars * num_output_tokens
     ) * AVGERAGE_CHARS_PER_TOEKN / 1000
 
+  @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/google/'
+        f'models/{self.model}:generateContent'
+    )
+
+  def request(
+      self, prompt: lf.Message, sampling_options: lf.LMSamplingOptions
+  ) -> dict[str, Any]:
+    request = dict(
+        generationConfig=self._generation_config(prompt, sampling_options)
+    )
+    request['contents'] = [self._content_from_message(prompt)]
+    return request
+
   def _generation_config(
       self, prompt: lf.Message, options: lf.LMSamplingOptions
-  ) -> Any:  # generative_models.GenerationConfig
-    """Creates generation config from langfun sampling options."""
-    assert generative_models is not None
-    # Users could use `metadata_json_schema` to pass additional
-    # request arguments.
-    json_schema = prompt.metadata.get('json_schema')
-    response_mime_type = None
-    if json_schema is not None:
+  ) -> dict[str, Any]:
+    """Returns a dict as generation config for prompt and LMSamplingOptions."""
+    config = dict(
+        temperature=options.temperature,
+        maxOutputTokens=options.max_tokens,
+        candidateCount=options.n,
+        topK=options.top_k,
+        topP=options.top_p,
+        stopSequences=options.stop,
+        seed=options.random_seed,
+        responseLogprobs=options.logprobs,
+        logprobs=options.top_logprobs,
+    )
+
+    if json_schema := prompt.metadata.get('json_schema'):
       if not isinstance(json_schema, dict):
         raise ValueError(
             f'`json_schema` must be a dict, got {json_schema!r}.'
         )
-      response_mime_type = 'application/json'
+      json_schema = pg.to_json(json_schema)
+      config['responseSchema'] = json_schema
+      config['responseMimeType'] = 'application/json'
       prompt.metadata.formatted_text = (
           prompt.text
           + '\n\n [RESPONSE FORMAT (not part of prompt)]\n'
           + pg.to_json_str(json_schema, json_indent=2)
       )
+    return config
 
-    return generative_models.GenerationConfig(
-        temperature=options.temperature,
-        top_p=options.top_p,
-        top_k=options.top_k,
-        max_output_tokens=options.max_tokens,
-        stop_sequences=options.stop,
-        response_mime_type=response_mime_type,
-        response_schema=json_schema,
-    )
-
-  def _content_from_message(
-      self, prompt: lf.Message
-  ) -> list[str | Any]:
-    """Gets generation input from langfun message."""
-    assert generative_models is not None
-    chunks = []
-
+  def _content_from_message(self, prompt: lf.Message) -> dict[str, Any]:
+    """Gets generation content from langfun message."""
+    parts = []
     for lf_chunk in prompt.chunk():
       if isinstance(lf_chunk, str):
-        chunks.append(lf_chunk)
+        parts.append({'text': lf_chunk})
       elif isinstance(lf_chunk, lf_modalities.Mime):
         try:
           modalities = lf_chunk.make_compatible(
@@ -312,174 +295,52 @@ class VertexAI(lf.LanguageModel):
             modalities = [modalities]
           for modality in modalities:
             if modality.is_text:
-              chunk = modality.to_text()
+              parts.append({'text': modality.to_text()})
             else:
-              chunk = generative_models.Part.from_data(
-                  modality.to_bytes(), modality.mime_type
-              )
-            chunks.append(chunk)
+              parts.append({
+                  'inlineData': {
+                      'data': base64.b64encode(modality.to_bytes()).decode(),
+                      'mimeType': modality.mime_type,
+                  }
+              })
         except lf.ModalityError as e:
           raise lf.ModalityError(f'Unsupported modality: {lf_chunk!r}') from e
       else:
         raise lf.ModalityError(f'Unsupported modality: {lf_chunk!r}')
-    return chunks
-
-  def _generation_response_to_message(
-      self,
-      response: Any,  # generative_models.GenerationResponse
-  ) -> lf.Message:
-    """Parses generative response into message."""
-    return lf.AIMessage(response.text)
+    return dict(role='user', parts=parts)
 
-  def _generation_endpoint_response_to_message(
-      self,
-      response: Any,  # google.cloud.aiplatform.aiplatform.models.Prediction
-  ) -> lf.Message:
-    """Parses Endpoint response into message."""
-    return lf.AIMessage(response.predictions[0])
-
-  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    assert self._api_initialized, 'Vertex AI API is not initialized.'
-    # TODO(yifenglu): It seems this exception is due to the instability of the
-    # API. We should revisit this later.
-    retry_on_errors = [
-        (Exception, 'InternalServerError'),
-        (Exception, 'ResourceExhausted'),
-        (Exception, '_InactiveRpcError'),
-        (Exception, 'ValueError'),
+  def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
+    messages = [
+        self._message_from_content_parts(candidate['content']['parts'])
+        for candidate in json['candidates']
     ]
-
-    return self._parallel_execute_with_currency_control(
-        self._sample_single,
-        prompts,
-        retry_on_errors=retry_on_errors,
-    )
-
-  def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
-    if self.sampling_options.n > 1:
-      raise ValueError(
-          f'`n` greater than 1 is not supported: {self.sampling_options.n}.'
-      )
-    api = SUPPORTED_MODELS_AND_SETTINGS[self.model].api
-    match api:
-      case 'gemini':
-        return self._sample_generative_model(prompt)
-      case 'palm':
-        return self._sample_text_generation_model(prompt)
-      case 'endpoint':
-        return self._sample_endpoint_model(prompt)
-      case _:
-        raise ValueError(f'Unsupported API: {api}')
-
-  def _sample_generative_model(self, prompt: lf.Message) -> lf.LMSamplingResult:
-    """Samples a generative model."""
-    model = _VERTEXAI_MODEL_HUB.get_generative_model(self.model)
-    input_content = self._content_from_message(prompt)
-    response = model.generate_content(
-        input_content,
-        generation_config=self._generation_config(
-            prompt, self.sampling_options
-        ),
-    )
-    usage_metadata = response.usage_metadata
-    usage = lf.LMSamplingUsage(
-        prompt_tokens=usage_metadata.prompt_token_count,
-        completion_tokens=usage_metadata.candidates_token_count,
-        total_tokens=usage_metadata.total_token_count,
-        estimated_cost=self.estimate_cost(
-            num_input_tokens=usage_metadata.prompt_token_count,
-            num_output_tokens=usage_metadata.candidates_token_count,
-        ),
-    )
+    usage = json['usageMetadata']
+    input_tokens = usage['promptTokenCount']
+    output_tokens = usage['candidatesTokenCount']
     return lf.LMSamplingResult(
-        [
-            # Scoring is not supported.
-            lf.LMSample(
-                self._generation_response_to_message(response), score=0.0
+        [lf.LMSample(message) for message in messages],
+        usage=lf.LMSamplingUsage(
+            prompt_tokens=input_tokens,
+            completion_tokens=output_tokens,
+            total_tokens=input_tokens + output_tokens,
+            estimated_cost=self.estimate_cost(
+                num_input_tokens=input_tokens,
+                num_output_tokens=output_tokens,
             ),
-        ],
-        usage=usage,
-    )
-
-  def _sample_text_generation_model(
-      self, prompt: lf.Message
-  ) -> lf.LMSamplingResult:
-    """Samples a text generation model."""
-    model = _VERTEXAI_MODEL_HUB.get_text_generation_model(self.model)
-    predict_options = dict(
-        temperature=self.sampling_options.temperature,
-        top_k=self.sampling_options.top_k,
-        top_p=self.sampling_options.top_p,
-        max_output_tokens=self.sampling_options.max_tokens,
-        stop_sequences=self.sampling_options.stop,
-    )
-    response = model.predict(prompt.text, **predict_options)
-    return lf.LMSamplingResult([
-        # Scoring is not supported.
-        lf.LMSample(lf.AIMessage(response.text), score=0.0)
-    ])
-
-  def _sample_endpoint_model(self, prompt: lf.Message) -> lf.LMSamplingResult:
-    """Samples a text generation model."""
-    assert aiplatform_models is not None
-    model = aiplatform_models.Endpoint(self.endpoint_name)
-    # TODO(chengrun): Add support for stop_sequences.
-    predict_options = dict(
-        temperature=self.sampling_options.temperature
-        if self.sampling_options.temperature is not None
-        else 1.0,
-        top_k=self.sampling_options.top_k
-        if self.sampling_options.top_k is not None
-        else 32,
-        top_p=self.sampling_options.top_p
-        if self.sampling_options.top_p is not None
-        else 1,
-        max_tokens=self.sampling_options.max_tokens
-        if self.sampling_options.max_tokens is not None
-        else 8192,
+        ),
     )
-    instances = [{'prompt': prompt.text, **predict_options}]
-    response = model.predict(instances=instances)
-
-    return lf.LMSamplingResult([
-        # Scoring is not supported.
-        lf.LMSample(
-            self._generation_endpoint_response_to_message(response), score=0.0
-        )
-    ])
-
-
-class _ModelHub:
-  """Vertex AI model hub."""
-
-  def __init__(self):
-    self._generative_model_cache = {}
-    self._text_generation_model_cache = {}
-
-  def get_generative_model(
-      self, model_id: str
-  ) -> Any:  # generative_models.GenerativeModel:
-    """Gets a generative model by model id."""
-    model = self._generative_model_cache.get(model_id, None)
-    if model is None:
-      assert generative_models is not None
-      model = generative_models.GenerativeModel(model_id)
-      self._generative_model_cache[model_id] = model
-    return model
 
-  def get_text_generation_model(
-      self, model_id: str
-  ) -> Any:  # language_models.TextGenerationModel
-    """Gets a text generation model by model id."""
-    model = self._text_generation_model_cache.get(model_id, None)
-    if model is None:
-      assert language_models is not None
-      model = language_models.TextGenerationModel.from_pretrained(model_id)
-      self._text_generation_model_cache[model_id] = model
-    return model
-
-
-_VERTEXAI_MODEL_HUB = _ModelHub()
+  def _message_from_content_parts(
+      self, parts: list[dict[str, Any]]
+  ) -> lf.Message:
+    """Converts Vertex AI's content parts protocol to message."""
+    chunks = []
+    for part in parts:
+      if text_part := part.get('text'):
+        chunks.append(text_part)
+      else:
+        raise ValueError(f'Unsupported part: {part}')
+    return lf.AIMessage.from_chunks(chunks)
 
 
 _IMAGE_TYPES = [
@@ -535,12 +396,6 @@ class VertexAIGemini1_5(VertexAI):  # pylint: disable=invalid-name
   )
 
 
-class VertexAIGeminiPro1_5_Latest(VertexAIGemini1_5):  # pylint: disable=invalid-name
-  """Vertex AI Gemini 1.5 Pro model."""
-
-  model = 'gemini-1.5-pro-latest'
-
-
 class VertexAIGeminiPro1_5(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
@@ -571,12 +426,6 @@ class VertexAIGeminiPro1_5_0409(VertexAIGemini1_5):  # pylint: disable=invalid-n
   model = 'gemini-1.5-pro-preview-0409'
 
 
-class VertexAIGeminiFlash1_5_Latest(VertexAIGemini1_5):  # pylint: disable=invalid-name
-  """Vertex AI Gemini 1.5 Flash model."""
-
-  model = 'gemini-1.5-flash-latest'
-
-
 class VertexAIGeminiFlash1_5(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
@@ -608,7 +457,7 @@ class VertexAIGeminiPro1(VertexAI):  # pylint: disable=invalid-name
 
 
 class VertexAIGeminiPro1Vision(VertexAI):  # pylint: disable=invalid-name
-  """Vertex AI Gemini 1.0 Pro model."""
+  """Vertex AI Gemini 1.0 Pro Vision model."""
 
   model = 'gemini-1.0-pro-vision'
   supported_modalities: pg.typing.List(str).freeze(   # pytype: disable=invalid-annotation
@@ -616,19 +465,17 @@ class VertexAIGeminiPro1Vision(VertexAI):  # pylint: disable=invalid-name
   )
 
 
-class VertexAIPalm2(VertexAI):  # pylint: disable=invalid-name
-  """Vertex AI PaLM2 text generation model."""
-
-  model = 'text-bison'
-
+class VertexAIEndpoint(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Endpoint model."""
 
-class VertexAIPalm2_32K(VertexAI):  # pylint: disable=invalid-name
-  """Vertex AI PaLM2 text generation model (32K context length)."""
+  model = 'vertexai-endpoint'
 
-  model = 'text-bison-32k'
+  endpoint: Annotated[str, 'Vertex AI Endpoint ID.']
 
-
-class VertexAICustom(VertexAI):  # pylint: disable=invalid-name
-  """Vertex AI Custom model (Endpoint)."""
-
-  model = 'custom'
+  @property
+  def api_endpoint(self) -> str:
+    return (
+        f'https://{self.location}-aiplatform.googleapis.com/v1/projects/'
+        f'{self.project}/locations/{self.location}/'
+        f'endpoints/{self.endpoint}:generateContent'
+    )