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

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(

langfun/core/structured/querying_test.py

@@ -249,35 +249,60 @@ class QueryTest(unittest.TestCase):
 
   def test_root_modality_to_structure_render(self):
     lm = fake.StaticResponse('1')
+    image = modalities.Image.from_bytes(b'mock_image')
     self.assert_render(
-        modalities.Image.from_bytes(b'mock_image'),
+        image,
         int,
         lm=lm,
-        expected_snippet='\n\nREQUEST:\n  <<[[input]]>>\n\n',
+        expected_snippet=f'\n\nREQUEST:\n  <<[[{image.id}]]>>\n\n',
         expected_modalities=1,
     )
 
   def test_root_modality_to_str_render(self):
     lm = fake.StaticResponse('1')
+    modality = modalities.Image.from_bytes(b'mock_image')
     self.assert_render(
-        modalities.Image.from_bytes(b'mock_image'),
+        modality,
         None,
         lm=lm,
-        expected_snippet='<<[[input]]>>',
+        expected_snippet=f'<<[[{modality.id}]]>>',
         exact_match=True,
         expected_modalities=1,
     )
 
   def test_str_with_modality_to_str_render(self):
     lm = fake.StaticResponse('A cat and a mouse.')
+    cat_image = modalities.Image.from_bytes(b'cat_image')
+    mouse_image = modalities.Image.from_bytes(b'mouse_image')
     self.assert_render(
         'What are these? {{this_image}} and {{that_image}}',
         None,
-        this_image=modalities.Image.from_bytes(b'cat_image'),
-        that_image=modalities.Image.from_bytes(b'mouse_image'),
+        this_image=cat_image,
+        that_image=mouse_image,
         lm=lm,
         expected_snippet=(
-            'What are these? <<[[this_image]]>> and <<[[that_image]]>>'
+            f'What are these? <<[[{cat_image.id}]]>> and '
+            f'<<[[{mouse_image.id}]]>>'
+        ),
+        exact_match=True,
+        expected_modalities=2,
+    )
+
+  def test_message_with_modality_to_str_render(self):
+    lm = fake.StaticResponse('A cat and a mouse.')
+    cat_image = modalities.Image.from_bytes(b'cat_image')
+    mouse_image = modalities.Image.from_bytes(b'mouse_image')
+    self.assert_render(
+        lf.Template(
+            'What are these? {{this_image}} and {{that_image}}',
+            this_image=cat_image,
+            that_image=mouse_image,
+        ).render(),
+        None,
+        lm=lm,
+        expected_snippet=(
+            f'What are these? <<[[{cat_image.id}]]>> and '
+            f'<<[[{mouse_image.id}]]>>'
         ),
         exact_match=True,
         expected_modalities=2,
@@ -285,33 +310,33 @@ class QueryTest(unittest.TestCase):
 
   def test_structure_with_modality_to_str_render(self):
     lm = fake.StaticResponse('A cat and a mouse.')
+    cat_image = modalities.Image.from_bytes(b'cat_image')
+    mouse_image = modalities.Image.from_bytes(b'mouse_image')
     self.assert_render(
-        [
-            modalities.Image.from_bytes(b'cat_image'),
-            modalities.Image.from_bytes(b'mouse_image'),
-        ],
+        [cat_image, mouse_image],
         None,
         lm=lm,
-        expected_snippet='`[<<[[input[0]]]>>, <<[[input[1]]]>>]`',
+        expected_snippet=(
+            f'`[<<[[{cat_image.id}]]>>, <<[[{mouse_image.id}]]>>]`'
+        ),
         exact_match=True,
         expected_modalities=2,
     )
 
   def test_structure_with_modality_to_structure_render(self):
     lm = fake.StaticResponse('["cat", "mouse"]')
+    cat_image = modalities.Image.from_bytes(b'cat_image')
+    mouse_image = modalities.Image.from_bytes(b'mouse_image')
     self.assert_render(
-        [
-            modalities.Image.from_bytes(b'cat_image'),
-            modalities.Image.from_bytes(b'mouse_image'),
-        ],
+        [cat_image, mouse_image],
         list[str],
         lm=lm,
-        expected_snippet=inspect.cleandoc("""
+        expected_snippet=inspect.cleandoc(f"""
             REQUEST:
               ```python
               [
-                <<[[input[0]]]>>,
-                <<[[input[1]]]>>
+                <<[[{cat_image.id}]]>>,
+                <<[[{mouse_image.id}]]>>
               ]
               ```
             """),
@@ -320,25 +345,25 @@ class QueryTest(unittest.TestCase):
 
   def test_structure_with_modality_and_examples_to_structure_render(self):
     lm = fake.StaticResponse('["cat", "mouse"]')
+    cat_image = modalities.Image.from_bytes(b'cat_image')
+    mouse_image = modalities.Image.from_bytes(b'mouse_image')
+    dog_image = modalities.Image.from_bytes(b'dog_image')
     self.assert_render(
-        [
-            modalities.Image.from_bytes(b'cat_image'),
-            modalities.Image.from_bytes(b'mouse_image'),
-        ],
+        [cat_image, mouse_image],
         list[str],
         examples=[
             mapping.MappingExample(
-                input=[modalities.Image.from_bytes(b'dog_image')],
+                input=[dog_image],
                 schema=list[str],
                 output=['dog'],
             ),
         ],
         lm=lm,
-        expected_snippet=inspect.cleandoc("""
+        expected_snippet=inspect.cleandoc(f"""
             REQUEST:
               ```python
               [
-                <<[[examples[0].input[0]]]>>
+                <<[[{dog_image.id}]]>>
               ]
               ```
 
@@ -356,8 +381,8 @@ class QueryTest(unittest.TestCase):
             REQUEST:
               ```python
               [
-                <<[[input[0]]]>>,
-                <<[[input[1]]]>>
+                <<[[{cat_image.id}]]>>,
+                <<[[{mouse_image.id}]]>>
               ]
               ```
 
@@ -369,6 +394,17 @@ class QueryTest(unittest.TestCase):
         expected_modalities=3,
     )
 
+  def test_query_with_modality_output(self):
+    cat_image = modalities.Image.from_bytes(b'cat_image')
+    lm = fake.StaticResponse(
+        lf.Template('Here you go: {{image}}', image=cat_image).render(
+            message_cls=lf.AIMessage
+        )
+    )
+    response = querying.query('Generate a cat image', lm=lm)
+    self.assertIsInstance(response, lf.AIMessage)
+    self.assertEqual(response.modalities(), [cat_image])
+
   def test_multiple_queries(self):
     self.assertEqual(
         querying.query(
@@ -545,7 +581,7 @@ class QueryTest(unittest.TestCase):
             )
         ).input,
     )
-    self.assertIsNotNone(output.get_modality('image'))
+    self.assertEqual(len(output.referred_modalities), 1)
 
   def test_query_and_reduce(self):
     self.assertEqual(

langfun/core/template.py

@@ -171,6 +171,7 @@ class Template(
 
     # Last render output.
     self._cached_render_output = None
+    self._referred_modalities = None
 
   @property
   def render_output(self) -> message_lib.Message | None:
@@ -322,24 +323,46 @@ class Template(
                 compact=True,
                 python_format=True,
             ):
-              # Natural language formattable objects will be returned in natural
-              # language when they are directly returned as rendering elements
+
+              # Capture the modality objects whose references are being
+              # rendered
               # in the template.
-              with modality.format_modality_as_ref():
-                rendered_text = self._template.render(**inputs)
+              with modality.capture_rendered_modalities() as modality_refs:
+
+                # Natural language formattable objects will be returned in
+                # natural language when they are directly returned as rendering
+                # elements in the template.
+                with modality.format_modality_as_ref():
+                  rendered_text = self._template.render(**inputs)
 
-            # Carry additional metadata.
-            metadata = self.additional_metadata()
+              # Carry the modality references passed from the constructor.
+              # This is to support modality objects that is already rendered
+              # in the template string.
+              if self._referred_modalities:
+                modality_refs.update(self._referred_modalities)
 
         if self.clean:
           rendered_text = rendered_text.strip()
 
-        metadata.update(
-            {k: pg.Ref(v) for k, v in inputs.items() if not inspect.ismethod(v)}
-        )
+        # Fill message metadata.
+        metadata = {
+            '__template_input__': {
+                k: pg.Ref(v) for k, v in inputs.items()
+                if not inspect.ismethod(v)
+            },
+        }
+
+        # Carry additional metadata.
+        # TODO(daiyip): Consider to put additional metadata into a separate
+        # key under `metadata`.
+        metadata.update(self.additional_metadata())
 
         # Fill the variables for rendering the template as metadata.
-        message = message_cls(text=rendered_text, metadata=metadata)
+        message = message_cls(
+            text=rendered_text,
+            referred_modalities=modality_refs,
+            metadata=metadata
+        )
 
         # Tag input as rendered message.
         message.tag(message_lib.Message.TAG_RENDERED)
@@ -518,10 +541,13 @@ class Template(
     if isinstance(value, str):
       return cls(template_str=value, **kwargs)
     if isinstance(value, message_lib.Message):
-      kwargs.update(value.metadata)
-      return cls(template_str=value.text, **kwargs)
+      for k, v in value.metadata.sym_items():  # pylint: disable=attribute-error
+        kwargs[_ADDITIONAL_METADATA_PREFIX + k] = v
+      t = cls(template_str=value.text, **kwargs)
+      t._referred_modalities = value.referred_modalities
+      return t
     if isinstance(value, Template):
-      lfun = cls(template_str=value.template_str, **kwargs)
+      lfun = cls(template_str=value.template_str, **kwargs)  # pylint: disable=attribute-error
       # So lfun could acccess all attributes from value.
       lfun.sym_setparent(value)
       return lfun