langfun 0.1.2.dev202504290805__py3-none-any.whl → 0.1.2.dev202504300804__py3-none-any.whl

langfun/__init__.py

@@ -37,6 +37,9 @@ generate_class = structured.generate_class
 
 track_queries = structured.track_queries
 
+# Context manager for setting the query protocol for the scope.
+query_protocol = structured.query_protocol
+
 # Helper function for map-reduce style querying.
 query_and_reduce = structured.query_and_reduce
 

langfun/core/agentic/action_eval.py

@@ -35,13 +35,18 @@ class ActionEval(lf.eval.v2.Evaluation):
     example_input = example.input
     action = example_input.action
     session = action_lib.Session(id=f'{self.id}#example-{example.id}')
+
+    # NOTE(daiyip): Setting session as metadata before action execution, so we
+    # could use `Evaluation.state.in_progress_examples` to access the session
+    # for status reporting from other threads.
+    example.metadata['session'] = session
+
     with lf.logging.use_log_level('fatal'):
       kwargs = self.action_args.copy()
       kwargs.update(verbose=True)
       action(session=session, **kwargs)
     return session.final_result, dict(session=session)
 
-
 #
 # TODO(daiyip): Remove V1 once V2 is fully launched.
 #

langfun/core/eval/base_test.py

@@ -195,7 +195,7 @@ class EvaluationTest(unittest.TestCase):
             score=1.0,
             logprobs=None,
             is_cached=False,
-            usage=lf.LMSamplingUsage(387, 24, 411),
+            usage=lf.LMSamplingUsage(428, 24, 452),
             tags=['lm-response', 'lm-output', 'transformed'],
         ),
     )
@@ -234,12 +234,12 @@ class EvaluationTest(unittest.TestCase):
                 }
             ),
             usage=dict(
-                total_prompt_tokens=774,
+                total_prompt_tokens=856,
                 total_completion_tokens=25,
                 num_usages=2,
-                average_prompt_tokens=387,
+                average_prompt_tokens=428,
                 average_completion_tokens=12,
-                average_total_tokens=399,
+                average_total_tokens=440,
             ),
         ),
     )

langfun/core/eval/v2/evaluation.py

@@ -167,6 +167,8 @@ class Evaluation(experiment_lib.Experiment):
       example.input = self.example_input_by_id(example.id)
 
     checkpointed = self._state.ckpt_example(example.id)
+    self._state.update(example, in_progress=True)
+
     with pg.timeit('evaluate') as timeit, lf.track_usages() as usage_summary:
       if checkpointed is None or checkpointed.has_error:
         if checkpointed is None:
@@ -221,7 +223,7 @@ class Evaluation(experiment_lib.Experiment):
     if example.newly_processed:
       example.end_time = time.time()
 
-    self._state.update(example)
+    self._state.update(example, in_progress=False)
     return example
 
   def _process(
@@ -501,6 +503,21 @@ class Evaluation(experiment_lib.Experiment):
           )
       )
 
+    def _in_progress_tab() -> pg.views.html.controls.Tab | None:
+      """Renders a tab for the in progress examples."""
+      if not self.state.in_progress_examples:
+        return None
+      return pg.views.html.controls.Tab(
+          label='In Progress',
+          content=pg.Html.element(
+              'div', [
+                  self._in_progress_view(
+                      list(self.state.in_progress_examples.values())
+                  )
+              ]
+          )
+      )
+
     def _metric_tab(metric: metrics_lib.Metric) -> pg.views.html.controls.Tab:
       """Renders a tab for a metric (group)."""
       return pg.views.html.controls.Tab(
@@ -571,10 +588,9 @@ class Evaluation(experiment_lib.Experiment):
               pg.views.html.controls.TabControl(
                   [
                       _definition_tab(),
-                  ] + [
-                      _metric_tab(m) for m in self.metrics
-                  ] + [
-                      _logs_tab()
+                      [_metric_tab(m) for m in self.metrics],
+                      _in_progress_tab(),
+                      _logs_tab(),
                   ],
                   selected=1,
               )
@@ -598,6 +614,27 @@ class Evaluation(experiment_lib.Experiment):
         css_classes=['eval-details'],
     )
 
+  def _in_progress_view(
+      self, in_progress_examples: list[example_lib.Example]
+  ) -> pg.Html:
+    """Renders a HTML view for the in-progress examples."""
+    current_time = time.time()
+    logs = [f'(Total {len(in_progress_examples)} examples in progress)']
+    for example in in_progress_examples:
+      if example.newly_processed:
+        logs.append(
+            f'Example {example.id}: In progress for '
+            f'{current_time - example.start_time:.2f} seconds.'
+        )
+      else:
+        logs.append(f'Example {example.id}: Recomputing metrics...')
+    return pg.Html.element(
+        'textarea',
+        [pg.Html.escape('\n'.join(logs))],
+        readonly=True,
+        css_classes=['logs-textarea'],
+    )
+
   def _html_tree_view_config(self) -> dict[str, Any]:
     return dict(
         css_classes=['eval-card'] if self.is_leaf else None
@@ -716,14 +753,27 @@ class EvaluationState:
         'Whether the example is evaluated.'
     ] = False
 
+    in_progress: Annotated[
+        bool,
+        (
+            'Whether the example is in progress. '
+        )
+    ] = False
+
     newly_processed: Annotated[
         bool,
-        'Whether the example is newly processed.'
+        (
+            'Whether the example is newly processed. '
+            'Applicable only when evaluated is True.'
+        )
     ] = False
 
     has_error: Annotated[
         bool,
-        'Whether the example has error.'
+        (
+            'Whether the example has error. '
+            'Applicable only when evaluated is True.'
+        )
     ] = False
 
   def __init__(self):
@@ -732,6 +782,7 @@ class EvaluationState:
     self._evaluation_status: dict[
         int, EvaluationState.ExampleStatus
     ] = {}
