Expose default examples for structured methods.

- `lf.structured.DEFAULT_PARSE_EXAMPLES` - `lf.structured.DEFAULT_QUERY_EXAMPLES` - `lf.structured.DEFAULT_COMPLETE_EXAMPLES` - `lf.structured.DEFAULT_DESCRIBE_EXAMPLES` PiperOrigin-RevId: 570101645
google · Oct 2, 2023 · e32b131 · e32b131
1 parent 487857e
commit e32b131
Show file tree

Hide file tree

Showing 6 changed files with 107 additions and 102 deletions.
diff --git a/langfun/core/structured/__init__.py b/langfun/core/structured/__init__.py
@@ -68,6 +68,15 @@
 from langfun.core.structured.completion import CompleteStructure
 from langfun.core.structured.completion import complete
 
+# Expose default examples for structured operations so users could refer to
+# them.
+from langfun.core.structured.parsing import DEFAULT_PARSE_EXAMPLES
+from langfun.core.structured.prompting import DEFAULT_QUERY_EXAMPLES
+from langfun.core.structured.completion import DEFAULT_COMPLETE_EXAMPLES
+from langfun.core.structured.description import DEFAULT_DESCRIBE_EXAMPLES
+
+# Default examples.
+
 
 # pylint: enable=g-importing-member
 # pylint: enable=g-bad-import-order
diff --git a/langfun/core/structured/completion.py b/langfun/core/structured/completion.py
@@ -59,21 +59,20 @@ class _Country(pg.Object):
   hobby: str
 
 
