langfun 0.0.2.dev20240530__py3-none-any.whl → 0.0.2.dev20240601__py3-none-any.whl

langfun/core/modalities/ms_office_test.py

@@ -12,11 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Video tests."""
+import base64
 import io
 import unittest
 from unittest import mock
 
 from langfun.core.modalities import ms_office as ms_office_lib
+from langfun.core.modalities import pdf as pdf_lib
 import pyglove as pg
 
 
@@ -243,23 +245,72 @@ def pptx_mock_request(*args, **kwargs):
   return pg.Dict(content=pptx_bytes)
 
 
+pdf_bytes = (
+    b'%PDF-1.1\n%\xc2\xa5\xc2\xb1\xc3\xab\n\n1 0 obj\n'
+    b'<< /Type /Catalog\n     /Pages 2 0 R\n  >>\nendobj\n\n2 0 obj\n '
+    b'<< /Type /Pages\n     /Kids [3 0 R]\n     '
+    b'/Count 1\n     /MediaBox [0 0 300 144]\n  '
+    b'>>\nendobj\n\n3 0 obj\n  '
+    b'<<  /Type /Page\n      /Parent 2 0 R\n      /Resources\n       '
+    b'<< /Font\n'
+    b'<< /F1\n'
+    b'<< /Type /Font\n'
+    b'/Subtype /Type1\n'
+    b'/BaseFont /Times-Roman\n'
+    b'>>\n>>\n>>\n      '
+    b'/Contents 4 0 R\n  >>\nendobj\n\n4 0 obj\n  '
+    b'<< /Length 55 >>\nstream\n  BT\n    /F1 18 Tf\n    0 0 Td\n    '
+    b'(Hello World) Tj\n  ET\nendstream\nendobj\n\nxref\n0 5\n0000000000 '
+    b'65535 f \n0000000018 00000 n \n0000000077 00000 n \n0000000178 00000 n '
+    b'\n0000000457 00000 n \ntrailer\n  <<  /Root 1 0 R\n      /Size 5\n  '
+    b'>>\nstartxref\n565\n%%EOF\n'
+)
+
+
+def convert_mock_request(*args, **kwargs):
+  del args, kwargs
+
+  class Result:
+    def json(self):
+      return {
+          'Files': [
+              {
+                  'FileData': base64.b64encode(pdf_bytes).decode()
+              }
+          ]
+      }
+  return Result()
+
+
 class DocxTest(unittest.TestCase):
 
-  def test_content(self):
+  def test_from_bytes(self):
     content = ms_office_lib.Docx.from_bytes(docx_bytes)
-    self.assertEqual(
+    self.assertIn(
         content.mime_type,
-        'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
+        (
+            'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
+            'application/octet-stream',
+        ),
     )
     self.assertEqual(content.to_bytes(), docx_bytes)
+    self.assertTrue(content.is_compatible('text/plain'))
+    self.assertFalse(content.is_compatible('application/pdf'))
+    self.assertEqual(
+        content.make_compatible(['image/png', 'text/plain']).mime_type,
+        'text/plain'
+    )
 
-  def test_file(self):
+  def test_from_uri(self):
     content = ms_office_lib.Docx.from_uri('http://mock/web/a.docx')
     with mock.patch('requests.get') as mock_requests_get:
       mock_requests_get.side_effect = docx_mock_request
-      self.assertEqual(
+      self.assertIn(
           content.mime_type,
-          'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
+          (
+              'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
+              'application/octet-stream',
+          ),
       )
       self.assertEqual(content.to_bytes(), docx_bytes)
       self.assertEqual(content.to_xml(), expected_docx_xml)
@@ -267,21 +318,33 @@ class DocxTest(unittest.TestCase):
 
 class XlsxTest(unittest.TestCase):
 
-  def test_content(self):
+  def test_from_bytes(self):
     content = ms_office_lib.Xlsx.from_bytes(xlsx_bytes)
-    self.assertEqual(
+    self.assertIn(
         content.mime_type,
-        'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
+        (
+            'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
+            'application/octet-stream',
+        ),
     )
     self.assertEqual(content.to_bytes(), xlsx_bytes)
+    self.assertTrue(content.is_compatible('text/plain'))
+    self.assertFalse(content.is_compatible('application/pdf'))
+    self.assertEqual(
+        content.make_compatible('text/plain').mime_type,
+        'text/html'
+    )
 
-  def test_file(self):
+  def test_from_uri(self):
     content = ms_office_lib.Xlsx.from_uri('http://mock/web/a.xlsx')
     with mock.patch('requests.get') as mock_requests_get:
       mock_requests_get.side_effect = xlsx_mock_request
-      self.assertEqual(
+      self.assertIn(
           content.mime_type,
-          'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
+          (
+              'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
+              'application/octet-stream',
+          ),
       )
       self.assertEqual(content.to_bytes(), xlsx_bytes)
       self.assertEqual(content.to_html(), expected_xlsx_html)
@@ -291,22 +354,36 @@ class PptxTest(unittest.TestCase):
 
   def test_content(self):
     content = ms_office_lib.Pptx.from_bytes(pptx_bytes)
-    self.assertEqual(
+    self.assertIn(
         content.mime_type,
-        'application/vnd.openxmlformats-officedocument.presentationml.presentation',
+        (
+            'application/vnd.openxmlformats-officedocument.presentationml.presentation',
+            'application/octet-stream',
+        ),
     )
     self.assertEqual(content.to_bytes(), pptx_bytes)
 
   def test_file(self):
     content = ms_office_lib.Pptx.from_uri('http://mock/web/a.pptx')
+    self.assertFalse(content.is_compatible('text/plain'))
+    self.assertTrue(content.is_compatible('application/pdf'))
     with mock.patch('requests.get') as mock_requests_get:
       mock_requests_get.side_effect = pptx_mock_request
-      self.assertEqual(
+      self.assertIn(
           content.mime_type,
-          'application/vnd.openxmlformats-officedocument.presentationml.presentation',
+          (
+              'application/vnd.openxmlformats-officedocument.presentationml.presentation',
+              'application/octet-stream',
+          ),
       )
       self.assertEqual(content.to_bytes(), pptx_bytes)
 
+    with mock.patch('requests.post') as mock_requests_post:
+      mock_requests_post.side_effect = convert_mock_request
+      self.assertIsInstance(
+          content.make_compatible('application/pdf'), pdf_lib.PDF
+      )
+
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/modalities/pdf.py

@@ -16,7 +16,7 @@
 from langfun.core.modalities import mime
 
 
-class PDF(mime.MimeType):
+class PDF(mime.Mime):
   """PDF document."""
 
   MIME_PREFIX = 'application/pdf'

langfun/core/modalities/video.py

@@ -17,7 +17,7 @@ import functools
 from langfun.core.modalities import mime
 
 
-class Video(mime.MimeType):
+class Video(mime.Mime):
   """Video."""
 
   MIME_PREFIX = 'video'

langfun/core/modality.py

@@ -108,3 +108,7 @@ class ModalityRef(pg.Object, pg.typing.CustomTyping):
         return ModalityRef(name=value.sym_path + k)
       return v
     return value.clone().rebind(_placehold, raise_on_no_change=False)
+
+
+class ModalityError(RuntimeError):  # pylint: disable=g-bad-exception-name
+  """Exception raised when modality is not supported."""

langfun/core/structured/__init__.py

@@ -16,6 +16,8 @@
 # pylint: disable=g-bad-import-order
 # pylint: disable=g-importing-member
 
+from langfun.core.structured.schema import include_method_in_prompt
+
 from langfun.core.structured.schema import Missing
 from langfun.core.structured.schema import MISSING
 from langfun.core.structured.schema import Unknown

langfun/core/structured/completion.py

@@ -107,7 +107,9 @@ class CompleteStructure(mapping.Mapping):
 
   def class_defs_repr(self, value: Any) -> str | None:
     return schema_lib.class_definitions(
-        self.missing_type_dependencies(value), markdown=True
+        self.missing_type_dependencies(value),
+        markdown=True,
+        allowed_dependencies=set()
     )
 
   def postprocess_result(self, result: Any) -> Any:

langfun/core/structured/mapping.py

@@ -251,7 +251,7 @@ class Mapping(lf.LangFunc):
 
       {%- if example.schema -%}
       {{ schema_title }}:
-      {{ example.schema_repr(protocol, include_methods=include_methods) | indent(2, True) }}
+      {{ example.schema_repr(protocol) | indent(2, True) }}
 
       {% endif -%}
 
@@ -279,10 +279,6 @@ 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

@@ -114,7 +114,6 @@ 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,
@@ -174,8 +173,6 @@ 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.
@@ -225,7 +222,6 @@ 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

@@ -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