langfun 0.0.2.dev20240330__py3-none-any.whl → 0.1.2.dev202501140804__py3-none-any.whl

langfun/core/llms/openai_test.py

@@ -11,73 +11,22 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for openai models."""
+"""Tests for OpenAI models."""
 
 import unittest
-from unittest import mock
-
 import langfun.core as lf
-from langfun.core import modalities as lf_modalities
 from langfun.core.llms import openai
-import pyglove as pg
-
-
-def mock_completion_query(prompt, *, n=1, **kwargs):
-  del kwargs
-  choices = []
-  for i, _ in enumerate(prompt):
-    for k in range(n):
-      choices.append(pg.Dict(
-          index=i,
-          text=f'Sample {k} for prompt {i}.',
-          logprobs=k / 10,
-      ))
-  return pg.Dict(choices=choices, usage=openai.Usage(
-      prompt_tokens=100,
-      completion_tokens=100,
-      total_tokens=200,
-  ))
-
-
-def mock_chat_completion_query(messages, *, n=1, **kwargs):
-  del messages, kwargs
-  choices = []
-  for k in range(n):
-    choices.append(pg.Dict(
-        message=pg.Dict(
-            content=f'Sample {k} for message.'
-        ),
-        logprobs=None,
-    ))
-  return pg.Dict(choices=choices, usage=openai.Usage(
-      prompt_tokens=100,
-      completion_tokens=100,
-      total_tokens=200,
-  ))
 
 
-def mock_chat_completion_query_vision(messages, *, n=1, **kwargs):
-  del kwargs
-  choices = []
-  urls = [
-      c['image_url'] for c in messages[0]['content'] if c['type'] == 'image_url'
-  ]
-  for k in range(n):
-    choices.append(pg.Dict(
-        message=pg.Dict(
-            content=f'Sample {k} for message: {"".join(urls)}'
-        ),
-        logprobs=None,
-    ))
-  return pg.Dict(choices=choices, usage=openai.Usage(
-      prompt_tokens=100,
-      completion_tokens=100,
-      total_tokens=200,
-  ))
+class OpenAITest(unittest.TestCase):
+  """Tests for OpenAI language model."""
 
+  def test_dir(self):
+    self.assertIn('gpt-4-turbo', openai.OpenAI.dir())
 
-class OpenaiTest(unittest.TestCase):
-  """Tests for OpenAI language model."""
+  def test_key(self):
+    with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
+      openai.Gpt4()('hi')
 
   def test_model_id(self):
     self.assertEqual(
@@ -88,136 +37,48 @@ class OpenaiTest(unittest.TestCase):
         openai.Gpt35(api_key='test_key').resource_id, 'OpenAI(text-davinci-003)'
     )
 
+  def test_headers(self):
+    self.assertEqual(
+        openai.Gpt35(api_key='test_key').headers,
+        {
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer test_key',
+        },
+    )
+
   def test_max_concurrency(self):
-    self.assertEqual(openai.Gpt35(api_key='test_key').max_concurrency, 8)
+    self.assertGreater(openai.Gpt35(api_key='test_key').max_concurrency, 0)
 
-  def test_get_request_args(self):
+  def test_request_args(self):
     self.assertEqual(
-        openai.Gpt35(api_key='test_key', timeout=90.0)._get_request_args(
+        openai.Gpt4(api_key='test_key')._request_args(
             lf.LMSamplingOptions(
-                temperature=2.0,
-                n=2,
-                max_tokens=4096,
-                top_p=1.0)),
-        dict(
-            engine='text-davinci-003',
-            logprobs=False,
-            top_logprobs=None,
-            n=2,
-            temperature=2.0,
-            max_tokens=4096,
-            stream=False,
-            timeout=90.0,
-            top_p=1.0,
-        )
-    )
-    self.assertEqual(
-        openai.Gpt4(api_key='test_key')._get_request_args(
-            lf.LMSamplingOptions(temperature=1.0, stop=['\n'], n=1)
+                temperature=1.0, stop=['\n'], n=1, random_seed=123
+            )
         ),
         dict(
             model='gpt-4',
-            logprobs=False,
             top_logprobs=None,
             n=1,
             temperature=1.0,
-            max_tokens=1024,
-            stream=False,
-            timeout=120.0,
             stop=['\n'],
+            seed=123,
         ),
     )
-
-  def test_call_completion(self):
-    with mock.patch('openai.Completion.create') as mock_completion:
-      mock_completion.side_effect = mock_completion_query
-      lm = openai.OpenAI(api_key='test_key', model='text-davinci-003')
-      self.assertEqual(
-          lm('hello', sampling_options=lf.LMSamplingOptions(n=2)),
-          'Sample 0 for prompt 0.',
-      )
-
-  def test_call_chat_completion(self):
-    with mock.patch('openai.ChatCompletion.create') as mock_chat_completion:
-      mock_chat_completion.side_effect = mock_chat_completion_query
-      lm = openai.OpenAI(api_key='test_key', model='gpt-4')
-      self.assertEqual(
-          lm('hello', sampling_options=lf.LMSamplingOptions(n=2)),
-          'Sample 0 for message.',
-      )
-
-  def test_call_chat_completion_vision(self):
-    with mock.patch('openai.ChatCompletion.create') as mock_chat_completion:
-      mock_chat_completion.side_effect = mock_chat_completion_query_vision
-      lm = openai.Gpt4TurboVision(api_key='test_key')
-      self.assertEqual(
-          lm(
-              lf.UserMessage(
-                  'hello {{image}}',
-                  image=lf_modalities.Image.from_uri('https://fake/image')
-              ),
-              sampling_options=lf.LMSamplingOptions(n=2)
-          ),
-          'Sample 0 for message: https://fake/image',
-      )
-
-  def test_sample_completion(self):
-    with mock.patch('openai.Completion.create') as mock_completion:
-      mock_completion.side_effect = mock_completion_query
-      lm = openai.OpenAI(api_key='test_key', model='text-davinci-003')
-      results = lm.sample(
-          ['hello', 'bye'], sampling_options=lf.LMSamplingOptions(n=3)
-      )
-
-    self.assertEqual(len(results), 2)
-    self.assertEqual(results[0], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for prompt 0.', score=0.0),
-        lf.LMSample('Sample 1 for prompt 0.', score=0.1),
-        lf.LMSample('Sample 2 for prompt 0.', score=0.2),
-    ], usage=openai.Usage(
-        prompt_tokens=100, completion_tokens=100, total_tokens=200)))
-
-    self.assertEqual(results[1], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for prompt 1.', score=0.0),
-        lf.LMSample('Sample 1 for prompt 1.', score=0.1),
-        lf.LMSample('Sample 2 for prompt 1.', score=0.2),
-    ]))
-
-  def test_sample_chat_completion(self):
-    with mock.patch('openai.ChatCompletion.create') as mock_chat_completion:
-      mock_chat_completion.side_effect = mock_chat_completion_query
-      lm = openai.OpenAI(api_key='test_key', model='gpt-4')
-      results = lm.sample(
-          ['hello', 'bye'], sampling_options=lf.LMSamplingOptions(n=3)
+    with self.assertRaisesRegex(RuntimeError, '`logprobs` is not supported.*'):
+      openai.GptO1Preview(api_key='test_key')._request_args(
+          lf.LMSamplingOptions(
+              temperature=1.0, logprobs=True
+          )
       )
 
-    self.assertEqual(len(results), 2)
-    self.assertEqual(results[0], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for message.', score=0.0),
-        lf.LMSample('Sample 1 for message.', score=0.0),
-        lf.LMSample('Sample 2 for message.', score=0.0),
-    ], usage=openai.Usage(
-        prompt_tokens=100, completion_tokens=100, total_tokens=200)))
-    self.assertEqual(results[1], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for message.', score=0.0),
-        lf.LMSample('Sample 1 for message.', score=0.0),
-        lf.LMSample('Sample 2 for message.', score=0.0),
-    ], usage=openai.Usage(
-        prompt_tokens=100, completion_tokens=100, total_tokens=200)))
-
-  def test_sample_with_contextual_options(self):
-    with mock.patch('openai.Completion.create') as mock_completion:
-      mock_completion.side_effect = mock_completion_query
-      lm = openai.OpenAI(api_key='test_key', model='text-davinci-003')
-      with lf.use_settings(sampling_options=lf.LMSamplingOptions(n=2)):
-        results = lm.sample(['hello'])
-
-    self.assertEqual(len(results), 1)
-    self.assertEqual(results[0], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for prompt 0.', score=0.0),
-        lf.LMSample('Sample 1 for prompt 0.', score=0.1),
-    ], usage=openai.Usage(
-        prompt_tokens=100, completion_tokens=100, total_tokens=200)))
+  def test_estimate_cost(self):
+    self.assertEqual(
+        openai.Gpt4(api_key='test_key').estimate_cost(
+            num_input_tokens=100, num_output_tokens=100
+        ),
+        0.009
+    )
 
 
 if __name__ == '__main__':

langfun/core/llms/rest.py

@@ -0,0 +1,113 @@
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Base class for language models through REST APIs."""
+
+import functools
+from typing import Annotated, Any, Callable
+
+import langfun.core as lf
+import requests
+
+
+class REST(lf.LanguageModel):
+  """REST-based language model."""
+
+  api_endpoint: Annotated[
+      str,
+      'The endpoint of the REST API.'
+  ]
+
+  request: Annotated[
+      Callable[[lf.Message, lf.LMSamplingOptions], dict[str, Any]],
+      'A function to convert a Langfun message to a JSON request.'
+  ]
+
+  result: Annotated[
+      Callable[[dict[str, Any]], lf.LMSamplingResult],
+      'A function to convert a JSON response to an LMSamplingResult.'
+  ]
+
+  model: Annotated[
+      str | None,
+      'Model ID.'
+  ] = None
+
+  headers: Annotated[
+      dict[str, Any] | None,
+      'The headers for the REST API.'
+  ] = None
+
+  @property
+  def model_id(self) -> str:
+    """Returns a string to identify the model."""
+    return self.model or 'unknown'
+
+  @functools.cached_property
+  def _api_initialized(self) -> bool:
+    """Returns whether the API is initialized."""
+    self._initialize()
+    return True
+
+  def _initialize(self) -> None:
+    """Initializes the API. Subclasses can override."""
+
+  @functools.cached_property
+  def _session(self) -> requests.Session:
+    assert self._api_initialized
+    s = requests.Session()
+    s.headers.update(self.headers or {})
+    return s
+
+  def _on_bound(self):
+    super()._on_bound()
+    self.__dict__.pop('_session', None)
+    self.__dict__.pop('_api_initialized', None)
+
+  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
+    assert self._api_initialized
+    return self._parallel_execute_with_currency_control(
+        self._sample_single, prompts
+    )
+
+  def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
+    try:
+      response = self._session.post(
+          self.api_endpoint,
+          json=self.request(prompt, self.sampling_options),
+          timeout=self.timeout,
+      )
+      return self._parse_response(response)
+    except ConnectionError as e:
+      raise lf.LMError(str(e)) from e
+
+  def _error(self, status_code: int, content: str) -> lf.LMError:
+    if status_code == 429:
+      error_cls = lf.RateLimitError
+    elif status_code in (
+        500,  # Server side issue (might be bug).
+        502,  # Bad gateway (upstream issue, might retry).
+        503,  # Servers currently under load, retry after a brief wait.
+        529,  # Overloaded, retry after a brief wait.
+    ):
+      error_cls = lf.TemporaryLMError
+    else:
+      error_cls = lf.LMError
+    return error_cls(f'{status_code}: {content}')
+
+  def _parse_response(self, response: requests.Response) -> lf.LMSamplingResult:
+    """Parses Anthropic's response."""
+    if response.status_code == 200:
+      return self.result(response.json())
+    else:
+      raise self._error(response.status_code, response.content)

