langfun 0.0.2.dev20240314__py3-none-any.whl → 0.0.2.dev20240316__py3-none-any.whl

langfun/__init__.py

@@ -31,6 +31,9 @@ query = structured.query
 describe = structured.describe
 complete = structured.complete
 score = structured.score
+generate_class = structured.generate_class
+
+source_form = structured.source_form
 
 from langfun.core import eval  # pylint: disable=redefined-builtin
 from langfun.core import templates

langfun/core/langfunc_test.py

@@ -95,8 +95,8 @@ class LangFuncCallTest(unittest.TestCase):
         ' lm=ExcitedEchoer(sampling_options=LMSamplingOptions(temperature=0.0,'
         ' max_tokens=1024, n=1, top_k=40, top_p=None, stop=None,'
         ' random_seed=None, logprobs=False, top_logprobs=None), cache=None,'
-        ' timeout=120.0, max_attempts=5, retry_interval=(5, 60),'
-        ' exponential_backoff=True, debug=False))',
+        ' max_concurrency=None, timeout=120.0, max_attempts=5,'
+        ' retry_interval=(5, 60), exponential_backoff=True, debug=False))',
     )
 
     l = LangFunc('Hello')

langfun/core/language_model.py

@@ -17,8 +17,9 @@ import abc
 import dataclasses
 import enum
 import time
-from typing import Annotated, Any
+from typing import Annotated, Any, Callable, Sequence, Tuple, Type, Union
 from langfun.core import component
+from langfun.core import concurrent
 from langfun.core import console
 from langfun.core import message as message_lib
 import pyglove as pg
@@ -209,6 +210,22 @@ class LanguageModel(component.Component):
       )
   ] = component.contextual(default=None)
 
+  max_concurrency: Annotated[
+      int | None,
+      (
+          'Max concurrent requests being sent to the server. '
+          'If None, there is no limit. '
+          'Please note that the concurrency control is based on the '
+          '`resource_id` property, meaning that model instances shared '
+          'the same resource ID will be accounted under the same concurrency '
+          'control key. This allows a process-level concurrency control '
+          'for specific models regardless the number of LM (client) instances '
+          'created by the program. Subclasses could override this number or '
+          'replace it with a `max_concurrency` property to allow dynamic '
+          'concurrency control.'
+      ),
+  ] = None
+
   timeout: Annotated[
       float | None, 'Timeout in seconds. If None, there is no timeout.'
   ] = 120.0
@@ -284,11 +301,6 @@ class LanguageModel(component.Component):
     """Resource ID for performing request parallism control."""
     return self.model_id
 
-  @property
-  def max_concurrency(self) -> int:
-    """Max concurrent requests."""
-    return 32
-
   def sample(
       self,
       prompts: list[str | message_lib.Message],
@@ -355,6 +367,28 @@ class LanguageModel(component.Component):
   ) -> list[LMSamplingResult]:
     """Subclass should override."""
 
+  def _parallel_execute_with_currency_control(
+      self,
+      action: Callable[..., Any],
+      inputs: Sequence[Any],
+      retry_on_errors: Union[
+          None,
+          Union[Type[Exception], Tuple[Type[Exception], str]],
+          Sequence[Union[Type[Exception], Tuple[Type[Exception], str]]],
+      ] = None,
+  ) -> Any:
+    """Helper method for subclasses for implementing _sample."""
+    return concurrent.concurrent_execute(
+        action,
+        inputs,
+        executor=self.resource_id if self.max_concurrency else None,
+        max_workers=self.max_concurrency or len(inputs),
+        retry_on_errors=retry_on_errors,
+        max_attempts=self.max_attempts,
+        retry_interval=self.retry_interval,
+        exponential_backoff=self.exponential_backoff,
+    )
+
   def __call__(
       self, prompt: message_lib.Message, *, cache_seed: int = 0, **kwargs
   ) -> message_lib.Message:

langfun/core/language_model_test.py

@@ -89,7 +89,7 @@ class LanguageModelTest(unittest.TestCase):
     lm = MockModel(1, temperature=0.5, top_k=2, max_attempts=2)
     self.assertEqual(lm.model_id, 'MockModel')
     self.assertEqual(lm.resource_id, 'MockModel')
-    self.assertEqual(lm.max_concurrency, 32)
+    self.assertIsNone(lm.max_concurrency)
     self.assertEqual(lm.failures_before_attempt, 1)
     self.assertEqual(lm.sampling_options.temperature, 0.5)
     self.assertEqual(lm.sampling_options.top_k, 2)

langfun/core/llms/gemini.py

@@ -133,14 +133,9 @@ class Gemini(lf.LanguageModel):
 
   def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
     assert self._api_initialized, 'Vertex AI API is not initialized.'
-    return lf.concurrent_execute(
+    return self._parallel_execute_with_currency_control(
         self._sample_single,
         prompts,
-        executor=self.resource_id,
-        max_workers=self.max_concurrency,
-        # NOTE(daiyip): Vertex has its own policy on handling
-        # with rate limit, so we do not retry on errors.
-        retry_on_errors=None,
     )
 
   def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:

langfun/core/llms/llama_cpp.py

@@ -67,13 +67,6 @@ class LlamaCppRemote(lf.LanguageModel):
         results.append(result)
       return results
 
-    return lf.concurrent_execute(
-        _complete_fn,
-        [prompts],
-        executor=self.resource_id,
-        max_workers=self.max_concurrency,
-        retry_on_errors=(),
-        max_attempts=self.max_attempts,
-        retry_interval=self.retry_interval,
-        exponential_backoff=self.exponential_backoff,
+    return self._parallel_execute_with_currency_control(
+        _complete_fn, [prompts]
     )[0]

langfun/core/llms/openai.py

@@ -214,18 +214,13 @@ class OpenAI(lf.LanguageModel):
           for index in sorted(samples_by_index.keys())
       ]
 
-    return lf.concurrent_execute(
+    return self._parallel_execute_with_currency_control(
         _open_ai_completion,
         [prompts],
-        executor=self.resource_id,
-        max_workers=self.max_concurrency,
         retry_on_errors=(
             openai_error.ServiceUnavailableError,
             openai_error.RateLimitError,
         ),
-        max_attempts=self.max_attempts,
-        retry_interval=self.retry_interval,
-        exponential_backoff=self.exponential_backoff,
     )[0]
 
   def _chat_complete_batch(
@@ -280,18 +275,13 @@ class OpenAI(lf.LanguageModel):
           ),
       )
 
-    return lf.concurrent_execute(
+    return self._parallel_execute_with_currency_control(
         _open_ai_chat_completion,
         prompts,
-        executor=self.resource_id,
-        max_workers=self.max_concurrency,
         retry_on_errors=(
             openai_error.ServiceUnavailableError,
             openai_error.RateLimitError,
         ),
-        max_attempts=self.max_attempts,
-        retry_interval=self.retry_interval,
-        exponential_backoff=self.exponential_backoff,
     )
 
 

langfun/core/structured/__init__.py

@@ -41,8 +41,12 @@ 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
 
+from langfun.core.structured.schema_generation import generate_class
+from langfun.core.structured.schema_generation import classgen_example
+from langfun.core.structured.schema_generation import default_classgen_examples
 
 from langfun.core.structured.mapping import Mapping
 from langfun.core.structured.mapping import MappingExample
@@ -68,8 +72,8 @@ from langfun.core.structured.scoring import score
 
 # 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.description import DEFAULT_DESCRIBE_EXAMPLES
+from langfun.core.structured.parsing import default_parse_examples
+from langfun.core.structured.description import default_describe_examples
 
 # Default examples.
 

langfun/core/structured/description.py

@@ -106,58 +106,61 @@ def describe(
   Returns:
     The parsed result based on the schema.
   """
-  if examples is None:
-    examples = DEFAULT_DESCRIBE_EXAMPLES
   return DescribeStructure(
-      input=value, context=context, examples=examples, **kwargs
+      input=value,
+      context=context,
+      examples=examples or default_describe_examples(),
+      **kwargs,
   )(lm=lm, cache_seed=cache_seed).text
 
 
-class _Country(pg.Object):
-  """A example dataclass for structured mapping."""
-
-  name: str
-  continents: list[
-      Literal[
-          'Africa',
-          'Asia',
-          'Europe',
-          'Oceania',
-          'North America',
-          'South America',
-      ]
+def default_describe_examples() -> list[mapping.MappingExample]:
+  """Default describe examples."""
+
+  class Country(pg.Object):
+    """A example dataclass for structured mapping."""
+
+    name: str
+    continents: list[
+        Literal[
+            'Africa',
+            'Asia',
+            'Europe',
+            'Oceania',
+            'North America',
+            'South America',
+        ]
+    ]
+    num_states: int
+    neighbor_countries: list[str]
+    population: int
+    capital: str | None
+    president: str | None
+
+  return [
+      mapping.MappingExample(
+          context='Brief intro to United States',
+          input=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,
+          ),
+          output=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.
+              """),
+      ),
   ]
