langfun 0.1.2.dev202509120804__py3-none-any.whl → 0.1.2.dev202512040805__py3-none-any.whl

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[

langfun/core/structured/function_generation.py

@@ -26,10 +26,10 @@ import pyglove as pg
 
 
 def unittest_gen(signature, lm, num_retries=1):
-  """Generates unit tests for a python function signature."""
+  """Generates unit tests for a Python function signature."""
 
   class UnitTest(pg.Object):
-    """A valid unit test for a python function."""
+    """A valid unit test for a Python function."""
 
     input: dict[str, Any]
     expected_output: Any
@@ -55,7 +55,7 @@ def unittest_gen(signature, lm, num_retries=1):
 
 
 def unittest_with_test_cases(f, unittests):
-  """Applies unit tests to a python function to be tested."""
+  """Applies unit tests to a Python function to be tested."""
   if not unittests:
     raise ValueError(f"No unit tests provided: {unittests}")
 
@@ -87,10 +87,10 @@ def _function_gen(
     ] = None,
     unittest_num_retries: int = 1,
 ):
-  """Generates a python function with LLM and verify its quality with unit testing."""
+  """Generates a Python function with LLM and verifies it with unit testing."""
 
   class PythonFunctionPrompt(template.Template):
-    r"""A template for a python function generation.
+    r"""A template for a Python function generation.
 
     Please reply to the last PYTHON_FUNCTION_SIGNATURE with a self-sufficient,
     error-free, and efficiently coded PYTHON_FUNCTION, crafted to the standards
@@ -195,11 +195,28 @@ def function_gen(
     ] = None,
     unittest_num_retries: int = 1,
 ):
-  """A decorator for automating function generation using a language model.
+  r"""Decorator for generating function implementations using an LLM.
 
-  This decorator should be applied to functions that are not yet implemented. It
-  facilitates the implementation via the specified LLM, ensuring
-  quality through unit tests.
+  `lf.function_gen` is a decorator that automatically generates the
+  implementation of a Python function based on its signature and docstring,
+  using the specified language model. This is useful for quickly prototyping
+  functions or generating boilerplate code.
+
+  The decorator can also automatically generate and run unit tests to verify
+  the correctness of the generated implementation.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+
+  @lf.function_gen(lm=lf.llms.Gemini25Flash())
+  def product(a: int, b: int) -> int:
+    \"\"\"Returns product of a and b.\"\"\"
+
+  print(product(2, 3))
+  # Output: 6
+  ```
 
   Args:
       lm (lf.LanguageModel): The language model used for generating function
@@ -212,10 +229,10 @@ def function_gen(
         tests. You can either provide a list of test cases as tuples of inputs
         and outputs, or a function that throws an error if a test fails, or let
         LLM automatically create the unit test cases. If a generated function is
-        and returned, it should pass all the unittests.
+        returned, it should pass all the unit tests.
       unittest_num_retries: If unittest is set to "auto", this parameter
-        specifies the number of times the LLM's attempts to generate unit test
-        cases.
+        specifies the number of times the LLM should attempt to generate unit
+        test cases.
 
   Returns:
       The implemented function object.