langfun 0.1.2.dev202509120804__py3-none-any.whl → 0.1.2.dev202512150805__py3-none-any.whl

langfun/core/structured/querying_test.py

@@ -249,35 +249,60 @@ class QueryTest(unittest.TestCase):
 
   def test_root_modality_to_structure_render(self):
     lm = fake.StaticResponse('1')
+    image = modalities.Image.from_bytes(b'mock_image')
     self.assert_render(
-        modalities.Image.from_bytes(b'mock_image'),
+        image,
         int,
         lm=lm,
-        expected_snippet='\n\nREQUEST:\n  <<[[input]]>>\n\n',
+        expected_snippet=f'\n\nREQUEST:\n  <<[[{image.id}]]>>\n\n',
         expected_modalities=1,
     )
 
   def test_root_modality_to_str_render(self):
     lm = fake.StaticResponse('1')
+    modality = modalities.Image.from_bytes(b'mock_image')
     self.assert_render(
-        modalities.Image.from_bytes(b'mock_image'),
+        modality,
         None,
         lm=lm,
-        expected_snippet='<<[[input]]>>',
+        expected_snippet=f'<<[[{modality.id}]]>>',
         exact_match=True,
         expected_modalities=1,
     )
 
   def test_str_with_modality_to_str_render(self):
     lm = fake.StaticResponse('A cat and a mouse.')
+    cat_image = modalities.Image.from_bytes(b'cat_image')
+    mouse_image = modalities.Image.from_bytes(b'mouse_image')
     self.assert_render(
         'What are these? {{this_image}} and {{that_image}}',
         None,
-        this_image=modalities.Image.from_bytes(b'cat_image'),
-        that_image=modalities.Image.from_bytes(b'mouse_image'),
+        this_image=cat_image,
+        that_image=mouse_image,
         lm=lm,
         expected_snippet=(
-            'What are these? <<[[this_image]]>> and <<[[that_image]]>>'
+            f'What are these? <<[[{cat_image.id}]]>> and '
+            f'<<[[{mouse_image.id}]]>>'
+        ),
+        exact_match=True,
+        expected_modalities=2,
+    )
+
+  def test_message_with_modality_to_str_render(self):
+    lm = fake.StaticResponse('A cat and a mouse.')
+    cat_image = modalities.Image.from_bytes(b'cat_image')
+    mouse_image = modalities.Image.from_bytes(b'mouse_image')
+    self.assert_render(
+        lf.Template(
+            'What are these? {{this_image}} and {{that_image}}',
+            this_image=cat_image,
+            that_image=mouse_image,
+        ).render(),
+        None,
+        lm=lm,
+        expected_snippet=(
+            f'What are these? <<[[{cat_image.id}]]>> and '
+            f'<<[[{mouse_image.id}]]>>'
         ),
         exact_match=True,
         expected_modalities=2,
@@ -285,33 +310,33 @@ class QueryTest(unittest.TestCase):
 
   def test_structure_with_modality_to_str_render(self):
     lm = fake.StaticResponse('A cat and a mouse.')
+    cat_image = modalities.Image.from_bytes(b'cat_image')
+    mouse_image = modalities.Image.from_bytes(b'mouse_image')
     self.assert_render(
-        [
-            modalities.Image.from_bytes(b'cat_image'),
-            modalities.Image.from_bytes(b'mouse_image'),
-        ],
+        [cat_image, mouse_image],
         None,
         lm=lm,
-        expected_snippet='`[<<[[input[0]]]>>, <<[[input[1]]]>>]`',
+        expected_snippet=(
+            f'`[<<[[{cat_image.id}]]>>, <<[[{mouse_image.id}]]>>]`'
+        ),
         exact_match=True,
         expected_modalities=2,
     )
 
   def test_structure_with_modality_to_structure_render(self):
     lm = fake.StaticResponse('["cat", "mouse"]')
+    cat_image = modalities.Image.from_bytes(b'cat_image')
+    mouse_image = modalities.Image.from_bytes(b'mouse_image')
     self.assert_render(
-        [
-            modalities.Image.from_bytes(b'cat_image'),
-            modalities.Image.from_bytes(b'mouse_image'),
-        ],
+        [cat_image, mouse_image],
         list[str],
         lm=lm,
-        expected_snippet=inspect.cleandoc("""
+        expected_snippet=inspect.cleandoc(f"""
             REQUEST:
               ```python
               [
-                <<[[input[0]]]>>,
-                <<[[input[1]]]>>
+                <<[[{cat_image.id}]]>>,
+                <<[[{mouse_image.id}]]>>
               ]
               ```
             """),
@@ -320,25 +345,25 @@ class QueryTest(unittest.TestCase):
 
   def test_structure_with_modality_and_examples_to_structure_render(self):
     lm = fake.StaticResponse('["cat", "mouse"]')
+    cat_image = modalities.Image.from_bytes(b'cat_image')
+    mouse_image = modalities.Image.from_bytes(b'mouse_image')
+    dog_image = modalities.Image.from_bytes(b'dog_image')
     self.assert_render(
-        [
-            modalities.Image.from_bytes(b'cat_image'),
-            modalities.Image.from_bytes(b'mouse_image'),
-        ],
+        [cat_image, mouse_image],
         list[str],
         examples=[
             mapping.MappingExample(
-                input=[modalities.Image.from_bytes(b'dog_image')],
+                input=[dog_image],
                 schema=list[str],
                 output=['dog'],
             ),
         ],
         lm=lm,
-        expected_snippet=inspect.cleandoc("""
+        expected_snippet=inspect.cleandoc(f"""
             REQUEST:
               ```python
               [
-                <<[[examples[0].input[0]]]>>
+                <<[[{dog_image.id}]]>>
               ]
               ```
 
@@ -356,8 +381,8 @@ class QueryTest(unittest.TestCase):
             REQUEST:
               ```python
               [
-                <<[[input[0]]]>>,
-                <<[[input[1]]]>>
+                <<[[{cat_image.id}]]>>,
+                <<[[{mouse_image.id}]]>>
               ]
               ```
 
@@ -369,6 +394,17 @@ class QueryTest(unittest.TestCase):
         expected_modalities=3,
     )
 
+  def test_query_with_modality_output(self):
+    cat_image = modalities.Image.from_bytes(b'cat_image')
+    lm = fake.StaticResponse(
+        lf.Template('Here you go: {{image}}', image=cat_image).render(
+            message_cls=lf.AIMessage
+        )
+    )
+    response = querying.query('Generate a cat image', lm=lm)
+    self.assertIsInstance(response, lf.AIMessage)
+    self.assertEqual(response.modalities(), [cat_image])
+
   def test_multiple_queries(self):
     self.assertEqual(
         querying.query(
@@ -545,7 +581,7 @@ class QueryTest(unittest.TestCase):
             )
         ).input,
     )
-    self.assertIsNotNone(output.get_modality('image'))
+    self.assertEqual(len(output.referred_modalities), 1)
 
   def test_query_and_reduce(self):
     self.assertEqual(
@@ -1590,7 +1626,7 @@ class TrackQueriesTest(unittest.TestCase):
         'bar',
     ])
     state = {}
-    def start_callabck(query):
+    def start_callback(query):
       self.assertFalse(query.is_completed)
       self.assertIsNone(query.end_time)
       elapse1 = query.elapse
@@ -1615,7 +1651,7 @@ class TrackQueriesTest(unittest.TestCase):
       state['end'] = query
 
     with querying.track_queries(
-        start_callabck=start_callabck, end_callabck=end_callback
+        start_callback=start_callback, end_callback=end_callback
     ) as queries:
       querying.query('foo', lm=lm)
     self.assertIs(state['start'], queries[0])
@@ -1626,7 +1662,7 @@ class TrackQueriesTest(unittest.TestCase):
         'bar',
     ])
     state = {}
-    def start_callabck(query):
+    def start_callback(query):
       self.assertFalse(query.is_completed)
       self.assertIsNone(query.end_time)
       self.assertIsNotNone(query.usage_summary)
@@ -1648,7 +1684,7 @@ class TrackQueriesTest(unittest.TestCase):
 
     with self.assertRaises(mapping.MappingError):
       with querying.track_queries(
-          start_callabck=start_callabck, end_callabck=end_callback
+          start_callback=start_callback, end_callback=end_callback
       ) as queries:
         querying.query('foo', int, lm=lm)
     self.assertIs(state['start'], queries[0])

langfun/core/structured/schema/__init__.py

@@ -0,0 +1,49 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Schema and prompting protocols for structured data."""
+
+# pylint: disable=g-bad-import-order
+# pylint: disable=g-importing-member
+from langfun.core.structured.schema.base import Schema
+from langfun.core.structured.schema.base import SchemaError
+from langfun.core.structured.schema.base import SchemaType
+
+from langfun.core.structured.schema.base import class_dependencies
+from langfun.core.structured.schema.base import schema_spec
+from langfun.core.structured.schema.base import annotation
+
+from langfun.core.structured.schema.base import PromptingProtocol
+from langfun.core.structured.schema.base import schema_repr
+from langfun.core.structured.schema.base import value_repr
+from langfun.core.structured.schema.base import parse_value
+
+from langfun.core.structured.schema.base import mark_missing
+from langfun.core.structured.schema.base import Missing
+from langfun.core.structured.schema.base import MISSING
+from langfun.core.structured.schema.base import Unknown
+from langfun.core.structured.schema.base import UNKNOWN
+
+# JSON protocol.
+from langfun.core.structured.schema.json import JsonError
+from langfun.core.structured.schema.json import JsonPromptingProtocol
+from langfun.core.structured.schema.json import cleanup_json
+
+# Python protocol.
+from langfun.core.structured.schema.python import PythonPromptingProtocol
+from langfun.core.structured.schema.python import source_form
+from langfun.core.structured.schema.python import structure_from_python
+
+from langfun.core.structured.schema.python import class_definition
+from langfun.core.structured.schema.python import class_definitions
+from langfun.core.structured.schema.python import include_method_in_prompt