-  num_states: int
-  neighbor_countries: list[str]
-  population: int
-  capital: str | None
-  president: str | None
-
-
-DEFAULT_DESCRIBE_EXAMPLES: list[mapping.MappingExample] = [
-    mapping.MappingExample(
-        context='Brief intro to United States',
-        input=_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,
-        ),
-        output=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.
-            """),
-    ),
-]

langfun/core/structured/mapping.py

@@ -293,25 +293,27 @@ class Mapping(lf.LangFunc):
 
   def transform_output(self, lm_output: lf.Message) -> lf.Message:
     """Transforms LM response into structure if schema is present."""
-    schema = self.mapping_request.schema
-    if schema is None:
-      return lm_output
-
     try:
-      result = schema.parse(
-          lm_output.text,
-          protocol=self.protocol,
-          additional_context=self.globals(),
-          autofix=self.autofix,
-          autofix_lm=self.autofix_lm or self.lm,
-      )
-      lm_output.result = self.postprocess_result(result)
+      lm_output.result = self.postprocess_result(self.parse_result(lm_output))
     except Exception as e:  # pylint: disable=broad-exception-caught
       if self.default == lf.RAISE_IF_HAS_ERROR:
         raise e
       lm_output.result = self.default
     return lm_output
 
+  def parse_result(self, lm_output: lf.Message) -> Any:
+    """Parse result from LLM response."""
+    schema = self.mapping_request.schema
+    if schema is None:
+      return None
+    return schema.parse(
+        lm_output.text,
+        protocol=self.protocol,
+        additional_context=self.globals(),
+        autofix=self.autofix,
+        autofix_lm=self.autofix_lm or self.lm,
+    )
+
   def postprocess_result(self, result: Any) -> Any:
     """Post process structured output."""
     return result

langfun/core/structured/parsing.py

@@ -162,11 +162,11 @@ def parse(
     message.source = lf.UserMessage(user_prompt, tags=['lm-input'])
   context = getattr(message.lm_input, 'text', None) if include_context else None
 
-  if examples is None:
-    examples = DEFAULT_PARSE_EXAMPLES
-
   t = _parse_structure_cls(protocol)(
-      schema=schema, context=context, default=default, examples=examples
+      schema=schema,
+      context=context,
+      default=default,
+      examples=examples or default_parse_examples(),
   )
 
   # Setting up context.
@@ -296,17 +296,19 @@ def _parse_structure_cls(
     raise ValueError(f'Unknown protocol: {protocol!r}.')
 
 
-class _AdditionResults(pg.Object):
-  one_plus_one_equals: int | None
-  two_plus_two_equals: int | None
+def default_parse_examples() -> list[mapping.MappingExample]:
+  """Default parsing examples."""
 
+  class AdditionResults(pg.Object):
+    one_plus_one_equals: int | None
+    two_plus_two_equals: int | None
 
-DEFAULT_PARSE_EXAMPLES: list[mapping.MappingExample] = [
-    mapping.MappingExample(
-        input='Two plus two equals four. Three plus three equals six.',
-        schema=_AdditionResults,
-        output=_AdditionResults(
-            one_plus_one_equals=None, two_plus_two_equals=4
-        ),
-    ),
-]
+  return [
+      mapping.MappingExample(
+          input='Two plus two equals four. Three plus three equals six.',
+          schema=AdditionResults,
+          output=AdditionResults(
+              one_plus_one_equals=None, two_plus_two_equals=4
+          ),
+      ),
+  ]

langfun/core/structured/schema.py

@@ -301,23 +301,43 @@ class SchemaRepr(metaclass=abc.ABCMeta):
 class SchemaPythonRepr(SchemaRepr):
   """Python-representation for a schema."""
 
-  def repr(self, schema: Schema) -> str:
-    ret = self.result_definition(schema)
-    class_definition_str = self.class_definitions(schema)
+  def repr(
+      self,
+      schema: Schema,
+      *,
+      include_result_definition: bool = True,
+      markdown: bool = True,
+      **kwargs,
+  ) -> str:
+    ret = ''
+    if include_result_definition:
+      ret += self.result_definition(schema)
+    class_definition_str = self.class_definitions(
+        schema, markdown=markdown, **kwargs
+    )
     if class_definition_str:
-      ret += f'\n\n```python\n{class_definition_str}```'
-    return ret
+      ret += f'\n\n{class_definition_str}'
+    return ret.strip()
 
-  def class_definitions(self, schema: Schema) -> str | None:
+  def class_definitions(self, schema: Schema, **kwargs) -> str | None:
     deps = schema.class_dependencies(include_subclasses=True)
-    return class_definitions(deps)
+    return class_definitions(deps, **kwargs)
 
   def result_definition(self, schema: Schema) -> str:
     return annotation(schema.spec)
 
 
+def source_form(value, markdown: bool = False) -> str:
+  """Returns the source code form of an object."""
+  return ValuePythonRepr().repr(value, markdown=markdown)
+
+
 def class_definitions(
-    classes: Sequence[Type[Any]], strict: bool = False, markdown: bool = False
+    classes: Sequence[Type[Any]],
+    *,
+    include_pg_object_as_base: bool = False,
+    strict: bool = False,
+    markdown: bool = False,
 ) -> str | None:
   """Returns a str for class definitions."""
   if not classes:
@@ -326,14 +346,22 @@ def class_definitions(
   for i, cls in enumerate(classes):
     if i > 0:
       def_str.write('\n')
-    def_str.write(class_definition(cls, strict))
+    def_str.write(
+        class_definition(
+            cls,
+            strict=strict,
+            include_pg_object_as_base=include_pg_object_as_base,
+        )
+    )
   ret = def_str.getvalue()
   if markdown and ret:
     ret = f'```python\n{ret}```'
   return ret
 
 
-def class_definition(cls, strict: bool = False) -> str:
+def class_definition(
+    cls, strict: bool = False, include_pg_object_as_base: bool = False
+) -> str:
   """Returns the Python class definition."""
   out = io.StringIO()
   if not issubclass(cls, pg.Object):
@@ -344,10 +372,9 @@ def class_definition(cls, strict: bool = False) -> str:
   schema = cls.__schema__
   eligible_bases = []
   for base_cls in cls.__bases__:
-    if issubclass(base_cls, pg.Symbolic) and not base_cls.__module__.startswith(
-        'pyglove'
-    ):
-      eligible_bases.append(base_cls.__name__)
+    if issubclass(base_cls, pg.Object):
+      if include_pg_object_as_base or base_cls is not pg.Object:
+        eligible_bases.append(base_cls.__name__)
   if eligible_bases:
     base_cls_str = ', '.join(eligible_bases)
     out.write(f'class {cls.__name__}({base_cls_str}):\n')
@@ -547,8 +574,20 @@ class ValuePythonRepr(ValueRepr):
            markdown: bool = True,
            **kwargs) -> str:
     del schema
-    object_code = pg.format(
-        value, compact=compact, verbose=verbose, python_format=True)
+    if inspect.isclass(value):
+      cls_schema = Schema.from_value(value)
+      if isinstance(cls_schema.spec, pg.typing.Object):
+        object_code = SchemaPythonRepr().class_definitions(
+            cls_schema, markdown=markdown, include_pg_object_as_base=True
+        )
+        assert object_code is not None
+        return object_code
+      else:
+        object_code = SchemaPythonRepr().result_definition(cls_schema)
+    else:
+      object_code = pg.format(
+          value, compact=compact, verbose=verbose, python_format=True
+      )
     if markdown:
       return f'```python\n{ object_code }\n```'
     return object_code
@@ -588,6 +627,7 @@ def structure_from_python(
   global_vars = global_vars or {}
   global_vars.update({
       'pg': pg,
+      'Object': pg.Object,
       'Any': typing.Any,
       'List': typing.List,
       'Tuple': typing.Tuple,

langfun/core/structured/schema_generation.py

@@ -0,0 +1,175 @@
+# Copyright 2023 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.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""LLM-based class generation."""
+
+import typing
+from typing import Any, Type
+import langfun.core as lf
+from langfun.core.coding.python import correction
+from langfun.core.structured import mapping
+import pyglove as pg
+
+
+class GenerateClass(mapping.Mapping):
+  """Python class generation."""
+
+  input_title = 'GENERATION_CONTEXT'
+  context_title = 'CLASS_NAME'
+  output_title = 'OUTPUT_CLASS'
+
+  preamble: lf.Template = lf.Template("""
+      Help generate a class based on the last {{ context_title }} and {{ input_title }}.
+
+      Instructions:
+      - Use `Object` as the base class for all generated classes
+      - Create auxillary classes for composition if needed.
+      - Use Python type annotation for declaraing fields:
+        (e.g. bool, str, int, float, Optional[str], List[int], Union[str, int])
+      - Do not use types that need import.
+      - Avoid self-referential types. e.g:
+        ```
+        class Node(Object):
+          children: list[Node]
+        ```
+      - Do not generate methods.
+      """)
+
+  def parse_result(self, lm_output: lf.Message) -> Type[Any]:
+    output_vars, final_code = correction.run_with_correction(
+        lm_output.text,
+        global_vars=self.allowed_annotation_types,
+        sandbox=False,
+        max_attempts=self.autofix,
+        lm=self.autofix_lm,
+        returns_code=True,
+        outputs_intermediate=True,
+    )
+    class_name = self.context
+    cls = output_vars.get(class_name, None)
+    if cls is None:
+      raise correction.errors.CodeError(
+          final_code,
+          TypeError(f'Class {class_name} is absent from LLM output.'),
+      )
+    return cls
+
+  @property
+  def allowed_annotation_types(self):
+    return dict(
+        pg=pg,
+        Any=typing.Any,
+        Object=pg.Object,
+        List=typing.List,
+        Dict=typing.Tuple,
+        Tuple=typing.Tuple,
+        Sequence=typing.Sequence,
+        Optional=typing.Optional,
+        Union=typing.Union,
+    )
+
+
+def generate_class(
+    name: str,
+    prompt: str | pg.Symbolic,
+    *,
+    lm: lf.LanguageModel | None = None,
+    examples: list[mapping.MappingExample] | None = None,
+    returns_message: bool = False,
+    skip_lm: bool = False,
+    **kwargs,
+) -> Type[Any] | lf.Message:
+  """Generate a class with specified name based on the prompt.
+
+  Example:
+    ```
+    trip_cls = lf.classgen(
+        'Trip',
+        'A trip plan to visit {{ city }}, city='San Francisco',
+        lm=lf.llms.GeminiPro()
+    )
+    ```
+
+  Args:
+    name: Class name to be generated.
+    prompt: A str (may contain {{}} as template) as natural language input, or a
+      `pg.Symbolic` object as structured input as prompt to LLM.
+    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 class generation.
+      If None, a default single shot example will be used. Use
+      `lf.structured.classgen_example` to generate example.
+    returns_message: If True, returns `lf.Message` as the output, instead of
+      returning the structured `message.result`.
+    skip_lm: If True, returns the rendered prompt as a UserMessage object.
+      otherwise return the LLM response based on the rendered prompt.
+    **kwargs: Template variables passed to `prompt` and keyword arguments passed
+      to `lf.structured.GenerateClass`.
+
+  Returns:
+    Generated class.
+
+  Raises:
+    CodeError: if generation failed.
+  """
+  if isinstance(prompt, str):
+    prompt = lf.Template(prompt, **kwargs)
+  elif isinstance(prompt, lf.Template):
+    prompt = prompt.rebind(**kwargs, raise_on_no_change=False)
+
+  if isinstance(prompt, lf.Template):
+    prompt = prompt.render(lm=lm)
+
+  call_kwargs = dict(skip_lm=skip_lm)
+  if lm is not None:
+    call_kwargs['lm'] = lm
+  message = GenerateClass(
+      input=prompt,
+      context=name,
+      examples=examples or default_classgen_examples(),
+      **kwargs,
+  )(**call_kwargs)
+  return message if returns_message else message.result
+
+
+def classgen_example(
+    class_name: str, prompt: str | pg.Symbolic, cls: Type[Any]
+) -> mapping.MappingExample:
+  """Creates a class generation example."""
+  if isinstance(prompt, lf.Template):
+    prompt = prompt.render()
+  return mapping.MappingExample(
+      input=prompt,
+      context=class_name,
+      output=cls,
+  )
+
+
+def default_classgen_examples() -> list[mapping.MappingExample]:
+  """Default examples for class generation."""
+
+  class Step(pg.Object):
+    description: str
+    output: float
+
+  class Solution(pg.Object):
+    steps: list[Step]  # pytype: disable=invalid-annotation
+    result: float
+
+  return [
+      classgen_example(
+          'Solution',
+          'How to evaluate an arithmetic expression?',
+          Solution,
+      )
+  ]

langfun/core/structured/schema_generation_test.py

@@ -0,0 +1,104 @@
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+import inspect
+import unittest
+
+import langfun.core.coding as lf_coding
+from langfun.core.llms import fake
+from langfun.core.structured import schema_generation
+
+
+class GenerateClassTest(unittest.TestCase):
+
+  def test_generate_class_prompt(self):
+    input_message = schema_generation.generate_class(
+        'Trip',
+        'Generate a trip class',
+        skip_lm=True,
+        returns_message=True,
+    )
+    self.maxDiff = None
+    self.assertEqual(
+        input_message.text,
+        inspect.cleandoc("""
+            Help generate a class based on the last CLASS_NAME and GENERATION_CONTEXT.
+
+            Instructions:
+            - Use `Object` as the base class for all generated classes
+            - Create auxillary classes for composition if needed.
+            - Use Python type annotation for declaraing fields:
+              (e.g. bool, str, int, float, Optional[str], List[int], Union[str, int])
+            - Do not use types that need import.
+            - Avoid self-referential types. e.g:
+              ```
+              class Node(Object):
+                children: list[Node]
+              ```
+            - Do not generate methods.
+
+            CLASS_NAME:
+              Solution
+
+            GENERATION_CONTEXT:
+              How to evaluate an arithmetic expression?
+
+            OUTPUT_CLASS:
+              ```python
+              class Step(Object):
+                description: str
+                output: float
+
+              class Solution(Object):
+                steps: list[Step]
+                result: float
+              ```
+
+
+            CLASS_NAME:
+              Trip
+
+            GENERATION_CONTEXT:
+              Generate a trip class
+
+            OUTPUT_CLASS:
+            """),
+    )
+
+  def test_generate_class(self):
+    lm = fake.StaticResponse("""
+        ```python
+        class A(Object):
+          x: int
+
+        class B(Object):
+          a: A
+        ```
+        """)
+    cls = schema_generation.generate_class(
+        'B',
+        'Generate a B class with a field pointing to another class A',
+        lm=lm,
+    )
+    self.assertIs(cls.__name__, 'B')
+
+    with self.assertRaises(lf_coding.CodeError):
+      schema_generation.generate_class(
+          'Foo',
+          'Generate a Foo class with a field pointing to another class A',
+          lm=lm,
+      )
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/structured/schema_test.py

@@ -435,6 +435,10 @@ class SchemaPythonReprTest(unittest.TestCase):
         schema_lib.class_definition(A),
         'class A:\n  pass\n',
     )
