langfun 0.1.2.dev202510230805__py3-none-any.whl → 0.1.2.dev202510250803__py3-none-any.whl

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/vertexai.py

@@ -369,6 +369,16 @@ class VertexAIAnthropic(VertexAI, anthropic.Anthropic):
 # pylint: disable=invalid-name
 
 
+class VertexAIClaude45Haiku_20251001(VertexAIAnthropic):
+  """Anthropic's Claude 4.5 Haiku model on VertexAI."""
+  model = 'claude-haiku-4-5@20251001'
+
+
+class VertexAIClaude45Sonnet_20250929(VertexAIAnthropic):
+  """Anthropic's Claude 4.5 Sonnet model on VertexAI."""
+  model = 'claude-sonnet-4-5@20250929'
+
+
 class VertexAIClaude4Opus_20250514(VertexAIAnthropic):
   """Anthropic's Claude 4 Opus model on VertexAI."""
   model = 'claude-opus-4@20250514'
@@ -487,7 +497,7 @@ _LLAMA_MODELS_BY_MODEL_ID = {m.model_id: m for m in LLAMA_MODELS}
 
 @pg.use_init_args(['model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
-class VertexAILlama(VertexAI, openai_compatible.OpenAICompatible):
+class VertexAILlama(VertexAI, openai_compatible.OpenAIChatCompletionAPI):
   """Llama models on VertexAI."""
 
   model: pg.typing.Annotated[
@@ -600,7 +610,7 @@ _MISTRAL_MODELS_BY_MODEL_ID = {m.model_id: m for m in MISTRAL_MODELS}
 
 @pg.use_init_args(['model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
-class VertexAIMistral(VertexAI, openai_compatible.OpenAICompatible):
+class VertexAIMistral(VertexAI, openai_compatible.OpenAIChatCompletionAPI):
   """Mistral AI models on VertexAI."""
 
   model: pg.typing.Annotated[

langfun/core/message.py

@@ -20,7 +20,7 @@ import contextlib
 import functools
 import inspect
 import io
-from typing import Annotated, Any, ClassVar, Optional, Type, Union
+from typing import Annotated, Any, Callable, ClassVar, Optional, Type, Union
 
 from langfun.core import modality
 from langfun.core import natural_language
@@ -86,6 +86,11 @@ class Message(
 
   sender: Annotated[str, 'The sender of the message.']
 
+  referred_modalities: Annotated[
+      dict[str, pg.Ref[modality.Modality]],
+      'The modality objects referred in the message.'
+  ] = pg.Dict()
+
   metadata: Annotated[
       dict[str, Any],
       (
@@ -111,6 +116,11 @@ class Message(
       *,
       # Default sender is specified in subclasses.
       sender: str | pg.object_utils.MissingValue = pg.MISSING_VALUE,
+      referred_modalities: (
+          list[modality.Modality]
+          | dict[str, modality.Modality]
+          | None
+      ) = None,
       metadata: dict[str, Any] | None = None,
       tags: list[str] | None = None,
       source: Optional['Message'] = None,
@@ -125,6 +135,7 @@ class Message(
     Args:
       text: The text in the message.
       sender: The sender name of the message.
+      referred_modalities: The modality objects referred in the message.
       metadata: Structured meta-data associated with this message.
       tags: Tags for the message.
       source: The source message of the current message.
@@ -138,9 +149,13 @@ class Message(
     """
     metadata = metadata or {}
     metadata.update(kwargs)
+    if isinstance(referred_modalities, list):
+      referred_modalities = {m.id: pg.Ref(m) for m in referred_modalities}
+
     super().__init__(
         text=text,
         metadata=metadata,
+        referred_modalities=referred_modalities or {},
         tags=tags or [],
         sender=sender,
         allow_partial=allow_partial,
@@ -186,7 +201,7 @@ class Message(
       A message created from the value.
     """
     if isinstance(value, modality.Modality):
-      return cls('<<[[object]]>>', object=value)
+      return cls(f'<<[[{value.id}]]>>', referred_modalities=[value])
     if isinstance(value, Message):
       return value
     if isinstance(value, str):
@@ -280,8 +295,7 @@ class Message(
     if key_path == Message.PATH_TEXT:
       return self.text
     else:
-      v = self.metadata.sym_get(key_path, default, use_inferred=True)
-      return v.value if isinstance(v, pg.Ref) else v
+      return self.metadata.sym_get(key_path, default, use_inferred=True)
 
   #
   # API for accessing the structured result and error.
@@ -361,43 +375,53 @@ class Message(
   # API for supporting modalities.
   #
 
+  def modalities(
+      self,
+      filter: (  # pylint: disable=redefined-builtin
+          Type[modality.Modality]
+          | Callable[[modality.Modality], bool]
+          | None
+      ) = None  # pylint: disable=bad-whitespace
+  ) -> list[modality.Modality]:
+    """Returns the modality objects referred in the message."""
+    if inspect.isclass(filter) and issubclass(filter, modality.Modality):
+      filter_fn = lambda v: isinstance(v, filter)  # pytype: disable=wrong-arg-types
+    elif filter is None:
+      filter_fn = lambda v: True
+    else:
+      filter_fn = filter
+    return [v for v in self.referred_modalities.values() if filter_fn(v)]
+
   @property
-  def text_with_modality_hash(self) -> str:
-    """Returns text with modality object placeheld by their 8-byte MD5 hash."""
-    parts = [self.text]
-    for name, modality_obj in self.referred_modalities().items():
-      parts.append(
-          f'<{name}>{modality_obj.hash}</{name}>'
-      )
-    return ''.join(parts)
+  def images(self) -> list[modality.Modality]:
+    """Returns the image objects referred in the message."""
+    assert False, 'Overridden in core/modalities/__init__.py'
+
+  @property
+  def videos(self) -> list[modality.Modality]:
+    """Returns the video objects referred in the message."""
+    assert False, 'Overridden in core/modalities/__init__.py'
+
+  @property
+  def audios(self) -> list[modality.Modality]:
+    """Returns the audio objects referred in the message."""
+    assert False, 'Overridden in core/modalities/__init__.py'
 
   def get_modality(
-      self, var_name: str, default: Any = None, from_message_chain: bool = True
+      self,
+      var_name: str,
+      default: Any = None
   ) -> modality.Modality | None:
     """Gets the modality object referred in the message.
 
     Args:
       var_name: The referred variable name for the modality object.
       default: default value.
-      from_message_chain: If True, the look up will be performed from the
-        message chain. Otherwise it will be performed in current message.
 
     Returns:
       A modality object if found, otherwise None.
     """
-    obj = self.get(var_name, None)
-    if isinstance(obj, modality.Modality):
-      return obj
-    elif obj is None and self.source is not None:
-      return self.source.get_modality(var_name, default, from_message_chain)
-    return default
-
-  def referred_modalities(self) -> dict[str, modality.Modality]:
-    """Returns modality objects attached on this message."""
-    chunks = self.chunk()
-    return {
-        m.referred_name: m for m in chunks if isinstance(m, modality.Modality)
-    }
+    return self.referred_modalities.get(var_name, default)
 
   def chunk(self, text: str | None = None) -> list[str | modality.Modality]:
     """Chunk a message into a list of str or modality objects."""
@@ -425,10 +449,15 @@ class Message(
 
       var_name = text[var_start:ref_end].strip()
       var_value = self.get_modality(var_name)
-      if var_value is not None:
-        add_text_chunk(text[chunk_start:ref_start].strip(' '))
-        chunks.append(var_value)
-        chunk_start = ref_end + len(modality.Modality.REF_END)
+      if var_value is None:
+        raise ValueError(
+            f'Unknown modality reference: {var_name!r}. '
+            'Please make sure the modality object is present in '
+            f'`referred_modalities` when creating {self.__class__.__name__}.'
+        )
+      add_text_chunk(text[chunk_start:ref_start].strip(' '))
+      chunks.append(var_value)
+      chunk_start = ref_end + len(modality.Modality.REF_END)
     return chunks
 
   @classmethod
@@ -437,8 +466,8 @@ class Message(
   ) -> 'Message':
     """Assembly a message from a list of string or modality objects."""
     fused_text = io.StringIO()
-    ref_index = 0
     metadata = dict()
+    referred_modalities = dict()
     last_char = None
     for i, chunk in enumerate(chunks):
       if i > 0 and last_char not in ('\t', ' ', '\n', None):
@@ -451,14 +480,16 @@ class Message(
           last_char = None
       else:
         assert isinstance(chunk, modality.Modality), chunk
-        var_name = f'obj{ref_index}'
-        fused_text.write(modality.Modality.text_marker(var_name))
+        fused_text.write(modality.Modality.text_marker(chunk.id))
         last_char = modality.Modality.REF_END[-1]
         # Make a reference if the chunk is already owned by another object
         # to avoid copy.
-        metadata[var_name] = pg.maybe_ref(chunk)
-        ref_index += 1
-    return cls(fused_text.getvalue().strip(), metadata=metadata)
+        referred_modalities[chunk.id] = pg.Ref(chunk)
+    return cls(
+        fused_text.getvalue().strip(),
+        referred_modalities=referred_modalities,
+        metadata=metadata,
+    )
 
   #
   # Tagging
@@ -551,6 +582,11 @@ class Message(
   #
 
   def natural_language_format(self) -> str:
+    """Returns the natural language format representation."""
+    # Propagate the modality references to parent context if any.
+    if capture_context := modality.get_modality_capture_context():
+      for v in self.referred_modalities.values():
+        capture_context.capture(v)
     return self.text
 
   def __eq__(self, other: Any) -> bool:
@@ -568,8 +604,7 @@ class Message(
   def __getattr__(self, key: str) -> Any:
     if key not in self.metadata:
       raise AttributeError(key)
-    v = self.metadata[key]
-    return v.value if isinstance(v, pg.Ref) else v
+    return self.metadata[key]
 
   def _html_tree_view_content(
       self,
@@ -646,15 +681,14 @@ class Message(
           s.write(s.escape(chunk))
         else:
           assert isinstance(chunk, modality.Modality), chunk
-          child_path = pg.KeyPath(['metadata', chunk.referred_name], root_path)
           s.write(
               pg.Html.element(
                   'div',
                   [
                       view.render(
                           chunk,
-                          name=chunk.referred_name,
-                          root_path=child_path,
+                          name=chunk.id,
+                          root_path=chunk.sym_path,
                           collapse_level=(
                               0 if collapse_modalities_in_text else 1
                           ),
@@ -667,7 +701,7 @@ class Message(
                   css_classes=['modality-in-text'],
               )
           )
-          referred_chunks[chunk.referred_name] = chunk
+          referred_chunks[chunk.id] = chunk
       s.write('</div>')
       return s