langfun 0.0.2.dev20240506__py3-none-any.whl → 0.0.2.dev20240507__py3-none-any.whl

langfun/core/structured/mapping.py

@@ -251,7 +251,7 @@ class Mapping(lf.LangFunc):
 
       {%- if example.schema -%}
       {{ schema_title }}:
-      {{ example.schema_repr(protocol) | indent(2, True) }}
+      {{ example.schema_repr(protocol, include_methods=include_methods) | indent(2, True) }}
 
       {% endif -%}
 
@@ -279,6 +279,10 @@ class Mapping(lf.LangFunc):
       'The protocol for representing the schema and value.',
   ] = 'python'
 
+  include_methods: Annotated[
+      bool, 'If True, include method definitions in the schema.'
+  ] = False
+
   #
   # Other user-provided flags.
   #

langfun/core/structured/prompting.py

@@ -113,6 +113,7 @@ def query(
     autofix: int = 0,
     autofix_lm: lf.LanguageModel | None = None,
     protocol: schema_lib.SchemaProtocol = 'python',
+    include_methods: bool = False,
     returns_message: bool = False,
     skip_lm: bool = False,
     **kwargs,
@@ -160,11 +161,11 @@ def query(
     examples: An optional list of fewshot examples for helping parsing. If None,
       the default one-shot example will be added.
     cache_seed: Seed for computing cache key. The cache key is determined by a
-      tuple of (lm, prompt, cache seed). If None, cache will be disabled for
-      the query even cache is configured by the LM.
+      tuple of (lm, prompt, cache seed). If None, cache will be disabled for the
+      query even cache is configured by the LM.
     response_postprocess: An optional callable object to process the raw LM
-      response before parsing it into the final output object. If None, the
-      raw LM response will not be processed.
+      response before parsing it into the final output object. If None, the raw
+      LM response will not be processed.
     autofix: Number of attempts to auto fix the generated code. If 0, autofix is
       disabled. Auto-fix is not supported for 'json' protocol.
     autofix_lm: The language model to use for autofix. If not specified, the
@@ -172,6 +173,8 @@ def query(
       will use `lm`.
     protocol: The protocol for schema/value representation. Applicable values
       are 'json' and 'python'. By default `python` will be used.
+    include_methods: If True, include method definitions in the output type
+      during prompting.
     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.
@@ -226,6 +229,7 @@ def query(
       schema=schema,
       default=default,
       examples=examples,
+      include_methods=include_methods,
       response_postprocess=response_postprocess,
       autofix=autofix if protocol == 'python' else 0,
       **kwargs,

langfun/core/structured/schema.py

@@ -301,6 +301,7 @@ class SchemaPythonRepr(SchemaRepr):
       schema: Schema,
       *,
       include_result_definition: bool = True,
+      include_methods: bool = False,
       markdown: bool = True,
       **kwargs,
   ) -> str:
@@ -308,7 +309,7 @@ class SchemaPythonRepr(SchemaRepr):
     if include_result_definition:
       ret += self.result_definition(schema)
     class_definition_str = self.class_definitions(
-        schema, markdown=markdown, **kwargs
+        schema, markdown=markdown, include_methods=include_methods, **kwargs
     )
     if class_definition_str:
       ret += f'\n\n{class_definition_str}'
@@ -331,6 +332,7 @@ def class_definitions(
     classes: Sequence[Type[Any]],
     *,
     include_pg_object_as_base: bool = False,
+    include_methods: bool = False,
     strict: bool = False,
     markdown: bool = False,
 ) -> str | None:
@@ -346,6 +348,7 @@ def class_definitions(
             cls,
             strict=strict,
             include_pg_object_as_base=include_pg_object_as_base,
+            include_methods=include_methods,
         )
     )
   ret = def_str.getvalue()
@@ -355,7 +358,10 @@ def class_definitions(
 
 
 def class_definition(
-    cls, strict: bool = False, include_pg_object_as_base: bool = False
+    cls,
+    strict: bool = False,
+    include_pg_object_as_base: bool = False,
+    include_methods: bool = False,
 ) -> str:
   """Returns the Python class definition."""
   out = io.StringIO()
@@ -383,6 +389,7 @@ def class_definition(
         out.write('\n')
       out.write('  """\n')
 
+  empty_class = True
   if schema.fields:
     for key, field in schema.items():
       if not isinstance(key, pg.typing.ConstStrKey):
@@ -401,11 +408,33 @@ def class_definition(
             out.write('\n')
       out.write(f'  {field.key}: {annotation(field.value, strict=strict)}')
       out.write('\n')
-  else:
+      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)
+      )
+      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))
+  for base in cls.__bases__:
+    names -= set(dir(base))
+  for name in names:
+    attr = getattr(cls, name)
+    if callable(attr):
+      yield attr
+
+
 def annotation(
     vs: pg.typing.ValueSpec,
     annotate_optional: bool = True,
@@ -493,7 +522,8 @@ def annotation(
 class SchemaJsonRepr(SchemaRepr):
   """JSON-representation for a schema."""
 
-  def repr(self, schema: Schema) -> str:
+  def repr(self, schema: Schema, **kwargs) -> str:
+    del kwargs
     out = io.StringIO()
     def _visit(node: Any) -> None:
       if isinstance(node, str):

langfun/core/structured/schema_test.py

@@ -461,6 +461,23 @@ class SchemaPythonReprTest(unittest.TestCase):
 
     self.assertEqual(schema_lib.class_definition(C), 'class C:\n  x: str\n')
 
+    class D(pg.Object):
+      x: str
+      def __call__(self, y: int) -> int:
+        return len(self.x) + y
+
+    self.assertEqual(
+        schema_lib.class_definition(D, include_methods=True),
+        inspect.cleandoc(
+            """
+            class D:
+              x: str
+
+              def __call__(self, y: int) -> int:
+                return len(self.x) + y
+            """) + '\n'
+    )
+
   def test_repr(self):
     class Foo(pg.Object):
       x: int
@@ -477,13 +494,21 @@ class SchemaPythonReprTest(unittest.TestCase):
     class A(pg.Object):
       foo: Foo
 
+      def foo_value(self) -> int:
+        return self.foo.x
+
     class B(A):
       bar: Bar
       foo2: Foo
 
+      def bar_value(self) -> str:
+        return self.bar.y
+
     schema = schema_lib.Schema([B])
     self.assertEqual(
-        schema_lib.SchemaPythonRepr().class_definitions(schema),
+        schema_lib.SchemaPythonRepr().class_definitions(
+            schema, include_methods=True
+        ),
         inspect.cleandoc('''
             class Foo:
               x: int
@@ -491,6 +516,9 @@ class SchemaPythonReprTest(unittest.TestCase):
             class A:
               foo: Foo
 
+              def foo_value(self) -> int:
+                return self.foo.x
+
             class Bar:
               """Class Bar."""
               y: str
@@ -503,6 +531,9 @@ class SchemaPythonReprTest(unittest.TestCase):
               foo: Foo
               bar: Bar
               foo2: Foo
+
+              def bar_value(self) -> str:
+                return self.bar.y
             ''') + '\n',
     )
 

{langfun-0.0.2.dev20240506.dist-info → langfun-0.0.2.dev20240507.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.0.2.dev20240506
+Version: 0.0.2.dev20240507
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.0.2.dev20240506.dist-info → langfun-0.0.2.dev20240507.dist-info}/RECORD

@@ -82,16 +82,16 @@ langfun/core/structured/description.py,sha256=SXW4MJvshFjbR-0gw6rE21o6WXq12UlRXa
 langfun/core/structured/description_test.py,sha256=UtZGjSFUaQ6130t1E5tcL7ODu0xIefkapb53TbnqsK8,7362
 langfun/core/structured/function_generation.py,sha256=pFgS3vcRAWiuFBol2x5Eeip3XqoudONsOpeJpWyjT3s,7479
 langfun/core/structured/function_generation_test.py,sha256=ZJI-aaGgWWszn92u7h5IZ9Pl70N2DgAGGJrIxPzsvwg,10065
-langfun/core/structured/mapping.py,sha256=Vq3bQZWi4iYjcVn8D2kvPXTAm9jrQ-_1ueHLbXtGRNQ,12112
+langfun/core/structured/mapping.py,sha256=V2EI53KwhXxqcoH2ouhuei8aYWny0ml_FwMTiS7Zr9M,12253
 langfun/core/structured/mapping_test.py,sha256=PiXklMeIa8L6KtMi3ju7J9Y39gZy0hIGz-Oeq4A_7XE,3835
 langfun/core/structured/parsing.py,sha256=keoVqEfzAbdULh6GawWFsTQzU91MzJXYFZjXGXLaD8g,11492
 langfun/core/structured/parsing_test.py,sha256=34wDrXaQ-EYhJLfDL8mX9K53oQMSzh5pVYdKjnESmK8,20895
-langfun/core/structured/prompting.py,sha256=WICdX_figP2u98Hvg3BFSTF83nNKxM4x1STxkWe2_9Y,7925
+langfun/core/structured/prompting.py,sha256=mzQwZBQi2RU3-W8RNdlr_tr_9G6UA2Z9QuowyzcC9ww,8099
 langfun/core/structured/prompting_test.py,sha256=vslaCAUikfwOvqsKzqs_oyEacrefFsr2SWSqu6OHi3w,20813
-langfun/core/structured/schema.py,sha256=mJXirgqx3N7SA9zBO_ISHrzcV-ZRshLhnMJyCcSjGjY,25057
+langfun/core/structured/schema.py,sha256=Zy9y6Vq9DrFwcuP5o5VL_PvMCmzavF-nuDqyviBnaxk,25818
 langfun/core/structured/schema_generation.py,sha256=U3nRQsqmMZg_qIVDh2fiY3K4JLfsAL1LcKzIFP1iXFg,5316
 langfun/core/structured/schema_generation_test.py,sha256=RM9s71kMNg2jTePwInkiW9fK1ACN37eyPeF8OII-0zw,2950
-langfun/core/structured/schema_test.py,sha256=yw_Uo3xJC3JA9dBDjZdkQdBGPf04e7t1oT9SZTAiSdg,22360
+langfun/core/structured/schema_test.py,sha256=NgQK1zGSliZVx_Af6gDBTqQxXRHvmAvGARv4dUs8IbI,23078
 langfun/core/structured/scoring.py,sha256=5DsMNrWKf98ZYCEkxA4-HvA62nMSNBs9DC5m8dYL7Cs,2442
 langfun/core/structured/scoring_test.py,sha256=39_dw6p_FkoqeUccO67yIqos-MccAWezoozS21i8mi0,1732
 langfun/core/templates/__init__.py,sha256=bO0eMsVJbi7sxEB2YlInKRQ2EVP-RyyKUwcD-8msuN4,927
@@ -103,8 +103,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=DYVrkk7uNKCqJGEHH31HssU2BPuMItU1vJLzfcXIlYg,2156
-langfun-0.0.2.dev20240506.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.0.2.dev20240506.dist-info/METADATA,sha256=lzwASgqmdYxjw1o3ivVL4au5jvKsvY6dQXdfBxLqcec,3405
-langfun-0.0.2.dev20240506.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-langfun-0.0.2.dev20240506.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.0.2.dev20240506.dist-info/RECORD,,
+langfun-0.0.2.dev20240507.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.0.2.dev20240507.dist-info/METADATA,sha256=kz0C5CzRBOml2nNCsUkvDY36x_49qo3HvwS-ifZY_sU,3405
+langfun-0.0.2.dev20240507.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+langfun-0.0.2.dev20240507.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.0.2.dev20240507.dist-info/RECORD,,