langfun 0.0.2.dev20240403__tar.gz → 0.0.2.dev20240404__tar.gz

{langfun-0.0.2.dev20240403 → langfun-0.0.2.dev20240404}/PKG-INFO

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

{langfun-0.0.2.dev20240403 → langfun-0.0.2.dev20240404}/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-0.0.2.dev20240403 → langfun-0.0.2.dev20240404}/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-0.0.2.dev20240403 → langfun-0.0.2.dev20240404}/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-0.0.2.dev20240403 → langfun-0.0.2.dev20240404}/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-0.0.2.dev20240403 → langfun-0.0.2.dev20240404}/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-0.0.2.dev20240403 → langfun-0.0.2.dev20240404}/langfun.egg-info/PKG-INFO

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