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

langfun/core/llms/openai_test.py

@@ -14,80 +14,20 @@
 """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=lf.LMSamplingUsage(
-          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=lf.LMSamplingUsage(
-          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=lf.LMSamplingUsage(
-          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())
+
+  def test_key(self):
+    with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
+      openai.Gpt4()('hi')
+
   def test_model_id(self):
     self.assertEqual(
         openai.Gpt35(api_key='test_key').model_id, 'OpenAI(text-davinci-003)')
@@ -97,352 +37,47 @@ 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.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,
-            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.',
+    with self.assertRaisesRegex(RuntimeError, '`logprobs` is not supported.*'):
+      openai.GptO1Preview(api_key='test_key')._request_args(
+          lf.LMSamplingOptions(
+              temperature=1.0, logprobs=True
+          )
       )
 
-  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_1 = openai.Gpt4Turbo(api_key='test_key')
-      lm_2 = openai.Gpt4VisionPreview(api_key='test_key')
-      for lm in (lm_1, lm_2):
-        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],
-        lf.LMSamplingResult(
-            [
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 0 for prompt 0.',
-                        score=0.0,
-                        logprobs=None,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 1 for prompt 0.',
-                        score=0.1,
-                        logprobs=None,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.1,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 2 for prompt 0.',
-                        score=0.2,
-                        logprobs=None,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.2,
-                    logprobs=None,
-                ),
-            ],
-            usage=lf.LMSamplingUsage(
-                prompt_tokens=100, completion_tokens=100, total_tokens=200
-            ),
-        ),
-    )
-    self.assertEqual(
-        results[1],
-        lf.LMSamplingResult(
-            [
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 0 for prompt 1.',
-                        score=0.0,
-                        logprobs=None,
-                        usage=None,
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 1 for prompt 1.',
-                        score=0.1,
-                        logprobs=None,
-                        usage=None,
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.1,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 2 for prompt 1.',
-                        score=0.2,
-                        logprobs=None,
-                        usage=None,
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.2,
-                    logprobs=None,
-                ),
-            ],
-        ),
-    )
-
-  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)
-      )
-
-    self.assertEqual(len(results), 2)
-    self.assertEqual(
-        results[0],
-        lf.LMSamplingResult(
-            [
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 0 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 1 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 2 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-            ],
-            usage=lf.LMSamplingUsage(
-                prompt_tokens=100, completion_tokens=100, total_tokens=200
-            ),
-        ),
-    )
-    self.assertEqual(
-        results[1],
-        lf.LMSamplingResult(
-            [
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 0 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 1 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 2 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-            ],
-            usage=lf.LMSamplingUsage(
-                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)
+  def test_estimate_cost(self):
     self.assertEqual(
-        results[0],
-        lf.LMSamplingResult(
-            [
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 0 for prompt 0.',
-                        score=0.0,
-                        logprobs=None,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=50,
-                            completion_tokens=50,
-                            total_tokens=100,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 1 for prompt 0.',
-                        score=0.1,
-                        logprobs=None,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=50,
-                            completion_tokens=50,
-                            total_tokens=100,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.1,
-                    logprobs=None,
-                ),
-            ],
-            usage=lf.LMSamplingUsage(
-                prompt_tokens=100, completion_tokens=100, total_tokens=200
-            ),
+        openai.Gpt4(api_key='test_key').estimate_cost(
+            num_input_tokens=100, num_output_tokens=100
         ),
+        0.009
     )
 
 

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()