langfun 0.1.2.dev202504290805__tar.gz → 0.1.2.dev202504300804__tar.gz

{langfun-0.1.2.dev202504290805 → langfun-0.1.2.dev202504300804}/PKG-INFO

@@ -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 → langfun-0.1.2.dev202504300804}/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-0.1.2.dev202504290805 → langfun-0.1.2.dev202504300804}/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-0.1.2.dev202504290805 → langfun-0.1.2.dev202504300804}/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-0.1.2.dev202504290805 → langfun-0.1.2.dev202504300804}/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-0.1.2.dev202504290805 → langfun-0.1.2.dev202504300804}/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-0.1.2.dev202504290805 → langfun-0.1.2.dev202504300804}/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-0.1.2.dev202504290805 → langfun-0.1.2.dev202504300804}/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-0.1.2.dev202504290805 → langfun-0.1.2.dev202504300804}/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-0.1.2.dev202504290805 → langfun-0.1.2.dev202504300804}/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.
 #