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

langfun/core/modalities/mime.py

@@ -15,6 +15,7 @@
 
 import base64
 import functools
+import hashlib
 from typing import Annotated, Any, Iterable, Type, Union
 import langfun.core as lf
 # Placeholder for Google-internal internet access import.
@@ -87,6 +88,14 @@ class Mime(lf.Modality):
     """Returns True if the MIME type is a binary type."""
     return not self.is_text
 
+  @property
+  def hash(self) -> str:
+    """Returns the hash of the MIME content."""
+    # Hash the URI to avoid downloading the content.
+    if self.uri is not None:
+      return hashlib.md5(self.uri.encode()).hexdigest()[:8]
+    return super().hash
+
   def to_text(self) -> str:
     """Returns the text content of the MIME type."""
     if not self.is_text:

langfun/core/modality.py

@@ -14,23 +14,15 @@
 """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__'
-
-
-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 Modality(component.Component, pg.views.HtmlTreeView.Extension):
   """Base class for multimodal object."""
 
@@ -39,15 +31,18 @@ class Modality(component.Component, pg.views.HtmlTreeView.Extension):
 
   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 +65,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 +78,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
 
@@ -102,7 +94,7 @@ class ModalityRef(pg.Object, pg.typing.CustomTyping):
   structure.
   """
 