-def _default_complete_examples() -> list[mapping.MappingExample]:
-  return [
-      mapping.MappingExample(
-          value=mapping.Pair(
-              left=_Country.partial(name='United States of America'),
-              right=_Country(
-                  name='United States of America',
-                  founding_date=_Date(year=1776, month=7, day=4),
-                  continent='North America',
-                  population=33_000_000,
-                  hobby=schema_lib.UNKNOWN,
-              ),
-          )
-      )
-  ]
+DEFAULT_COMPLETE_EXAMPLES: list[mapping.MappingExample] = [
+    mapping.MappingExample(
+        value=mapping.Pair(
+            left=_Country.partial(name='United States of America'),
+            right=_Country(
+                name='United States of America',
+                founding_date=_Date(year=1776, month=7, day=4),
+                continent='North America',
+                population=33_000_000,
+                hobby=schema_lib.UNKNOWN,
+            ),
+        )
+    ),
+]
 
 
 def complete(
@@ -128,6 +127,6 @@ class Flight(pg.Object):
     The result based on the schema.
   """
   if examples is None:
-    examples = _default_complete_examples()
+    examples = DEFAULT_COMPLETE_EXAMPLES
   t = CompleteStructure(default=default, examples=examples, **kwargs)
   return t.transform(message=lf.UserMessage(text='', result=value)).result
diff --git a/langfun/core/structured/completion_test.py b/langfun/core/structured/completion_test.py
@@ -172,7 +172,7 @@ def test_render_no_class_definitions(self):
 
   def test_render_with_examples(self):
     l = completion.CompleteStructure(
-        examples=completion._default_complete_examples()
+        examples=completion.DEFAULT_COMPLETE_EXAMPLES
     )
     m = lf.UserMessage(
         '',

diff --git a/langfun/core/structured/description.py b/langfun/core/structured/description.py
@@ -104,7 +104,7 @@ class Flight(pg.Object):
     )
 
   if examples is None:
-    examples = _default_describe_examples()
+    examples = DEFAULT_DESCRIBE_EXAMPLES
   return DescribeStructure(examples, **kwargs)(message=message).text
 
 
@@ -129,32 +129,31 @@ class _Country(pg.Object):
   president: str | None
 
 
-def _default_describe_examples() -> list[mapping.MappingExample]:
-  return [
-      mapping.MappingExample(
-          nl_context='Brief intro to United States',
-          nl_text=inspect.cleandoc("""
-              The United States of America is a country primarily located in North America
-              consisting of fifty states. It shares land borders with Canada to its north
-              and with Mexico to its south and has maritime borders with the Bahamas, Cuba,
-              Russia, and other nations. With a population of over 333 million. The national
-              capital of the United States is Washington, D.C.
-              """),
-          schema=None,
-          value=_Country(
-              name='The United States of America',
-              continents=['North America'],
-              num_states=50,
-              neighbor_countries=[
-                  'Canada',
-                  'Mexico',
-                  'Bahamas',
-                  'Cuba',
-                  'Russia',
-              ],
-              population=333000000,
-              capital='Washington, D.C',
-              president=None,
-          ),
-      )
-  ]
+DEFAULT_DESCRIBE_EXAMPLES: list[mapping.MappingExample] = [
+    mapping.MappingExample(
+        nl_context='Brief intro to United States',
+        nl_text=inspect.cleandoc("""
+            The United States of America is a country primarily located in North America
+            consisting of fifty states. It shares land borders with Canada to its north
+            and with Mexico to its south and has maritime borders with the Bahamas, Cuba,
+            Russia, and other nations. With a population of over 333 million. The national
+            capital of the United States is Washington, D.C.
+            """),
+        schema=None,
+        value=_Country(
+            name='The United States of America',
+            continents=['North America'],
+            num_states=50,
+            neighbor_countries=[
+                'Canada',
+                'Mexico',
+                'Bahamas',
+                'Cuba',
+                'Russia',
+            ],
+            population=333000000,
+            capital='Washington, D.C',
+            president=None,
+        ),
+    ),
+]
diff --git a/langfun/core/structured/parsing.py b/langfun/core/structured/parsing.py
@@ -156,7 +156,7 @@ class Flight(pg.Object):
     The parsed result based on the schema.
   """
   if examples is None:
-    examples = _default_parsing_examples()
+    examples = DEFAULT_PARSE_EXAMPLES
 
   t = _parse_structure_cls(protocol)(
       schema, default=default, examples=examples, **kwargs
@@ -211,7 +211,7 @@ def as_structured(
   if annotation is None:
     return self
   if examples is None:
-    examples = _default_parsing_examples()
+    examples = DEFAULT_PARSE_EXAMPLES
   return self >> _parse_structure_cls(protocol)(
       schema=annotation,
       default=default,
@@ -238,36 +238,35 @@ class _Country(pg.Object):
   president: str | None
 
 
-def _default_parsing_examples() -> list[mapping.MappingExample]:
-  return [
-      mapping.MappingExample(
-          nl_context='Brief introduction of the U.S.A.',
-          nl_text=inspect.cleandoc("""
-              The United States of America is a country primarily located in North America
-              consisting of fifty states, a federal district, five major unincorporated territories,
-              nine Minor Outlying Islands, and 326 Indian reservations. It shares land borders
-              with Canada to its north and with Mexico to its south and has maritime borders
-              with the Bahamas, Cuba, Russia, and other nations. With a population of over 333
-              million. The national capital of the United States is Washington, D.C.
-              """),
-          schema=_Country,
-          value=_Country(
-              name='The United States of America',
-              continents=['North America'],
-              num_states=50,
-              neighbor_countries=[
-                  'Canada',
-                  'Mexico',
-                  'Bahamas',
-                  'Cuba',
-                  'Russia',
-              ],
-              population=333000000,
-              capital='Washington, D.C',
-              president=None,
-          ),
-      )
-  ]
+DEFAULT_PARSE_EXAMPLES: list[mapping.MappingExample] = [
+    mapping.MappingExample(
+        nl_context='Brief introduction of the U.S.A.',
+        nl_text=inspect.cleandoc("""
+            The United States of America is a country primarily located in North America
+            consisting of fifty states, a federal district, five major unincorporated territories,
+            nine Minor Outlying Islands, and 326 Indian reservations. It shares land borders
+            with Canada to its north and with Mexico to its south and has maritime borders
+            with the Bahamas, Cuba, Russia, and other nations. With a population of over 333
+            million. The national capital of the United States is Washington, D.C.
+            """),
+        schema=_Country,
+        value=_Country(
+            name='The United States of America',
+            continents=['North America'],
+            num_states=50,
+            neighbor_countries=[
+                'Canada',
+                'Mexico',
+                'Bahamas',
+                'Cuba',
+                'Russia',
+            ],
+            population=333000000,
+            capital='Washington, D.C',
+            president=None,
+        ),
+    ),
+]
 
 
 lf.MessageTransform.as_structured = as_structured
diff --git a/langfun/core/structured/prompting.py b/langfun/core/structured/prompting.py
@@ -92,28 +92,27 @@ class _Country(pg.Object):
   president: str | None
 
 
-def _default_query_examples() -> list[mapping.MappingExample]:
-  return [
-      mapping.MappingExample(
-          nl_context='Brief introduction of the U.S.A.',
-          schema=_Country,
-          value=_Country(
-              name='The United States of America',
-              continents=['North America'],
-              num_states=50,
-              neighbor_countries=[
-                  'Canada',
-                  'Mexico',
-                  'Bahamas',
-                  'Cuba',
-                  'Russia',
-              ],
-              population=333000000,
-              capital='Washington, D.C',
-              president=None,
-          ),
-      )
-  ]
+DEFAULT_QUERY_EXAMPLES: list[mapping.MappingExample] = [
+    mapping.MappingExample(
+        nl_context='Brief introduction of the U.S.A.',
+        schema=_Country,
+        value=_Country(
+            name='The United States of America',
+            continents=['North America'],
+            num_states=50,
+            neighbor_countries=[
+                'Canada',
+                'Mexico',
+                'Bahamas',
+                'Cuba',
+                'Russia',
+            ],
+            population=333000000,
+            capital='Washington, D.C',
+            president=None,
+        ),
+    ),
+]
 
 
 def query(
@@ -175,7 +174,7 @@ class Flight(pg.Object):
     The result based on the schema.
   """
   if examples is None:
-    examples = _default_query_examples()
+    examples = DEFAULT_QUERY_EXAMPLES
   t = _query_structure_cls(protocol)(
       schema, default=default, examples=examples, **kwargs
   )