+    self._in_progress_examples: dict[int, example_lib.Example] = {}
 
   def load(
       self,
@@ -758,6 +809,11 @@ class EvaluationState:
     """Returns the evaluation status of the examples."""
     return self._evaluation_status
 
+  @property
+  def in_progress_examples(self) -> dict[int, example_lib.Example]:
+    """Returns the in-progress examples."""
+    return self._in_progress_examples
+
   @property
   def ckpt_examples(self) -> dict[int, example_lib.Example]:
     """Returns the unevaluated examples from checkpoints."""
@@ -773,17 +829,27 @@ class EvaluationState:
         example_id, EvaluationState.ExampleStatus()
     )
 
-  def update(self, example: example_lib.Example) -> None:
+  def update(self, example: example_lib.Example, in_progress: bool) -> None:
     """Updates the state with the given example."""
-    self._update_status(example)
-    # Processed examples will be removed once it's done.
-    self._ckpt_examples.pop(example.id, None)
+    self._update_status(example, in_progress)
+
+    if in_progress:
+      self._in_progress_examples[example.id] = example
+    else:
+      self._in_progress_examples.pop(example.id, None)
+      # Processed examples will be removed once it's done.
+      self._ckpt_examples.pop(example.id, None)
 
-  def _update_status(self, example: example_lib.Example) -> None:
+  def _update_status(
+      self,
+      example: example_lib.Example,
+      in_progress: bool
+  ) -> None:
     """Updates the evaluation status of the example."""
     self._evaluation_status[example.id] = (
         EvaluationState.ExampleStatus(
             evaluated=example.output != pg.MISSING_VALUE,
+            in_progress=in_progress,
             newly_processed=example.newly_processed,
             has_error=example.has_error,
         )

langfun/core/eval/v2/evaluation_test.py

@@ -79,8 +79,10 @@ class EvaluationTest(unittest.TestCase):
     exp = eval_test_helper.TestEvaluation()
     example = exp.evaluate(Example(id=3))
     self.assertTrue(exp.state.get_status(3).evaluated)
+    self.assertFalse(exp.state.get_status(3).in_progress)
     self.assertTrue(exp.state.get_status(3).newly_processed)
     self.assertFalse(exp.state.get_status(3).has_error)
+    self.assertEqual(exp.state.in_progress_examples, {})
     self.assertTrue(example.newly_processed)
     self.assertEqual(example.input, pg.Dict(x=2, y=4, groundtruth=6))
     self.assertEqual(example.output, 6)

langfun/core/structured/__init__.py

@@ -55,8 +55,11 @@ from langfun.core.structured.parsing import call
 
 from langfun.core.structured.querying import track_queries
 from langfun.core.structured.querying import QueryInvocation
+
+from langfun.core.structured.querying import LfQuery
 from langfun.core.structured.querying import query
 from langfun.core.structured.querying import query_and_reduce
+from langfun.core.structured.querying import query_protocol
 
 from langfun.core.structured.querying import query_prompt
 from langfun.core.structured.querying import query_output

langfun/core/structured/mapping.py

@@ -340,8 +340,11 @@ class Mapping(lf.LangFunc):
   schema_title: Annotated[str, 'The section title for schema.'] = 'SCHEMA'
 
   protocol: Annotated[
-      schema_lib.SchemaProtocol,
-      'The protocol for representing the schema and value.',
+      str,
+      (
+          'A string representing the protocol for formatting the prompt. '
+          'Built-in Langfun protocols are: `python` and `json`.'
+      ),
   ] = 'python'
 
   #

langfun/core/structured/parsing_test.py

@@ -646,7 +646,7 @@ class CallTest(unittest.TestCase):
             score=1.0,
             logprobs=None,
             is_cached=False,
-            usage=lf.LMSamplingUsage(315, 1, 316),
+            usage=lf.LMSamplingUsage(356, 1, 357),
             tags=['lm-response', 'lm-output', 'transformed']
         ),
     )

langfun/core/structured/querying.py

@@ -15,8 +15,9 @@
 
 import contextlib
 import functools
+import inspect
 import time
-from typing import Annotated, Any, Callable, Iterator, Type, Union
+from typing import Annotated, Any, Callable, ClassVar, Iterator, Type, Union
 import uuid
 
 import langfun.core as lf
@@ -26,8 +27,35 @@ import pyglove as pg
 
 
 @lf.use_init_args(['schema', 'default', 'examples'])
-class _QueryStructure(mapping.Mapping):
-  """Query an object out from a natural language text."""
+class LfQuery(mapping.Mapping):
+  """Base class for different implementations of `lf.query`.
+
+  By subclassing this class, users could create variations of prompts for
+  `lf.query` and associated them with specific protocols and versions.
+
+  For example:
+
+  ```
+  class _MyLfQuery(LFQuery):
+    protocol = 'my_format'
+    version = '1.0'
+
+    template_str = inspect.cleandoc(
+      '''
+      ...
+      '''
+    )
+    mapping_template = lf.Template(
+      '''
+      ...
+      '''
+    )
+
+  lf.query(..., protocol='my_format:1.0')
+  ```
+
+  (THIS IS NOT A TEMPLATE)
+  """
 
   context_title = 'CONTEXT'
   input_title = 'INPUT_OBJECT'
@@ -37,8 +65,81 @@ class _QueryStructure(mapping.Mapping):
       schema_lib.schema_spec(), 'Required schema for parsing.'
   ]
 
