langfun 0.1.2.dev202509120804__py3-none-any.whl → 0.1.2.dev202512150805__py3-none-any.whl

langfun/core/llms/openai_compatible.py

@@ -23,8 +23,18 @@ import pyglove as pg
 
 
 @lf.use_init_args(['api_endpoint', 'model'])
-class OpenAICompatible(rest.REST):
-  """Base for OpenAI compatible models."""
+class OpenAIChatCompletionAPI(rest.REST):
+  """Base class for models compatible with OpenAI's Chat Completion API.
+
+  This class provides a common interface for language models that adhere to
+  the OpenAI Chat Completion API format, which is used by providers like
+  Groq, DeepSeek, and others. It standardizes request formatting and
+  response parsing for these models.
+
+  **References:**
+
+  *   https://platform.openai.com/docs/api-reference/chat
+  """
 
   model: Annotated[
       str, 'The name of the model to use.',
@@ -42,12 +52,14 @@ class OpenAICompatible(rest.REST):
     # Reference:
     # https://platform.openai.com/docs/api-reference/completions/create
     # NOTE(daiyip): options.top_k is not applicable.
-    args = dict(
-        n=options.n,
-        top_logprobs=options.top_logprobs,
-    )
+    args = {}
+
     if self.model:
       args['model'] = self.model
+    if options.n != 1:
+      args['n'] = options.n
+    if options.top_logprobs is not None:
+      args['top_logprobs'] = options.top_logprobs
     if options.logprobs:
       args['logprobs'] = options.logprobs
     if options.temperature is not None:
@@ -62,6 +74,8 @@ class OpenAICompatible(rest.REST):
       args['seed'] = options.random_seed
     if options.reasoning_effort is not None:
       args['reasoning_effort'] = options.reasoning_effort
+    if options.extras:
+      args.update(options.extras)
     return args
 
   def request(
@@ -72,27 +86,13 @@ class OpenAICompatible(rest.REST):
     """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}.'
-        )
-      if 'title' not in json_schema:
-        raise ValueError(
-            f'The root of `json_schema` must have a `title` field, '
-            f'got {json_schema!r}.'
-        )
+    # Handle structured output.
+    output_schema = self._structure_output_schema(prompt)
+    if output_schema is not None:
       request_args.update(
           response_format=dict(
               type='json_schema',
-              json_schema=dict(
-                  schema=json_schema,
-                  name=json_schema['title'],
-                  strict=True,
-              )
+              json_schema=output_schema,
           )
       )
       prompt.metadata.formatted_text = (
@@ -118,17 +118,43 @@ class OpenAICompatible(rest.REST):
       assert isinstance(system_message, lf.SystemMessage), type(system_message)
       messages.append(
           system_message.as_format(
-              'openai', chunk_preprocessor=modality_check
+              'openai_chat_completion_api', chunk_preprocessor=modality_check
           )
       )
     messages.append(
-        prompt.as_format('openai', chunk_preprocessor=modality_check)
+        prompt.as_format(
+            'openai_chat_completion_api',
+            chunk_preprocessor=modality_check
+        )
     )
     request = dict()
     request.update(request_args)
     request['messages'] = messages
     return request
 
+  def _structure_output_schema(
+      self, prompt: lf.Message
+  ) -> dict[str, Any] | None:
+    # 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}.'
+        )
+      return dict(
+          schema=json_schema,
+          name=json_schema['title'],
+          strict=True,
+      )
+    return None
+
   def _parse_choice(self, choice: dict[str, Any]) -> lf.LMSample:
     # Reference:
     # https://platform.openai.com/docs/api-reference/chat/object
@@ -144,7 +170,10 @@ class OpenAICompatible(rest.REST):
           for t in choice_logprobs['content']
       ]
     return lf.LMSample(
-        lf.Message.from_value(choice['message'], format='openai'),
+        lf.Message.from_value(
+            choice['message'],
+            format='openai_chat_completion_api'
+        ),
         score=0.0,
         logprobs=logprobs,
     )
@@ -169,3 +198,95 @@ class OpenAICompatible(rest.REST):
         or (status_code == 400 and b'string_above_max_length' in content)):
       return lf.ContextLimitError(f'{status_code}: {content}')
     return super()._error(status_code, content)
+
+
+class OpenAIResponsesAPI(OpenAIChatCompletionAPI):
+  """Base class for models compatible with OpenAI's Responses API.
+
+  This class provides a common interface for language models that adhere to
+  the new OpenAI Responses API format. It standardizes request formatting
+  and response parsing for these models, including handling instructions
+  (system messages) and structured outputs.
+
+  **References:**
+
+  *   https://platform.openai.com/docs/api-reference/responses
+  """
+
+  def _request_args(
+      self, options: lf.LMSamplingOptions) -> dict[str, Any]:
+    """Returns a dict as request arguments."""
+    if options.logprobs:
+      raise ValueError('logprobs is not supported on Responses API.')
+    if options.n != 1:
+      raise ValueError('n must be 1 for Responses API.')
+    return super()._request_args(options)
+
+  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)
+
+    # Handle structured output.
+    output_schema = self._structure_output_schema(prompt)
+    if output_schema is not None:
+      output_schema['type'] = 'json_schema'
+      request_args.update(text=dict(format=output_schema))
+      prompt.metadata.formatted_text = (
+          prompt.text
+          + '\n\n [RESPONSE FORMAT (not part of prompt)]\n'
+          + pg.to_json_str(request_args['text'], json_indent=2)
+      )
+
+    request = dict()
+    request.update(request_args)
+
+    # Users could use `metadata_system_message` to pass system message.
+    system_message = prompt.metadata.get('system_message')
+    if system_message:
+      assert isinstance(system_message, lf.SystemMessage), type(system_message)
+      request['instructions'] = system_message.text
+
+    # Prepare input.
+    def modality_check(chunk: str | lf.Modality) -> Any:
+      if (isinstance(chunk, lf_modalities.Mime)
+          and not self.supports_input(chunk.mime_type)):
+        raise ValueError(
+            f'Unsupported modality: {chunk!r}.'
+        )
+      return chunk
+
+    request['input'] = [
+        prompt.as_format(
+            'openai_responses_api',
+            chunk_preprocessor=modality_check
+        )
+    ]
+    return request
+
+  def _parse_output(self, output: dict[str, Any]) -> lf.LMSample:
+    for item in output:
+      if isinstance(item, dict) and item.get('type') == 'message':
+        return lf.LMSample(
+            lf.Message.from_value(item, format='openai_responses_api'),
+            score=0.0,
+        )
+    raise ValueError('No message found in output.')
+
+  def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
+    """Returns a LMSamplingResult from a JSON response."""
+    usage = json['usage']
+    return lf.LMSamplingResult(
+        samples=[self._parse_output(json['output'])],
+        usage=lf.LMSamplingUsage(
+            prompt_tokens=usage['input_tokens'],
+            completion_tokens=usage['output_tokens'],
+            total_tokens=usage['total_tokens'],
+            completion_tokens_details=usage.get(
+                'output_tokens_details', None
+            ),
+        ),
+    )

langfun/core/llms/openai_compatible_test.py

@@ -38,7 +38,7 @@ def mock_chat_completion_request(url: str, json: dict[str, Any], **kwargs):
     response_format = ''
 
   choices = []
-  for k in range(json['n']):
+  for k in range(json.get('n', 1)):
     if json.get('logprobs'):
       logprobs = dict(
           content=[
@@ -89,7 +89,7 @@ def mock_chat_completion_request_vision(
       c['image_url']['url']
       for c in json['messages'][0]['content'] if c['type'] == 'image_url'
   ]
-  for k in range(json['n']):
+  for k in range(json.get('n', 1)):
     choices.append(pg.Dict(
         message=pg.Dict(
             content=f'Sample {k} for message: {"".join(urls)}'
@@ -111,12 +111,88 @@ def mock_chat_completion_request_vision(
   return response
 
 
-class OpenAIComptibleTest(unittest.TestCase):
+def mock_responses_request(url: str, json: dict[str, Any], **kwargs):
+  del url, kwargs
+  _ = json['input']
+
+  system_message = ''
+  if 'instructions' in json:
+    system_message = f' system={json["instructions"]}'
+
+  response_format = ''
+  if 'text' in json and 'format' in json['text']:
+    response_format = f' format={json["text"]["format"]["type"]}'
+
+  output = [
+      dict(
+          type='message',
+          content=[
+              dict(
+                  type='output_text',
+                  text=(
+                      f'Sample 0 for message.{system_message}{response_format}'
+                  )
+              )
+          ],
+      )
+  ]
+
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str(
+      dict(
+          output=output,
+          usage=dict(
+              input_tokens=100,
+              output_tokens=100,
+              total_tokens=200,
+          ),
+      )
+  ).encode()
+  return response
+
+
+def mock_responses_request_vision(
+    url: str, json: dict[str, Any], **kwargs
+):
+  del url, kwargs
+  urls = [
+      c['image_url']
+      for c in json['input'][0]['content']
+      if c['type'] == 'input_image'
+  ]
+  output = [
+      pg.Dict(
+          type='message',
+          content=[
+              pg.Dict(
+                  type='output_text',
+                  text=f'Sample 0 for message: {"".join(urls)}',
+              )
+          ],
+      )
+  ]
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str(
+      dict(
+          output=output,
+          usage=dict(
+              input_tokens=100,
+              output_tokens=100,
+              total_tokens=200,
+          ),
+      )
+  ).encode()
+  return response
+
+
+class OpenAIChatCompletionAPITest(unittest.TestCase):
   """Tests for OpenAI compatible language model."""
 
   def test_request_args(self):
     self.assertEqual(
-        openai_compatible.OpenAICompatible(
+        openai_compatible.OpenAIChatCompletionAPI(
             api_endpoint='https://test-server',
             model='test-model'
         )._request_args(
@@ -126,8 +202,6 @@ class OpenAIComptibleTest(unittest.TestCase):
         ),
         dict(
             model='test-model',
-            top_logprobs=None,
-            n=1,
             temperature=1.0,
             stop=['\n'],
             seed=123,
@@ -137,7 +211,7 @@ class OpenAIComptibleTest(unittest.TestCase):
   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_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model',
       )
       self.assertEqual(
@@ -148,7 +222,7 @@ class OpenAIComptibleTest(unittest.TestCase):
   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_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model',
       )
       results = lm.sample(['hello'], logprobs=True)
@@ -214,13 +288,14 @@ class OpenAIComptibleTest(unittest.TestCase):
       def mime_type(self) -> str:
         return 'image/png'
 
+    image = FakeImage.from_uri('https://fake/image')
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_chat_completion_request_vision
-      lm_1 = openai_compatible.OpenAICompatible(
+      lm_1 = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server',
           model='test-model1',
       )
-      lm_2 = openai_compatible.OpenAICompatible(
+      lm_2 = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server',
           model='test-model2',
       )
@@ -228,15 +303,15 @@ class OpenAIComptibleTest(unittest.TestCase):
         self.assertEqual(
             lm(
                 lf.UserMessage(
-                    'hello <<[[image]]>>',
-                    image=FakeImage.from_uri('https://fake/image')
+                    f'hello <<[[{image.id}]]>>',
+                    referred_modalities=[image],
                 ),
                 sampling_options=lf.LMSamplingOptions(n=2)
             ),
             'Sample 0 for message: https://fake/image',
         )
 
-    class TextOnlyModel(openai_compatible.OpenAICompatible):
+    class TextOnlyModel(openai_compatible.OpenAIChatCompletionAPI):
 
       class ModelInfo(lf.ModelInfo):
         input_modalities: list[str] = lf.ModelInfo.TEXT_INPUT_ONLY
@@ -251,15 +326,15 @@ class OpenAIComptibleTest(unittest.TestCase):
     with self.assertRaisesRegex(ValueError, 'Unsupported modality'):
       lm_3(
           lf.UserMessage(
-              'hello <<[[image]]>>',
-              image=FakeImage.from_uri('https://fake/image')
+              f'hello <<[[{image.id}]]>>',
+              referred_modalities=[image],
           ),
       )
 
   def test_sample_chat_completion(self):
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_chat_completion_request
-      lm = openai_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model'
       )
       results = lm.sample(
@@ -400,7 +475,7 @@ class OpenAIComptibleTest(unittest.TestCase):
   def test_sample_with_contextual_options(self):
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_chat_completion_request
-      lm = openai_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model'
       )
       with lf.use_settings(sampling_options=lf.LMSamplingOptions(n=2)):
@@ -458,7 +533,7 @@ class OpenAIComptibleTest(unittest.TestCase):
   def test_call_with_system_message(self):
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_chat_completion_request
-      lm = openai_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model'
       )
       self.assertEqual(
@@ -475,7 +550,7 @@ class OpenAIComptibleTest(unittest.TestCase):
   def test_call_with_json_schema(self):
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_chat_completion_request
-      lm = openai_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model'
       )
       self.assertEqual(
@@ -515,7 +590,7 @@ class OpenAIComptibleTest(unittest.TestCase):
 
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_context_limit_error
-      lm = openai_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model'
       )
       with self.assertRaisesRegex(
@@ -524,5 +599,117 @@ class OpenAIComptibleTest(unittest.TestCase):
         lm(lf.UserMessage('hello'))
 
 
+class OpenAIResponsesAPITest(unittest.TestCase):
+  """Tests for OpenAI compatible language model on Responses API."""
+
+  def test_request_args(self):
+    lm = openai_compatible.OpenAIResponsesAPI(
+        api_endpoint='https://test-server', model='test-model'
+    )
+    # Test valid args.
+    self.assertEqual(
+        lm._request_args(
+            lf.LMSamplingOptions(
+                temperature=1.0, stop=['\n'], n=1, random_seed=123
+            )
+        ),
+        dict(
+            model='test-model',
+            temperature=1.0,
+            stop=['\n'],
+            seed=123,
+        ),
+    )
+    # Test unsupported n.
+    with self.assertRaisesRegex(ValueError, 'n must be 1 for Responses API.'):
+      lm._request_args(lf.LMSamplingOptions(n=2))
+
+    # Test unsupported logprobs.
+    with self.assertRaisesRegex(
+        ValueError, 'logprobs is not supported on Responses API.'
+    ):
+      lm._request_args(lf.LMSamplingOptions(logprobs=True))
+
+  def test_call_responses(self):
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_responses_request
+      lm = openai_compatible.OpenAIResponsesAPI(
+          api_endpoint='https://test-server',
+          model='test-model',
+      )
+      self.assertEqual(lm('hello'), 'Sample 0 for message.')
+
+  def test_call_responses_vision(self):
+    class FakeImage(lf_modalities.Image):
+      @property
+      def mime_type(self) -> str:
+        return 'image/png'
+
+    image = FakeImage.from_uri('https://fake/image')
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_responses_request_vision
+      lm = openai_compatible.OpenAIResponsesAPI(
+          api_endpoint='https://test-server',
+          model='test-model1',
+      )
+      self.assertEqual(
+          lm(
+              lf.UserMessage(
+                  f'hello <<[[{image.id}]]>>',
+                  referred_modalities=[image],
+              )
+          ),
+          'Sample 0 for message: https://fake/image',
+      )
+
+  def test_call_with_system_message(self):
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_responses_request
+      lm = openai_compatible.OpenAIResponsesAPI(
+          api_endpoint='https://test-server', model='test-model'
+      )
+      self.assertEqual(
+          lm(
+              lf.UserMessage(
+                  'hello',
+                  system_message=lf.SystemMessage('hi'),
+              )
+          ),
+          'Sample 0 for message. system=hi',
+      )
+
+  def test_call_with_json_schema(self):
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_responses_request
+      lm = openai_compatible.OpenAIResponsesAPI(
+          api_endpoint='https://test-server', model='test-model'
+      )
+      self.assertEqual(
+          lm(
+              lf.UserMessage(
+                  'hello',
+                  json_schema={
+                      'type': 'object',
+                      'properties': {
+                          'name': {'type': 'string'},
+                      },
+                      'required': ['name'],
+                      'title': 'Person',
+                  },
+              )
+          ),
+          'Sample 0 for message. format=json_schema',
+      )
+
+    # Test bad json schema.
+    with self.assertRaisesRegex(ValueError, '`json_schema` must be a dict'):
+      lm(lf.UserMessage('hello', json_schema='foo'))
+
+    with self.assertRaisesRegex(
+        ValueError, 'The root of `json_schema` must have a `title` field'
+    ):
+      lm(lf.UserMessage('hello', json_schema={}))
+
+
 if __name__ == '__main__':
   unittest.main()

langfun/core/llms/openai_test.py

@@ -61,8 +61,6 @@ class OpenAITest(unittest.TestCase):
         ),
         dict(
             model='gpt-4',
-            top_logprobs=None,
-            n=1,
             temperature=1.0,
             stop=['\n'],
             seed=123,

langfun/core/llms/rest.py

@@ -22,7 +22,18 @@ import requests
 
 
 class REST(lf.LanguageModel):
-  """REST-based language model."""
+  """Base class for language models accessed via REST APIs.
+
+  The `REST` class provides a foundation for implementing language models
+  that are accessed through RESTful endpoints. It handles the details of
+  making HTTP requests, managing sessions, and handling common errors like
+  timeouts and connection issues.
+
+  Subclasses need to implement the `request` and `result` methods to
+  convert Langfun messages to API-specific request formats and to parse
+  API responses back into `LMSamplingResult` objects. They also need to
+  provide the `api_endpoint` and can override `headers` for authentication.
+  """
 
   api_endpoint: Annotated[
       str,
@@ -98,7 +109,9 @@ class REST(lf.LanguageModel):
       raise lf.TemporaryLMError(str(e)) from e
     except (
         requests.exceptions.ConnectionError,
+        requests.exceptions.ChunkedEncodingError,
         ConnectionError,
+        ConnectionResetError,
     ) as e:
       error_message = str(e)
       if 'REJECTED_CLIENT_THROTTLED' in error_message:
@@ -107,6 +120,8 @@ class REST(lf.LanguageModel):
         raise lf.TemporaryLMError(error_message) from e
       if 'UNREACHABLE_ERROR' in error_message:
         raise lf.TemporaryLMError(error_message) from e
+      if 'Connection reset by peer' in error_message:
+        raise lf.TemporaryLMError(error_message) from e
       raise lf.LMError(error_message) from e
 
   def _error(self, status_code: int, content: str) -> lf.LMError: