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

langfun/core/modality.py

@@ -14,40 +14,63 @@
 """Interface for modality (e.g. Image, Video, etc.)."""
 
 import abc
+import contextlib
 import functools
 import hashlib
-from typing import Any, ContextManager
+import re
+from typing import Any, ContextManager, Iterator
 from langfun.core import component
 import pyglove as pg
 
 
-_TLS_MODALITY_AS_REF = '__format_modality_as_ref__'
+class Modality(component.Component, pg.views.HtmlTreeView.Extension):
+  """Base class for representing non-text content in prompts.
 
+  `lf.Modality` is the base class for multimodal objects such as `lf.Image`,
+  `lf.Audio`, and `lf.Video`. It allows these non-text inputs to be
+  seamlessly embedded within text prompts for processing by multimodal
+  language models.
 
-def format_modality_as_ref(enabled: bool = True) -> ContextManager[None]:
-  """A context manager that formats modality objects as references."""
-  return pg.object_utils.thread_local_value_scope(
-      _TLS_MODALITY_AS_REF, enabled, False
-  )
+  When a `Modality` object is rendered within an `lf.Template`, it is
+  replaced by a text marker (e.g., `<<[[image:b10a8db1]]>>`), and the
+  modality object itself is stored in the `referred_modalities` field of
+  the resulting `lf.Message`. This allows language models to associate
+  the placeholder with its content during processing.
 
+  **Example:**
 
-class Modality(component.Component, pg.views.HtmlTreeView.Extension):
-  """Base class for multimodal object."""
+  ```python
+  import langfun as lf
+
+  image = lf.Image.from_path('/path/to/image.png')
+  prompt = lf.Template('What is in this image? {{image}}', image=image)
+
+  message = prompt.render()
+  print(message.text)
+  # Output: What is in this image? <<[[image:b10a8db1]]>>
+
+  print(message.modalities())
+  # Output: [<Image object>]
+  ```
+  """
 
   REF_START = '<<[['
   REF_END = ']]>>'
 
   def _on_bound(self):
     super()._on_bound()
-    # Invalidate cached hash if modality member is changed.
+    # Invalidate cached hash and id if modality member is changed.
     self.__dict__.pop('hash', None)
+    self.__dict__.pop('id', None)
 
   def format(self, *args, **kwargs) -> str:
-    if self.referred_name is None or not pg.object_utils.thread_local_get(
-        _TLS_MODALITY_AS_REF, False
-    ):
+    if not pg.object_utils.thread_local_get(_TLS_MODALITY_AS_REF, False):
       return super().format(*args, **kwargs)
-    return Modality.text_marker(self.referred_name)
+
+    capture_scope = get_modality_capture_context()
+    if capture_scope is not None:
+      capture_scope.capture(self)
+    return Modality.text_marker(self.id)
 
   def __str_kwargs__(self) -> dict[str, Any]:
     # For modality objects, we don't want to use markdown format when they
@@ -70,14 +93,11 @@ class Modality(component.Component, pg.views.HtmlTreeView.Extension):
     """Returns a marker in the text for this object."""
     return Modality.REF_START + var_name + Modality.REF_END
 
-  @property
-  def referred_name(self) -> str | None:
+  @functools.cached_property
+  def id(self) -> str | None:
     """Returns the referred name of this object in its template."""
-    if not self.sym_path:
-      return None
-    # Strip the metadata prefix under message.
-    path = str(self.sym_path)
-    return path[9:] if path.startswith('metadata.') else path
+    modality_type = _camel_to_snake(self.__class__.__name__)
+    return f'{modality_type}:{self.hash}'
 
   @classmethod
   def from_value(cls, value: pg.Symbolic) -> dict[str, 'Modality']:
@@ -86,7 +106,7 @@ class Modality(component.Component, pg.views.HtmlTreeView.Extension):
     def _visit(k, v, p):
       del k, p
       if isinstance(v, Modality):
-        modalities[v.referred_name] = v
+        modalities[v.id] = v
         return pg.TraverseAction.CONTINUE
       return pg.TraverseAction.ENTER
 
@@ -95,14 +115,47 @@ class Modality(component.Component, pg.views.HtmlTreeView.Extension):
 
 
 class ModalityRef(pg.Object, pg.typing.CustomTyping):
-  """References of modality objects in a symbolic tree.
+  """Lightweight placeholder for a `lf.Modality` object in a symbolic tree.
 