+  # A map from (protocol, version) to the query structure class.
+  # This is used to map different protocols/versions to different templates.
+  # So users can use `lf.query(..., protocol='<protocol>:<version>')` to use
+  # a specific version of the prompt. We use this feature to support variations
+  # of prompts and maintain backward compatibility.
+  _OOP_PROMPT_MAP: ClassVar[
+      dict[
+          str,        # protocol.
+          dict[
+              str,    # version.
+              Type['LfQuery']
+          ]
+      ]
+  ] = {}
+
+  # This the flag to update default protocol version.
+  _DEFAULT_PROTOCOL_VERSIONS: ClassVar[dict[str, str]] = {
+      'python': '2.0',
+      'json': '1.0',
+  }
+
+  def __init_subclass__(cls) -> Any:
+    super().__init_subclass__()
+    if not inspect.isabstract(cls):
+      protocol = cls.__schema__['protocol'].default_value
+      version_dict = cls._OOP_PROMPT_MAP.get(protocol)
+      if version_dict is None:
+        version_dict = {}
+        cls._OOP_PROMPT_MAP[protocol] = version_dict
+      dest_cls = version_dict.get(cls.version)
+      if dest_cls is not None and dest_cls.__type_name__ != cls.__type_name__:
+        raise ValueError(
+            f'Version {cls.version} is already registered for {dest_cls!r} '
+            f'under protocol {protocol!r}. Please use a different version.'
+        )
+      version_dict[cls.version] = cls
+
+  @classmethod
+  def from_protocol(cls, protocol: str) -> Type['LfQuery']:
+    """Returns a query structure from the given protocol and version."""
+    if ':' in protocol:
+      protocol, version = protocol.split(':')
+    else:
+      version = cls._DEFAULT_PROTOCOL_VERSIONS.get(protocol)
+      if version is None:
+        version_dict = cls._OOP_PROMPT_MAP.get(protocol)
+        if version_dict is None:
+          raise ValueError(
+              f'Protocol {protocol!r} is not supported. Available protocols: '
+              f'{sorted(cls._OOP_PROMPT_MAP.keys())}.'
+          )
+        elif len(version_dict) == 1:
+          version = list(version_dict.keys())[0]
+        else:
+          raise ValueError(
+              f'Multiple versions found for protocol {protocol!r}, please '
+              f'specify a version with "{protocol}:<version>".'
+          )
+
+    version_dict = cls._OOP_PROMPT_MAP.get(protocol)
+    if version_dict is None:
+      raise ValueError(
+          f'Protocol {protocol!r} is not supported. Available protocols: '
+          f'{sorted(cls._OOP_PROMPT_MAP.keys())}.'
+      )
+    dest_cls = version_dict.get(version)
+    if dest_cls is None:
+      raise ValueError(
+          f'Version {version!r} is not supported for protocol {protocol!r}. '
+          f'Available versions: {sorted(version_dict.keys())}.'
+      )
+    return dest_cls
 
-class _QueryStructureJson(_QueryStructure):
+
+class _LfQueryJsonV1(LfQuery):
   """Query a structured value using JSON as the protocol."""
 
   preamble = """
@@ -58,12 +159,13 @@ class _QueryStructureJson(_QueryStructure):
         {"result": {"_type": "langfun.core.structured.query.Answer", "final_answer": 2}}
       """
 
+  version = '1.0'
   protocol = 'json'
   schema_title = 'SCHEMA'
   output_title = 'JSON'
 
 
-class _QueryStructurePython(_QueryStructure):
+class _LfQueryPythonV1(LfQuery):
   """Query a structured value using Python as the protocol."""
 
   preamble = """
@@ -87,20 +189,87 @@ class _QueryStructurePython(_QueryStructure):
         )
         ```
       """
+  version = '1.0'
   protocol = 'python'
   schema_title = 'OUTPUT_TYPE'
   output_title = 'OUTPUT_OBJECT'
