langfun 0.0.2.dev20240527__py3-none-any.whl → 0.0.2.dev20240531__py3-none-any.whl

langfun/core/structured/schema.py

@@ -16,6 +16,7 @@
 import abc
 import inspect
 import io
+import re
 import textwrap
 import typing
 from typing import Any, Literal, Sequence, Type, Union
@@ -24,6 +25,17 @@ from langfun.core.coding.python import correction
 import pyglove as pg
 
 
+def include_method_in_prompt(method):
+  """Decorator to include a method in the class definition of the prompt."""
+  setattr(method, '__show_in_prompt__', True)
+  return method
+
+
+def should_include_method_in_prompt(method):
+  """Returns true if the method should be shown in the prompt."""
+  return getattr(method, '__show_in_prompt__', False)
+
+
 def parse_value_spec(value) -> pg.typing.ValueSpec:
   """Parses a PyGlove ValueSpec equivalence into a ValueSpec."""
   if isinstance(value, pg.typing.ValueSpec):
@@ -163,9 +175,12 @@ class Schema(lf.NaturalLanguageFormattable, pg.Object):
 
   def class_dependencies(
       self,
+      include_base_classes: bool = True,
       include_subclasses: bool = True) -> list[Type[Any]]:
     """Returns a list of class dependencies for current schema."""
-    return class_dependencies(self.spec, include_subclasses)
+    return class_dependencies(
+        self.spec, include_base_classes, include_subclasses
+    )
 
   @classmethod
   def from_value(cls, value) -> 'Schema':
@@ -198,11 +213,12 @@ def class_dependencies(
         Type[pg.Object],
         tuple[Union[pg.typing.ValueSpec, Type[pg.Object]], ...],
     ],
+    include_base_classes: bool = True,
     include_subclasses: bool = True,
 ) -> list[Type[Any]]:
   """Returns a list of class dependencies from a value or specs."""
   if isinstance(value_or_spec, Schema):