-  name: str
+  id: str
 
   def custom_apply(
       self, path: pg.KeyPath, value_spec: pg.ValueSpec, *args, **kwargs
@@ -122,12 +114,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/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/completion.py

@@ -118,13 +118,8 @@ class _CompleteStructure(mapping.Mapping):
   def postprocess_result(self, result: Any) -> Any:
     """Postprocess 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):

langfun/core/structured/completion_test.py

@@ -407,22 +407,17 @@ class CompleteStructureTest(unittest.TestCase):
       image: modalities.Image
       name: str
 
+    image_elephant = modalities.Image.from_bytes(b'image_of_elephant')
+    image_rabbit = modalities.Image.from_bytes(b'image_of_rabbit')
     input_value = schema_lib.mark_missing(
-        Animal.partial(
-            modalities.Image.from_bytes(b'image_of_elephant'),
-        )
+        Animal.partial(image_elephant)
     )
     l = completion._CompleteStructure(
         input=input_value,
         examples=[
             mapping.MappingExample(
-                input=Animal.partial(
-                    modalities.Image.from_bytes(b'image_of_rabbit')
-                ),
-                output=Animal(
-                    modalities.Image.from_bytes(b'image_of_rabbit'),
-                    'rabbit',
-                ),
+                input=Animal.partial(image_rabbit),
+                output=Animal(image_rabbit, 'rabbit'),
             )
         ],
     )
@@ -430,7 +425,7 @@ class CompleteStructureTest(unittest.TestCase):
     self.maxDiff = None
     self.assertEqual(
         lm_input.text,
-        inspect.cleandoc("""
+        inspect.cleandoc(f"""
             Please generate the OUTPUT_OBJECT by completing the MISSING fields from the last INPUT_OBJECT.
 
             INSTRUCTIONS:
@@ -457,22 +452,22 @@ class CompleteStructureTest(unittest.TestCase):
               ```python
               Animal(
                 image=ModalityRef(
-                  name='examples[0].input.image'
+                  id='{image_rabbit.id}'
                 ),
                 name=MISSING(str)
               )
               ```
 
             MODALITY_REFERENCES:
-              {
-                'examples[0].input.image': <<[[examples[0].input.image]]>>
-              }
+              {{
+                '{image_rabbit.id}': <<[[{image_rabbit.id}]]>>
+              }}
 
             OUTPUT_OBJECT:
               ```python
               Animal(
                 image=ModalityRef(
-                  name='examples[0].output.image'
+                  id='{image_rabbit.id}'
                 ),
                 name='rabbit'
               )
@@ -483,16 +478,16 @@ class CompleteStructureTest(unittest.TestCase):
               ```python
               Animal(
                 image=ModalityRef(
-                  name='input.image'
+                  id='{image_elephant.id}'
                 ),
                 name=MISSING(str)
               )
               ```
 
             MODALITY_REFERENCES:
-              {
-                'input.image': <<[[input.image]]>>
-              }
+              {{
+                '{image_elephant.id}': <<[[{image_elephant.id}]]>>
+              }}
 
             OUTPUT_OBJECT:
             """),
@@ -500,39 +495,27 @@ class CompleteStructureTest(unittest.TestCase):
     self.assertTrue(
         pg.eq(
             {
-                'examples': lm_input.get('examples'),
-                'input': lm_input.get('input'),
+                'examples': lm_input.__template_input__.examples,
+                'input': lm_input.__template_input__.mapping_request.input,
             },
             {
                 'examples': [
                     mapping.MappingExample(
-                        input=Animal.partial(
-                            image=modalities.Image.from_bytes(
-                                b'image_of_rabbit'
-                            )
-                        ),
-                        output=Animal.partial(
-                            image=modalities.Image.from_bytes(
-                                b'image_of_rabbit'
-                            ),
-                            name='rabbit',
-                        ),
+                        input=Animal.partial(image_rabbit),
+                        output=Animal.partial(image_rabbit, 'rabbit'),
                     )
                 ],
-                'input': Animal(
-                    image=modalities.Image.from_bytes(b'image_of_elephant'),
-                    name=schema_lib.MISSING,
-                ),
+                'input': Animal(image_elephant, name=schema_lib.MISSING),
             },
         )
     )
     lm_output = l(
         input=input_value,
-        lm=fake.StaticResponse(inspect.cleandoc("""
+        lm=fake.StaticResponse(inspect.cleandoc(f"""
             ```python
             Animal(
               image=ModalityRef(
-                name='input.image'
+                id='{image_elephant.id}'
               ),
               name='elephant'
             )
@@ -542,10 +525,7 @@ class CompleteStructureTest(unittest.TestCase):
     self.assertTrue(
         pg.eq(
             lm_output.result,
-            Animal(
-                image=modalities.Image.from_bytes(b'image_of_elephant'),
-                name='elephant',
-            ),
+            Animal(image=image_elephant, name='elephant'),
         )
     )
 

langfun/core/structured/mapping.py

@@ -127,6 +127,8 @@ class MappingExample(lf.NaturalLanguageFormattable,
   ) -> str:
     if isinstance(value, str):
       return value
+    if isinstance(value, lf.Message):
+      return str(value)
     if isinstance(value, lf.Modality):
       with lf.modality.format_modality_as_ref():
         return str(value)
@@ -192,9 +194,7 @@ class MappingExample(lf.NaturalLanguageFormattable,
 
     def render_value(view, *, value, **kwargs):
       if isinstance(value, lf.Template):
-        # Make a shallow copy to make sure modalities are rooted by
-        # the input.
-        value = value.clone().render()
+        value = value.render()
       if value is None:
         return None
       return view.render(value, **kwargs)
@@ -286,12 +286,8 @@ class Mapping(lf.LangFunc):
   @property
   def mapping_request(self) -> MappingExample:
     """Returns a MappingExample as the mapping request."""
-    if isinstance(self.input, lf.Message):
-      input_value = self.input.text
-    else:
-      input_value = pg.Ref(self.input)
     return MappingExample(
-        input=input_value,
+        input=pg.Ref(self.input),
         schema=pg.Ref(self.schema),
         context=self.context,
     )
@@ -402,11 +398,6 @@ class Mapping(lf.LangFunc):
 
   def transform_input(self, lm_input: lf.Message) -> lf.Message:
     # Find modalities to fill the input message.
-    lm_input.metadata.update(
-        examples=pg.Ref(self.examples),
-        input=pg.Ref(self.input),
-        schema=pg.Ref(self.schema) if self.schema is not None else None,
-    )
     if isinstance(self.input, lf.Message):
       lm_input.source = self.input
     return lm_input

langfun/core/structured/querying.py

@@ -529,24 +529,22 @@ def query(
     ).render(message_cls=lf.SystemMessage)
 
   # Normalize query input.
-  if isinstance(prompt, (lf.Message, str)):
+  if isinstance(prompt, str):
     # Query with structured output.
     prompt_kwargs = kwargs.copy()
     prompt_kwargs.pop('template_str', None)
     query_input = lf.Template.from_value(prompt, **prompt_kwargs)
+  elif isinstance(prompt, lf.Message):
+    query_input = prompt
   elif isinstance(prompt, lf.Template):
-    # Create a copy of the prompt if it has a parent object, so all child
-    # modality objects could be referred by path relative to the prompt.
-    query_input = prompt.clone() if prompt.sym_parent is not None else prompt
-
     # Attach template metadata from kwargs. This is used to pass through fields
     # from kwargs to the rendered message.
-    template_metadata = {
-        k: v for k, v in kwargs.items() if k.startswith('metadata_')
-    }
-    query_input.rebind(
-        template_metadata, skip_notification=True, raise_on_no_change=False
+    prompt.rebind(
+        {k: v for k, v in kwargs.items() if k.startswith('metadata_')},
+        skip_notification=True,
+        raise_on_no_change=False
     )
+    query_input = prompt
   elif pg.MISSING_VALUE == prompt:
     query_input = lf.UserMessage('')
   else:
@@ -665,7 +663,11 @@ def query(
 
   if returns_message:
     return output_message
-  return output_message.text if schema in (None, str) else output_message.result
+  if schema not in (None, str):
+    return output_message.result
+  if returns_message or output_message.referred_modalities:
+    return output_message
+  return output_message.text
 
 
 async def aquery(