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

langfun/core/llms/vertexai.py

@@ -13,17 +13,21 @@
 # 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
+  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
@@ -32,6 +36,8 @@ try:
 
   Credentials = credentials_lib.Credentials
 except ImportError:
+  google_auth = None
+  auth_requests = None
   credentials_lib = None   # pylint: disable=invalid-name
   vertexai = None
   generative_models = None
@@ -449,6 +455,238 @@ class VertexAI(lf.LanguageModel):
     ])
 
 
+@lf.use_init_args(['model'])
+@pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
+class VertexAIRest(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 with REST API support. See '
+          'https://cloud.google.com/vertex-ai/generative-ai/docs/'
+          'model-reference/inference#supported-models'
+          ' for details.'
+      ),
+  ]
+
+  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()
+    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.'
+      )
+
+    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.'
+      )
+
+    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
+
+  @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
+
+  @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
+  ) -> 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}.'
+        )
+      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
+
+  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):
+        parts.append({'text': 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:
+              parts.append({'text': modality.to_text()})
+            else:
+              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 dict(role='user', parts=parts)
+
+  def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
+    messages = [
+        self._message_from_content_parts(candidate['content']['parts'])
+        for candidate in json['candidates']
+    ]
+    usage = json['usageMetadata']
+    input_tokens = usage['promptTokenCount']
+    output_tokens = usage['candidatesTokenCount']
+    return lf.LMSamplingResult(
+        [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,
+            ),
+        ),
+    )
+
+  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)
+
+
 class _ModelHub:
   """Vertex AI model hub."""
 
@@ -547,13 +785,21 @@ class VertexAIGeminiPro1_5(VertexAIGemini1_5):  # pylint: disable=invalid-name
   model = 'gemini-1.5-pro'
 
 
