Clean up lf.structured.

1) Rename prompting.py to querying.py as it's easier for users to locate the source code of `lf.query`.
2) Make helper classes without backward compatibility guarentee private.

PiperOrigin-RevId: 706800418

docs/notebooks/langfun101.ipynb

@@ -10,7 +10,7 @@
         "\n",
         "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/langfun/blob/main/docs/notebooks/langfun101.ipynb)\n",
         "\n",
-        "Effective programming of Large Language Models (LLMs) demands a seamless integration of natural language text with structured data. Langfun, leveraging [PyGlove](https://github.com/google/pyglove)'s symbolic objects, provides a simple yet powerful interface for mapping between Python objects with the assistance of LLMs. The input/output objects may include natural language texts (in string form), structured data (objects of a specific class), modalities (such as images), and more. The unified API for accomplishing all conceivable mappings is [`lf.query`](https://github.com/google/langfun/blob/145f4a48603d2e6e19f9025e032ac3c86dfd3e35/langfun/core/structured/prompting.py#L103).\n"
+        "Effective programming of Large Language Models (LLMs) demands a seamless integration of natural language text with structured data. Langfun, leveraging [PyGlove](https://github.com/google/pyglove)'s symbolic objects, provides a simple yet powerful interface for mapping between Python objects with the assistance of LLMs. The input/output objects may include natural language texts (in string form), structured data (objects of a specific class), modalities (such as images), and more. The unified API for accomplishing all conceivable mappings is [`lf.query`](https://github.com/google/langfun/blob/145f4a48603d2e6e19f9025e032ac3c86dfd3e35/langfun/core/structured/querying.py#L103).\n"
       ]
     },
     {

langfun/core/coding/python/correction.py

@@ -76,7 +76,7 @@ def run_with_correction(
   # Delay import at runtime to avoid circular depenency.
   # pylint: disable=g-import-not-at-top
   # pytype: disable=import-error
-  from langfun.core.structured import prompting
+  from langfun.core.structured import querying
   # pytype: enable=import-error
   # pylint: enable=g-import-not-at-top
 
@@ -119,7 +119,7 @@ def result_and_error(code: str) -> tuple[Any, str | None]:
     # structure.
     try:
       # Disable autofix for code correction to avoid recursion.
-      correction = prompting.query(
+      correction = querying.query(
           CodeWithError(code=code, error=error), CorrectedCode, lm=lm, autofix=0
       )
     except errors.CodeError:

langfun/core/structured/__init__.py

@@ -36,12 +36,6 @@
 from langfun.core.structured.schema import annotation
 from langfun.core.structured.schema import structure_from_python
 
-from langfun.core.structured.schema import SchemaRepr
-from langfun.core.structured.schema import SchemaJsonRepr
-from langfun.core.structured.schema import SchemaPythonRepr
-from langfun.core.structured.schema import ValueRepr
-from langfun.core.structured.schema import ValueJsonRepr
-from langfun.core.structured.schema import ValuePythonRepr
 from langfun.core.structured.schema import schema_repr
 from langfun.core.structured.schema import source_form
 from langfun.core.structured.schema import value_repr
@@ -56,26 +50,17 @@
 from langfun.core.structured.mapping import MappingError
 from langfun.core.structured.mapping import MappingExample
 
-from langfun.core.structured.parsing import ParseStructure
-from langfun.core.structured.parsing import ParseStructureJson
-from langfun.core.structured.parsing import ParseStructurePython
 from langfun.core.structured.parsing import parse
 from langfun.core.structured.parsing import call
 
-from langfun.core.structured.prompting import QueryStructure
-from langfun.core.structured.prompting import QueryStructureJson
-from langfun.core.structured.prompting import QueryStructurePython
-from langfun.core.structured.prompting import query
-from langfun.core.structured.prompting import query_prompt
-from langfun.core.structured.prompting import query_output
-from langfun.core.structured.prompting import query_reward
-from langfun.core.structured.prompting import QueryInvocation
-from langfun.core.structured.prompting import track_queries
-
-from langfun.core.structured.description import DescribeStructure
-from langfun.core.structured.description import describe
+from langfun.core.structured.querying import track_queries
+from langfun.core.structured.querying import QueryInvocation
+from langfun.core.structured.querying import query
+from langfun.core.structured.querying import query_prompt
+from langfun.core.structured.querying import query_output
+from langfun.core.structured.querying import query_reward
 
-from langfun.core.structured.completion import CompleteStructure
+from langfun.core.structured.description import describe
 from langfun.core.structured.completion import complete
 
 from langfun.core.structured.scoring import score

langfun/core/structured/completion.py

@@ -21,7 +21,7 @@
 import pyglove as pg
 
 
-class CompleteStructure(mapping.Mapping):
+class _CompleteStructure(mapping.Mapping):
   """Complete structure by filling the missing fields."""
 
   input: Annotated[
@@ -241,7 +241,7 @@ class Flight(pg.Object):
   Returns:
     The result based on the schema.
   """
-  t = CompleteStructure(
+  t = _CompleteStructure(
       input=schema_lib.mark_missing(input_value),
       default=default,
       examples=examples,

langfun/core/structured/completion_test.py

@@ -46,7 +46,7 @@ class TripPlan(pg.Object):
 class CompleteStructureTest(unittest.TestCase):
 
   def test_render_no_examples(self):
-    l = completion.CompleteStructure()
+    l = completion._CompleteStructure()
     input_value = schema_lib.mark_missing(
         TripPlan.partial(
             place='San Francisco',
@@ -120,7 +120,7 @@ class Activity:
     )
 
   def test_render_no_class_definitions(self):
-    l = completion.CompleteStructure()
+    l = completion._CompleteStructure()
     input_value = schema_lib.mark_missing(
         TripPlan.partial(
             place='San Francisco',
@@ -200,7 +200,7 @@ def test_render_no_class_definitions(self):
     )
 
   def test_render_with_examples(self):
-    l = completion.CompleteStructure()
+    l = completion._CompleteStructure()
     input_value = schema_lib.mark_missing(
         TripPlan.partial(
             place='San Francisco',
@@ -411,7 +411,7 @@ class Animal(pg.Object):
             modalities.Image.from_bytes(b'image_of_elephant'),
         )
     )
-    l = completion.CompleteStructure(
+    l = completion._CompleteStructure(
         input=input_value,
         examples=[
             mapping.MappingExample(

langfun/core/structured/description.py

@@ -22,7 +22,7 @@
 
 
 @pg.use_init_args(['examples'])
-class DescribeStructure(mapping.Mapping):
+class _DescribeStructure(mapping.Mapping):
   """Describe a structured value in natural language."""
 
   input_title = 'PYTHON_OBJECT'
@@ -106,7 +106,7 @@ class Flight(pg.Object):
   Returns:
     The parsed result based on the schema.
   """
-  return DescribeStructure(
+  return _DescribeStructure(
       input=value,
       context=context,
       examples=examples or default_describe_examples(),

langfun/core/structured/description_test.py

@@ -36,7 +36,7 @@ class Itinerary(pg.Object):
 class DescribeStructureTest(unittest.TestCase):
 
   def test_render(self):
-    l = description_lib.DescribeStructure(
+    l = description_lib._DescribeStructure(
         input=Itinerary(
             day=1,
             type='daytime',
@@ -137,7 +137,7 @@ def test_render_no_examples(self):
         ],
         hotel=None,
     )
-    l = description_lib.DescribeStructure(
+    l = description_lib._DescribeStructure(
         input=value, context='1 day itinerary to SF'
     )
     self.assertEqual(
@@ -187,7 +187,7 @@ def test_render_no_context(self):
         ],
         hotel=None,
     )
-    l = description_lib.DescribeStructure(input=value)
+    l = description_lib._DescribeStructure(input=value)
     self.assertEqual(
         l.render().text,
         inspect.cleandoc("""

langfun/core/structured/function_generation.py

@@ -21,7 +21,7 @@
 from langfun.core import language_model
 from langfun.core import template
 from langfun.core.coding import python
-from langfun.core.structured import prompting
+from langfun.core.structured import querying
 import pyglove as pg
 
 
@@ -39,7 +39,7 @@ class PythonFunctionSignature(pg.Object):
 
   unittest_examples = None
   for _ in range(num_retries):
-    r = prompting.query(
+    r = querying.query(
         PythonFunctionSignature(signature=signature),
         list[UnitTest],
         lm=lm,
@@ -145,7 +145,7 @@ def calculate_area_circle(radius: float) -> float:
   last_error = None
   for _ in range(num_retries):
     try:
-      source_code = prompting.query(
+      source_code = querying.query(
           PythonFunctionPrompt(signature=signature), lm=lm
       )
       f = python.evaluate(source_code, global_vars=context)

langfun/core/structured/parsing.py

@@ -16,13 +16,13 @@
 
 import langfun.core as lf
 from langfun.core.structured import mapping
-from langfun.core.structured import prompting
+from langfun.core.structured import querying
 from langfun.core.structured import schema as schema_lib
 import pyglove as pg
 
 
 @lf.use_init_args(['schema', 'default', 'examples'])
-class ParseStructure(mapping.Mapping):
+class _ParseStructure(mapping.Mapping):
   """Parse an object out from a natural language text."""
 
   context_title = 'USER_REQUEST'
@@ -37,7 +37,7 @@ class ParseStructure(mapping.Mapping):
   ]
 
 
-class ParseStructureJson(ParseStructure):
+class _ParseStructureJson(_ParseStructure):
   """Parse an object out from a NL text using JSON as the protocol."""
 
   preamble = """
@@ -53,7 +53,7 @@ class ParseStructureJson(ParseStructure):
   output_title = 'JSON'
 
 
-class ParseStructurePython(ParseStructure):
+class _ParseStructurePython(_ParseStructure):
   """Parse an object out from a NL text using Python as the protocol."""
 
   preamble = """
@@ -87,7 +87,7 @@ def parse(
     returns_message: bool = False,
     **kwargs,
 ) -> Any:
-  """Parse a natural langugage message based on schema.
+  """Parse a natural language message based on schema.
 
   Examples:
 
@@ -271,7 +271,7 @@ def call(
     return lm_output if returns_message else lm_output.text
 
   # Call `parsing_lm` for structured parsing.
-  parsing_message = prompting.query(
+  parsing_message = querying.query(
       lm_output.text,
       schema,
       examples=parsing_examples,
@@ -293,11 +293,11 @@ def call(
 
 def _parse_structure_cls(
     protocol: schema_lib.SchemaProtocol,
-) -> Type[ParseStructure]:
+) -> Type[_ParseStructure]:
   if protocol == 'json':
-    return ParseStructureJson
+    return _ParseStructureJson
   elif protocol == 'python':
-    return ParseStructurePython
+    return _ParseStructurePython
   else:
     raise ValueError(f'Unknown protocol: {protocol!r}.')
 

langfun/core/structured/parsing_test.py

@@ -37,7 +37,7 @@ class Itinerary(pg.Object):
 class ParseStructurePythonTest(unittest.TestCase):
 
   def test_render_no_examples(self):
-    l = parsing.ParseStructurePython(int)
+    l = parsing._ParseStructurePython(int)
     m = lf.AIMessage('12 / 6 + 2 = 4')
     self.assertEqual(
         l.render(input=m, context='Compute 12 / 6 + 2.').text,
@@ -62,7 +62,7 @@ def test_render_no_examples(self):
     )
 
   def test_render_no_context(self):
-    l = parsing.ParseStructurePython(int)
+    l = parsing._ParseStructurePython(int)
     m = lf.AIMessage('12 / 6 + 2 = 4')
 
     self.assertEqual(
@@ -85,7 +85,7 @@ def test_render_no_context(self):
     )
 
   def test_render(self):
-    l = parsing.ParseStructurePython(
+    l = parsing._ParseStructurePython(
         int,
         examples=[
             mapping.MappingExample(
@@ -212,7 +212,7 @@ def test_invocation(self):
         ),
         override_attrs=True,
     ):
-      l = parsing.ParseStructurePython(
+      l = parsing._ParseStructurePython(
           [Itinerary],
           examples=[
               mapping.MappingExample(
@@ -295,7 +295,7 @@ def test_parse(self):
 class ParseStructureJsonTest(unittest.TestCase):
 
   def test_render_no_examples(self):
-    l = parsing.ParseStructureJson(int)
+    l = parsing._ParseStructureJson(int)
     m = lf.AIMessage('12 / 6 + 2 = 4')
     self.assertEqual(
         l.render(input=m, context='Compute 12 / 6 + 2.').text,
@@ -320,7 +320,7 @@ def test_render_no_examples(self):
     )
 
   def test_render_no_context(self):
-    l = parsing.ParseStructureJson(int)
+    l = parsing._ParseStructureJson(int)
     m = lf.AIMessage('12 / 6 + 2 = 4')
 
     self.assertEqual(
@@ -343,7 +343,7 @@ def test_render_no_context(self):
     )
 
   def test_render(self):
-    l = parsing.ParseStructureJson(
+    l = parsing._ParseStructureJson(
         int,
         examples=[
             mapping.MappingExample(
@@ -504,7 +504,7 @@ def test_invocation(self):
         override_attrs=True,
     ):
       message = lf.LangFunc(lm_input)()
-      l = parsing.ParseStructureJson(
+      l = parsing._ParseStructureJson(
           [Itinerary],
           examples=[
               mapping.MappingExample(