langfun 0.1.2.dev202505080804__py3-none-any.whl → 0.1.2.dev202505100803__py3-none-any.whl

langfun/core/llms/vertexai_test.py

@@ -17,6 +17,7 @@ import os
 import unittest
 from unittest import mock
 
+from google import auth
 from google.auth import exceptions
 import langfun.core as lf
 from langfun.core.llms import rest
@@ -61,7 +62,11 @@ class VertexAITest(unittest.TestCase):
         lf.concurrent.RetryError,
         'Failed to refresh Google authentication credentials'
     ):
-      with mock.patch.object(rest.REST, '_sample_single') as mock_sample_single:
+      with (
+          mock.patch.object(auth, 'default') as mock_auth,
+          mock.patch.object(rest.REST, '_sample_single') as mock_sample_single
+      ):
+        mock_auth.return_value = mock.MagicMock(), None
         mock_sample_single.side_effect = _auth_refresh_error
         model = vertexai.VertexAIGemini15Pro(
             project='abc', location='us-central1', max_attempts=1

langfun/core/structured/querying.py

@@ -325,7 +325,7 @@ def query(
 
     ```
     lf.query('1 + 1 = ?', int, lm=lf.llms.Gpt4Turbo())
-    
+
     # Output: 2
     ```
 
@@ -349,7 +349,7 @@ def query(
 
     class Dog(Animal):
       pass
-    
+
     class Entity(pg.Object):
       name: str
 
@@ -562,7 +562,10 @@ def query(
           output_message = lf.AIMessage(processed_text, source=output_message)
     else:
       # Query with structured output.
-      output_message = LfQuery.from_protocol(protocol)(
+      query_cls = LfQuery.from_protocol(protocol)
+      if ':' not in protocol:
+        protocol = f'{protocol}:{query_cls.version}'
+      output_message = query_cls(
           input=(
               query_input.render(lm=lm)
               if isinstance(query_input, lf.Template)
@@ -572,7 +575,7 @@ def query(
           default=default,
           examples=examples,
           response_postprocess=response_postprocess,
-          autofix=autofix if protocol == 'python' else 0,
+          autofix=autofix if protocol.startswith('python:') else 0,
           **kwargs,
       )(
           lm=lm,
@@ -605,6 +608,8 @@ def query(
           ),
           lm=pg.Ref(lm),
           examples=pg.Ref(examples) if examples else [],
+          protocol=protocol,
+          kwargs={k: pg.Ref(v) for k, v in kwargs.items()},
           lm_response=lf.AIMessage(output_message.text, metadata=metadata),
           usage_summary=usage_summary,
           start_time=start_time,
@@ -788,6 +793,14 @@ class QueryInvocation(pg.Object, pg.views.HtmlTreeView.Extension):
       list[mapping.MappingExample],
       'Fewshot exemplars for `lf.query`.'
   ]
+  protocol: Annotated[
+      str,
+      'Protocol of `lf.query`.'
+  ] = 'python'
+  kwargs: Annotated[
+      dict[str, Any],
+      'Kwargs of `lf.query`.'
+  ] = {}
   usage_summary: Annotated[
       lf.UsageSummary,
       'Usage summary for `lf.query`.'
@@ -803,13 +816,17 @@ class QueryInvocation(pg.Object, pg.views.HtmlTreeView.Extension):
 
   @functools.cached_property
   def lm_request(self) -> lf.Message:
-    return query_prompt(self.input, self.schema, examples=self.examples or None)
+    return query_prompt(
+        self.input, self.schema, examples=self.examples or None,
+        protocol=self.protocol,
+        **self.kwargs
+    )
 
   @functools.cached_property
   def output(self) -> Any:
     """The output of `lf.query`. If it failed, returns the `MappingError`."""
     try:
-      return query_output(self.lm_response, self.schema)
+      return query_output(self.lm_response, self.schema, protocol=self.protocol)
     except mapping.MappingError as e:
       return e
 

langfun/core/structured/querying_test.py

@@ -1321,6 +1321,26 @@ class QueryInvocationTest(unittest.TestCase):
     self.assertTrue(queries[0].has_error)
     self.assertIsInstance(queries[0].output, mapping.MappingError)
 
+  def test_kwargs(self):
+    lm = fake.StaticSequence([
+        'Activity(description="hi")',
+    ])
+    with querying.track_queries() as queries:
+      querying.query(
+          'foo {{x}}',
+          Activity,
+          lm=lm,
+          system_message='system message',
+          x=1,
+      )
+    self.assertTrue(queries[0].protocol.startswith('python:'))
+    self.assertEqual(
+        list(queries[0].kwargs.keys()),
+        ['x', 'metadata_system_message']
+    )
+    self.assertIn('foo 1', queries[0].lm_request.text)
+    self.assertEqual(queries[0].lm_request.system_message, 'system message')
+
   def test_to_html(self):
     lm = fake.StaticSequence([
         'Activity(description="hi")',

{langfun-0.1.2.dev202505080804.dist-info → langfun-0.1.2.dev202505100803.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langfun
-Version: 0.1.2.dev202505080804
+Version: 0.1.2.dev202505100803
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.2.dev202505080804.dist-info → langfun-0.1.2.dev202505100803.dist-info}/RECORD

@@ -107,7 +107,7 @@ langfun/core/llms/openai_test.py,sha256=gwuO6aoa296iM2welWV9ua4KF8gEVGsEPakgbtkW
 langfun/core/llms/rest.py,sha256=MCybcHApJcf49lubLnDzScN9Oc2IWY_JnMHIGdbDOuU,4474
 langfun/core/llms/rest_test.py,sha256=_zM7nV8DEVyoXNiQOnuwJ917mWjki0614H88rNmDboE,5020
 langfun/core/llms/vertexai.py,sha256=nPWwv1C-wVpWpE4Vyl4vJoAEMB519-XNzPjyLY4N3tM,18816
-langfun/core/llms/vertexai_test.py,sha256=0M4jsPOXGagdzPfEdJixmyLdhmmERePZWSFfTwnaYCQ,4875
+langfun/core/llms/vertexai_test.py,sha256=_e-acnNBAf9C3WO6i1b2J_mhRzdDdYQTorD9hIVZKOg,5034
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=rt3zwmyw0y9jsSGW-ZbV1vAfLxQ7_3AVk0l2EySlse4,3918
 langfun/core/llms/cache/in_memory.py,sha256=i58oiQL28RDsq37dwqgVpC2mBETJjIEFS20yHiV5MKU,5185
@@ -137,8 +137,8 @@ langfun/core/structured/mapping.py,sha256=gxdcYQP9yqbDRtiJQ1RRAOrKHiCr0h6xBYLCRK
 langfun/core/structured/mapping_test.py,sha256=OntYvfDitAf0tAnzQty3YS90vyEn6FY1Mi93r_ViEk8,9594
 langfun/core/structured/parsing.py,sha256=MGvI7ypXlwfzr5XB8_TFU9Ei0_5reYqkWkv64eAy0EA,12015
 langfun/core/structured/parsing_test.py,sha256=V8Cj1tJK4Lxv_b0YQj6-2hzXZgnYNBa2JR7rOLRBKoQ,22346
-langfun/core/structured/querying.py,sha256=lHSPQ4UKRfhLXRjofG1geIGpQ_ImZf0dIvjXIwr_jNk,30480
-langfun/core/structured/querying_test.py,sha256=_npZ3ztaZc6VerP7nU_QTJscWGgBqiwTE02z_S3Ahd4,40197
+langfun/core/structured/querying.py,sha256=v-C8cU4LNFPwVSVgfO5PFyYxKCTHHh6SdbaR5V5MKKs,30938
+langfun/core/structured/querying_test.py,sha256=6jN3oOSpVEBIlISrCUJtZfY6sBHahF_EBRUEhNNlzW8,40797
 langfun/core/structured/schema.py,sha256=UTonddBx3hVr0Zhm_38jqd8khTSXWy-bkk8l1YOUdLA,28797
 langfun/core/structured/schema_generation.py,sha256=3AcuKvv3VOtKY5zMVqODrxfOuDxzoZtGeBxHlOWDOWw,5308
 langfun/core/structured/schema_generation_test.py,sha256=RM9s71kMNg2jTePwInkiW9fK1ACN37eyPeF8OII-0zw,2950
@@ -156,8 +156,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=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
-langfun-0.1.2.dev202505080804.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202505080804.dist-info/METADATA,sha256=CAqwJ8wOEhPMNBg-5-9XUOB8qWL3wuZJkycpxgcKugg,8178
-langfun-0.1.2.dev202505080804.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-langfun-0.1.2.dev202505080804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202505080804.dist-info/RECORD,,
+langfun-0.1.2.dev202505100803.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202505100803.dist-info/METADATA,sha256=M52b_XjFPeZNK9ZFRbXrVjBPlT9LyI3NThMyeLjKaQ4,8178
+langfun-0.1.2.dev202505100803.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+langfun-0.1.2.dev202505100803.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202505100803.dist-info/RECORD,,

{langfun-0.1.2.dev202505080804.dist-info → langfun-0.1.2.dev202505100803.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.3.1)
+Generator: setuptools (80.4.0)
 Root-Is-Purelib: true
 Tag: py3-none-any