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

langfun/core/llms/vertexai.py

@@ -28,21 +28,13 @@ try:
   from google import auth as google_auth
   from google.auth import credentials as credentials_lib
   from google.auth.transport import requests as auth_requests
-  import vertexai
-  from google.cloud.aiplatform import models as aiplatform_models
-  from vertexai import generative_models
-  from vertexai import language_models
   # pylint: enable=g-import-not-at-top
 
   Credentials = credentials_lib.Credentials
 except ImportError:
   google_auth = None
+  credentials_lib = None
   auth_requests = None
-  credentials_lib = None   # pylint: disable=invalid-name
-  vertexai = None
-  generative_models = None
-  language_models = None
-  aiplatform_models = None
   Credentials = Any
 
 
@@ -56,408 +48,72 @@ 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."""
-
-  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.'
-      ),
-  ]
-
-  endpoint_name: pg.typing.Annotated[
-      str | None,
-      'Vertex Endpoint name or ID.',
-  ]
-
-  project: Annotated[
-      str | None,
-      (
-          'Vertex AI project ID. Or set from environment variable '
-          'VERTEXAI_PROJECT.'
-      ),
-  ] = None
-
-  location: Annotated[
-      str | None,
-      (
-          'Vertex AI service location. Or set from environment variable '
-          'VERTEXAI_LOCATION.'
-      ),
-  ] = None
-
-  credentials: Annotated[
-      Credentials | None,
-      (
-          'Credentials to use. If None, the default credentials to the '
-          'environment will be used.'
-      ),
-  ] = None
-
-  supported_modalities: Annotated[
-      list[str],
-      'A list of MIME types for supported modalities'
-  ] = []
-
-  def _on_bound(self):
-    super()._on_bound()
-    self.__dict__.pop('_api_initialized', None)
-    if generative_models is None:
-      raise RuntimeError(
-          'Please install "langfun[llm-google-vertex]" to use Vertex AI models.'
-      )
-
-  @functools.cached_property
-  def _api_initialized(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.'
-      )
-
-    location = self.location or os.environ.get('VERTEXAI_LOCATION', None)
-    if not location:
-      raise ValueError(
-          'Please specify `location` during `__init__` or set environment '
-          'variable `VERTEXAI_LOCATION` with your Vertex AI service location.'
-      )
-
-    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
-
-  @property
-  def max_concurrency(self) -> int:
-    """Returns the maximum number of concurrent requests."""
-    return self.rate_to_max_concurrency(
-        requests_per_min=SUPPORTED_MODELS_AND_SETTINGS[self.model].rpm,
-        tokens_per_min=0,
-    )
-
-  def estimate_cost(
-      self,
-      num_input_tokens: int,
-      num_output_tokens: int
-  ) -> float | None:
-    """Estimate the cost based on usage."""
-    cost_per_1k_input_chars = SUPPORTED_MODELS_AND_SETTINGS[self.model].get(
-        'cost_per_1k_input_chars', None
-    )
-    cost_per_1k_output_chars = SUPPORTED_MODELS_AND_SETTINGS[self.model].get(
-        'cost_per_1k_output_chars', None
-    )
-    if cost_per_1k_output_chars is None or cost_per_1k_input_chars is None:
-      return None
-    return (
-        cost_per_1k_input_chars * num_input_tokens
-        + cost_per_1k_output_chars * num_output_tokens
-    ) * AVGERAGE_CHARS_PER_TOEKN / 1000
-
-  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:
-      if not isinstance(json_schema, dict):
-        raise ValueError(
-            f'`json_schema` must be a dict, got {json_schema!r}.'
-        )
-      response_mime_type = '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 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 = []
-
-    for lf_chunk in prompt.chunk():
-      if isinstance(lf_chunk, str):
-        chunks.append(lf_chunk)
-      elif isinstance(lf_chunk, lf_modalities.Mime):
-        try:
-          modalities = lf_chunk.make_compatible(
-              self.supported_modalities + ['text/plain']
-          )
-          if isinstance(modalities, lf_modalities.Mime):
-            modalities = [modalities]
-          for modality in modalities:
-            if modality.is_text:
-              chunk = modality.to_text()
-            else:
-              chunk = generative_models.Part.from_data(
-                  modality.to_bytes(), modality.mime_type
-              )
-            chunks.append(chunk)
-        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)
-
-  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'),
-    ]
-
-    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,
-        ),
-    )
-    return lf.LMSamplingResult(
-        [
-            # Scoring is not supported.
-            lf.LMSample(
-                self._generation_response_to_message(response), score=0.0
-            ),
-        ],
-        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
-        )
-    ])
-
-
 @lf.use_init_args(['model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
-class VertexAIRest(rest.REST):
+class VertexAI(rest.REST):
   """Language model served on VertexAI with REST API."""
 
   model: pg.typing.Annotated[
@@ -687,39 +343,6 @@ class VertexAIRest(rest.REST):
     return lf.AIMessage.from_chunks(chunks)
 
 
-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()
-
-
 _IMAGE_TYPES = [
     'image/png',
     'image/jpeg',
@@ -773,33 +396,19 @@ 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."""
 
   model = 'gemini-1.5-pro'
 
 
-class VertexAIRestGemini1_5(VertexAIRest):  # pylint: disable=invalid-name
-  """Vertex AI Gemini 1.5 model with REST API."""
-
-  supported_modalities: pg.typing.List(str).freeze(   # pytype: disable=invalid-annotation
-      _DOCUMENT_TYPES + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
-  )
-
-
-class VertexAIGeminiPro1_5_002(VertexAIRestGemini1_5):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1_5_002(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-002'
 
 
-class VertexAIGeminiPro1_5_001(VertexAIRestGemini1_5):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1_5_001(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-001'
@@ -817,25 +426,19 @@ 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."""
 
   model = 'gemini-1.5-flash'
 
 
-class VertexAIGeminiFlash1_5_002(VertexAIRestGemini1_5):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash1_5_002(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-002'
 
 
-class VertexAIGeminiFlash1_5_001(VertexAIRestGemini1_5):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash1_5_001(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-001'
@@ -847,7 +450,7 @@ class VertexAIGeminiFlash1_5_0514(VertexAIGemini1_5):  # pylint: disable=invalid
   model = 'gemini-1.5-flash-preview-0514'
 
 
-class VertexAIGeminiPro1(VertexAIRest):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.0 Pro model."""
 
   model = 'gemini-1.0-pro'
@@ -862,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 VertexAIPalm2_32K(VertexAI):  # pylint: disable=invalid-name
-  """Vertex AI PaLM2 text generation model (32K context length)."""
+class VertexAIEndpoint(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Endpoint model."""
 
-  model = 'text-bison-32k'
+  model = 'vertexai-endpoint'
 
+  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'
+    )