+    self.assertEqual(
+        schema_lib.class_definition(A, include_pg_object_as_base=True),
+        'class A(Object):\n  pass\n',
+    )
 
     class B:
       pass
@@ -520,6 +524,32 @@ class SchemaPythonReprTest(unittest.TestCase):
             ```
             """),
     )
+    self.assertEqual(
+        schema_lib.SchemaPythonRepr().repr(
+            schema,
+            include_result_definition=False,
+            include_pg_object_as_base=True,
+            markdown=False,
+        ),
+        inspect.cleandoc("""
+            class Foo(Object):
+              x: int
+
+            class A(Object):
+              foo: Foo
+
+            class Bar(Object):
+              y: str
+
+            class Baz(Bar):
+              y: str
+
+            class B(A):
+              foo: Foo
+              bar: Bar
+              foo2: Foo
+            """),
+    )
 
 
 class SchemaJsonReprTest(unittest.TestCase):
@@ -559,6 +589,20 @@ class ValuePythonReprTest(unittest.TestCase):
         ),
         "A(foo=[Foo(x=1), Foo(x=2)], y='bar')",
     )
+    self.assertEqual(
+        schema_lib.ValuePythonRepr().repr(A),
+        inspect.cleandoc("""
+            ```python
+            class Foo(Object):
+              x: int
+            
+            class A(Object):
+              foo: list[Foo]
+              y: str | None
+            ```
+            """),
+    )
+    self.assertEqual(schema_lib.source_form(int), 'int')
 
   def test_parse(self):
     class Foo(pg.Object):

langfun/core/template.py

@@ -124,9 +124,12 @@ class Template(
 
   @classmethod
   def resolve_vars(cls, template_str: str) -> Set[str]:
-    env = jinja2.Environment()
-    ast = env.parse(template_str)
-    return jinja2_meta.find_undeclared_variables(ast)
+    try:
+      env = jinja2.Environment()
+      ast = env.parse(template_str)
+      return jinja2_meta.find_undeclared_variables(ast)
+    except jinja2.TemplateSyntaxError as e:
+      raise ValueError(f'Bad template string:\n\n{template_str}') from e
 
   def _on_bound(self) -> None:
     super()._on_bound()

langfun/core/template_test.py

@@ -154,6 +154,10 @@ class DefinitionTest(unittest.TestCase):
     with self.assertRaisesRegex(TypeError, '.* missing 1 required argument'):
       MyPrompt(y=2)()
 
+  def test_bad_template(self):
+    with self.assertRaisesRegex(ValueError, 'Bad template string.*'):
+      Template('{{x=1')
+
 
 class VarsTest(unittest.TestCase):
 

{langfun-0.0.2.dev20240314.dist-info → langfun-0.0.2.dev20240316.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.0.2.dev20240314
+Version: 0.0.2.dev20240316
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors
@@ -24,7 +24,7 @@ License-File: LICENSE
 Requires-Dist: google-generativeai >=0.3.2
 Requires-Dist: jinja2 >=3.1.2
 Requires-Dist: openai ==0.27.2
-Requires-Dist: pyglove >=0.4.5.dev20240201
+Requires-Dist: pyglove >=0.4.5.dev20240314
 Requires-Dist: python-magic >=0.4.27
 Requires-Dist: requests >=2.31.0
 Requires-Dist: termcolor ==1.1.0

{langfun-0.0.2.dev20240314.dist-info → langfun-0.0.2.dev20240316.dist-info}/RECORD

@@ -1,4 +1,4 @@
-langfun/__init__.py,sha256=8H9dYWG6gM3SlQfTH9BNiwVtIbe7Wz8XUIK5uolY1Z0,1760
+langfun/__init__.py,sha256=PqX3u18BC0szYIMu00j-RKxvwkNPwXtAFZ-96oxrQ0M,1841
 langfun/core/__init__.py,sha256=sVcPl89lWYHQ1cUoaLaM8dErCovugJo5e2F3A_94Q3Y,4192
 langfun/core/component.py,sha256=VRPfDB_2jEnxcB3-HoiVjG4ID-SMenNPIsytb0uXMPg,9674
 langfun/core/component_test.py,sha256=VAPd6V_-odAe8rBvesW3ogYDd6OSqRq4FaPhfgOM4Zg,7949
@@ -7,9 +7,9 @@ langfun/core/concurrent_test.py,sha256=mwFMZhDUdppnDr7vDSTwcbMHwrdsIoKJwRYNtl4ZW
 langfun/core/console.py,sha256=bk5rNPNm9rMGW5YT2HixxU04p2umnoabn5SDz6Dqe88,2317
 langfun/core/console_test.py,sha256=5SYJdxpJGLgdSSQqqMPoA1X6jpsLD8rgcyk-EgI65oE,1077
 langfun/core/langfunc.py,sha256=WXdTc3QsmGD_n80KD9dFRr5MHpGZ9E_y_Rhtk4t9-3w,11852
-langfun/core/langfunc_test.py,sha256=8WeiyNauZPkbAA3HiLjVw5-pRSmiLlz-77lB_fjHGdA,8317
-langfun/core/language_model.py,sha256=Qbm7wxgxW26bCVwtgpp-4aV3BKYAsb4IJrJuzhf3Q6o,15507
-langfun/core/language_model_test.py,sha256=h5MWooOb9HubvOzxaBnH6WuDYBdxTetu7JZSWDzva3M,11368
+langfun/core/langfunc_test.py,sha256=dFNJoEXExIkrAJ9_PSWh_iRQoR4Gmp2VOZ_ve61DSHM,8339
+langfun/core/language_model.py,sha256=jPuFfjnRCnbT8po-CBPgmXoa09Yfk5_21snCXURqaKU,17011
+langfun/core/language_model_test.py,sha256=q7pNdirVWfkQXPA3taCGnyLB2NNs1KqX4JjjnoJvFOQ,11365
 langfun/core/memory.py,sha256=f-asN1F7Vehgdn_fK84v73GrEUOxRtaW934keutTKjk,2416
 langfun/core/message.py,sha256=QhvV9t5qaryPcruyxxcXi3gm9QDInkSldwTtK6sVJ3c,15734
 langfun/core/message_test.py,sha256=Z23pUM5vPnDrYkIIibe2KL73D5HKur_awI0ut_EQFQA,9501
@@ -21,8 +21,8 @@ langfun/core/sampling.py,sha256=vygWvgC8MFw0_AKNSmz-ywMXJYWf8cl0tI8QycvAmyI,5795
 langfun/core/sampling_test.py,sha256=U7PANpMsl9E_pa4_Y4FzesSjcwg-u-LKHGCWSgv-8FY,3663
 langfun/core/subscription.py,sha256=euawEuSZP-BHydaT-AQpfYFL0m5pWPGcW0upFhrojqc,10930
 langfun/core/subscription_test.py,sha256=Y4ZdbZEwm83YNZBxHff0QR4QUa4rdaNXA3_jfIcArBo,8717
-langfun/core/template.py,sha256=1UnXgVuqdGRes0vSMamIQ8KpG0NgTXV1iCSWVbCybN4,17530
-langfun/core/template_test.py,sha256=WNf7O45V5BZz7IaAAcZhGANyaKMwpWjBpcpEIASfh-Q,13446
+langfun/core/template.py,sha256=zVD8dAsXFfgF25aKh2WqSuCEHVqriCC-4tLbQqTMa2w,17662
+langfun/core/template_test.py,sha256=1hDdYfvXJVoslTUudh3WhxU7VnDSiIz6MkxPfmuHKAY,13572
 langfun/core/text_formatting.py,sha256=ytjj7opnRJ6w-pkglL2CZUyfYDXLpNf65E42LBb31gc,5158
 langfun/core/text_formatting_test.py,sha256=nyKC6tn2L4hPJiqQHgxcbQsJJi4A4Nbj8FiO8iT6B80,1514
 langfun/core/coding/__init__.py,sha256=5utju_fwEsImaiftx4oXKl9FAM8p281k8-Esdh_-m1w,835
@@ -49,11 +49,11 @@ langfun/core/eval/scoring_test.py,sha256=_L_B40VZkyI2_PJce-jVKYC4llrO4jGUR5j86Gu
 langfun/core/llms/__init__.py,sha256=T4mgT091BLA4mHrOjAvEGhZPHf0tiYgqD88l_JTp1dQ,2386
 langfun/core/llms/fake.py,sha256=dVzOrW27RZ1p3DdQoRCRZs_vfoQcTcNrlWxia7oqmvw,2499
 langfun/core/llms/fake_test.py,sha256=Qk_Yoi4Z7P9o6f8Q_BZkaSlvxH89ZVsDxnVIbSBRBXk,3555
-langfun/core/llms/gemini.py,sha256=9HxrTvac_dMbDytNCEo6YcYqs8vsePtScfI_EygpI5Y,5677
+langfun/core/llms/gemini.py,sha256=E7JGewkgjMzuDQxAn8CBbwWsDDZH4jcmNbzmO3OvdsY,5474
 langfun/core/llms/gemini_test.py,sha256=ybNNCn3JW3hYpMe0wT5ILGDrMPaYYU8PN2kSookM0jk,5433
-langfun/core/llms/llama_cpp.py,sha256=EIjJa1-Tg4_VaIxVR88oDWSWc_axc1r2KwSPpl4PSp0,2549
+langfun/core/llms/llama_cpp.py,sha256=sJ9TOismqwGJ7QhgdYknWTEkqrbeZpWYc_nClOh36NU,2320
 langfun/core/llms/llama_cpp_test.py,sha256=ZxC6defGd_HX9SFRU9U4cJiQnBKundbOrchbXuC1Z2M,1683
-langfun/core/llms/openai.py,sha256=-PnJ8OICuPTzr-unIys4HftNVZ6seAhV5nXny4OfVYc,11715
+langfun/core/llms/openai.py,sha256=BV8NWjB1b6A1X4Kff8Pub5AECodsngZnXqeBvRIHFM0,11331
 langfun/core/llms/openai_test.py,sha256=yfw7A-4Zo9u1cIkAMk39evE-tO7z6isNYTXiSnJXDQw,7599
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=cFfYvOIUae842pncqCAsRvqXCk2AnAsRYVx0mcIoAeY,3338
@@ -69,19 +69,21 @@ langfun/core/modalities/mime.py,sha256=wVfaYflhGz1W4v3m972rAplW3OGOFtjFpHDYIaUD5
 langfun/core/modalities/mime_test.py,sha256=cVHxRvJ1QXC1SVhBmWkJdWGpL9Xl0UNfTQq6j0OGGL4,1881
 langfun/core/modalities/video.py,sha256=5-sIlzXb_ZY84RMFcpVD9ysP9GbcwbdKaZOEm3jECtc,1469
 langfun/core/modalities/video_test.py,sha256=jYuI2m8S8zDCAVBPEUbbpP205dXAht90A2_PHWo4-r8,2039
-langfun/core/structured/__init__.py,sha256=LZ5BFLX6VXy1zH17yChWCdg8bvIDrhtL2lqtSCwtZ-M,3187
+langfun/core/structured/__init__.py,sha256=SpObW-HKpyKvkLlX8FV5ixz7CRm098j2aGfOguM3AUI,3462
 langfun/core/structured/completion.py,sha256=skBxt6V_fv2TBUKnzFgnPMbVY8HSYn8sY04MLok2yvs,7299
 langfun/core/structured/completion_test.py,sha256=98UCgA4gzfp6H6HgP2s2kcKs25YH3k4Nxj1rgAvmVBw,19249
-langfun/core/structured/description.py,sha256=3MLTbpTpeiBqMRe3WfDNIxtrE6WQsKJsJdkbfcyPlsg,5088
+langfun/core/structured/description.py,sha256=SXW4MJvshFjbR-0gw6rE21o6WXq12UlRXawvDBXMZFA,5211
 langfun/core/structured/description_test.py,sha256=UtZGjSFUaQ6130t1E5tcL7ODu0xIefkapb53TbnqsK8,7362
-langfun/core/structured/mapping.py,sha256=lGkjhmvVdhBGgJmc5KbfT2xQjC1MuU4OCcCfsAYJjaQ,10192
+langfun/core/structured/mapping.py,sha256=tahkaAB-L6yKbYb7qjVI301-FfIARdw4w8nP3wqS2-k,10291
 langfun/core/structured/mapping_test.py,sha256=07DDCGbwytQHSMm7fCi5-Ly-JNgdV4ubHZq0wthX4A4,3338
-langfun/core/structured/parsing.py,sha256=YKWl9ZQ2uFkt78SXiRISWHg8_cDMGMwAN3SeK-OqWt4,11382
+langfun/core/structured/parsing.py,sha256=yTKuezai5i-X9W-jU0DeEZzqHHbCFom0plj-D0bhp98,11436
 langfun/core/structured/parsing_test.py,sha256=2_Uf3LYNRON1-5ysEr75xiG_cAxR3ZiixSfvUQu6mOQ,20846
 langfun/core/structured/prompting.py,sha256=0xRPC0K_RaFRv-j52x8_-1n1eRFSomJEpdZApVXsCV0,6902
 langfun/core/structured/prompting_test.py,sha256=SwoYbPyKhUT1H2QbqHvl93biCiE9Ttn1aWixoHH-v9Y,19129
-langfun/core/structured/schema.py,sha256=5DKba0LrvXCJFRY-NVfER3p54BLOB7M3Yi2-u5IAJTw,24115
-langfun/core/structured/schema_test.py,sha256=LEtCST5Bfwoke59I6Q1mnOJLf2cFXQwKwTeAkI2hgqM,20912
+langfun/core/structured/schema.py,sha256=60griJ-yC1SExX6g-aOcAOo8yFh53CdwMV4EVK3ivug,25207
+langfun/core/structured/schema_generation.py,sha256=Yv9flJ4GTtLw-bDB8S7A93G-z4gXsFMkMASkbiduT3E,5353
+langfun/core/structured/schema_generation_test.py,sha256=cfZyP0gHno2fXy_c9vsVdvHmqKQSfuyUsCtfO3JFmYQ,2945
+langfun/core/structured/schema_test.py,sha256=kMIgnAzm3f2O5ofn0pPKjT6H8hny4cWVaUVDOZuyjOQ,21987
 langfun/core/structured/scoring.py,sha256=a3vfGnqf-DOWjD07MF54GCZTO_R1RTxTDVPzerXnU0s,2325
 langfun/core/structured/scoring_test.py,sha256=TznLMl0x9QxzmhHz_3Vr44VOXuvFnUSeRQVhu33W5cA,1437
 langfun/core/templates/__init__.py,sha256=bO0eMsVJbi7sxEB2YlInKRQ2EVP-RyyKUwcD-8msuN4,927
@@ -93,8 +95,8 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=IB5rWbjK_9CTkqEo1BclQPzFAKcIiusJckH8J19HFgI,2096
-langfun-0.0.2.dev20240314.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.0.2.dev20240314.dist-info/METADATA,sha256=zA8mV-vbd9FftlkgYSeTTvvUqkuTmjewdbc4NSDhuas,3405
-langfun-0.0.2.dev20240314.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-langfun-0.0.2.dev20240314.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.0.2.dev20240314.dist-info/RECORD,,
+langfun-0.0.2.dev20240316.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.0.2.dev20240316.dist-info/METADATA,sha256=rvpQMtNiFs55Okrd1TNlJOS8szUWshlCt5NFB_2vPfs,3405
+langfun-0.0.2.dev20240316.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+langfun-0.0.2.dev20240316.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.0.2.dev20240316.dist-info/RECORD,,