+  mapping_template = lf.Template(
+      """
+      {%- if example.context -%}
+      {{ context_title}}:
+      {{ example.context | indent(2, True)}}
 
+      {% endif -%}
 
-def _query_structure_cls(
-    protocol: schema_lib.SchemaProtocol,
-) -> Type[_QueryStructure]:
-  if protocol == 'json':
-    return _QueryStructureJson
-  elif protocol == 'python':
-    return _QueryStructurePython
-  else:
-    raise ValueError(f'Unknown protocol: {protocol!r}.')
+      {{ input_title }}:
+      {{ example.input_repr(protocol, compact=False) | indent(2, True) }}
+
+      {% if example.schema -%}
+      {{ schema_title }}:
+      {{ example.schema_repr(protocol) | indent(2, True) }}
+
+      {% endif -%}
+
+      {{ output_title }}:
+      {%- if example.has_output %}
+      {{ example.output_repr(protocol, compact=False) | indent(2, True) }}
+      {% endif -%}
+      """
+  )
+
+
+class _LfQueryPythonV2(LfQuery):
+  """Query a structured value using Python as the protocol."""
+
+  preamble = """
+      Please respond to the last {{ input_title }} with {{ output_title }} only according to {{ schema_title }}.
+
+      {{ input_title }}:
+        1 + 1 =
+
+      {{ schema_title }}:
+        Answer
+
+        ```python
+        class Answer:
+          final_answer: int
+        ```
+
+      {{ output_title }}:
+        ```python
+        output = Answer(
+          final_answer=2
+        )
+        ```
+      """
+  version = '2.0'
+  protocol = 'python'
+  input_title = 'REQUEST'
+  schema_title = 'OUTPUT PYTHON TYPE'
+  output_title = 'OUTPUT PYTHON OBJECT'
+  mapping_template = lf.Template(
+      """
+      {%- if example.context -%}
+      {{ context_title}}:
+      {{ example.context | indent(2, True)}}
+
+      {% endif -%}
+
+      {{ input_title }}:
+      {{ example.input_repr(protocol, compact=False) | indent(2, True) }}
+
+      {% if example.schema -%}
+      {{ schema_title }}:
+      {{ example.schema_repr(protocol) | indent(2, True) }}
+
+      {% endif -%}
+
+      {{ output_title }}:
+      {%- if example.has_output %}
+      {{ example.output_repr(protocol, compact=False, assign_to_var='output') | indent(2, True) }}
+      {% endif -%}
+      """
+  )
 
 
 def query(
@@ -116,7 +285,7 @@ def query(
     response_postprocess: Callable[[str], str] | None = None,
     autofix: int = 0,
     autofix_lm: lf.LanguageModel | None = None,
-    protocol: schema_lib.SchemaProtocol = 'python',
+    protocol: str | None = None,
     returns_message: bool = False,
     skip_lm: bool = False,
     invocation_id: str | None = None,
@@ -259,8 +428,14 @@ def query(
       disable auto-fixing. Not supported with the `'json'` protocol.
     autofix_lm: The LM to use for auto-fixing. Defaults to the `autofix_lm`
       from `lf.context` or the main `lm`.
-    protocol: Format for schema representation. Choices are `'json'` or
-      `'python'`. Default is `'python'`.
+    protocol: Format for schema representation. Builtin choices are `'json'` or
+      `'python'`, users could extend with their own protocols by subclassing
+      `lf.structured.LfQuery'. Also protocol could be specified with a version
+      in the format of 'protocol:version', e.g., 'python:1.0', so users could
+      use a specific version of the prompt based on the protocol. Please see the
+      documentation of `LfQuery` for more details. If None, the protocol from
+      context manager `lf.query_protocol` will be used, or 'python' if not
+      specified.
     returns_message:  If `True`, returns an `lf.Message` object instead of
       the final parsed result.
     skip_lm: If `True`, skips the LLM call and returns the rendered 
@@ -280,6 +455,9 @@ def query(
   """
     # Internal usage logging.
 
+  if protocol is None:
+    protocol = lf.context_value('__query_protocol__', 'python')
+
   invocation_id = invocation_id or f'query@{uuid.uuid4().hex[-7:]}'
   # Multiple quries will be issued when `lm` is a list or `num_samples` is
   # greater than 1.
@@ -382,7 +560,7 @@ def query(
           output_message = lf.AIMessage(processed_text, source=output_message)
     else:
       # Query with structured output.
-      output_message = _query_structure_cls(protocol)(
+      output_message = LfQuery.from_protocol(protocol)(
           input=(
               query_input.render(lm=lm)
               if isinstance(query_input, lf.Template)
@@ -436,6 +614,15 @@ def query(
   return output_message if returns_message else _result(output_message)
 
 
+@contextlib.contextmanager
+def query_protocol(protocol: str) -> Iterator[None]:
+  """Context manager for setting the query protocol for the scope."""
+  with lf.context(__query_protocol__=protocol):
+    try:
+      yield
+    finally:
+      pass
+
 #
 # Helper function for map-reduce style querying.
 #

langfun/core/structured/querying_test.py

@@ -90,7 +90,7 @@ class QueryTest(unittest.TestCase):
             score=1.0,
             logprobs=None,
             is_cached=False,
-            usage=lf.LMSamplingUsage(323, 1, 324),
+            usage=lf.LMSamplingUsage(364, 1, 365),
             tags=['lm-response', 'lm-output', 'transformed'],
         ),
     )
@@ -143,26 +143,26 @@ class QueryTest(unittest.TestCase):
         y=2,
         lm=lm.clone(),
         expected_snippet=(
-            'Please respond to the last INPUT_OBJECT with OUTPUT_OBJECT '
-            'according to OUTPUT_TYPE.\n\n'
-            'INPUT_OBJECT:\n  1 + 1 =\n\n'
-            'OUTPUT_TYPE:\n'
+            'Please respond to the last REQUEST with OUTPUT PYTHON OBJECT '
+            'only according to OUTPUT PYTHON TYPE.\n\n'
+            'REQUEST:\n  1 + 1 =\n\n'
+            'OUTPUT PYTHON TYPE:\n'
             '  Answer\n\n'
             '  ```python\n'
             '  class Answer:\n'
             '    final_answer: int\n'
             '  ```\n\n'
-            'OUTPUT_OBJECT:\n'
+            'OUTPUT PYTHON OBJECT:\n'
             '  ```python\n'
-            '  Answer(\n'
+            '  output = Answer(\n'
             '    final_answer=2\n'
             '  )\n'
             '  ```\n\n'
-            'INPUT_OBJECT:\n'
+            'REQUEST:\n'
             '  What is 1 + 2?\n\n'
-            'OUTPUT_TYPE:\n'
+            'OUTPUT PYTHON TYPE:\n'
             '  int\n\n'
-            'OUTPUT_OBJECT:'
+            'OUTPUT PYTHON OBJECT:'
         ),
     )
 
@@ -176,26 +176,26 @@ class QueryTest(unittest.TestCase):
         lm=lm.clone(),
         template_str='!!{{ DEFAULT }}!!',
         expected_snippet=(
-            '!!Please respond to the last INPUT_OBJECT with OUTPUT_OBJECT '
-            'according to OUTPUT_TYPE.\n\n'
-            'INPUT_OBJECT:\n  1 + 1 =\n\n'
-            'OUTPUT_TYPE:\n'
+            '!!Please respond to the last REQUEST with OUTPUT PYTHON OBJECT '
+            'only according to OUTPUT PYTHON TYPE.\n\n'
+            'REQUEST:\n  1 + 1 =\n\n'
+            'OUTPUT PYTHON TYPE:\n'
             '  Answer\n\n'
             '  ```python\n'
             '  class Answer:\n'
             '    final_answer: int\n'
             '  ```\n\n'
-            'OUTPUT_OBJECT:\n'
+            'OUTPUT PYTHON OBJECT:\n'
             '  ```python\n'
-            '  Answer(\n'
+            '  output = Answer(\n'
             '    final_answer=2\n'
             '  )\n'
             '  ```\n\n'
-            'INPUT_OBJECT:\n'
+            'REQUEST:\n'
             '  What is 1 + 2?\n\n'
-            'OUTPUT_TYPE:\n'
+            'OUTPUT PYTHON TYPE:\n'
             '  int\n\n'
-            'OUTPUT_OBJECT:!!'
+            'OUTPUT PYTHON OBJECT:!!'
         ),
     )
 