-  `ModalityRef` was introduced to placehold modality objects in a symbolic
-  tree, to prevent message from being chunked in the middle of a Python
-  structure.
+  `ModalityRef` acts as a reference to a `Modality` object (like `lf.Image`
+  or `lf.Audio`) within a structured object hierarchy (e.g., a `pg.Object`).
+  Instead of embedding potentially large modality data directly, `ModalityRef`
+  stores only the ID of the modality object.
+
+  This is useful in scenarios where structured objects are serialized or
+  manipulated, and it's more efficient to refer to modalities by ID rather
+  than copying their content. The `lf.ModalityRef.placehold()` class method
+  can be used to replace `Modality` instances in a symbolic object with
+  `ModalityRef` placeholders, while `lf.ModalityRef.restore()` can reinstate
+  the original `Modality` objects using a lookup table.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+  import pyglove as pg
+
+  class ImagePair(pg.Object):
+    image1: lf.Image
+    image2: lf.Image
+
+  pair = ImagePair(
+      image1=lf.Image(content=b'abc'), image2=lf.Image(content=b'def')
+  )
+  modalities = lf.Modality.from_value(pair)
+
+  # Replace Image objects with ModalityRef placeholders
+  pair_with_refs = lf.ModalityRef.placehold(pair)
+  print(pair_with_refs.image1)
+  # Output: ModalityRef(id='image:d81e5a68')
+
+  # Restore Image objects from ModalityRef placeholders
+  pair_restored = lf.ModalityRef.restore(pair_with_refs, modalities)
+  assert pair_restored.image1.content == b'abc'
+  ```
   """
 
-  name: str
+  id: str
 
   def custom_apply(
       self, path: pg.KeyPath, value_spec: pg.ValueSpec, *args, **kwargs
@@ -122,12 +175,97 @@ class ModalityRef(pg.Object, pg.typing.CustomTyping):
     """
 
     def _placehold(k, v, p):
-      del p
+      del k, p
       if isinstance(v, Modality):
-        return ModalityRef(name=value.sym_path + k)
+        return ModalityRef(id=v.id)
       return v
     return value.clone().rebind(_placehold, raise_on_no_change=False)
 
+  @classmethod
+  def restore(cls, value: pg.Symbolic, modalities: dict[str, Modality]) -> Any:
+    """Returns a copy of value by replacing refs with modality objects."""
+    def _restore(k, v, p):
+      del k, p
+      if isinstance(v, ModalityRef):
+        modality_object = modalities.get(v.id)
+        if modality_object is None:
+          raise ValueError(
+              f'Modality {v.id} not found in modalities {modalities.keys()}'
+          )
+        return modality_object
+      return v
+    return value.rebind(_restore, raise_on_no_change=False)
+
 
 class ModalityError(RuntimeError):  # pylint: disable=g-bad-exception-name
   """Exception raised when modality is not supported."""
+
+
+#
+# Context managers to deal with modality objects.
+#
+
+
+_TLS_MODALITY_CAPTURE_SCOPE = '__modality_capture_scope__'
+_TLS_MODALITY_AS_REF = '__format_modality_as_ref__'
+
+
+def format_modality_as_ref(enabled: bool = True) -> ContextManager[None]:
+  """A context manager that formats modality objects as references."""
+  return pg.object_utils.thread_local_value_scope(
+      _TLS_MODALITY_AS_REF, enabled, False
+  )
+
+
+class _ModalityCaptureContext:
+  """A context to capture modality objects when being rendered."""
+
+  def __init__(self):
+    self._references: dict[str, pg.Ref[Modality]] = {}
+
+  def capture(self, modality: Modality) -> None:
+    """Captures the modality object."""
+    self._references[modality.id] = pg.Ref(modality)
+
+  @property
+  def references(self) -> dict[str, pg.Ref[Modality]]:
+    """Returns the modality references captured in this context."""
+    return self._references
+
+
+@contextlib.contextmanager
+def capture_rendered_modalities() -> Iterator[dict[str, pg.Ref[Modality]]]:
+  """Capture modality objects whose references is being rendered.
+
+  Example:
+    ```
+    image = lf.Image.from_url(...)
+    with lf.modality.capture_rendered_modalities() as rendered_modalities:
+      with lf.modality.format_modality_as_ref():
+        print(f'Hello {image}')
+    self.assertEqual(rendered_modalities, {'image:<hash>': pg.Ref(image)})
+    ```
+  """
+  context = get_modality_capture_context()
+  top_level = context is None
+  if top_level:
+    context = _ModalityCaptureContext()
+    pg.object_utils.thread_local_set(_TLS_MODALITY_CAPTURE_SCOPE, context)
+
+  try:
+    yield context.references  # pylint: disable=attribute-error
+  finally:
+    if top_level:
+      pg.object_utils.thread_local_del(_TLS_MODALITY_CAPTURE_SCOPE)
+
+
+def get_modality_capture_context() -> _ModalityCaptureContext | None:
+  """Returns the current modality capture context."""
+  return pg.object_utils.thread_local_get(_TLS_MODALITY_CAPTURE_SCOPE, None)
+
+
+def _camel_to_snake(name: str) -> str:
+  """Converts a camelCase name to snake_case."""
+  return re.sub(
+      pattern=r'([A-Z]+)', repl=r'_\1', string=name
+  ).lower().lstrip('_')

langfun/core/modality_test.py

@@ -29,34 +29,64 @@ class ModalityTest(unittest.TestCase):
 
   def test_basic(self):
     v = CustomModality('a')
-    self.assertIsNone(v.referred_name)
+    self.assertEqual(v.id, 'custom_modality:0cc175b9')
     self.assertEqual(str(v), "CustomModality(\n  content = 'a'\n)")
     self.assertEqual(v.hash, '0cc175b9')
 
     _ = pg.Dict(metadata=pg.Dict(x=pg.Dict(metadata=pg.Dict(y=v))))
-    self.assertEqual(v.referred_name, 'x.metadata.y')
+    self.assertEqual(v.id, 'custom_modality:0cc175b9')
     self.assertEqual(str(v), "CustomModality(\n  content = 'a'\n)")
     with modality.format_modality_as_ref():
-      self.assertEqual(str(v), '<<[[x.metadata.y]]>>')
+      self.assertEqual(str(v), '<<[[custom_modality:0cc175b9]]>>')
+
+  def test_capture_rendered_modalities(self):
+    x = CustomModality('a')
+    y = CustomModality('b')
+    z = CustomModality('b')
+
+    with modality.capture_rendered_modalities() as rendered_modalities:
+      with modality.format_modality_as_ref():
+        self.assertEqual(
+            f'Hello {x} {y} {z}',
+            (
+                'Hello <<[[custom_modality:0cc175b9]]>> '
+                '<<[[custom_modality:92eb5ffe]]>> '
+                '<<[[custom_modality:92eb5ffe]]>>'
+            )
+        )
+    self.assertEqual(len(rendered_modalities), 2)
+    self.assertIs(rendered_modalities['custom_modality:0cc175b9'].value, x)
+    # y and z share the same content will be treated as the same object.
+    self.assertIs(rendered_modalities['custom_modality:92eb5ffe'].value, z)
 
 
 class ModalityRefTest(unittest.TestCase):
 
-  def test_placehold(self):
+  def test_placehold_and_restore(self):
     class A(pg.Object):
       x: Any
       y: Any
 
-    a = A(x=dict(z=CustomModality('a')), y=CustomModality('b'))
+    image_a = CustomModality('a')
+    image_b = CustomModality('b')
+    a = A(x=dict(z=image_a), y=image_b)
+    a_placehold = modality.ModalityRef.placehold(a)
     self.assertEqual(
-        modality.ModalityRef.placehold(a),
-        A(x=dict(z=modality.ModalityRef('x.z')), y=modality.ModalityRef('y')),
+        a_placehold,
+        A(x=dict(z=modality.ModalityRef(image_a.id)),
+          y=modality.ModalityRef(image_b.id)),
+    )
+    a_restore = modality.ModalityRef.restore(
+        a_placehold.clone(),
+        {image_a.id: image_a, image_b.id: image_b},
     )
+    self.assertTrue(pg.eq(a_restore, a))
     self.assertEqual(
         modality.ModalityRef.placehold(a.x),
-        # The prefix 'x' of referred name is preserved.
-        dict(z=modality.ModalityRef('x.z')),
+        dict(z=modality.ModalityRef(image_a.id)),
     )
+    with self.assertRaisesRegex(ValueError, 'Modality .* not found'):
+      modality.ModalityRef.restore(a_placehold, {image_a.id: image_a})
 
   def test_from_value(self):
     class A(pg.Object):
@@ -68,8 +98,8 @@ class ModalityRefTest(unittest.TestCase):
         pg.eq(
             modality.Modality.from_value(a),
             {
-                'x.z': CustomModality('a'),
-                'y': CustomModality('b'),
+                'custom_modality:0cc175b9': CustomModality('a'),
+                'custom_modality:92eb5ffe': CustomModality('b'),
             },
         )
     )
@@ -77,7 +107,7 @@ class ModalityRefTest(unittest.TestCase):
         pg.eq(
             modality.Modality.from_value(a.x.z),
             {
-                'x.z': CustomModality('a'),
+                'custom_modality:0cc175b9': CustomModality('a'),
             },
         )
     )

langfun/core/natural_language.py

@@ -11,7 +11,7 @@
 # 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.
-"""Natural language utilities."""
+"""Natural language formatting."""
 
 import abc
 import pyglove as pg

langfun/core/sampling.py

@@ -38,10 +38,10 @@ def sweep(
         Union[message_lib.Message, BaseException, None],  # LM output.
     ],
 ]:
-  """Sweeps the input/output of this LangFunc concurrently.
+  """Sweeps the input/output of a LangFunc search space concurrently.
 
   Args:
-    lfun: An LangFunc object that contains `pg.oneof` as the search space 
+    lfun: An LangFunc object that contains `pg.oneof` as the search space
       for sampling.
     num_examples: Number of examples to sample.
     max_workers: Max number of concurrent workers to do sampling.
@@ -84,10 +84,10 @@ def random_sample(
         Union[message_lib.Message, BaseException, None],  # LM output.
     ],
 ]:
-  """Random samples the input/output of this LangFunc concurrently.
+  """Random samples the input/output of a LangFunc search space concurrently.
 
   Args:
-    lfun: An LangFunc object that contains `pg.oneof` as the search space 
+    lfun: An LangFunc object that contains `pg.oneof` as the search space
       for sampling.
     num_examples: Number of examples to sample.
     max_workers: Max number of concurrent workers to do sampling.

langfun/core/sampling_test.py

@@ -39,8 +39,13 @@ class SamplingTest(unittest.TestCase):
     l = LangFunc('Compute {{x}} and {{y}}', x=pg.oneof([1, 2]))
     with component.context(lm=ExcitedEchoer()):
       samples = list(sampling.sweep(l, y=pg.oneof([3, 4])))
-      samples = sorted(samples, key=lambda x: (x[0].x, x[0].y))
-
+      samples = sorted(
+          samples,
+          key=lambda x: (
+              x[0].__template_input__.x,
+              x[0].__template_input__.y
+          )
+      )
     self.assertEqual(
         samples,
         [
@@ -57,7 +62,12 @@ class SamplingTest(unittest.TestCase):
       samples = list(
           sampling.random_sample(l, y=pg.oneof([2, 4]), num_examples=3, seed=1)
       )
-      samples = sorted(samples, key=lambda x: (x[0].x, x[0].y))
+      samples = sorted(
+          samples, key=lambda x: (
+              x[0].__template_input__.x,
+              x[0].__template_input__.y
+          )
+      )
 
     self.assertEqual(
         samples,
@@ -97,7 +107,13 @@ class SamplingTest(unittest.TestCase):
         silence_on_errors=(AttributeError,),
         ignore_examples_with_errors=False))
 
-    samples = sorted(samples, key=lambda x: (x[0].x, x[0].y))
+    samples = sorted(
+        samples,
+        key=lambda x: (
+            x[0].__template_input__.x,
+            x[0].__template_input__.y
+        )
+    )
     self.assertEqual(
         [x[0] for x in samples],
         [

langfun/core/structured/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The Langfun Authors
+# Copyright 2025 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.
@@ -16,29 +16,7 @@
 # pylint: disable=g-bad-import-order
 # pylint: disable=g-importing-member
 
-from langfun.core.structured.schema import include_method_in_prompt
-
-from langfun.core.structured.schema import Missing
-from langfun.core.structured.schema import MISSING
-from langfun.core.structured.schema import Unknown
-from langfun.core.structured.schema import UNKNOWN
-
-from langfun.core.structured.schema import Schema
-from langfun.core.structured.schema import SchemaProtocol
-from langfun.core.structured.schema import schema_spec
-
-from langfun.core.structured.schema import SchemaError
-from langfun.core.structured.schema import JsonError
-
-from langfun.core.structured.schema import class_dependencies
-from langfun.core.structured.schema import class_definition
-from langfun.core.structured.schema import class_definitions
-from langfun.core.structured.schema import annotation
-from langfun.core.structured.schema import structure_from_python
-
-from langfun.core.structured.schema import schema_repr
-from langfun.core.structured.schema import source_form
-from langfun.core.structured.schema import value_repr
+from langfun.core.structured.schema import *
 
 from langfun.core.structured.schema_generation import generate_class
 from langfun.core.structured.schema_generation import classgen_example

langfun/core/structured/completion.py

@@ -116,15 +116,10 @@ class _CompleteStructure(mapping.Mapping):
     )
 
   def postprocess_result(self, result: Any) -> Any:
-    """Postprocess result."""
+    """Postprocesses result."""
     # Try restore modality objects from the input value to output value.
-    modalities = self.modalities(self.input)
-    if modalities:
-      # Remove the `input` prefix for all entries.
-      modalities = pg.object_utils.flatten(
-          pg.object_utils.canonicalize(modalities)['input']
-      )
-      result.rebind(modalities)
+    if modalities := self.modalities(self.input):
+      result = lf.ModalityRef.restore(result, modalities)
     return result
 
   def globals(self):
@@ -156,7 +151,7 @@ class _CompleteStructure(mapping.Mapping):
   #
 
   def has_modality_refs(self, value: Any) -> bool:
-    """Returns true if the value has modalities."""
+    """Returns True if the value has modalities."""
     return not isinstance(value, lf.Modality) and pg.contains(
         value, type=lf.Modality
     )
@@ -186,41 +181,36 @@ def complete(
     returns_message: bool = False,
     **kwargs,
 ) -> Any:
-  """Complete a symbolic value by filling its missing fields.
-
-  Examples:
-
-    ```
-    class FlightDuration:
-      hours: int
-      minutes: int
-
-    class Flight(pg.Object):
-      airline: str
-      flight_number: str
-      departure_airport_code: str
-      arrival_airport_code: str
-      departure_time: str
-      arrival_time: str
-      duration: FlightDuration
-      stops: int
-      price: float
-
-    prompt = '''
-      Information about flight UA2631.
-      '''
-
-    r = lf.query(prompt, Flight)
-    assert isinstance(r, Flight)
-    assert r.airline == 'United Airlines'
-    assert r.departure_airport_code == 'SFO'
-    assert r.duration.hour = 7
-    ```
+  """Completes a symbolic value by filling its missing fields using an LLM.
+
+  `lf.complete` is used to fill in missing information in structured
+  data. It takes a partially defined `pg.Object` instance where some fields
+  are marked as `lf.MISSING`, and uses a language model to infer and
+  populate those fields based on the provided values.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+  import pyglove as pg
+
+  class Country(pg.Object):
+    name: str
+    capital: str = lf.MISSING
+    population: int = lf.MISSING
+
+  # Filling missing fields of Country(name='France')
+  country = lf.complete(Country(name='France'), lm=lf.llms.Gemini25Flash())
+  print(country)
+  # Output: Country(name='France', capital='Paris', population=67000000)
+  ```
 
   Args:
-    input_value: A symbolic value that may contain missing values.
-    default: The default value if parsing failed. If not specified, error will
-      be raised.
+    input_value: A symbolic value that may contain missing values marked
+      by `lf.MISSING`.
+    default: The default value to return if parsing fails. If
+      `lf.RAISE_IF_HAS_ERROR` is used (default), an error will be raised
+      instead.
     lm: The language model to use. If not specified, the language model from
       `lf.context` context manager will be used.
     examples: An optional list of fewshot examples for helping parsing. If None,
@@ -236,10 +226,10 @@ def complete(
     returns_message: If True, returns `lf.Message` as the output, instead of
       returning the structured `message.result`.
     **kwargs: Keyword arguments passed to the
-      `lf.structured.NaturalLanguageToStructureed` transform.
+      `lf.structured.Mapping` transform.
 
   Returns:
-    The result based on the schema.
+    The input object with missing fields completed by LLM.
   """
   t = _CompleteStructure(
       input=schema_lib.mark_missing(input_value),