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

langfun/core/llms/openai.py

@@ -13,34 +13,14 @@
 # limitations under the License.
 """Language models from OpenAI."""
 
-import collections
-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:
-  import openai  # pylint: disable=g-import-not-at-top
-
-  if hasattr(openai, 'error'):
-    # For lower versions.
-    ServiceUnavailableError = openai.error.ServiceUnavailableError
-    RateLimitError = openai.error.RateLimitError
-    APITimeoutError = (
-        openai.error.APIError,
-        '.*The server had an error processing your request'
-    )
-  else:
-    # For higher versions.
-    ServiceUnavailableError = getattr(openai, 'InternalServerError')
-    RateLimitError = getattr(openai, 'RateLimitError')
-    APITimeoutError = getattr(openai, 'APITimeoutError')
-except ImportError:
-  openai = None
-
 
 # From https://platform.openai.com/settings/organization/limits
 _DEFAULT_TPM = 250000
@@ -289,7 +269,7 @@ SUPPORTED_MODELS_AND_SETTINGS = {
         rpm=_DEFAULT_RPM,
         tpm=_DEFAULT_TPM
     ),
-    # GPT-3 instruction-tuned models
+    # GPT-3 instruction-tuned models (Deprecated)
     'text-curie-001': pg.Dict(
         in_service=False,
         rpm=_DEFAULT_RPM,
@@ -325,9 +305,9 @@ SUPPORTED_MODELS_AND_SETTINGS = {
         rpm=_DEFAULT_RPM,
         tpm=_DEFAULT_TPM
     ),
-    # GPT-3 base models
+    # GPT-3 base models that are still in service.
     'babbage-002': pg.Dict(
-        in_service=False,
+        in_service=True,
         rpm=_DEFAULT_RPM,
         tpm=_DEFAULT_TPM
     ),
@@ -340,7 +320,7 @@ SUPPORTED_MODELS_AND_SETTINGS = {
 
 
 @lf.use_init_args(['model'])
-class OpenAI(lf.LanguageModel):
+class OpenAI(rest.REST):
   """OpenAI model."""
 
   model: pg.typing.Annotated[
@@ -348,7 +328,9 @@ class OpenAI(lf.LanguageModel):
           pg.MISSING_VALUE, list(SUPPORTED_MODELS_AND_SETTINGS.keys())
       ),
       'The name of the model to use.',
-  ] = 'gpt-3.5-turbo'
+  ]
+
+  api_endpoint: str = 'https://api.openai.com/v1/chat/completions'
 
   multimodal: Annotated[
       bool,
@@ -372,27 +354,45 @@ class OpenAI(lf.LanguageModel):
       ),
   ] = None
 
+  project: Annotated[
+      str | None,
+      (
+          'Project. If None, the key will be read from environment '
+          "variable 'OPENAI_PROJECT'. Based on the value, usages from "
+          "these API requests will count against the project's quota. "
+      ),
+  ] = None
+
   def _on_bound(self):
     super()._on_bound()
-    self.__dict__.pop('_api_initialized', None)
-    if openai is None:
-      raise RuntimeError(
-          'Please install "langfun[llm-openai]" to use OpenAI models.'
-      )
+    self._api_key = None
+    self._organization = None
+    self._project = None
 
-  @functools.cached_property
-  def _api_initialized(self):
+  def _initialize(self):
     api_key = self.api_key or os.environ.get('OPENAI_API_KEY', None)
     if not api_key:
       raise ValueError(
           'Please specify `api_key` during `__init__` or set environment '
           'variable `OPENAI_API_KEY` with your OpenAI API key.'
       )
-    openai.api_key = api_key
-    org = self.organization or os.environ.get('OPENAI_ORGANIZATION', None)
-    if org:
-      openai.organization = org
-    return True
+    self._api_key = api_key
+    self._organization = self.organization or os.environ.get(
+        'OPENAI_ORGANIZATION', None
+    )
+    self._project = self.project or os.environ.get('OPENAI_PROJECT', None)
+
+  @property
+  def headers(self) -> dict[str, Any]:
+    headers = {
+        'Content-Type': 'application/json',
+        'Authorization': f'Bearer {self._api_key}',
+    }
+    if self._organization:
+      headers['OpenAI-Organization'] = self._organization
+    if self._project:
+      headers['OpenAI-Project'] = self._project
+    return headers
 
   @property
   def model_id(self) -> str:
@@ -428,23 +428,16 @@ class OpenAI(lf.LanguageModel):
 
   @classmethod
   def dir(cls):
-    assert openai is not None
-    return openai.Model.list()
+    return [k for k, v in SUPPORTED_MODELS_AND_SETTINGS.items() if v.in_service]
 
-  @property
-  def is_chat_model(self):
-    """Returns True if the model is a chat model."""
-    return self.model.startswith(('o1', 'gpt-4', 'gpt-3.5-turbo'))
-
-  def _get_request_args(
+  def _request_args(
       self, options: lf.LMSamplingOptions) -> dict[str, Any]:
     # Reference:
     # https://platform.openai.com/docs/api-reference/completions/create
     # NOTE(daiyip): options.top_k is not applicable.
     args = dict(
+        model=self.model,
         n=options.n,
-        stream=False,
-        timeout=self.timeout,
         top_logprobs=options.top_logprobs,
     )
     if options.logprobs:
@@ -453,13 +446,10 @@ class OpenAI(lf.LanguageModel):
         raise RuntimeError('`logprobs` is not supported on {self.model!r}.')
       args['logprobs'] = options.logprobs
 
-    # Completion and ChatCompletion uses different parameter name for model.
-    args['model' if self.is_chat_model else 'engine'] = self.model
-
     if options.temperature is not None:
       args['temperature'] = options.temperature
     if options.max_tokens is not None:
-      args['max_tokens'] = options.max_tokens
+      args['max_completion_tokens'] = options.max_tokens
     if options.top_p is not None:
       args['top_p'] = options.top_p
     if options.stop:
@@ -468,168 +458,113 @@ class OpenAI(lf.LanguageModel):
       args['seed'] = options.random_seed
     return args
 
-  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    assert self._api_initialized
-    if self.is_chat_model:
-      return self._chat_complete_batch(prompts)
-    else:
-      return self._complete_batch(prompts)
-
-  def _complete_batch(
-      self, prompts: list[lf.Message]
-  ) -> list[lf.LMSamplingResult]:
-
-    def _open_ai_completion(prompts):
-      assert openai is not None
-      response = openai.Completion.create(
-          prompt=[p.text for p in prompts],
-          **self._get_request_args(self.sampling_options),
-      )
-      # Parse response.
-      samples_by_index = collections.defaultdict(list)
-      for choice in response.choices:
-        samples_by_index[choice.index].append(
-            lf.LMSample(choice.text.strip(), score=choice.logprobs or 0.0)
-        )
-
-      n = len(samples_by_index)
-      estimated_cost = self.estimate_cost(
-          num_input_tokens=response.usage.prompt_tokens,
-          num_output_tokens=response.usage.completion_tokens,
-      )
-      usage = lf.LMSamplingUsage(
-          prompt_tokens=response.usage.prompt_tokens // n,
-          completion_tokens=response.usage.completion_tokens // n,
-          total_tokens=response.usage.total_tokens // n,
-          estimated_cost=(
-              None if estimated_cost is None else (estimated_cost // n)
-          )
-      )
-      return [
-          lf.LMSamplingResult(samples_by_index[index], usage=usage)
-          for index in sorted(samples_by_index.keys())
-      ]
-
-    return self._parallel_execute_with_currency_control(
-        _open_ai_completion,
-        [prompts],
-        retry_on_errors=(
-            ServiceUnavailableError,
-            RateLimitError,
-            APITimeoutError,
-        ),
-    )[0]
-
-  def _chat_complete_batch(
-      self, prompts: list[lf.Message]
-  ) -> list[lf.LMSamplingResult]:
-    def _content_from_message(message: lf.Message):
-      if self.multimodal:
-        content = []
-        for chunk in message.chunk():
-          if isinstance(chunk, str):
-            item = dict(type='text', text=chunk)
-          elif isinstance(chunk, lf_modalities.Image):
-            if chunk.uri and chunk.uri.lower().startswith(
-                ('http:', 'https:', 'ftp:')
-            ):
-              uri = chunk.uri
-            else:
-              uri = chunk.content_uri
-            item = dict(type='image_url', image_url=dict(url=uri))
-          else:
-            raise ValueError(f'Unsupported modality object: {chunk!r}.')
-          content.append(item)
+  def _content_from_message(self, message: lf.Message):
+    """Returns a OpenAI content object from a Langfun message."""
+    def _uri_from(chunk: lf.Modality) -> str:
+      if chunk.uri and chunk.uri.lower().startswith(
+          ('http:', 'https:', 'ftp:')
+      ):
+        return chunk.uri
+      return chunk.content_uri
+
+    content = []
+    for chunk in message.chunk():
+      if isinstance(chunk, str):
+        item = dict(type='text', text=chunk)
+      elif isinstance(chunk, lf_modalities.Image) and self.multimodal:
+        item = dict(type='image_url', image_url=dict(url=_uri_from(chunk)))
       else:
-        content = message.text
-      return content
-
-    def _open_ai_chat_completion(prompt: lf.Message):
-      request_args = self._get_request_args(self.sampling_options)
-      # Users could use `metadata_json_schema` to pass additional
-      # request arguments.
-      json_schema = prompt.metadata.get('json_schema')
-      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}.'
-          )
-        if 'title' not in json_schema:
-          raise ValueError(
-              f'The root of `json_schema` must have a `title` field, '
-              f'got {json_schema!r}.'
-          )
-        request_args.update(
-            response_format=dict(
-                type='json_schema',
-                json_schema=dict(
-                    schema=json_schema,
-                    name=json_schema['title'],
-                    strict=True,
-                )
-            )
-        )
-        prompt.metadata.formatted_text = (
-            prompt.text
-            + '\n\n [RESPONSE FORMAT (not part of prompt)]\n'
-            + pg.to_json_str(request_args['response_format'], json_indent=2)
-        )
+        raise ValueError(f'Unsupported modality: {chunk!r}.')
+      content.append(item)
+    return content
 