-class VertexAIGeminiPro1_5_002(VertexAIGemini1_5):  # pylint: disable=invalid-name
+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
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-002'
 
 
-class VertexAIGeminiPro1_5_001(VertexAIGemini1_5):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1_5_001(VertexAIRestGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-001'
@@ -583,13 +829,13 @@ class VertexAIGeminiFlash1_5(VertexAIGemini1_5):  # pylint: disable=invalid-name
   model = 'gemini-1.5-flash'
 
 
-class VertexAIGeminiFlash1_5_002(VertexAIGemini1_5):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash1_5_002(VertexAIRestGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-002'
 
 
-class VertexAIGeminiFlash1_5_001(VertexAIGemini1_5):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash1_5_001(VertexAIRestGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-001'
@@ -601,14 +847,14 @@ class VertexAIGeminiFlash1_5_0514(VertexAIGemini1_5):  # pylint: disable=invalid
   model = 'gemini-1.5-flash-preview-0514'
 
 
-class VertexAIGeminiPro1(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1(VertexAIRest):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.0 Pro model."""
 
   model = 'gemini-1.0-pro'
 
 
 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

langfun/core/llms/vertexai_test.py

@@ -14,6 +14,7 @@
 """Tests for Gemini models."""
 
 import os
+from typing import Any
 import unittest
 from unittest import mock
 
@@ -23,6 +24,7 @@ import langfun.core as lf
 from langfun.core import modalities as lf_modalities
 from langfun.core.llms import vertexai
 import pyglove as pg
+import requests
 
 
 example_image = (
@@ -64,6 +66,40 @@ def mock_generate_content(content, generation_config, **kwargs):
   })
 
 
+def mock_requests_post(url: str, json: dict[str, Any], **kwargs):
+  del url, kwargs
+  c = pg.Dict(json['generationConfig'])
+  content = json['contents'][0]['parts'][0]['text']
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str({
+      'candidates': [
+          {
+              'content': {
+                  'role': 'model',
+                  'parts': [
+                      {
+                          'text': (
+                              f'This is a response to {content} with '
+                              f'temperature={c.temperature}, '
+                              f'top_p={c.topP}, '
+                              f'top_k={c.topK}, '
+                              f'max_tokens={c.maxOutputTokens}, '
+                              f'stop={"".join(c.stopSequences)}.'
+                          )
+                      },
+                  ],
+              },
+          },
+      ],
+      'usageMetadata': {
+          'promptTokenCount': 3,
+          'candidatesTokenCount': 4,
+      }
+  }).encode()
+  return response
+
+
 def mock_endpoint_predict(instances, **kwargs):
   del kwargs
   assert len(instances) == 1
@@ -83,7 +119,7 @@ class VertexAITest(unittest.TestCase):
 
   def test_content_from_message_text_only(self):
     text = 'This is a beautiful day'
-    model = vertexai.VertexAIGeminiPro1()
+    model = vertexai.VertexAIGeminiPro1Vision()
     chunks = model._content_from_message(lf.UserMessage(text))
     self.assertEqual(chunks, [text])
 
@@ -95,7 +131,7 @@ class VertexAITest(unittest.TestCase):
 
     # Non-multimodal model.
     with self.assertRaisesRegex(lf.ModalityError, 'Unsupported modality'):
-      vertexai.VertexAIGeminiPro1()._content_from_message(message)
+      vertexai.VertexAIPalm2()._content_from_message(message)
 
     model = vertexai.VertexAIGeminiPro1Vision()
     chunks = model._content_from_message(message)
@@ -119,7 +155,7 @@ class VertexAITest(unittest.TestCase):
             },
         ],
     })
-    model = vertexai.VertexAIGeminiPro1()
+    model = vertexai.VertexAIGeminiPro1Vision()
     message = model._generation_response_to_message(response)
     self.assertEqual(message, lf.AIMessage('hello world'))
 
@@ -158,25 +194,25 @@ class VertexAITest(unittest.TestCase):
 
   def test_project_and_location_check(self):
     with self.assertRaisesRegex(ValueError, 'Please specify `project`'):
-      _ = vertexai.VertexAIGeminiPro1()._api_initialized
+      _ = vertexai.VertexAIGeminiPro1Vision()._api_initialized
 
     with self.assertRaisesRegex(ValueError, 'Please specify `location`'):
-      _ = vertexai.VertexAIGeminiPro1(project='abc')._api_initialized
+      _ = vertexai.VertexAIGeminiPro1Vision(project='abc')._api_initialized
 
     self.assertTrue(
-        vertexai.VertexAIGeminiPro1(
+        vertexai.VertexAIGeminiPro1Vision(
             project='abc', location='us-central1'
         )._api_initialized
     )
 
     os.environ['VERTEXAI_PROJECT'] = 'abc'
     os.environ['VERTEXAI_LOCATION'] = 'us-central1'
-    self.assertTrue(vertexai.VertexAIGeminiPro1()._api_initialized)
+    self.assertTrue(vertexai.VertexAIGeminiPro1Vision()._api_initialized)
     del os.environ['VERTEXAI_PROJECT']
     del os.environ['VERTEXAI_LOCATION']
 
   def test_generation_config(self):
-    model = vertexai.VertexAIGeminiPro1()
+    model = vertexai.VertexAIGeminiPro1Vision()
     json_schema = {
         'type': 'object',
         'properties': {
@@ -245,7 +281,9 @@ class VertexAITest(unittest.TestCase):
       ) as mock_generate:
         mock_generate.side_effect = mock_generate_content
 
-        lm = vertexai.VertexAIGeminiPro1(project='abc', location='us-central1')
+        lm = vertexai.VertexAIGeminiPro1Vision(
+            project='abc', location='us-central1'
+        )
         self.assertEqual(
             lm(
                 'hello',
@@ -328,5 +366,163 @@ class VertexAITest(unittest.TestCase):
         )
 
 
+class VertexRestfulAITest(unittest.TestCase):
+  """Tests for Vertex model with REST API."""
+
+  def test_content_from_message_text_only(self):
+    text = 'This is a beautiful day'
+    model = vertexai.VertexAIGeminiPro1_5_002()
+    chunks = model._content_from_message(lf.UserMessage(text))
+    self.assertEqual(chunks, {'role': 'user', 'parts': [{'text': text}]})
+
+  def test_content_from_message_mm(self):
+    message = lf.UserMessage(
+        'This is an <<[[image]]>>, what is it?',
+        image=lf_modalities.Image.from_bytes(example_image),
+    )
+
+    # Non-multimodal model.
+    with self.assertRaisesRegex(lf.ModalityError, 'Unsupported modality'):
+      vertexai.VertexAIGeminiPro1()._content_from_message(message)
+
+    model = vertexai.VertexAIGeminiPro1Vision()
+    chunks = model._content_from_message(message)
+    self.maxDiff = None
+    self.assertEqual([chunks[0], chunks[2]], ['This is an', ', what is it?'])
+    self.assertIsInstance(chunks[1], generative_models.Part)
+
+  def test_generation_response_to_message_text_only(self):
+    response = generative_models.GenerationResponse.from_dict({
+        'candidates': [
+            {
+                'index': 0,
+                'content': {
+                    'role': 'model',
+                    'parts': [
+                        {
+                            'text': 'hello world',
+                        },
+                    ],
+                },
+            },
+        ],
+    })
+    model = vertexai.VertexAIGeminiPro1Vision()
+    message = model._generation_response_to_message(response)
+    self.assertEqual(message, lf.AIMessage('hello world'))
+
+  def test_model_hub(self):
+    with mock.patch(
+        'vertexai.generative_models.'
+        'GenerativeModel.__init__'
+    ) as mock_model_init:
+      mock_model_init.side_effect = lambda *args, **kwargs: None
+      model = vertexai._VERTEXAI_MODEL_HUB.get_generative_model(
+          'gemini-1.0-pro'
+      )
+      self.assertIsNotNone(model)
+      self.assertIs(
+          vertexai._VERTEXAI_MODEL_HUB.get_generative_model('gemini-1.0-pro'),
+          model,
+      )
+
+  @mock.patch.object(vertexai.VertexAIRest, 'credentials', new=True)
+  def test_project_and_location_check(self):
+    with self.assertRaisesRegex(ValueError, 'Please specify `project`'):
+      _ = vertexai.VertexAIGeminiPro1()._api_initialized
+
+    with self.assertRaisesRegex(ValueError, 'Please specify `location`'):
+      _ = vertexai.VertexAIGeminiPro1(project='abc')._api_initialized
+
+    self.assertTrue(
+        vertexai.VertexAIGeminiPro1(
+            project='abc', location='us-central1'
+        )._api_initialized
+    )
+
+    os.environ['VERTEXAI_PROJECT'] = 'abc'
+    os.environ['VERTEXAI_LOCATION'] = 'us-central1'
+    self.assertTrue(vertexai.VertexAIGeminiPro1()._api_initialized)
+    del os.environ['VERTEXAI_PROJECT']
+    del os.environ['VERTEXAI_LOCATION']
+
+  def test_generation_config(self):
+    model = vertexai.VertexAIGeminiPro1()
+    json_schema = {
+        'type': 'object',
+        'properties': {
+            'name': {'type': 'string'},
+        },
+        'required': ['name'],
+        'title': 'Person',
+    }
+    actual = model._generation_config(
+        lf.UserMessage('hi', json_schema=json_schema),
+        lf.LMSamplingOptions(
+            temperature=2.0,
+            top_p=1.0,
+            top_k=20,
+            max_tokens=1024,
+            stop=['\n'],
+        ),
+    )
+    self.assertEqual(
+        actual,
+        dict(
+            candidateCount=1,
+            temperature=2.0,
+            topP=1.0,
+            topK=20,
+            maxOutputTokens=1024,
+            stopSequences=['\n'],
+            responseLogprobs=False,
+            logprobs=None,
+            seed=None,
+            responseMimeType='application/json',
+            responseSchema={
+                'type': 'object',
+                'properties': {
+                    'name': {'type': 'string'}
+                },
+                'required': ['name'],
+                'title': 'Person',
+            }
+        ),
+    )
+    with self.assertRaisesRegex(
+        ValueError, '`json_schema` must be a dict, got'
+    ):
+      model._generation_config(
+          lf.UserMessage('hi', json_schema='not a dict'),
+          lf.LMSamplingOptions(),
+      )
+
+  @mock.patch.object(vertexai.VertexAIRest, 'credentials', new=True)
+  def test_call_model(self):
+    with mock.patch('requests.Session.post') as mock_generate:
+      mock_generate.side_effect = mock_requests_post
+
+      lm = vertexai.VertexAIGeminiPro1_5_002(
+          project='abc', location='us-central1'
+      )
+      r = lm(
+          'hello',
+          temperature=2.0,
+          top_p=1.0,
+          top_k=20,
+          max_tokens=1024,
+          stop='\n',
+      )
+      self.assertEqual(
+          r.text,
+          (
+              'This is a response to hello with temperature=2.0, '
+              'top_p=1.0, top_k=20, max_tokens=1024, stop=\n.'
+          ),
+      )
+      self.assertEqual(r.metadata.usage.prompt_tokens, 3)
+      self.assertEqual(r.metadata.usage.completion_tokens, 4)
+
+
 if __name__ == '__main__':
   unittest.main()

{langfun-0.1.2.dev202412010804.dist-info → langfun-0.1.2.dev202412030000.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.1.2.dev202412010804
+Version: 0.1.2.dev202412030000
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors
@@ -32,7 +32,6 @@ Requires-Dist: termcolor==1.1.0; extra == "all"
 Requires-Dist: tqdm>=4.64.1; extra == "all"
 Requires-Dist: google-cloud-aiplatform>=1.5.0; extra == "all"
 Requires-Dist: google-generativeai>=0.3.2; extra == "all"
-Requires-Dist: openai>=0.27.2; extra == "all"
 Requires-Dist: python-magic>=0.4.27; extra == "all"
 Requires-Dist: python-docx>=0.8.11; extra == "all"
 Requires-Dist: pillow>=10.0.0; extra == "all"
@@ -44,7 +43,6 @@ Requires-Dist: tqdm>=4.64.1; extra == "ui"
 Provides-Extra: llm
 Requires-Dist: google-cloud-aiplatform>=1.5.0; extra == "llm"
 Requires-Dist: google-generativeai>=0.3.2; extra == "llm"
-Requires-Dist: openai>=0.27.2; extra == "llm"
 Provides-Extra: llm-google
 Requires-Dist: google-cloud-aiplatform>=1.5.0; extra == "llm-google"
 Requires-Dist: google-generativeai>=0.3.2; extra == "llm-google"
@@ -52,8 +50,6 @@ Provides-Extra: llm-google-vertex
 Requires-Dist: google-cloud-aiplatform>=1.5.0; extra == "llm-google-vertex"
 Provides-Extra: llm-google-genai
 Requires-Dist: google-generativeai>=0.3.2; extra == "llm-google-genai"
-Provides-Extra: llm-openai
-Requires-Dist: openai>=0.27.2; extra == "llm-openai"
 Provides-Extra: mime
 Requires-Dist: python-magic>=0.4.27; extra == "mime"
 Requires-Dist: python-docx>=0.8.11; extra == "mime"
@@ -214,7 +210,6 @@ If you want to customize your installation, you can select specific features usi
 | llm-google          | All supported Google-powered LLMs.       |
 | llm-google-vertexai | LLMs powered by Google Cloud VertexAI    |
 | llm-google-genai    | LLMs powered by Google Generative AI API |
-| llm-openai          | LLMs powered by OpenAI                   |
 | mime                | All MIME supports.                       |
 | mime-auto           | Automatic MIME type detection.           |
 | mime-docx           | DocX format support.                     |