langfun/core/llms/rest_test.py

@@ -0,0 +1,111 @@
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for REST models."""
+
+from typing import Any
+import unittest
+from unittest import mock
+import langfun.core as lf
+from langfun.core.llms import rest
+import pyglove as pg
+import requests
+
+
+def mock_requests_post(url: str, json: dict[str, Any], **kwargs):
+  del url, kwargs
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str({
+      'content': [(
+          f'hello with temperature={json.get("temperature")}, '
+          f'top_k={json.get("top_k")}, '
+          f'top_p={json.get("top_p")}, '
+          f'max_tokens={json.get("max_tokens")}, '
+          f'stop={json.get("stop_sequences")}.'
+      )],
+  }).encode()
+  return response
+
+
+def mock_requests_post_error(status_code, error_type, error_message):
+  def _mock_requests(url: str, json: dict[str, Any], **kwargs):
+    del url, json, kwargs
+    response = requests.Response()
+    response.status_code = status_code
+    response._content = pg.to_json_str(
+        {
+            'error': {
+                'type': error_type,
+                'message': error_message,
+            }
+        }
+    ).encode()
+    return response
+
+  return _mock_requests
+
+
+class RestTest(unittest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self._lm = rest.REST(
+        api_endpoint='https://fake-api.com',
+        request=lambda x, o: dict(
+            model='test-model',
+            prompt=x.text,
+            temperature=0.0,
+            top_k=0.1,
+            top_p=0.2,
+            stop_sequences=['\n'],
+            max_tokens=4096,
+        ),
+        result=lambda x: lf.LMSamplingResult(
+            [lf.LMSample(c) for c in x['content']]),
+        headers=dict(api_key='fake_key'),
+    )
+
+  def test_call(self):
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_requests_post
+      self.assertEqual(self._lm.model_id, 'unknown')
+      response = self._lm(
+          'hello', temperature=0.0, top_k=0.1, top_p=0.2, stop=['\n'])
+      self.assertEqual(
+          response.text,
+          (
+              'hello with temperature=0.0, top_k=0.1, top_p=0.2, '
+              "max_tokens=4096, stop=['\\n']."
+          ),
+      )
+      self.assertIsInstance(response.usage, lf.UsageNotAvailable)
+
+  def test_call_errors(self):
+    for status_code, error_type, error_message in [
+        (429, 'rate_limit', 'Rate limit exceeded.'),
+        (529, 'service_unavailable', 'Service unavailable.'),
+        (500, 'bad_request', 'Bad request.'),
+    ]:
+      with mock.patch('requests.Session.post') as mock_mm_request:
+        mock_mm_request.side_effect = mock_requests_post_error(
+            status_code, error_type, error_message
+        )
+        with self.assertRaisesRegex(
+            Exception, f'.*{status_code}: .*{error_message}'
+        ):
+          self._lm('hello', max_attempts=1)
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/llms/vertexai.py

@@ -0,0 +1,192 @@
+# Copyright 2023 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.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Vertex AI generative models."""
+
+import functools
+import os
+from typing import Annotated, Any
+
+import langfun.core as lf
+from langfun.core.llms import gemini
+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
+  # pylint: enable=g-import-not-at-top
+
+  Credentials = credentials_lib.Credentials
+except ImportError:
+  google_auth = None
+  credentials_lib = None
+  auth_requests = None
+  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."""
+
+  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
+
+  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
+    self._location = location
+
+    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 model_id(self) -> str:
+    """Returns a string to identify the model."""
+    return f'VertexAI({self.model})'
+
+  @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 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/google/'
+        f'models/{self.model}:generateContent'
+    )
+
+
+class VertexAIGeminiFlash2_0ThinkingExp_20241219(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini Flash 2.0 Thinking model launched on 12/19/2024."""
+
+  api_version = 'v1alpha'
+  model = 'gemini-2.0-flash-thinking-exp-1219'
+  timeout = None
+
+
+class VertexAIGeminiFlash2_0Exp(VertexAI):  # 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
+  """Vertex AI Gemini Experimental model launched on 12/06/2024."""
+
+  model = 'gemini-exp-1206'
+
+
+class VertexAIGeminiExp_20241114(VertexAI):  # 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
+  """Vertex AI Gemini 1.5 Pro model."""
+
+  model = 'gemini-1.5-pro-latest'
+
+
+class VertexAIGeminiPro1_5_002(VertexAI):  # 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
+  """Vertex AI Gemini 1.5 Pro model."""
+
+  model = 'gemini-1.5-pro-001'
+
+
+class VertexAIGeminiFlash1_5(VertexAI):  # 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
+  """Vertex AI Gemini 1.5 Flash model."""
+
+  model = 'gemini-1.5-flash-002'
+
+
+class VertexAIGeminiFlash1_5_001(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.5 Flash model."""
+
+  model = 'gemini-1.5-flash-001'
+
+
+class VertexAIGeminiPro1(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.0 Pro model."""
+
+  model = 'gemini-1.0-pro'