diff --git a/langfun/core/structured/__init__.py b/langfun/core/structured/__init__.py index 8f03f60a..398457d7 100644 --- 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 index 900f5ec2..0fefe23d 100644 --- 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 index 345fb2af..88a87dc4 100644 --- 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 index 76aa63fa..70f3b8dd 100644 --- 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 index e9bfdf3c..bae1ac1a 100644 --- 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 index 84c363ae..1fdbe9f5 100644 --- 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 )