-      # Prepare messages.
-      messages = []
-      # Users could use `metadata_system_message` to pass system message.
-      system_message = prompt.metadata.get('system_message')
-      if system_message:
-        system_message = lf.SystemMessage.from_value(system_message)
-        messages.append(
-            dict(role='system', content=_content_from_message(system_message))
+  def request(
+      self,
+      prompt: lf.Message,
+      sampling_options: lf.LMSamplingOptions
+  ) -> dict[str, Any]:
+    """Returns the JSON input for a message."""
+    request_args = self._request_args(sampling_options)
+
+    # Users could use `metadata_json_schema` to pass additional
+    # request arguments.
+    json_schema = prompt.metadata.get('json_schema')
+    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}.'
         )
-      messages.append(dict(role='user', content=_content_from_message(prompt)))
-
-      assert openai is not None
-      response = openai.ChatCompletion.create(messages=messages, **request_args)
-
-      samples = []
-      for choice in response.choices:
-        logprobs = None
-        choice_logprobs = getattr(choice, 'logprobs', None)
-        if choice_logprobs:
-          logprobs = [
-              (
-                  t.token,
-                  t.logprob,
-                  [(tt.token, tt.logprob) for tt in t.top_logprobs],
-              )
-              for t in choice_logprobs.content
-          ]
-        samples.append(
-            lf.LMSample(
-                choice.message.content,
-                score=0.0,
-                logprobs=logprobs,
-            )
+      if 'title' not in json_schema:
+        raise ValueError(
+            f'The root of `json_schema` must have a `title` field, '
+            f'got {json_schema!r}.'
         )
-
-      return lf.LMSamplingResult(
-          samples=samples,
-          usage=lf.LMSamplingUsage(
-              prompt_tokens=response.usage.prompt_tokens,
-              completion_tokens=response.usage.completion_tokens,
-              total_tokens=response.usage.total_tokens,
-              estimated_cost=self.estimate_cost(
-                  num_input_tokens=response.usage.prompt_tokens,
-                  num_output_tokens=response.usage.completion_tokens,
+      request_args.update(
+          response_format=dict(
+              type='json_schema',
+              json_schema=dict(
+                  schema=json_schema,
+                  name=json_schema['title'],
+                  strict=True,
               )
-          ),
+          )
+      )
+      prompt.metadata.formatted_text = (
+          prompt.text
+          + '\n\n [RESPONSE FORMAT (not part of prompt)]\n'
+          + pg.to_json_str(request_args['response_format'], json_indent=2)
+      )
+
+    # Prepare messages.
+    messages = []
+    # Users could use `metadata_system_message` to pass system message.
+    system_message = prompt.metadata.get('system_message')
+    if system_message:
+      system_message = lf.SystemMessage.from_value(system_message)
+      messages.append(
+          dict(role='system',
+               content=self._content_from_message(system_message))
       )
+    messages.append(
+        dict(role='user', content=self._content_from_message(prompt))
+    )
+    request = dict()
+    request.update(request_args)
+    request['messages'] = messages
+    return request
 
-    return self._parallel_execute_with_currency_control(
-        _open_ai_chat_completion,
-        prompts,
-        retry_on_errors=(
-            ServiceUnavailableError,
-            RateLimitError,
-            APITimeoutError
+  def _parse_choice(self, choice: dict[str, Any]) -> lf.LMSample:
+    # Reference:
+    # https://platform.openai.com/docs/api-reference/chat/object
+    logprobs = None
+    choice_logprobs = choice.get('logprobs')
+    if choice_logprobs:
+      logprobs = [
+          (
+              t['token'],
+              t['logprob'],
+              [(tt['token'], tt['logprob']) for tt in t['top_logprobs']],
+          )
+          for t in choice_logprobs['content']
+      ]
+    return lf.LMSample(
+        choice['message']['content'],
+        score=0.0,
+        logprobs=logprobs,
+    )
+
+  def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
+    usage = json['usage']
+    return lf.LMSamplingResult(
+        samples=[self._parse_choice(choice) for choice in json['choices']],
+        usage=lf.LMSamplingUsage(
+            prompt_tokens=usage['prompt_tokens'],
+            completion_tokens=usage['completion_tokens'],
+            total_tokens=usage['total_tokens'],
+            estimated_cost=self.estimate_cost(
+                num_input_tokens=usage['prompt_tokens'],
+                num_output_tokens=usage['completion_tokens'],
+            )
         ),
     )
 

langfun/core/llms/openai_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Tests for OpenAI models."""
 
+from typing import Any
 import unittest
 from unittest import mock
 
@@ -20,86 +21,106 @@ import langfun.core as lf
 from langfun.core import modalities as lf_modalities
 from langfun.core.llms import openai
 import pyglove as pg
+import requests
 
 
-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):
+def mock_chat_completion_request(url: str, json: dict[str, Any], **kwargs):
+  del url, kwargs
+  messages = json['messages']
   if len(messages) > 1:
     system_message = f' system={messages[0]["content"]}'
   else:
     system_message = ''
 
-  if 'response_format' in kwargs:
-    response_format = f' format={kwargs["response_format"]["type"]}'
+  if 'response_format' in json:
+    response_format = f' format={json["response_format"]["type"]}'
   else:
     response_format = ''
 
   choices = []
-  for k in range(n):
-    choices.append(pg.Dict(
-        message=pg.Dict(
+  for k in range(json['n']):
+    if json.get('logprobs'):
+      logprobs = dict(
+          content=[
+              dict(
+                  token='chosen_token',
+                  logprob=0.5,
+                  top_logprobs=[
+                      dict(
+                          token=f'alternative_token_{i + 1}',
+                          logprob=0.1
+                      ) for i in range(3)
+                  ]
+              )
+          ]
+      )
+    else:
+      logprobs = None
+
+    choices.append(dict(
+        message=dict(
             content=(
                 f'Sample {k} for message.{system_message}{response_format}'
             )
         ),
-        logprobs=None,
+        logprobs=logprobs,
     ))
-  return pg.Dict(
-      choices=choices,
-      usage=lf.LMSamplingUsage(
-          prompt_tokens=100,
-          completion_tokens=100,
-          total_tokens=200,
-      ),
-  )
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str(
+      dict(
+          choices=choices,
+          usage=lf.LMSamplingUsage(
+              prompt_tokens=100,
+              completion_tokens=100,
+              total_tokens=200,
+          ),
+      )
+  ).encode()
+  return response
 
 
-def mock_chat_completion_query_vision(messages, *, n=1, **kwargs):
-  del kwargs
+def mock_chat_completion_request_vision(
+    url: str, json: dict[str, Any], **kwargs
+):
+  del url, kwargs
   choices = []
   urls = [
       c['image_url']['url']
-      for c in messages[0]['content'] if c['type'] == 'image_url'
+      for c in json['messages'][0]['content'] if c['type'] == 'image_url'
   ]
-  for k in range(n):
+  for k in range(json['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,
-      ),
-  )
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str(
+      dict(
+          choices=choices,
+          usage=lf.LMSamplingUsage(
+              prompt_tokens=100,
+              completion_tokens=100,
+              total_tokens=200,
+          ),
+      )
+  ).encode()
+  return response
 
 
 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)')
@@ -112,29 +133,9 @@ class OpenAITest(unittest.TestCase):
   def test_max_concurrency(self):
     self.assertGreater(openai.Gpt35(api_key='test_key').max_concurrency, 0)
 
-  def test_get_request_args(self):
-    self.assertEqual(
-        openai.Gpt35(api_key='test_key', timeout=90.0)._get_request_args(
-            lf.LMSamplingOptions(
-                temperature=2.0,
-                logprobs=True,
-                n=2,
-                max_tokens=4096,
-                top_p=1.0)),
-        dict(
-            engine='text-davinci-003',
-            logprobs=True,
-            top_logprobs=None,
-            n=2,
-            temperature=2.0,
-            max_tokens=4096,
-            stream=False,
-            timeout=90.0,
-            top_p=1.0,
-        )
-    )
+  def test_request_args(self):
     self.assertEqual(
-        openai.Gpt4(api_key='test_key')._get_request_args(
+        openai.Gpt4(api_key='test_key')._request_args(
             lf.LMSamplingOptions(
                 temperature=1.0, stop=['\n'], n=1, random_seed=123
             )
@@ -144,40 +145,93 @@ class OpenAITest(unittest.TestCase):
             top_logprobs=None,
             n=1,
             temperature=1.0,
-            stream=False,
-            timeout=120.0,
             stop=['\n'],
             seed=123,
         ),
     )
     with self.assertRaisesRegex(RuntimeError, '`logprobs` is not supported.*'):
-      openai.GptO1Preview(api_key='test_key')._get_request_args(
+      openai.GptO1Preview(api_key='test_key')._request_args(
           lf.LMSamplingOptions(
               temperature=1.0, logprobs=True
           )
       )
 
-  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')
+  def test_call_chat_completion(self):
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_chat_completion_request
+      lm = openai.OpenAI(
+          model='gpt-4',
+          api_key='test_key',
+          organization='my_org',
+          project='my_project'
+      )
       self.assertEqual(
           lm('hello', sampling_options=lf.LMSamplingOptions(n=2)),
-          'Sample 0 for prompt 0.',
+          'Sample 0 for message.',
       )
 
-  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')
+  def test_call_chat_completion_with_logprobs(self):
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_chat_completion_request
+      lm = openai.OpenAI(
+          model='gpt-4',
+          api_key='test_key',
+          organization='my_org',
+          project='my_project'
+      )
+      results = lm.sample(['hello'], logprobs=True)
+      self.assertEqual(len(results), 1)
       self.assertEqual(
-          lm('hello', sampling_options=lf.LMSamplingOptions(n=2)),
-          'Sample 0 for message.',
+          results[0],
+          lf.LMSamplingResult(
+              [
+                  lf.LMSample(
+                      response=lf.AIMessage(
+                          text='Sample 0 for message.',
+                          metadata={
+                              'score': 0.0,
+                              'logprobs': [(
+                                  'chosen_token',
+                                  0.5,
+                                  [
+                                      ('alternative_token_1', 0.1),
+                                      ('alternative_token_2', 0.1),
+                                      ('alternative_token_3', 0.1),
+                                  ],
+                              )],
+                              'is_cached': False,
+                              'usage': lf.LMSamplingUsage(
+                                  prompt_tokens=100,
+                                  completion_tokens=100,
+                                  total_tokens=200,
+                                  estimated_cost=0.009,
+                              ),
+                          },
+                          tags=['lm-response'],
+                      ),
+                      logprobs=[(
+                          'chosen_token',
+                          0.5,
+                          [
+                              ('alternative_token_1', 0.1),
+                              ('alternative_token_2', 0.1),
+                              ('alternative_token_3', 0.1),
+                          ],
+                      )],
+                  )
+              ],
+              usage=lf.LMSamplingUsage(
+                  prompt_tokens=100,
+                  completion_tokens=100,
+                  total_tokens=200,
+                  estimated_cost=0.009,
+              ),
+          ),
       )
 
   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
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_chat_completion_request_vision
       lm_1 = openai.Gpt4Turbo(api_key='test_key')
       lm_2 = openai.Gpt4VisionPreview(api_key='test_key')
       for lm in (lm_1, lm_2):
@@ -191,136 +245,18 @@ class OpenAITest(unittest.TestCase):
             ),
             '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)
+    lm_3 = openai.Gpt35Turbo(api_key='test_key')
+    with self.assertRaisesRegex(ValueError, 'Unsupported modality'):
+      lm_3(
+          lf.UserMessage(
+              'hello <<[[image]]>>',
+              image=lf_modalities.Image.from_uri('https://fake/image')
+          ),
       )
 
-    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,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=16,
-                            completion_tokens=16,
-                            total_tokens=33
-                        ),
-                        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,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=16,
-                            completion_tokens=16,
-                            total_tokens=33
-                        ),
-                        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,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=16,
-                            completion_tokens=16,
-                            total_tokens=33
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.2,
-                    logprobs=None,
-                ),
-            ],
-            usage=lf.LMSamplingUsage(
-                prompt_tokens=50, completion_tokens=50, total_tokens=100
-            ),
-        ),
-    )
-    self.assertEqual(
-        results[1],
-        lf.LMSamplingResult(
-            [
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 0 for prompt 1.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=16,
-                            completion_tokens=16,
-                            total_tokens=33
-                        ),
-                        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,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=16,
-                            completion_tokens=16,
-                            total_tokens=33
-                        ),
-                        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,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=16,
-                            completion_tokens=16,
-                            total_tokens=33
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.2,
-                    logprobs=None,
-                ),
-            ],
-            usage=lf.LMSamplingUsage(
-                prompt_tokens=50, completion_tokens=50, total_tokens=100
-            ),
-        ),
-    )
-
   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
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_chat_completion_request
       openai.SUPPORTED_MODELS_AND_SETTINGS['gpt-4'].update({
           'cost_per_1k_input_tokens': 1.0,
           'cost_per_1k_output_tokens': 1.0,
@@ -458,8 +394,8 @@ class OpenAITest(unittest.TestCase):
     )
 
   def test_sample_with_contextual_options(self):
-    with mock.patch('openai.Completion.create') as mock_completion:
-      mock_completion.side_effect = mock_completion_query
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_chat_completion_request
       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'])
@@ -471,7 +407,7 @@ class OpenAITest(unittest.TestCase):
             [
                 lf.LMSample(
                     lf.AIMessage(
-                        'Sample 0 for prompt 0.',
+                        'Sample 0 for message.',
                         score=0.0,
                         logprobs=None,
                         is_cached=False,
@@ -487,8 +423,8 @@ class OpenAITest(unittest.TestCase):
                 ),
                 lf.LMSample(
                     lf.AIMessage(
-                        'Sample 1 for prompt 0.',
-                        score=0.1,
+                        'Sample 1 for message.',
+                        score=0.0,
                         logprobs=None,
                         is_cached=False,
                         usage=lf.LMSamplingUsage(
@@ -498,19 +434,19 @@ class OpenAITest(unittest.TestCase):
                         ),
                         tags=[lf.Message.TAG_LM_RESPONSE],
                     ),
-                    score=0.1,
+                    score=0.0,
                     logprobs=None,
                 ),
             ],
             usage=lf.LMSamplingUsage(
                 prompt_tokens=100, completion_tokens=100, total_tokens=200
             ),
-        ),
+        )
     )
 
   def test_call_with_system_message(self):
-    with mock.patch('openai.ChatCompletion.create') as mock_chat_completion:
-      mock_chat_completion.side_effect = mock_chat_completion_query
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_chat_completion_request
       lm = openai.OpenAI(api_key='test_key', model='gpt-4')
       self.assertEqual(
           lm(
@@ -520,12 +456,12 @@ class OpenAITest(unittest.TestCase):
               ),
               sampling_options=lf.LMSamplingOptions(n=2)
           ),
-          'Sample 0 for message. system=hi',
+          '''Sample 0 for message. system=[{'type': 'text', 'text': 'hi'}]''',
       )
 
   def test_call_with_json_schema(self):
-    with mock.patch('openai.ChatCompletion.create') as mock_chat_completion:
-      mock_chat_completion.side_effect = mock_chat_completion_query
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_chat_completion_request
       lm = openai.OpenAI(api_key='test_key', model='gpt-4')
       self.assertEqual(
           lm(

langfun/core/llms/vertexai_test.py

@@ -426,6 +426,7 @@ class VertexRestfulAITest(unittest.TestCase):
           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
@@ -496,6 +497,7 @@ class VertexRestfulAITest(unittest.TestCase):
           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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.1.2.dev202412020805
+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.                     |

{langfun-0.1.2.dev202412020805.dist-info → langfun-0.1.2.dev202412030000.dist-info}/RECORD

@@ -92,12 +92,12 @@ langfun/core/llms/groq.py,sha256=dCnR3eAECEKuKKAAj-PDTs8NRHl6CQPdf57m1f6a79U,103
 langfun/core/llms/groq_test.py,sha256=GYF_Qtq5S1H1TrKH38t6_lkdroqT7v-joYLDKnmS9e0,5274
 langfun/core/llms/llama_cpp.py,sha256=9tXQntSCDtjTF3bnyJrAPCr4N6wycy5nXYvp9uduygE,2843
 langfun/core/llms/llama_cpp_test.py,sha256=MWO_qaOeKjRniGjcaWPDScd7HPaIJemqUZoslrt4FPs,1806
-langfun/core/llms/openai.py,sha256=_VwOSuDsyXDngUM2iiES0CW1aN0BzMjXNBMegLzm4J4,23209
-langfun/core/llms/openai_test.py,sha256=_8cd3VRNEUfE0-Ko1RiM6MlC5hjalRj7nYTJNhG1p3E,18907
+langfun/core/llms/openai.py,sha256=l49v6RubfInvV0iG114AymTKNogTX4u4N-UFCeSgIxw,20963
+langfun/core/llms/openai_test.py,sha256=kOWa1nf-nJvtYY10REUw5wojh3ZgfU8tRaCZ8wUgJbA,16623
 langfun/core/llms/rest.py,sha256=sWbYUV8S3SuOg9giq7xwD-xDRfaF7NP_ig7bI52-Rj4,3442
 langfun/core/llms/rest_test.py,sha256=NZ3Nf0XQVpT9kLP5cBVo_yBHLI7vWTYhWQxYEJVMGs4,3472
 langfun/core/llms/vertexai.py,sha256=EZhJrdN-SsZVV0KT3NHzaJLVKsNMxCT6M3W6f5fpIWQ,27068
-langfun/core/llms/vertexai_test.py,sha256=nGv59yE4xu1zUxqmP_U941QjSBrr_sW15Q2YakuxMv4,16982
+langfun/core/llms/vertexai_test.py,sha256=qapDa7fvLkHm3BhG12a-HopxGCn625r-eVud2QqRITo,17120
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=rt3zwmyw0y9jsSGW-ZbV1vAfLxQ7_3AVk0l2EySlse4,3918
 langfun/core/llms/cache/in_memory.py,sha256=l6b-iU9OTfTRo9Zmg4VrQIuArs4cCJDOpXiEpvNocjo,5004
@@ -148,8 +148,8 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
-langfun-0.1.2.dev202412020805.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202412020805.dist-info/METADATA,sha256=c3yjg186RyrDaIHGLMpmXsI7-Kqj4V1vLGxYsjJJN2Y,8890
-langfun-0.1.2.dev202412020805.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-langfun-0.1.2.dev202412020805.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202412020805.dist-info/RECORD,,
+langfun-0.1.2.dev202412030000.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202412030000.dist-info/METADATA,sha256=PoROaIMontFjWm5sPVdo3DpJATWFoFbO8IOr9t-3K2o,8651
+langfun-0.1.2.dev202412030000.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+langfun-0.1.2.dev202412030000.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202412030000.dist-info/RECORD,,