@@ -220,7 +220,7 @@ class QueryTest(unittest.TestCase):
         y=2,
         lm=lm.clone(),
         expected_snippet=(
-            '\n\nINPUT_OBJECT:\n  ```python\n  [\n    1\n  ]\n  ```\n\n'
+            '\n\nREQUEST:\n  ```python\n  [\n    1\n  ]\n  ```\n\n'
         ),
     )
 
@@ -236,7 +236,7 @@ class QueryTest(unittest.TestCase):
         modalities.Image.from_bytes(b'mock_image'),
         int,
         lm=lm,
-        expected_snippet='\n\nINPUT_OBJECT:\n  <<[[input]]>>\n\n',
+        expected_snippet='\n\nREQUEST:\n  <<[[input]]>>\n\n',
         expected_modalities=1,
     )
 
@@ -290,7 +290,7 @@ class QueryTest(unittest.TestCase):
         list[str],
         lm=lm,
         expected_snippet=inspect.cleandoc("""
-            INPUT_OBJECT:
+            REQUEST:
               ```python
               [
                 <<[[input[0]]]>>,
@@ -318,31 +318,36 @@ class QueryTest(unittest.TestCase):
         ],
         lm=lm,
         expected_snippet=inspect.cleandoc("""
-            INPUT_OBJECT:
+            REQUEST:
               ```python
               [
                 <<[[examples[0].input[0]]]>>
               ]
               ```
 
-            OUTPUT_TYPE:
+            OUTPUT PYTHON TYPE:
               list[str]
 
-            OUTPUT_OBJECT:
+            OUTPUT PYTHON OBJECT:
               ```python
-              [
+              output = [
                 'dog'
               ]
               ```
 
 
-            INPUT_OBJECT:
+            REQUEST:
               ```python
               [
                 <<[[input[0]]]>>,
                 <<[[input[1]]]>>
               ]
               ```
+
+            OUTPUT PYTHON TYPE:
+              list[str]
+
+            OUTPUT PYTHON OBJECT:
             """),
         expected_modalities=3,
     )
@@ -410,20 +415,69 @@ class QueryTest(unittest.TestCase):
     self.assertEqual([r.text for r in results], ['1', 'abc'])
     self.assertEqual([r.result for r in results], [1, 0])
 
-  def test_bad_protocol(self):
-    with self.assertRaisesRegex(ValueError, 'Unknown protocol'):
-      querying.query('what is 1 + 1', int, protocol='text')
+  def test_from_protocol(self):
+    self.assertIs(
+        querying.LfQuery.from_protocol('python'), querying._LfQueryPythonV2
+    )
+    self.assertIs(
+        querying.LfQuery.from_protocol('python:1.0'),
+        querying._LfQueryPythonV1
+    )
+
+    class MyLfQuery(querying.LfQuery):
+      protocol = 'yaml'
+      version = '1.0'
+
+    self.assertIs(
+        querying.LfQuery.from_protocol('yaml:1.0'),
+        MyLfQuery
+    )
+    self.assertIs(
+        querying.LfQuery.from_protocol('yaml'),
+        MyLfQuery
+    )
+
+    with self.assertRaisesRegex(
+        ValueError, 'Version .* is already registered'
+    ):
+      class MyLfQuery2(querying.LfQuery):  # pylint: disable=unused-variable
+        protocol = 'yaml'
+        version = '1.0'
+
+    with self.assertRaisesRegex(
+        ValueError, 'Version \'2.0\' is not supported for protocol \'yaml\''
+    ):
+      querying.LfQuery.from_protocol('yaml:2.0')
+
+    class MyLfQuery3(querying.LfQuery):  # pylint: disable=unused-variable
+      protocol = 'yaml'
+      version = '3.0'
+
+    with self.assertRaisesRegex(
+        ValueError, 'Multiple versions found for protocol \'yaml\''
+    ):
+      querying.LfQuery.from_protocol('yaml')
+
+    with self.assertRaisesRegex(
+        ValueError, 'Protocol \'text\' is not supported'
+    ):
+      querying.LfQuery.from_protocol('text')
+
+    with self.assertRaisesRegex(
+        ValueError, 'Protocol \'text\' is not supported'
+    ):
+      querying.LfQuery.from_protocol('text:1.0')
 
   def test_query_prompt(self):
     self.assertEqual(
         querying.query_prompt('what is this?', int),
         inspect.cleandoc("""
-            Please respond to the last INPUT_OBJECT with OUTPUT_OBJECT according to OUTPUT_TYPE.
+            Please respond to the last REQUEST with OUTPUT PYTHON OBJECT only according to OUTPUT PYTHON TYPE.
 
-            INPUT_OBJECT:
+            REQUEST:
               1 + 1 =
 
-            OUTPUT_TYPE:
+            OUTPUT PYTHON TYPE:
               Answer
 
               ```python
@@ -431,20 +485,20 @@ class QueryTest(unittest.TestCase):
                 final_answer: int
               ```
 
-            OUTPUT_OBJECT:
+            OUTPUT PYTHON OBJECT:
               ```python
-              Answer(
+              output = Answer(
                 final_answer=2
               )
               ```
 
-            INPUT_OBJECT:
+            REQUEST:
               what is this?
 
-            OUTPUT_TYPE:
+            OUTPUT PYTHON TYPE:
               int
 
-            OUTPUT_OBJECT:
+            OUTPUT PYTHON OBJECT:
             """),
     )
 
@@ -632,10 +686,10 @@ class QueryTest(unittest.TestCase):
       )
 
 
-class QueryStructurePythonTest(unittest.TestCase):
+class LfQueryPythonV1Test(unittest.TestCase):
 
   def test_render_no_examples(self):
-    l = querying._QueryStructurePython(
+    l = querying.LfQuery.from_protocol('python:1.0')(
         input=lf.AIMessage('Compute 12 / 6 + 2.'), schema=int
     )
     self.assertEqual(
@@ -672,7 +726,7 @@ class QueryStructurePythonTest(unittest.TestCase):
     )
 
   def test_render(self):
-    l = querying._QueryStructurePython(
+    l = querying.LfQuery.from_protocol('python:1.0')(
         input=lf.AIMessage('Compute 12 / 6 + 2.'),
         schema=int,
         examples=[
@@ -782,7 +836,7 @@ class QueryStructurePythonTest(unittest.TestCase):
         ),
         override_attrs=True,
     ):
-      l = querying._QueryStructurePython(
+      l = querying.LfQuery.from_protocol('python:1.0')(
           input=lm_input,
           schema=[Itinerary],
           examples=[
@@ -810,6 +864,115 @@ class QueryStructurePythonTest(unittest.TestCase):
       self.assertEqual(len(r.result[0].activities), 3)
       self.assertIsNone(r.result[0].hotel)
 
+
+class LfQueryPythonV2Test(unittest.TestCase):
+
+  def test_render_no_examples(self):
+    l = querying.LfQuery.from_protocol('python:2.0')(
+        input=lf.AIMessage('Compute 12 / 6 + 2.'), schema=int
+    )
+    self.assertEqual(
+        l.render().text,
+        inspect.cleandoc("""
+            Please respond to the last REQUEST with OUTPUT PYTHON OBJECT only according to OUTPUT PYTHON TYPE.
+
+            REQUEST:
+              1 + 1 =
+
+            OUTPUT PYTHON TYPE:
+              Answer
+
+              ```python
+              class Answer:
+                final_answer: int
+              ```
+
+            OUTPUT PYTHON OBJECT:
+              ```python
+              output = Answer(
+                final_answer=2
+              )
+              ```
+
+            REQUEST:
+              Compute 12 / 6 + 2.
+
+            OUTPUT PYTHON TYPE:
+              int
+
+            OUTPUT PYTHON OBJECT:
+            """),
+    )
+
+  def test_render_with_examples(self):
+    l = querying.LfQuery.from_protocol('python:2.0')(
+        input=lf.AIMessage('Compute 12 / 6 + 2.'),
+        schema=int,
+        examples=[
+            mapping.MappingExample(
+                input='What is the answer of 1 plus 1?', output=2
+            ),
+            mapping.MappingExample(
+                input='Compute the value of 3 + (2 * 6).', output=15
+            ),
+        ],
+    )
+    self.assertEqual(
+        l.render().text,
+        inspect.cleandoc("""
+            Please respond to the last REQUEST with OUTPUT PYTHON OBJECT only according to OUTPUT PYTHON TYPE.
+
+            REQUEST:
+              1 + 1 =
+
+            OUTPUT PYTHON TYPE:
+              Answer
+
+              ```python
+              class Answer:
+                final_answer: int
+              ```
+
+            OUTPUT PYTHON OBJECT:
+              ```python
+              output = Answer(
+                final_answer=2
+              )
+              ```
+
+            REQUEST:
+              What is the answer of 1 plus 1?
+
+            OUTPUT PYTHON TYPE:
+              int
+
+            OUTPUT PYTHON OBJECT:
+              ```python
+              output = 2
+              ```
+
+            REQUEST:
+              Compute the value of 3 + (2 * 6).
+
+            OUTPUT PYTHON TYPE:
+              int
+
+            OUTPUT PYTHON OBJECT:
+              ```python
+              output = 15
+              ```
+
+
+            REQUEST:
+              Compute 12 / 6 + 2.
+
+            OUTPUT PYTHON TYPE:
+              int
+
+            OUTPUT PYTHON OBJECT:
+            """),
+    )
+
   def test_bad_response(self):
     with lf.context(
         lm=fake.StaticSequence(['a2']),
@@ -849,11 +1012,80 @@ class QueryStructurePythonTest(unittest.TestCase):
           3
       )
 
+  def test_render(self):
+    l = querying.LfQuery.from_protocol('python:2.0')(
+        input=lf.AIMessage('Compute 12 / 6 + 2.'),
+        schema=int,
+        examples=[
+            mapping.MappingExample(
+                input='What is the answer of 1 plus 1?', output=2
+            ),
+            mapping.MappingExample(
+                input='Compute the value of 3 + (2 * 6).', output=15
+            ),
+        ],
+    )
+    self.assertEqual(
+        l.render().text,
+        inspect.cleandoc("""
+            Please respond to the last REQUEST with OUTPUT PYTHON OBJECT only according to OUTPUT PYTHON TYPE.
+
+            REQUEST:
+              1 + 1 =
+
+            OUTPUT PYTHON TYPE:
+              Answer
+
+              ```python
+              class Answer:
+                final_answer: int
+              ```
+
+            OUTPUT PYTHON OBJECT:
+              ```python
+              output = Answer(
+                final_answer=2
+              )
+              ```
+
+            REQUEST:
+              What is the answer of 1 plus 1?
+
+            OUTPUT PYTHON TYPE:
+              int
+
+            OUTPUT PYTHON OBJECT:
+              ```python
+              output = 2
+              ```
+
+            REQUEST:
+              Compute the value of 3 + (2 * 6).
+
+            OUTPUT PYTHON TYPE:
+              int
+
+            OUTPUT PYTHON OBJECT:
+              ```python
+              output = 15
+              ```
+
+
+            REQUEST:
+              Compute 12 / 6 + 2.
+
+            OUTPUT PYTHON TYPE:
+              int
+
+            OUTPUT PYTHON OBJECT:
+            """),
+    )
 
-class QueryStructureJsonTest(unittest.TestCase):
+
+class LfQueryJsonV1Test(unittest.TestCase):
 
   def test_render_no_examples(self):
-    l = querying._QueryStructureJson(
+    l = querying.LfQuery.from_protocol('json:1.0')(
         input=lf.AIMessage('Compute 12 / 6 + 2.'), schema=int
     )
     self.assertEqual(
@@ -885,7 +1117,7 @@ class QueryStructureJsonTest(unittest.TestCase):
     )
 
   def test_render(self):
-    l = querying._QueryStructureJson(
+    l = querying.LfQuery.from_protocol('json:1.0')(
         input=lf.AIMessage('Compute 12 / 6 + 2.'),
         schema=int,
         examples=[
@@ -1020,7 +1252,7 @@ class QueryStructureJsonTest(unittest.TestCase):
         ),
         override_attrs=True,
     ):
-      l = querying._QueryStructureJson(
+      l = querying.LfQuery.from_protocol('json:1.0')(
           input=lm_input,
           schema=[Itinerary],
           examples=[
@@ -1063,10 +1295,17 @@ class QueryStructureJsonTest(unittest.TestCase):
       self.assertEqual(len(cache), 0)
 
   def test_query(self):
-    lm = fake.StaticSequence(['{"result": 1}'])
+    lm = fake.StaticResponse('{"result": 1}')
     self.assertEqual(
         querying.query('what is 1 + 0', int, lm=lm, protocol='json'), 1
     )
+    self.assertEqual(
+        querying.query('what is 1 + 0', int, lm=lm, protocol='json:1.0'), 1
+    )
+    with querying.query_protocol('json'):
+      self.assertEqual(
+          querying.query('what is 1 + 0', int, lm=lm), 1
+      )
 
 
 class QueryInvocationTest(unittest.TestCase):

{langfun-0.1.2.dev202504290805.dist-info → langfun-0.1.2.dev202504300804.dist-info}/METADATA

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

{langfun-0.1.2.dev202504290805.dist-info → langfun-0.1.2.dev202504300804.dist-info}/RECORD

@@ -1,4 +1,4 @@
-langfun/__init__.py,sha256=LuB1difPJsI8V9H4kmUJM2DAEmMEV_WtaBLGYp3qEFM,2527
+langfun/__init__.py,sha256=krEJ1lyDkNARsacY6nBQpD3bQrFi4fifD-FwpwPbFPM,2635
 langfun/core/__init__.py,sha256=pW4prpiyWNkRbtWBGYF1thn7_0F_TgDVfAIZPvGn6HA,4758
 langfun/core/component.py,sha256=g1kQM0bryYYYWVDrSMnHfc74wIBbpfe5_B3s-UIP5GE,3028
 langfun/core/component_test.py,sha256=0CxTgjAud3aj8wBauFhG2FHDqrxCTl4OI4gzQTad-40,9254
@@ -27,7 +27,7 @@ langfun/core/template.py,sha256=jNhYSrbLIn9kZOa03w5QZbyjgfnzJzE_ZrrMvvWY4t4,2492
 langfun/core/template_test.py,sha256=AQv_m9qE93WxhEhSlm1xaBgB4hu0UVtA53dljngkUW0,17090
 langfun/core/agentic/__init__.py,sha256=qR3jlfUO4rhIoYdRDLz-d22YZf3FvU4FW88vsjiGDQQ,1224
 langfun/core/agentic/action.py,sha256=9P7xDiZVUV9MvJDfuAfLx-xa7qvS5F0EOGWDQnjAZBw,38931
-langfun/core/agentic/action_eval.py,sha256=8cdH-ubWUL1sscqagrQltPg19vTYNKrI3CKLvluGP24,4615
+langfun/core/agentic/action_eval.py,sha256=NwjQ5hR-7YT6mo2q0mbDOgmNCKzTMpEzslYtR3fjXJY,4862
 langfun/core/agentic/action_eval_test.py,sha256=tRUkWmOE9p0rpNOq19xAY2oDEnYsEEykjg6sUpAwJk0,2832
 langfun/core/agentic/action_test.py,sha256=9EZKgLaBrqTErSRoxtrSlzmCz_cbnwWu0ZqpwKLst-s,10224
 langfun/core/coding/__init__.py,sha256=5utju_fwEsImaiftx4oXKl9FAM8p281k8-Esdh_-m1w,835
@@ -52,7 +52,7 @@ langfun/core/data/conversion/openai.py,sha256=sSpkDSxMJWJ3I1dNICBCzvLsJv4iiLg8FP
 langfun/core/data/conversion/openai_test.py,sha256=38WV_3ofFZiUF10bTKnZp4VyuDP5-81aR3h3Q0HlBm0,5283
 langfun/core/eval/__init__.py,sha256=OEXr1ZRuvLuhJJfuQ1ZWQ-SvYzjyrtiAAEogYaB7E6o,1933
 langfun/core/eval/base.py,sha256=qIJnrO1jX5pzY8yoQTtcTn5lGdD9adz5U6C_jla1BV4,75806
-langfun/core/eval/base_test.py,sha256=rFWCIDzNGadXMxxM4Ofavz52SsjtlH4wqS1fnaA4ksI,27191
+langfun/core/eval/base_test.py,sha256=q4wEd2KDUxzUkeELwof0HXBKe9TMQYUq84ddA043VPg,27191
 langfun/core/eval/matching.py,sha256=AVKkGoc-BaHEzgSBamaAk3194TgqckDe_dinpS6LrXI,9323
 langfun/core/eval/matching_test.py,sha256=2xtwsTi-UzLTt0QnXl3u_eAG3fFjCG2tsae7YkcQTB0,5312
 langfun/core/eval/patching.py,sha256=R0s2eAd1m97exQt06dmUL0V_MBG0W2Hxg7fhNB7cXW0,3866
@@ -63,8 +63,8 @@ langfun/core/eval/v2/__init__.py,sha256=9lNKJwbvl0lcFblAXYT_OHI8fOubJsTOdSkxEqsP
 langfun/core/eval/v2/checkpointing.py,sha256=t47rBfzGZYgIqWW1N1Ak9yQnNtHd-IRbEO0cZjG2VRo,11755
 langfun/core/eval/v2/checkpointing_test.py,sha256=NggOSJ_6XSa4cNP6nGIu9wLsK59dUwe8SPWDiXtGGDE,9197
 langfun/core/eval/v2/eval_test_helper.py,sha256=sKFi_wPYCNmr96WyTduuXY0KnxjFxcJyEhXey-_nGX8,3962
-langfun/core/eval/v2/evaluation.py,sha256=07ebNx3bYlN25sg4Nam_QJRuikKVt9Pc0oXmSG6IS-k,24937
-langfun/core/eval/v2/evaluation_test.py,sha256=TKbeyeenoaUDH23E3TXXJ4otuq3VZBuWkWdV7c5vMk4,6804
+langfun/core/eval/v2/evaluation.py,sha256=4Kcpve_dlti3FWxd7KtZaapsva_C7sOGjyJ69xsXLY4,26992
+langfun/core/eval/v2/evaluation_test.py,sha256=QNp_HEvRTupvNuLEeYTvylykh1Ut2jpMqHQ-gCUZQ10,6919
 langfun/core/eval/v2/example.py,sha256=Jegt-viQSNYzPVkOZE_M19GON2TYGTct4Cp9HnJ7DGo,10861
 langfun/core/eval/v2/example_test.py,sha256=1DNm6EuyZOq827DKvf3oTRVFkMNM_qTnLUpvOjpgz5I,3419
 langfun/core/eval/v2/experiment.py,sha256=xlQvx-AgPCgHNaBoW1HxddA9wby-odADF0VJ3rQjw_M,32978
@@ -126,19 +126,19 @@ langfun/core/modalities/pdf.py,sha256=mfaeCbUA4JslFVTARiJh8hW7imvL4tLVw9gUhO5bAZ
 langfun/core/modalities/pdf_test.py,sha256=ulZ0FbnlsU0wkrdckJ4ONZPTYRyMPO9Aob1UO6FXygk,1950
 langfun/core/modalities/video.py,sha256=vI9apcHIHGyp90i34Srg7S3G6IBDtDCk8qiXhwRQmkw,967
 langfun/core/modalities/video_test.py,sha256=7OXZoohKMYjt7vrJUdPb553HLyl1oBOKRgzBePFv68Q,2042
-langfun/core/structured/__init__.py,sha256=3Wb4ks14D5E5z1nQ5trXSXiLqFN5E_3HvcEJQAfFRl0,3220
+langfun/core/structured/__init__.py,sha256=xWO1RO-gy12REgiPn4oWWALZxJhHtBOOdDsawKN_SfY,3334
 langfun/core/structured/completion.py,sha256=yW95Yd4wbt964I5wIyPUtIVeeqeZbA6HidLgN0ZpWm0,8110
 langfun/core/structured/completion_test.py,sha256=VtYfI3ciVSSWbi8x3l1WwpWK-Ofn2DMHYEEqm2uTzhw,19314
 langfun/core/structured/description.py,sha256=6BztYOiucPkF4CrTQtPLPJo1gN2dwnKmaJW83GBf4H0,5213
 langfun/core/structured/description_test.py,sha256=UxaXnKKP7TnyPDPUyf3U-zPE0TvLlIP6DGr8thjcePw,7365
 langfun/core/structured/function_generation.py,sha256=g7AOR_e8HxFU6n6Df750aGkgMgV1KExLZMAz0yd5Agg,8555
 langfun/core/structured/function_generation_test.py,sha256=LaXYDXf9GlqUrR6v_gtmK_H4kxzonmU7SYbn7XXMgjU,12128
-langfun/core/structured/mapping.py,sha256=of-EeBq0RgmkiUaSk2rVEDVCzgn_wXU8tRke7NCcC6E,13649
+langfun/core/structured/mapping.py,sha256=gxdcYQP9yqbDRtiJQ1RRAOrKHiCr0h6xBYLCRKE4I3c,13723
 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=kNPrhpdPY3iWhUld0TFYU-Zgn44wC0d6YuQ9XdVbQ8o,22346
-langfun/core/structured/querying.py,sha256=GGNtHtJcKh8rzLBNx_Df1ATvsPZzyfZuGkzSQVabdpo,24885
-langfun/core/structured/querying_test.py,sha256=vUuVUClYBFGEaO9KuD60huPE1dmP6RCRLeRnBv67NmQ,34263
+langfun/core/structured/parsing_test.py,sha256=V8Cj1tJK4Lxv_b0YQj6-2hzXZgnYNBa2JR7rOLRBKoQ,22346
+langfun/core/structured/querying.py,sha256=NA5LL97ny9aIml3wz4GPzFsvAH2XafZIGoZbIhuV8nA,30456
+langfun/core/structured/querying_test.py,sha256=_npZ3ztaZc6VerP7nU_QTJscWGgBqiwTE02z_S3Ahd4,40197
 langfun/core/structured/schema.py,sha256=pGiAjez-ON2nKLUSeAls27gJsMto5aJnCXLVwH3pUKM,28296
 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.dev202504290805.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202504290805.dist-info/METADATA,sha256=KWW7NslPxpSS2u0eVRcN022ylI5TeWFdswXMnve9144,8178
-langfun-0.1.2.dev202504290805.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
-langfun-0.1.2.dev202504290805.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202504290805.dist-info/RECORD,,
+langfun-0.1.2.dev202504300804.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202504300804.dist-info/METADATA,sha256=rlaBu8x--oQz7DAbnG3snNBBk7ltVc1RFvtOcE9cW3Q,8178
+langfun-0.1.2.dev202504300804.dist-info/WHEEL,sha256=ooBFpIzZCPdw3uqIQsOo4qqbA4ZRPxHnOH7peeONza0,91
+langfun-0.1.2.dev202504300804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202504300804.dist-info/RECORD,,

{langfun-0.1.2.dev202504290805.dist-info → langfun-0.1.2.dev202504300804.dist-info}/WHEEL

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