langfun 0.1.2.dev202509120804__py3-none-any.whl → 0.1.2.dev202512150805__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/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),

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

@@ -23,7 +23,7 @@ import pyglove as pg
 
 @pg.use_init_args(['examples'])
 class _DescribeStructure(mapping.Mapping):
-  """Describe a structured value in natural language."""
+  """Describes a structured value in natural language."""
 
   input_title = 'PYTHON_OBJECT'
   context_title = 'CONTEXT_FOR_DESCRIPTION'
@@ -47,64 +47,68 @@ def describe(
     cache_seed: int | None = 0,
     **kwargs,
 ) -> str:
-  """Describes a structured value using natural language.
-
-  Examples:
-
-    ```
-    class FlightDuration(pg.Object):
-      hours: int
-      minutes: int
-
-    class Flight(pg.Object):
-      airline: str
-      flight_number: str
-      departure_airport: str
-      arrival_airport: str
-      departure_time: str
-      arrival_time: str
-      duration: FlightDuration
-      stops: int
-      price: float
-
-    text = lf.describe(
-        Flight(
-            airline='United Airlines',
-            flight_number='UA2631',
-            depature_airport: 'SFO',
-            arrival_airport: 'JFK',
-            depature_time: '2023-09-07T05:15:00',
-            arrival_time: '2023-09-07T12:12:00',
-            duration: FlightDuration(
-                hours=7,
-                minutes=57
-            ),
-            stops=1,
-            price=227,
-        ))
-    print(text)
-
-    >> The flight is operated by United Airlines, has the flight number UA2631,
-    >> departs from San Francisco International Airport (SFO), arrives at John
-    >> F. Kennedy International Airport (JFK), It departs at
-    >> 2023-09-07T05:15:00, arrives at 2023-09-07T12:12:00, has a duration of 7
-    >> hours and 57 minutes, makes 1 stop, and costs $227.
-    ```
+  """Describes a structured value in natural language using an LLM.
+
+  `lf.describe` takes a Python object, often a `pg.Object` instance,
+  and uses a language model to generate a human-readable, natural language
+  description of its content. It is the inverse of `lf.parse`.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+  import pyglove as pg
+
+  class FlightDuration(pg.Object):
+    hours: int
+    minutes: int
+
+  class Flight(pg.Object):
+    airline: str
+    flight_number: str
+    departure_airport: str
+    arrival_airport: str
+    departure_time: str
+    arrival_time: str
+    duration: FlightDuration
+    stops: int
+    price: float
+
+  flight_info = Flight(
+      airline='United Airlines',
+      flight_number='UA2631',
+      departure_airport='SFO',
+      arrival_airport='JFK',
+      departure_time='2023-09-07T05:15:00',
+      arrival_time='2023-09-07T12:12:00',
+      duration=FlightDuration(hours=7, minutes=57),
+      stops=1,
+      price=227,
+  )
+
+  description = lf.describe(flight_info, lm=lf.llms.Gemini25Flash())
+  print(description)
+  # Possible output:
+  # The flight is operated by United Airlines, with the flight number UA2631,
+  # departing from SFO at 2023-09-07T05:15:00 and arriving at JFK at
+  # 2023-09-07T12:12:00. The flight duration is 7 hours and 57 minutes,
+  # with 1 stop, and costs $227.
+  ```
 
   Args:
     value: A structured value to be mapped.
     context: The context information for describing the structured value.
     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,
-      the default one-shot example will be added.
+    examples: An optional list of fewshot examples for guiding description.
+      If None, default examples will be used.
     cache_seed: Seed for computing cache key. The cache key is determined by a
       tuple of (lm, prompt, cache seed). If None, cache will be disabled for
       the query even cache is configured by the LM.
-    **kwargs: Keyword arguments passed to the `lf.structured.DescribeStructure`.
+    **kwargs: Keyword arguments passed to the `_DescribeStructure`.
 
   Returns:
-    The parsed result based on the schema.
+    A natural language description of the input value.
   """
   return _DescribeStructure(
       input=value,
@@ -115,10 +119,10 @@ def describe(
 
 
 def default_describe_examples() -> list[mapping.MappingExample]:
-  """Default describe examples."""
+  """Returns default examples for `lf.describe`."""
 
   class Country(pg.Object):
-    """A example dataclass for structured mapping."""
+    """An example dataclass for structured mapping."""
 
     name: str
     continents: list[