langfun 0.0.2.dev20240403__py3-none-any.whl → 0.0.2.dev20240407__py3-none-any.whl

langfun/core/coding/python/correction.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Python code error correction."""
-import re
 from typing import Any
 import langfun.core as lf
 from langfun.core.coding.python import errors
@@ -31,11 +30,6 @@ class CorrectedCode(pg.Object):
   corrected_code: str
 
 
-def remove_docstrings(code):
-  pattern = re.compile(r"(def .+?:\s*?)('''|\"\"\")((.|\s)*?)(\2)", re.DOTALL)
-  return pattern.sub(r"\1", code)
-
-
 def run_with_correction(
     code: str,
     error: str | None = None,
@@ -86,7 +80,6 @@ def run_with_correction(
   # pytype: enable=import-error
   # pylint: enable=g-import-not-at-top
 
-  code = remove_docstrings(code)
   if max_attempts == 0:
     result = execution.run(
         code,

langfun/core/eval/base.py

@@ -816,11 +816,6 @@ class Evaluation(Evaluable):
     if self.schema_fn is None:
       return None
 
-    kwargs = {}
-    # Allow schema to be a function based on current evaluation.
-    if 'evaluation' in self.schema_fn.__signature__.arg_names:
-      kwargs['evaluation'] = self
-
     schema = self._call_schema_fn()
     fewshot_examples = None
     if isinstance(schema, tuple):

langfun/core/modalities/video.py

@@ -15,9 +15,7 @@
 
 import base64
 from typing import cast
-
 from langfun.core.modalities import mime
-import magic
 
 
 class Video(mime.MimeType):
@@ -29,6 +27,11 @@ class Video(mime.MimeType):
 
   @property
   def mime_type(self) -> str:
+    # TODO(daiyip): after cl/619658455, LaunchPad binaries cannot import `magic`
+    # correctly. This is to mitigate the issue for major Langfun users who do
+    # not use Video. We shall move this import out once the issue is fixed.
+    import magic  # pylint: disable=g-import-not-at-top
+
     video_mime_type = magic.from_buffer(self.to_bytes(), mime=True)
     if 'video/' not in video_mime_type:
       raise ValueError(f'Not a video: {video_mime_type!r}.')

langfun/core/structured/mapping.py

@@ -14,7 +14,7 @@
 """The base of symbolic mapping methods."""
 
 import io
-from typing import Annotated, Any
+from typing import Annotated, Any, Callable
 import langfun.core as lf
 from langfun.core.structured import schema as schema_lib
 import pyglove as pg
@@ -278,6 +278,14 @@ class Mapping(lf.LangFunc):
       ),
   ] = lf.RAISE_IF_HAS_ERROR
 
+  response_postprocess: Annotated[
+      Callable[[str], str] | None,
+      (
+          'A callable object that post process the raw LLM response before '
+          'parsing it into the output Python object.'
+      )
+  ] = None
+
   #
   # Key methods for implementing specific mappings.
   #
@@ -296,6 +304,7 @@ class Mapping(lf.LangFunc):
   def transform_output(self, lm_output: lf.Message) -> lf.Message:
     """Transforms LM response into structure if schema is present."""
     try:
+      lm_output = self.postprocess_response(lm_output)
       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:
@@ -316,6 +325,14 @@ class Mapping(lf.LangFunc):
         autofix_lm=self.autofix_lm or self.lm,
     )
 
+  def postprocess_response(self, response: lf.Message) -> lf.Message:
+    """Post process LLM response."""
+    if self.response_postprocess is not None:
+      postprocessed_text = self.response_postprocess(response.text)
+      if postprocessed_text != response.text:
+        return lf.AIMessage(postprocessed_text, source=response)
+    return response
+
   def postprocess_result(self, result: Any) -> Any:
     """Post process structured output."""
     return result

langfun/core/structured/prompting.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Symbolic query."""
 
-from typing import Any, Type, Union
+from typing import Any, Callable, Type, Union
 
 import langfun.core as lf
 from langfun.core.structured import mapping
@@ -107,6 +107,7 @@ def query(
     lm: lf.LanguageModel | None = None,
     examples: list[mapping.MappingExample] | None = None,
     cache_seed: int | None = 0,
+    response_postprocess: Callable[[str], str] | None = None,
     autofix: int = 0,
     autofix_lm: lf.LanguageModel | None = None,
     protocol: schema_lib.SchemaProtocol = 'python',
@@ -159,6 +160,9 @@ def query(
     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.
+    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.
     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
@@ -188,6 +192,10 @@ def query(
     output = lf.LangFunc.from_value(prompt, **kwargs)(
         lm=lm, cache_seed=cache_seed, skip_lm=skip_lm
     )
+    if response_postprocess:
+      processed_text = response_postprocess(output.text)
+      if processed_text != output.text:
+        output = lf.AIMessage(processed_text, source=output)
     return output if returns_message else output.text
 
   # Query with structured output.
@@ -206,6 +214,7 @@ def query(
       schema=schema,
       default=default,
       examples=examples,
+      response_postprocess=response_postprocess,
       autofix=autofix if protocol == 'python' else 0,
       **kwargs,
   )(

langfun/core/structured/prompting_test.py

@@ -436,6 +436,23 @@ class QueryStructurePythonTest(unittest.TestCase):
     ])
     self.assertEqual(prompting.query('what is 1 + 0', int, lm=lm, autofix=3), 1)
 
+  def test_response_postprocess(self):
+    with lf.context(
+        lm=fake.StaticResponse('<!-- some comment-->\n3'),
+        override_attrs=True,
+    ):
+      self.assertEqual(
+          prompting.query(
+              'Compute 1 + 2', response_postprocess=lambda x: x.split('\n')[1]),
+          '3'
+      )
+      self.assertEqual(
+          prompting.query(
+              'Compute 1 + 2', int,
+              response_postprocess=lambda x: x.split('\n')[1]),
+          3
+      )
+
 
 class QueryStructureJsonTest(unittest.TestCase):
 

langfun/core/structured/schema.py

@@ -386,10 +386,12 @@ def class_definition(
   if schema.fields:
     for key, field in schema.items():
       if not isinstance(key, pg.typing.ConstStrKey):
-        raise TypeError(
+        pg.logging.warning(
             'Variable-length keyword arguments is not supported in '
-            f'structured parsing or query. Encountered: {field}'
+            f'structured parsing or query. Encountered: {cls}, Schema: {schema}'
         )
+        continue
+
       # Write field doc string as comments before the field definition.
       if field.description:
         for line in field.description.split('\n'):

langfun/core/structured/schema_test.py

@@ -459,9 +459,7 @@ class SchemaPythonReprTest(unittest.TestCase):
       x: str
       __kwargs__: typing.Any
 
-    with self.assertRaisesRegex(
-        TypeError, 'Variable-length keyword arguments is not supported'):
-      schema_lib.class_definition(C)
+    self.assertEqual(schema_lib.class_definition(C), 'class C:\n  x: str\n')
 
   def test_repr(self):
     class Foo(pg.Object):

{langfun-0.0.2.dev20240403.dist-info → langfun-0.0.2.dev20240407.dist-info}/METADATA

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

{langfun-0.0.2.dev20240403.dist-info → langfun-0.0.2.dev20240407.dist-info}/RECORD

@@ -27,7 +27,7 @@ langfun/core/text_formatting.py,sha256=ytjj7opnRJ6w-pkglL2CZUyfYDXLpNf65E42LBb31
 langfun/core/text_formatting_test.py,sha256=nyKC6tn2L4hPJiqQHgxcbQsJJi4A4Nbj8FiO8iT6B80,1514
 langfun/core/coding/__init__.py,sha256=5utju_fwEsImaiftx4oXKl9FAM8p281k8-Esdh_-m1w,835
 langfun/core/coding/python/__init__.py,sha256=MJ-vubliz-ebrZH3OBRKBwMi0S9-FrhGCp8YQLR6_I4,1776
-langfun/core/coding/python/correction.py,sha256=uuQmZCrAl0EA9etIUmn2-FZ-ge8iNcjOAAlm-WgkYfo,6776
+langfun/core/coding/python/correction.py,sha256=a2aFUt9ocbXTCR6Z6OGNjQZDI1LfU0PBkSe7hJB8dEM,6589
 langfun/core/coding/python/correction_test.py,sha256=yLqmQ9BPORsnREkrS10PnljEaLR3BoydTVeT3OGoqfU,3507
 langfun/core/coding/python/errors.py,sha256=fX3Du63uGm25YFXW9D-bV2gntTdTAX3hBFtAnRlmg14,3166
 langfun/core/coding/python/errors_test.py,sha256=_ZbWJCFIb-FkCK7K1zCuH8W3x_NFt-jNe3dfP8yqaD4,2323
@@ -40,7 +40,7 @@ langfun/core/coding/python/parsing_test.py,sha256=9vAWF484kWIm6JZq8NFiMgKUDhXV-d
 langfun/core/coding/python/permissions.py,sha256=1QWGHvzL8MM0Ok_auQ9tURqZHtdOfJaDpBzZ29GUE-c,2544
 langfun/core/coding/python/permissions_test.py,sha256=w5EDb8QxpxgJyZkojyzVWQvDfg366zn99-g__6TbPQ0,2699
 langfun/core/eval/__init__.py,sha256=iDA2OcJ3kR6ixZizXIY3N9LsjkaVrfTbSClTiSP8ekY,1291
-langfun/core/eval/base.py,sha256=dLDWAYLHnLX7CyvMTBcditLFIuq-tUvGoPX1vT65GJQ,54730
+langfun/core/eval/base.py,sha256=YpAPtWLeh3OpJuuPk7913xkB2OTCaGls1uWMohDA8SI,54551
 langfun/core/eval/base_test.py,sha256=8MOum0DWMEm2-NpwmFgcqmlqEmuWYF5MesrCXTySylg,21083
 langfun/core/eval/matching.py,sha256=g2yuBb4FeOlAlB10hqdWvaIg4QVQlJbiViRDcD2Y8go,9567
 langfun/core/eval/matching_test.py,sha256=jFrNOaHteNo7wxCwc6w_mGylM0VHwezAcvfaZANKKmA,4898
@@ -67,23 +67,23 @@ langfun/core/modalities/image.py,sha256=zNpLHwJi6PJMeeAcVQG6vn0oYIbilTuJq6xu-Tvr
 langfun/core/modalities/image_test.py,sha256=YxDRvC49Bjwyyndd_P7y6XjyS7dOft0Zewwxk-7q4kE,2301
 langfun/core/modalities/mime.py,sha256=wVfaYflhGz1W4v3m972rAplW3OGOFtjFpHDYIaUD5D0,2238
 langfun/core/modalities/mime_test.py,sha256=cVHxRvJ1QXC1SVhBmWkJdWGpL9Xl0UNfTQq6j0OGGL4,1881
-langfun/core/modalities/video.py,sha256=5-sIlzXb_ZY84RMFcpVD9ysP9GbcwbdKaZOEm3jECtc,1469
+langfun/core/modalities/video.py,sha256=25M4XsNG5XEWRy57LYT_a6_aMURMPAgC41B3weEXFsY,1747
 langfun/core/modalities/video_test.py,sha256=jYuI2m8S8zDCAVBPEUbbpP205dXAht90A2_PHWo4-r8,2039
 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=SXW4MJvshFjbR-0gw6rE21o6WXq12UlRXawvDBXMZFA,5211
 langfun/core/structured/description_test.py,sha256=UtZGjSFUaQ6130t1E5tcL7ODu0xIefkapb53TbnqsK8,7362
-langfun/core/structured/mapping.py,sha256=m7i80GU3vfDZXw4TTnidWTS3K-c1H8JNX9KcoMw4E4s,10373
+langfun/core/structured/mapping.py,sha256=7JInwZLmQdu7asHhC0vFLJNOCBnY-hrD6v5RQgf-xKk,11020
 langfun/core/structured/mapping_test.py,sha256=07DDCGbwytQHSMm7fCi5-Ly-JNgdV4ubHZq0wthX4A4,3338
 langfun/core/structured/parsing.py,sha256=keoVqEfzAbdULh6GawWFsTQzU91MzJXYFZjXGXLaD8g,11492
 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=CuRXOBjoK8rv5b281-w2o7nWQ7ox2YX5kkq2dOk0Ry8,25020
+langfun/core/structured/prompting.py,sha256=OAejcd7E-GTRI86i71Uel66te5zcTqyhsNtxe_Nt46U,7407
+langfun/core/structured/prompting_test.py,sha256=FdEQ0Sy0zGSFXsaOYkWLjyso37-ckIBZE74hmD7QaAA,19610
+langfun/core/structured/schema.py,sha256=mJXirgqx3N7SA9zBO_ISHrzcV-ZRshLhnMJyCcSjGjY,25057
 langfun/core/structured/schema_generation.py,sha256=U3nRQsqmMZg_qIVDh2fiY3K4JLfsAL1LcKzIFP1iXFg,5316
 langfun/core/structured/schema_generation_test.py,sha256=cfZyP0gHno2fXy_c9vsVdvHmqKQSfuyUsCtfO3JFmYQ,2945
-langfun/core/structured/schema_test.py,sha256=hLhwbUaBtTJMG4c-q21rFUksTtI-jP_oNIyr5_S-0yo,22472
+langfun/core/structured/schema_test.py,sha256=XsTACSWwFmqlj4pCKsKWcVJbOaeu2DP0jMMCR4EZ_Js,22405
 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
@@ -95,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.dev20240403.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.0.2.dev20240403.dist-info/METADATA,sha256=2NOTKH53MG8WNbkap8zD-6kQYRhDBOGbc-eywzJKf3g,3405
-langfun-0.0.2.dev20240403.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-langfun-0.0.2.dev20240403.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.0.2.dev20240403.dist-info/RECORD,,
+langfun-0.0.2.dev20240407.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.0.2.dev20240407.dist-info/METADATA,sha256=AstNjq_o1Rpd5C6aHkrQ78z_Fe38CMBXUlnwz4ArvwY,3405
+langfun-0.0.2.dev20240407.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+langfun-0.0.2.dev20240407.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.0.2.dev20240407.dist-info/RECORD,,