-    return value_or_spec.class_dependencies(include_subclasses)
+    value_or_spec = value_or_spec.spec
 
   if inspect.isclass(value_or_spec) or isinstance(
       value_or_spec, pg.typing.ValueSpec
@@ -236,13 +252,14 @@ def class_dependencies(
       if vs.cls not in seen:
         seen.add(vs.cls)
 
-        # Add base classes as dependencies.
-        for base_cls in vs.cls.__bases__:
-          # We only keep track of user-defined symbolic classes.
-          if base_cls is not object and base_cls is not pg.Object:
-            _fill_dependencies(
-                pg.typing.Object(base_cls), include_subclasses=False
-            )
+        if include_base_classes:
+          # Add base classes as dependencies.
+          for base_cls in vs.cls.__bases__:
+            # We only keep track of user-defined symbolic classes.
+            if base_cls is not object and base_cls is not pg.Object:
+              _fill_dependencies(
+                  pg.typing.Object(base_cls), include_subclasses=False
+              )
 
         # Add members as dependencies.
         for field in _pg_schema(vs.cls).values():
@@ -301,7 +318,6 @@ class SchemaPythonRepr(SchemaRepr):
       schema: Schema,
       *,
       include_result_definition: bool = True,
-      include_methods: bool = False,
       markdown: bool = True,
       **kwargs,
   ) -> str:
@@ -309,15 +325,27 @@ class SchemaPythonRepr(SchemaRepr):
     if include_result_definition:
       ret += self.result_definition(schema)
     class_definition_str = self.class_definitions(
-        schema, markdown=markdown, include_methods=include_methods, **kwargs
+        schema, markdown=markdown, **kwargs
     )
     if class_definition_str:
       ret += f'\n\n{class_definition_str}'
     return ret.strip()
 
-  def class_definitions(self, schema: Schema, **kwargs) -> str | None:
-    deps = schema.class_dependencies(include_subclasses=True)
-    return class_definitions(deps, **kwargs)
+  def class_definitions(
+      self,
+      schema: Schema,
+      additional_dependencies: list[Type[Any]] | None = None,
+      **kwargs
+  ) -> str | None:
+    """Returns a string containing of class definitions from a schema."""
+    deps = schema.class_dependencies(
+        include_base_classes=False, include_subclasses=True
+    )
+    allowed_dependencies = set(deps)
+    if additional_dependencies:
+      allowed_dependencies.update(additional_dependencies)
+    return class_definitions(
+        deps, allowed_dependencies=allowed_dependencies, **kwargs)
 
   def result_definition(self, schema: Schema) -> str:
     return annotation(schema.spec)
@@ -331,8 +359,7 @@ def source_form(value, markdown: bool = False) -> str:
 def class_definitions(
     classes: Sequence[Type[Any]],
     *,
-    include_pg_object_as_base: bool = False,
-    include_methods: bool = False,
+    allowed_dependencies: set[Type[Any]] | None = None,
     strict: bool = False,
     markdown: bool = False,
 ) -> str | None:
@@ -347,8 +374,7 @@ def class_definitions(
         class_definition(
             cls,
             strict=strict,
-            include_pg_object_as_base=include_pg_object_as_base,
-            include_methods=include_methods,
+            allowed_dependencies=allowed_dependencies,
         )
     )
   ret = def_str.getvalue()
@@ -360,8 +386,7 @@ def class_definitions(
 def class_definition(
     cls,
     strict: bool = False,
-    include_pg_object_as_base: bool = False,
-    include_methods: bool = False,
+    allowed_dependencies: set[Type[Any]] | None = None,
 ) -> str:
   """Returns the Python class definition."""
   out = io.StringIO()
@@ -369,7 +394,7 @@ def class_definition(
   eligible_bases = []
   for base_cls in cls.__bases__:
     if base_cls is not object:
-      if include_pg_object_as_base or base_cls is not pg.Object:
+      if allowed_dependencies is None or base_cls in allowed_dependencies:
         eligible_bases.append(base_cls.__name__)
 
   if eligible_bases:
@@ -406,32 +431,41 @@ def class_definition(
             out.write('  # ')
             out.write(line)
             out.write('\n')
-      out.write(f'  {field.key}: {annotation(field.value, strict=strict)}')
-      out.write('\n')
-      empty_class = False
 
-  if include_methods:
-    for method in _iter_newly_defined_methods(cls):
-      out.write('\n')
-      out.write(
-          textwrap.indent(
-              inspect.cleandoc('\n' + inspect.getsource(method)), ' ' * 2)
+      annotation_str = annotation(
+          field.value, strict=strict, allowed_dependencies=allowed_dependencies
       )
+      out.write(f'  {field.key}: {annotation_str}')
       out.write('\n')
       empty_class = False
 
+  for method in _iter_newly_defined_methods(cls, allowed_dependencies):
+    source = inspect.getsource(method)
+    # Remove decorators from the method definition.
+    source = re.sub(r'\s*@.*\.include_method_in_prompt.*\n', '', source)
+    out.write('\n')
+    out.write(
+        textwrap.indent(
+            inspect.cleandoc('\n' + source), ' ' * 2)
+    )
+    out.write('\n')
+    empty_class = False
+
   if empty_class:
     out.write('  pass\n')
   return out.getvalue()
 
 
-def _iter_newly_defined_methods(cls):
-  names = set(dir(cls))
+def _iter_newly_defined_methods(
+    cls, allowed_dependencies: set[Type[Any]] | None):
+  names = {attr_name: True for attr_name in dir(cls)}
   for base in cls.__bases__:
-    names -= set(dir(base))
-  for name in names:
+    if allowed_dependencies is None or base in allowed_dependencies:
+      for name in dir(base):
+        names.pop(name, None)
+  for name in names.keys():
     attr = getattr(cls, name)
-    if callable(attr):
+    if callable(attr) and should_include_method_in_prompt(attr):
       yield attr
 
 
@@ -439,8 +473,12 @@ def annotation(
     vs: pg.typing.ValueSpec,
     annotate_optional: bool = True,
     strict: bool = False,
+    allowed_dependencies: set[Type[Any]] | None = None,
 ) -> str:
   """Returns the annotation string for a value spec."""
+  child_annotation_kwargs = dict(
+      strict=strict, allowed_dependencies=allowed_dependencies
+  )
   if isinstance(vs, pg.typing.Any):
     return 'Any'
   elif isinstance(vs, pg.typing.Enum):
@@ -449,7 +487,7 @@ def annotation(
   elif isinstance(vs, pg.typing.Union):
     candidate_str = ', '.join(
         [
-            annotation(c, annotate_optional=False, strict=strict)
+            annotation(c, annotate_optional=False, **child_annotation_kwargs)
             for c in vs.candidates
         ]
     )
@@ -485,20 +523,23 @@ def annotation(
         )
       x += '(' + ', '.join(constraints) + ')'
   elif isinstance(vs, pg.typing.Object):
-    x = vs.cls.__name__
+    if allowed_dependencies is None or vs.cls in allowed_dependencies:
+      x = vs.cls.__name__
+    else:
+      x = 'Any'
   elif isinstance(vs, pg.typing.List):
-    item_str = annotation(vs.element.value, strict=strict)
+    item_str = annotation(vs.element.value, **child_annotation_kwargs)
     x = f'list[{item_str}]'
   elif isinstance(vs, pg.typing.Tuple):
     elem_str = ', '.join(
-        [annotation(el.value, strict=strict) for el in vs.elements]
+        [annotation(el.value, **child_annotation_kwargs) for el in vs.elements]
     )
     x = f'tuple[{elem_str}]'
   elif isinstance(vs, pg.typing.Dict):
     kv_pairs = None
     if vs.schema is not None:
       kv_pairs = [
-          (k, annotation(f.value, strict=strict))
+          (k, annotation(f.value, **child_annotation_kwargs))
           for k, f in vs.schema.items()
           if isinstance(k, pg.typing.ConstStrKey)
       ]
@@ -509,7 +550,7 @@ def annotation(
       if strict:
         x = f'pg.typing.Dict({x})'
     elif vs.schema and vs.schema.dynamic_field:
-      v = annotation(vs.schema.dynamic_field.value, strict=strict)
+      v = annotation(vs.schema.dynamic_field.value, **child_annotation_kwargs)
       x = f'dict[str, {v}]'
     else:
       x = 'dict[str, Any]'
@@ -604,7 +645,12 @@ class ValuePythonRepr(ValueRepr):
       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
+            cls_schema,
+            markdown=markdown,
+            # We add `pg.Object` as additional dependencies to the class
+            # definition so exemplars for class generation could show
+            # pg.Object as their bases.
+            additional_dependencies=[pg.Object]
         )
         assert object_code is not None
         return object_code

langfun/core/structured/schema_test.py

@@ -283,9 +283,10 @@ class SchemaPythonReprTest(unittest.TestCase):
       value_spec: pg.typing.ValueSpec,
       expected_annotation: str,
       strict: bool = False,
+      **kwargs,
   ) -> None:
     self.assertEqual(
-        schema_lib.annotation(value_spec, strict=strict),
+        schema_lib.annotation(value_spec, strict=strict, **kwargs),
         expected_annotation,
     )
 
@@ -361,11 +362,27 @@ class SchemaPythonReprTest(unittest.TestCase):
     self.assert_annotation(
         pg.typing.Object(Activity).noneable(), 'Activity | None'
     )
+    self.assert_annotation(
+        pg.typing.Object(Activity).noneable(), 'Activity | None',
+        allowed_dependencies=set([Activity]),
+    )
+    self.assert_annotation(
+        pg.typing.Object(Activity).noneable(), 'Any | None',
+        allowed_dependencies=set(),
+    )
 
     # List.
     self.assert_annotation(
         pg.typing.List(pg.typing.Object(Activity)), 'list[Activity]'
     )
+    self.assert_annotation(
+        pg.typing.List(pg.typing.Object(Activity)), 'list[Activity]',
+        allowed_dependencies=set([Activity]),
+    )
+    self.assert_annotation(
+        pg.typing.List(pg.typing.Object(Activity)), 'list[Any]',
+        allowed_dependencies=set(),
+    )
     self.assert_annotation(
         pg.typing.List(pg.typing.Object(Activity)).noneable(),
         'list[Activity] | None',
@@ -377,16 +394,35 @@ class SchemaPythonReprTest(unittest.TestCase):
 
     # Tuple.
     self.assert_annotation(
-        pg.typing.Tuple([pg.typing.Int(), pg.typing.Str()]), 'tuple[int, str]'
+        pg.typing.Tuple([Activity, pg.typing.Str()]), 'tuple[Activity, str]'
+    )
+    self.assert_annotation(
+        pg.typing.Tuple([Activity, pg.typing.Str()]), 'tuple[Activity, str]',
+        allowed_dependencies=set([Activity]),
     )
     self.assert_annotation(
-        pg.typing.Tuple([pg.typing.Int(), pg.typing.Str()]).noneable(),
-        'tuple[int, str] | None',
+        pg.typing.Tuple([Activity, pg.typing.Str()]), 'tuple[Any, str]',
+        allowed_dependencies=set(),
+    )
+    self.assert_annotation(
+        pg.typing.Tuple([Activity, pg.typing.Str()]).noneable(),
+        'tuple[Activity, str] | None',
     )
 
     # Dict.
     self.assert_annotation(
-        pg.typing.Dict({'x': int, 'y': str}), '{\'x\': int, \'y\': str}'
+        pg.typing.Dict({'x': Activity, 'y': str}),
+        '{\'x\': Activity, \'y\': str}'
+    )
+    self.assert_annotation(
+        pg.typing.Dict({'x': Activity, 'y': str}),
+        '{\'x\': Activity, \'y\': str}',
+        allowed_dependencies=set([Activity]),
+    )
+    self.assert_annotation(
+        pg.typing.Dict({'x': Activity, 'y': str}),
+        '{\'x\': Any, \'y\': str}',
+        allowed_dependencies=set(),
     )
     self.assert_annotation(
         pg.typing.Dict({'x': int, 'y': str}),
@@ -420,6 +456,13 @@ class SchemaPythonReprTest(unittest.TestCase):
         ).noneable(),
         'Union[Activity, Itinerary, None]',
     )
+    self.assert_annotation(
+        pg.typing.Union(
+            [pg.typing.Object(Activity), pg.typing.Object(Itinerary)]
+        ).noneable(),
+        'Union[Activity, Any, None]',
+        allowed_dependencies=set([Activity]),
+    )
 
     # Any.
     self.assert_annotation(pg.typing.Any(), 'Any')
@@ -427,13 +470,13 @@ class SchemaPythonReprTest(unittest.TestCase):
 
   def test_class_definition(self):
     self.assertEqual(
-        schema_lib.class_definition(Activity),
+        schema_lib.class_definition(Activity, allowed_dependencies=set()),
         'class Activity:\n  description: str\n',
     )
     self.assertEqual(
         schema_lib.class_definition(Itinerary),
         inspect.cleandoc("""
-            class Itinerary:
+            class Itinerary(Object):
               \"\"\"A travel itinerary for a day.\"\"\"
               day: int(min=1)
               type: Literal['daytime', 'nighttime']
@@ -443,7 +486,9 @@ class SchemaPythonReprTest(unittest.TestCase):
             """) + '\n',
     )
     self.assertEqual(
-        schema_lib.class_definition(PlaceOfInterest),
+        schema_lib.class_definition(
+            PlaceOfInterest, allowed_dependencies=set()
+        ),
         inspect.cleandoc("""
             class PlaceOfInterest:
               \"\"\"The name of a place of interest.
@@ -459,11 +504,11 @@ class SchemaPythonReprTest(unittest.TestCase):
       pass
 
     self.assertEqual(
-        schema_lib.class_definition(A),
+        schema_lib.class_definition(A, allowed_dependencies=set()),
         'class A:\n  pass\n',
     )
     self.assertEqual(
-        schema_lib.class_definition(A, include_pg_object_as_base=True),
+        schema_lib.class_definition(A),
         'class A(Object):\n  pass\n',
     )
 
@@ -471,18 +516,21 @@ class SchemaPythonReprTest(unittest.TestCase):
       x: str
       __kwargs__: typing.Any
 
-    self.assertEqual(schema_lib.class_definition(C), 'class C:\n  x: str\n')
+    self.assertEqual(
+        schema_lib.class_definition(C), 'class C(Object):\n  x: str\n'
+    )
 
     class D(pg.Object):
       x: str
+      @schema_lib.include_method_in_prompt
       def __call__(self, y: int) -> int:
         return len(self.x) + y
 
     self.assertEqual(
-        schema_lib.class_definition(D, include_methods=True),
+        schema_lib.class_definition(D),
         inspect.cleandoc(
             """
-            class D:
+            class D(Object):
               x: str
 
               def __call__(self, y: int) -> int:
@@ -506,31 +554,28 @@ class SchemaPythonReprTest(unittest.TestCase):
     class A(pg.Object):
       foo: Foo
 
+      @schema_lib.include_method_in_prompt
       def foo_value(self) -> int:
         return self.foo.x
 
+      def baz_value(self) -> str:
+        return 'baz'
+
     class B(A):
       bar: Bar
       foo2: Foo
 
+      @schema_lib.include_method_in_prompt
       def bar_value(self) -> str:
         return self.bar.y
 
     schema = schema_lib.Schema([B])
     self.assertEqual(
-        schema_lib.SchemaPythonRepr().class_definitions(
-            schema, include_methods=True
-        ),
+        schema_lib.SchemaPythonRepr().class_definitions(schema),
         inspect.cleandoc('''
             class Foo:
               x: int
 
-            class A:
-              foo: Foo
-
-              def foo_value(self) -> int:
-                return self.foo.x
-
             class Bar:
               """Class Bar."""
               y: str
@@ -539,13 +584,16 @@ class SchemaPythonReprTest(unittest.TestCase):
               """Baz(y: str)"""
               y: str
 
-            class B(A):
+            class B:
               foo: Foo
               bar: Bar
               foo2: Foo
 
               def bar_value(self) -> str:
                 return self.bar.y
+
+              def foo_value(self) -> int:
+                return self.foo.x
             ''') + '\n',
     )
 
@@ -562,9 +610,6 @@ class SchemaPythonReprTest(unittest.TestCase):
             class Foo:
               x: int
 
-            class A:
-              foo: Foo
-
             class Bar:
               """Class Bar."""
               y: str
@@ -573,10 +618,16 @@ class SchemaPythonReprTest(unittest.TestCase):
               """Baz(y: str)"""
               y: str
 
-            class B(A):
+            class B:
               foo: Foo
               bar: Bar
               foo2: Foo
+
+              def bar_value(self) -> str:
+                return self.bar.y
+
+              def foo_value(self) -> int:
+                return self.foo.x
             ```
             '''),
     )
@@ -584,16 +635,12 @@ class SchemaPythonReprTest(unittest.TestCase):
         schema_lib.SchemaPythonRepr().repr(
             schema,
             include_result_definition=False,
-            include_pg_object_as_base=True,
             markdown=False,
         ),
         inspect.cleandoc('''
-            class Foo(Object):
+            class Foo:
               x: int
 
-            class A(Object):
-              foo: Foo
-
             class Bar:
               """Class Bar."""
               y: str
@@ -602,10 +649,16 @@ class SchemaPythonReprTest(unittest.TestCase):
               """Baz(y: str)"""
               y: str
 
-            class B(A):
+            class B:
               foo: Foo
               bar: Bar
               foo2: Foo
+
+              def bar_value(self) -> str:
+                return self.bar.y
+
+              def foo_value(self) -> int:
+                return self.foo.x
             '''),
     )
 
@@ -653,7 +706,7 @@ class ValuePythonReprTest(unittest.TestCase):
             ```python
             class Foo(Object):
               x: int
-            
+
             class A(Object):
               foo: list[Foo]
               y: str | None

langfun/core/template_test.py

@@ -309,7 +309,7 @@ class RenderTest(unittest.TestCase):
         Template(
             'This is {{ x }} and {{ a }}', x=1, a=CustomModality('foo')
         ).render(),
-        'This is 1 and {{a}}',
+        'This is 1 and <<[[a]]>>',
     )
 
   def test_render_with_default(self):

{langfun-0.0.2.dev20240527.dist-info → langfun-0.0.2.dev20240531.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.0.2.dev20240527
+Version: 0.0.2.dev20240531
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors
@@ -30,6 +30,8 @@ Requires-Dist: python-magic >=0.4.27
 Requires-Dist: requests >=2.31.0
 Requires-Dist: termcolor ==1.1.0
 Requires-Dist: tqdm >=4.64.1
+Requires-Dist: python-docx >=0.8.11
+Requires-Dist: pandas >=2.1.4
 
 <div align="center">
 <img src="https://raw.githubusercontent.com/google/langfun/main/docs/_static/logo.svg" width="520px" alt="logo"></img>