langfun 0.1.2.dev202412170805__py3-none-any.whl → 0.1.2.dev202412190804__py3-none-any.whl

langfun/__init__.py

@@ -37,6 +37,9 @@ generate_class = structured.generate_class
 
 track_queries = structured.track_queries
 
+# Helper function for map-reduce style querying.
+query_and_reduce = structured.query_and_reduce
+
 # Helper functions for input/output transformations based on
 # `lf.query` (e.g. jax-on-beam could use these for batch processing)
 query_prompt = structured.query_prompt

langfun/core/eval/v2/checkpointing.py

@@ -92,7 +92,7 @@ class PerExampleCheckpointer(Checkpointer):
             )
         )
         writer.add(example)
-        del writer
+        writer.close()
       runner.background_run(save_state, example)
 
   def _file_prefix_and_ext(self, filename: str) -> tuple[str, str]:
@@ -128,6 +128,8 @@ class BulkCheckpointer(Checkpointer):
   ) -> None:
     with self._lock:
       if self._sequence_writer is not None:
+        for writer in self._sequence_writer.values():
+          writer.close()
         self._sequence_writer.clear()
 
   def on_run_complete(
@@ -174,6 +176,9 @@ class BulkCheckpointer(Checkpointer):
     assert experiment.id in self._sequence_writer
     with self._lock:
       if self._sequence_writer is not None:
+        # Make sure the writer is closed without delay so the file will be
+        # available immediately.
+        self._sequence_writer[experiment.id].close()
         del self._sequence_writer[experiment.id]
 
   def on_example_complete(
@@ -207,9 +212,13 @@ class SequenceWriter:
         return
       self._sequence_writer.add(example_blob)
 
-  def __del__(self):
+  def close(self):
     # Make sure there is no write in progress.
     with self._lock:
-      assert self._sequence_writer is not None
+      if self._sequence_writer is None:
+        return
       self._sequence_writer.close()
       self._sequence_writer = None
+
+  def __del__(self):
+    self.close()

langfun/core/eval/v2/evaluation.py

@@ -14,7 +14,9 @@
 """Base class for Langfun evaluation tasks."""
 
 import abc
+import datetime
 import functools
+import threading
 import time
 
 from typing import Annotated, Any, Callable, Iterable
@@ -63,6 +65,8 @@ class Evaluation(experiment_lib.Experiment):
     self.__dict__.pop('is_leaf', None)
     self.__dict__.pop('children', None)
     super()._on_bound()
+    self._log_entries = []
+    self._log_lock = threading.Lock()
 
   #
   # Handling evaluation hierarchy (materialized vs. hyper evaluations).
@@ -277,6 +281,41 @@ class Evaluation(experiment_lib.Experiment):
       for metric in self.metrics:
         metric.reset()
 
+  #
+  # Evaluation-level logging.
+  #
+
+  def _log(self, level: lf.logging.LogLevel, message: str, **kwargs):
+    with self._log_lock:
+      self._log_entries.append(
+          lf.logging.LogEntry(
+              level=level,
+              time=datetime.datetime.now(),
+              message=message,
+              metadata=kwargs,
+          )
+      )
+
+  def debug(self, message: str, **kwargs):
+    """Logs a debug message to the session."""
+    self._log('debug', message, **kwargs)
+
+  def info(self, message: str, **kwargs):
+    """Logs an info message to the session."""
+    self._log('info', message, **kwargs)
+
+  def warning(self, message: str, **kwargs):
+    """Logs a warning message to the session."""
+    self._log('warning', message, **kwargs)
+
+  def error(self, message: str, **kwargs):
+    """Logs an error message to the session."""
+    self._log('error', message, **kwargs)
+
+  def fatal(self, message: str, **kwargs):
+    """Logs a fatal message to the session."""
+    self._log('fatal', message, **kwargs)
+
   #
   # HTML views.
   #
@@ -465,6 +504,25 @@ class Evaluation(experiment_lib.Experiment):
           )
       )
 
+    def _logs_tab() -> pg.views.html.controls.Tab:
+      """Renders a tab for the logs of the evaluation."""
+      with self._log_lock:
+        log_history = '\n'.join(str(l) for l in self._log_entries)
+      return pg.views.html.controls.Tab(
+          label='Logs',
+          content=pg.Html.element(
+              'div',
+              [
+                  pg.Html.element(
+                      'textarea',
+                      [pg.Html.escape(log_history)],
+                      readonly=True,
+                      css_classes=['logs-textarea'],
+                  )
+              ]
+          )
+      )
+
     def _main_tabs() -> pg.Html:
       return pg.Html.element(
           'div',
@@ -474,6 +532,8 @@ class Evaluation(experiment_lib.Experiment):
                       _definition_tab(),
                   ] + [
                       _metric_tab(m) for m in self.metrics
+                  ] + [
+                      _logs_tab()
                   ],
                   selected=1,
               )
@@ -593,6 +653,14 @@ class Evaluation(experiment_lib.Experiment):
           width:100%;
           height:100%;
         }
+        .logs-textarea {
+          width: 100%;
+          height: 500px;
+          padding: 5px;
+          border: 1px solid #DDD;
+          background-color: #EEE;
+          resize: vertical;
+        }
         """
     ]
 
@@ -615,6 +683,11 @@ class EvaluationState:
         assert isinstance(example, example_lib.Example), example
         self._evaluated_examples[example.id] = example
 
+  @property
+  def evaluated_examples(self) -> dict[int, example_lib.Example]:
+    """Returns the examples in the state."""
+    return self._evaluated_examples
+
   def get(self, example_id: int) -> example_lib.Example | None:
     """Returns the example with the given ID."""
     return self._evaluated_examples.get(example_id)
@@ -622,9 +695,3 @@ class EvaluationState:
   def update(self, example: example_lib.Example) -> None:
     """Updates the state with the given example."""
     self._evaluated_examples[example.id] = example
-
-  @property
-  def evaluated_examples(self) -> dict[int, example_lib.Example]:
-    """Returns the examples in the state."""
-    return self._evaluated_examples
-

langfun/core/eval/v2/evaluation_test.py

@@ -133,6 +133,12 @@ class EvaluationTest(unittest.TestCase):
 
   def test_html_view(self):
     exp = test_helper.TestEvaluation()
+    exp.debug('debug message')
+    exp.info('info message')
+    exp.warning('warning message', x=1)
+    exp.error('error message', x=1)
+    exp.fatal('fatal message')
+
     self.assertIn(
         exp.id,
         exp.to_html(extra_flags=dict(card_view=True, current_run=None)).content

langfun/core/llms/vertexai.py

@@ -409,7 +409,7 @@ class VertexAIGemini2_0(VertexAI):  # pylint: disable=invalid-name
   )
 
 
-class VertexAIGeminiFlash2_0Exp(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash2_0Exp(VertexAIGemini2_0):  # pylint: disable=invalid-name
   """Vertex AI Gemini 2.0 Flash model."""
 
   model = 'gemini-2.0-flash-exp'

langfun/core/logging.py

@@ -54,6 +54,25 @@ class LogEntry(pg.Object, pg.views.HtmlTreeView.Extension):
   def should_output(self, min_log_level: LogLevel) -> bool:
     return _LOG_LEVELS.index(self.level) >= _LOG_LEVELS.index(min_log_level)
 
+  def format(self,
+             compact: bool = False,
+             verbose: bool = True,
+             root_indent: int = 0,
+             *,
+             text_format: bool = True,
+             **kwargs):
+    if text_format:
+      s = f"""{self.time.strftime('%H:%M:%S')} {self.level.upper()} - {self.message}"""
+      if self.metadata:
+        s += f' (metadata: {self.metadata!r})'
+      return s
+    return super().format(
+        compact=compact,
+        verbose=verbose,
+        root_indent=root_indent,
+        **kwargs
+    )
+
   def _html_tree_view_summary(
       self,
       view: pg.views.HtmlTreeView,

langfun/core/logging_test.py

@@ -61,6 +61,25 @@ class LoggingTest(unittest.TestCase):
       print(actual)
     self.assertEqual(actual, expected)
 
+  def test_format(self):
+    time = datetime.datetime(2024, 10, 10, 12, 30, 45)
+    self.assertEqual(
+        str(
+            logging.LogEntry(
+                level='info', message='hello\nworld',
+                time=time, metadata=dict(x=1),
+            )
+        ),
+        '12:30:45 INFO - hello\nworld (metadata: {x=1})',
+    )
+    self.assertIn(
+        'LogEntry(',
+        logging.LogEntry(
+            level='info', message='hello\nworld',
+            time=time, metadata=dict(x=1),
+        ).format(text_format=False),
+    )
+
   def test_html(self):
     time = datetime.datetime(2024, 10, 10, 12, 30, 45)
     self.assert_html_content(

langfun/core/structured/__init__.py

@@ -56,6 +56,8 @@ 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 query
+from langfun.core.structured.querying import query_and_reduce
+
 from langfun.core.structured.querying import query_prompt
 from langfun.core.structured.querying import query_output
 from langfun.core.structured.querying import query_reward

langfun/core/structured/parsing.py

@@ -270,24 +270,31 @@ def call(
   if schema in (str, None):
     return lm_output if returns_message else lm_output.text
 
+  def _chain_nl_output_message(parsing_message: lf.Message):
+    """Chain the source of the parsed output to the LM output."""
+    parsing_message.root.source = lm_output
+    parsing_message.tag('parsing-lm-output')
+    parsing_message.lm_input.tag('parsing-lm-input')
+
   # Call `parsing_lm` for structured parsing.
-  parsing_message = querying.query(
-      lm_output.text,
-      schema,
-      examples=parsing_examples,
-      lm=parsing_lm or lm,
-      include_context=parsing_include_context,
-      cache_seed=cache_seed,
-      autofix=autofix,
-      autofix_lm=autofix_lm or lm,
-      protocol=protocol,
-      returns_message=True,
-      **kwargs,
-  )
-  # Chain the source of the parsed output to the LM output.
-  parsing_message.root.source = lm_output
-  parsing_message.tag('parsing-lm-output')
-  parsing_message.lm_input.tag('parsing-lm-input')
+  try:
+    parsing_message = querying.query(
+        lm_output.text,
+        schema,
+        examples=parsing_examples,
+        lm=parsing_lm or lm,
+        include_context=parsing_include_context,
+        cache_seed=cache_seed,
+        autofix=autofix,
+        autofix_lm=autofix_lm or lm,
+        protocol=protocol,
+        returns_message=True,
+        **kwargs,
+    )
+    _chain_nl_output_message(parsing_message)
+  except mapping.MappingError as e:
+    _chain_nl_output_message(e.lm_response)
+    raise e
   return parsing_message if returns_message else parsing_message.result
 
 

langfun/core/structured/parsing_test.py

@@ -686,6 +686,31 @@ class CallTest(unittest.TestCase):
         ],
         returns_message=True,
     )
+    self.assertIn('parsing-lm-output', output.tags)
+    self.assertIn('parsing-lm-input', output.source.tags)
+    self.assertEqual(output.root.text, 'Compute 1 + 2')
+
+  def test_call_with_parsing_message_chaining_on_parsing_error(self):
+    try:
+      output = parsing.call(
+          'Compute 1 + 2',
+          int,
+          lm=fake.StaticSequence(['three']),
+          parsing_lm=fake.StaticSequence(['abc']),
+          parsing_examples=[
+              mapping.MappingExample(
+                  context='Multiple four and five',
+                  input='twenty',
+                  schema=int,
+                  output=20,
+              )
+          ],
+          returns_message=True,
+      )
+    except mapping.MappingError as e:
+      output = e.lm_response
+    self.assertIn('parsing-lm-output', output.tags)
+    self.assertIn('parsing-lm-input', output.source.tags)
     self.assertEqual(output.root.text, 'Compute 1 + 2')
 
   def test_call_with_autofix(self):

langfun/core/structured/querying.py

@@ -105,12 +105,11 @@ def _query_structure_cls(
 
 def query(
     prompt: Union[str, lf.Template, Any],
-    schema: Union[
-        schema_lib.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
-    ] = None,
+    schema: schema_lib.SchemaType |  None = None,
     default: Any = lf.RAISE_IF_HAS_ERROR,
     *,
-    lm: lf.LanguageModel | None = None,
+    lm: lf.LanguageModel | list[lf.LanguageModel] | None = None,
+    num_samples: int | list[int] = 1,
     examples: list[mapping.MappingExample] | None = None,
     cache_seed: int | None = 0,
     response_postprocess: Callable[[str], str] | None = None,
@@ -121,76 +120,207 @@ def query(
     skip_lm: bool = False,
     **kwargs,
 ) -> Any:
-  """Queries an language model for a (maybe) structured output.
+  """Query one or more language models for structured or unstructured outputs.
+
+  This is the primary API in Langfun for interacting with language models,
+  supporting natural language prompts, structured inputs, and multiple advanced
+  features.
+
+  Key Features:
+
+    - **Input**: Accepts natural language strings, structured inputs (e.g.,
+      `pg.Object`), and templates (`lf.Template`) with modality objects.
+
+    - **Output**: Returns structured outputs when `schema` is specified;
+      otherwise, outputs raw natural language (as a string).
+
+    - **Few-shot examples**: Supports structured few-shot examples with the
+      `examples` argument.
+
+    - **Multi-LM fan-out**: Sends queries to multiple language models with in
+      multiple samples in parallel,  returning a list of outputs.
 
   Examples:
 
+    Case 1: Regular natural language-based LLM query:
+
+    ```
+    lf.query('1 + 1 = ?', lm=lf.llms.Gpt4Turbo())
+
+    # Outptut: '2'
+    ```
+
+    Case 2: Query with structured output.
+
+    ```
+    lf.query('1 + 1 = ?', int, lm=lf.llms.Gpt4Turbo())
+    
+    # Output: 2
+    ```
+
+    Case 3: Query with structured input.
+
+    ```
+    class Sum(pg.Object):
+      a: int
+      b: int
+
+    lf.query(Sum(1, 1), int, lm=lf.llms.Gpt4Turbo())
+
+    # Output: 2
+    ```
+
+    Case 4: Query with input of mixed modalities.
+
     ```
-    class FlightDuration:
-      hours: int
-      minutes: int
-
-    class Flight(pg.Object):
-      airline: str
-      flight_number: str
-      departure_airport_code: str
-      arrival_airport_code: str
-      departure_time: str
-      arrival_time: str
-      duration: FlightDuration
-      stops: int
-      price: float
-
-    prompt = '''
-      Information about flight UA2631.
-      '''
-
-    r = lf.query(prompt, Flight)
-    assert isinstance(r, Flight)
-    assert r.airline == 'United Airlines'
-    assert r.departure_airport_code == 'SFO'
-    assert r.duration.hour = 7
+    class Animal(pg.Object):
+      pass
+
+    class Dog(Animal):
+      pass
+    
+    class Entity(pg.Object):
+      name: str
+
+    lf.query(
+        'What is in this {{image}} and {{objects}}?'
+        list[Entity],
+        lm=lf.llms.Gpt4Turbo()
+        image=lf.Image(path='/path/to/a/airplane.png'),
+        objects=[Dog()],
+    )
+
+    # Output: [Entity(name='airplane'), Entity(name='dog')]
+    ```
+
+    Case 5: Query with structured few-shot examples.
+    ```
+    lf.query(
+        'What is in this {{image}} and {{objects}}?'
+        list[Entity],
+        lm=lf.llms.Gpt4Turbo()
+        image=lf.Image(path='/path/to/a/dinasaur.png'),
+        objects=[Dog()],
+        examples=[
+            lf.MappingExample(
+                input=lf.Template(
+                    'What is the object near the house in this {{image}}?',
+                    image=lf.Image(path='/path/to/image.png'),
+                ),
+                schema=Entity,
+                output=Entity('cat'),
+            ),
+        ],
+    )
+
+    # Output: [Entity(name='dinasaur'), Entity(name='dog')]
+    ```
+
+    Case 6: Multiple queries to multiple models.
+    ```
+    lf.query(
+        '1 + 1 = ?',
+        int,
+        lm=[
+            lf.llms.Gpt4Turbo(),
+            lf.llms.Gemini1_5Pro(),
+        ],
+        num_samples=[1, 2],
+    )
+    # Output: [2, 2, 2]
     ```
 
   Args:
-    prompt: A str (may contain {{}} as template) as natural language input, or a
-      `pg.Symbolic` object as structured input as prompt to LLM.
-    schema: A type annotation as the schema for output object. If str (default),
-      the response will be a str in natural language.
-    default: The default value if parsing failed. If not specified, error will
-      be raised.
-    lm: The language model to use. If not specified, the language model from
-      `lf.context` context manager will be used.
-    examples: An optional list of fewshot examples for helping parsing. If None,
-      the default one-shot example will be added.
-    cache_seed: Seed for computing cache key. The cache key is determined by a
-      tuple of (lm, prompt, cache seed). If None, cache will be disabled for the
-      query even cache is configured by the LM.
-    response_postprocess: An optional callable object to process the raw LM
-      response before parsing it into the final output object. If None, the raw
-      LM response will not be processed.
-    autofix: Number of attempts to auto fix the generated code. If 0, autofix is
-      disabled. Auto-fix is not supported for 'json' protocol.
-    autofix_lm: The language model to use for autofix. If not specified, the
-      `autofix_lm` from `lf.context` context manager will be used. Otherwise it
-      will use `lm`.
-    protocol: The protocol for schema/value representation. Applicable values
-      are 'json' and 'python'. By default `python` will be used.
-    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.
-      otherwise return the LLM response based on the rendered prompt.
-    **kwargs: Keyword arguments passed to render the prompt or configure the 
-      `lf.structured.Mapping` class. Notable kwargs are:
-      - template_str: Change the root template for query.
-      - preamble: Change the preamble for query.
-      - mapping_template: Change the template for each mapping examle.
+    prompt: The input query. Can be:
+      - A natural language string (supports templating with `{{}}`),
+      - A `pg.Object` object for structured input,
+      - An `lf.Template` for mixed or template-based inputs.
+    schema: Type annotation or `lf.Schema` object for the expected output. 
+      If `None` (default), the response will be a natural language string.
+    default: Default value to return if parsing fails. If not specified, an
+      error will be raised.
+    lm: The language model(s) to query. Can be:
+      - A single `LanguageModel`,
+      - A list of `LanguageModel`s for multi-model fan-out.
+      If `None`, the LM from `lf.context` will be used.
+    num_samples: Number of samples to generate. If a list is provided, its
+      length must match the number of models in `lm`.
+    examples: Few-shot examples to guide the model output. Defaults to `None`.
+    cache_seed: Seed for caching the query. Queries with the same
+      `(lm, prompt, cache_seed)` will use cached responses. If `None`,
+      caching is disabled.
+    response_postprocess: A post-processing function for the raw LM response.
+      If `None`, no post-processing occurs.
+    autofix: Number of attempts for auto-fixing code errors. Set to `0` to
+      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'`.
+    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 
+      prompt as a `UserMessage` object.
+    **kwargs: Additional keyword arguments for:
+      - Rendering templates (e.g., `template_str`, `preamble`),
+      - Configuring `lf.structured.Mapping`.
 
   Returns:
-    The result based on the schema.
+    The result of the query:
+    - A single output or a list of outputs if multiple models/samples are used.
+    - Each output is a parsed object matching `schema`, an `lf.Message` (if 
+      `returns_message=True`), or a natural language string (default).
   """
     # Internal usage logging.
 
+  # Multiple quries will be issued when `lm` is a list or `num_samples` is
+  # greater than 1.
+  if isinstance(lm, list) or num_samples != 1:
+    def _single_query(inputs):
+      lm, example_i = inputs
+      return query(
+          prompt,
+          schema,
+          default=default,
+          lm=lm,
+          examples=examples,
+          # Usually num_examples should not be large, so we multiple the user
+          # provided cache seed by 100 to avoid collision.
+          cache_seed=(
+              None if cache_seed is None else cache_seed * 100 + example_i
+          ),
+          response_postprocess=response_postprocess,
+          autofix=autofix,
+          autofix_lm=autofix_lm,
+          protocol=protocol,
+          returns_message=returns_message,
+          skip_lm=skip_lm,
+          **kwargs,
+      )
+    lm_list = lm if isinstance(lm, list) else [lm]
+    num_samples_list = (
+        num_samples if isinstance(num_samples, list)
+        else [num_samples] * len(lm_list)
+    )
+    assert len(lm_list) == len(num_samples_list), (
+        'Expect the length of `num_samples` to be the same as the '
+        f'the length of `lm`. Got {num_samples} and {lm_list}.'
+    )
+    query_inputs = []
+    total_queries = 0
+    for lm, num_samples in zip(lm_list, num_samples_list):
+      query_inputs.extend([(lm, i) for i in range(num_samples)])
+      total_queries += num_samples
+
+    samples = []
+    for _, output, error in lf.concurrent_map(
+        _single_query, query_inputs, max_workers=max(64, total_queries),
+        ordered=True,
+    ):
+      if error is None:
+        samples.append(output)
+    return samples
+
   # Normalize query schema.
   # When `lf.query` is used for symbolic completion, schema is automatically
   # inferred when it is None.
@@ -280,11 +410,52 @@ def query(
   return output_message if returns_message else _result(output_message)
 
 
+#
+# Helper function for map-reduce style querying.
+#
+
+
+def query_and_reduce(
+    prompt: Union[str, lf.Template, Any],
+    schema: schema_lib.SchemaType | None = None,
+    *,
+    reduce: Callable[[list[Any]], Any],
+    lm: lf.LanguageModel | list[lf.LanguageModel] | None = None,
+    num_samples: int | list[int] = 1,
+    **kwargs,
+) -> Any:
+  """Issues multiple `lf.query` calls in parallel and reduce the outputs.
+  
+  Args:
+    prompt: A str (may contain {{}} as template) as natural language input, or a
+      `pg.Symbolic` object as structured input as prompt to LLM.
+    schema: A type annotation as the schema for output object. If str (default),
+      the response will be a str in natural language.
+    reduce: A function to reduce the outputs of multiple `lf.query` calls. It
+      takes a list of outputs and returns the final object.
+    lm: The language model to use. If not specified, the language model from
+      `lf.context` context manager will be used.
+    num_samples: The number of samples to obtain from each language model being
+      requested. If a list is provided, it should have the same length as `lm`.
+    **kwargs: Additional arguments to pass to `lf.query`.
+
+  Returns:
+    The reduced output from multiple `lf.query` calls.
+  """
+  results = query(prompt, schema, lm=lm, num_samples=num_samples, **kwargs)
+  if isinstance(results, list):
+    results = reduce(results)
+  return results
+
+
+#
+# Functions for decomposing `lf.query` into pre-llm and post-llm operations.
+#
+
+
 def query_prompt(
     prompt: Union[str, lf.Template, Any],
-    schema: Union[
-        schema_lib.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
-    ] = None,
+    schema: schema_lib.SchemaType | None = None,
     **kwargs,
 ) -> lf.Message:
   """Returns the final prompt sent to LLM for `lf.query`."""
@@ -295,9 +466,7 @@ def query_prompt(
 
 def query_output(
     response: Union[str, lf.Message],
-    schema: Union[
-        schema_lib.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
-    ],
+    schema: schema_lib.SchemaType | None = None,
     **kwargs,
 ) -> Any:
   """Returns the final output of `lf.query` from a provided LLM response."""
@@ -308,6 +477,11 @@ def query_output(
   )
 
 
+#
+# Functions for computing reward of an LLM response based on a mapping example.
+#
+
+
 def query_reward(
     mapping_example: Union[str, mapping.MappingExample],
     response: Union[str, lf.Message],
@@ -362,6 +536,11 @@ def _reward_fn(cls) -> Callable[
   return _reward
 
 
+#
+# Functions for tracking `lf.query` invocations.
+#
+
+
 class QueryInvocation(pg.Object, pg.views.HtmlTreeView.Extension):
   """A class to represent the invocation of `lf.query`."""
 

langfun/core/structured/querying_test.py

@@ -327,6 +327,69 @@ class QueryTest(unittest.TestCase):
         expected_modalities=3,
     )
 
+  def test_multiple_queries(self):
+    self.assertEqual(
+        querying.query(
+            'Compute 1 + 2',
+            int,
+            lm=[
+                fake.StaticResponse('1'),
+                fake.StaticResponse('2'),
+            ],
+            num_samples=[1, 2],
+        ),
+        [1, 2, 2]
+    )
+    self.assertEqual(
+        querying.query(
+            'Compute 1 + 2',
+            int,
+            lm=[
+                fake.StaticResponse('1'),
+                fake.StaticResponse('2'),
+            ],
+            num_samples=2,
+        ),
+        [1, 1, 2, 2]
+    )
+    self.assertEqual(
+        querying.query(
+            'Compute 1 + 2',
+            int,
+            lm=[
+                fake.StaticResponse('1'),
+                fake.StaticResponse('abc'),
+            ],
+            num_samples=[1, 2],
+        ),
+        [1]
+    )
+    self.assertEqual(
+        querying.query(
+            'Compute 1 + 2',
+            int,
+            default=0,
+            lm=[
+                fake.StaticResponse('1'),
+                fake.StaticResponse('abc'),
+            ],
+            num_samples=[1, 2],
+        ),
+        [1, 0, 0]
+    )
+    results = querying.query(
+        'Compute 1 + 2',
+        int,
+        default=0,
+        lm=[
+            fake.StaticResponse('1'),
+            fake.StaticResponse('abc'),
+        ],
+        returns_message=True,
+    )
+    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')
@@ -393,6 +456,30 @@ class QueryTest(unittest.TestCase):
     )
     self.assertIsNotNone(output.get_modality('image'))
 
+  def test_query_and_reduce(self):
+    self.assertEqual(
+        querying.query_and_reduce(
+            'Compute 1 + 1',
+            int,
+            reduce=sum,
+            lm=[
+                fake.StaticResponse('1'),
+                fake.StaticResponse('2'),
+            ],
+            num_samples=[1, 2],
+        ),
+        5
+    )
+    self.assertEqual(
+        querying.query_and_reduce(
+            'Compute 1 + 1',
+            int,
+            reduce=sum,
+            lm=fake.StaticResponse('2'),
+        ),
+        2
+    )
+
   def test_query_output(self):
     self.assertEqual(
         querying.query_output(

langfun/core/structured/schema.py

@@ -213,18 +213,8 @@ class Schema(
         """
     )
 
-  def _html_tree_view_tooltip(
-      self,
-      *,
-      view: pg.views.HtmlTreeView,
-      content: pg.Html | str | None = None,
-      **kwargs,
-  ):
-    return view.tooltip(
-        self,
-        content=content or pg.Html.escape(self.schema_str(protocol='python')),
-        **kwargs
-    )
+
+SchemaType = Union[Schema, Type[Any], list[Type[Any]], dict[str, Any]]
 
 
 def _top_level_object_specs_from_value(value: pg.Symbolic) -> list[Type[Any]]:

{langfun-0.1.2.dev202412170805.dist-info → langfun-0.1.2.dev202412190804.dist-info}/METADATA

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

{langfun-0.1.2.dev202412170805.dist-info → langfun-0.1.2.dev202412190804.dist-info}/RECORD

@@ -1,4 +1,4 @@
-langfun/__init__.py,sha256=Xkzi8VV93jI7iLSrypzyHV9FtRbxHQpmoUCIZJSGHdA,2400
+langfun/__init__.py,sha256=fhfPXpHN7GoGqixpFfqhQkYxFs_siP_LhbjZhd3lhio,2497
 langfun/core/__init__.py,sha256=xlvFTXc7IKUTs8aCFRFhzOLTmmeuhXgk9yx2InBLNiA,4937
 langfun/core/component.py,sha256=HVrEoTL1Y01iqOHC3FYdbAOnffqfHHtGJXoK1vkdEwo,11583
 langfun/core/component_test.py,sha256=sG-T2wpvBfHqWGZE7sc4NayJj2aj5QFBzSwFiwrGEIc,10376
@@ -10,8 +10,8 @@ langfun/core/langfunc.py,sha256=G50YgoVZ0y1GFw2ev41MlOqr6qa8YakbvNC0h_E0PiA,1114
 langfun/core/langfunc_test.py,sha256=fKIAqcSNI_7M6nwoZW77HEam8Oa6vcWhsCNgVJanzb4,8822
 langfun/core/language_model.py,sha256=b15MZ_qbydnz5vQ09t7sf9tc3C7qWvMSxUrGfT0p99I,33827
 langfun/core/language_model_test.py,sha256=hnYhtw7GM_TbhgsJzHNYTaoDewUlPHpOVlI7xEkCFuI,31783
-langfun/core/logging.py,sha256=uslllP0RTGN223oro1m4nZZ0bFppcL07OwbFKm2iG6k,7519
-langfun/core/logging_test.py,sha256=b5bPTSUoYeICATaO6I8dOVumodwRbxSp1Oz96Sf3KcE,6104
+langfun/core/logging.py,sha256=W3mLEMXdo210Q5OX3a1ZTc4nU-xMy73-IfNKnsA-RFo,8051
+langfun/core/logging_test.py,sha256=N7-YvSXC8zvnr2SNwWHOykn1CFmqvIuTLDgn41Ku9JU,6642
 langfun/core/memory.py,sha256=f-asN1F7Vehgdn_fK84v73GrEUOxRtaW934keutTKjk,2416
 langfun/core/message.py,sha256=16oiMpg9O9VKrgpfrvJrfvga3n3FzUuD_zdWb9nvSWA,25686
 langfun/core/message_test.py,sha256=jtZoNBNbA99i2fjoKg5vTRgoUe84J4MH8ZMGakGmTHs,32577
@@ -58,10 +58,10 @@ langfun/core/eval/patching_test.py,sha256=8kCd54Egjju22FMgtJuxEsrXkW8ifs-UUBHtrC
 langfun/core/eval/scoring.py,sha256=B69IsIxiPs1xZcOBFIhZF70YmDue2Siik-CPL2bh33s,6254
 langfun/core/eval/scoring_test.py,sha256=O8olHbrUEg60gMxwOkWzKBJZpZoUlmVnBANX5Se2SXM,4546
 langfun/core/eval/v2/__init__.py,sha256=qoa6zKdFXOFyCX6vay6OdgPf1eUhYGoHYAxe35qECGk,1628
-langfun/core/eval/v2/checkpointing.py,sha256=ks0az5IJCceZG8V8pcqpkaaXcO9IHnbLpjklmBM1_uQ,6257
+langfun/core/eval/v2/checkpointing.py,sha256=8vxH3AfIBS8dxA0IiOZBUxAHXIx5m2tSWSSumDLpzp8,6546
 langfun/core/eval/v2/checkpointing_test.py,sha256=dAERKQTW_PM1B0oUauB0YVQkMEI-cgJq0q-wAVlGYpU,4383
-langfun/core/eval/v2/evaluation.py,sha256=h_AWRUSKhEs-bHLBgqo-GeBYXluD5bPbAqypRW0ajfA,19441
-langfun/core/eval/v2/evaluation_test.py,sha256=hh6L2HhQPQ6NBv1pXKcNkYraNcV9MLuJ--69t9jbmaI,5846
+langfun/core/eval/v2/evaluation.py,sha256=7PC-npbEQjwwv0pWbv8vGi_OkzZ7QpJrEpYoixFBlno,21429
+langfun/core/eval/v2/evaluation_test.py,sha256=ld8oBOjsfN-LNLL2eViSTu17wAq90GcsfURXX6oVlFo,6014
 langfun/core/eval/v2/example.py,sha256=fURrvdNmMsVMqoEErcsmLmC6Xq3ny16dYsnLH8HVlcY,9626
 langfun/core/eval/v2/example_test.py,sha256=WcJmU7IQQXvjFia63mokySC4CqxzVL9Wso1sC5F0YK8,3032
 langfun/core/eval/v2/experiment.py,sha256=0JBGckJ93aqSdffpJPDVPy_I5T2BXscghTxiglHzJWo,29556
@@ -96,7 +96,7 @@ langfun/core/llms/openai.py,sha256=l49v6RubfInvV0iG114AymTKNogTX4u4N-UFCeSgIxw,2
 langfun/core/llms/openai_test.py,sha256=kOWa1nf-nJvtYY10REUw5wojh3ZgfU8tRaCZ8wUgJbA,16623
 langfun/core/llms/rest.py,sha256=sWbYUV8S3SuOg9giq7xwD-xDRfaF7NP_ig7bI52-Rj4,3442
 langfun/core/llms/rest_test.py,sha256=NZ3Nf0XQVpT9kLP5cBVo_yBHLI7vWTYhWQxYEJVMGs4,3472
-langfun/core/llms/vertexai.py,sha256=adUTByiuiTHBQ31tM_EXPUWIyUwo3zqyYIe9UILAFDE,14981
+langfun/core/llms/vertexai.py,sha256=oEd665IBwzCTlHuLEMrCdwgQzrFB5ERcnxw6nrYNSyk,14990
 langfun/core/llms/vertexai_test.py,sha256=ffcA5yPecnQy_rhkuYAw_6o1iLW8AR8FgswmHt6aAys,6725
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=rt3zwmyw0y9jsSGW-ZbV1vAfLxQ7_3AVk0l2EySlse4,3918
@@ -118,7 +118,7 @@ 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=i5Plug6041OtEltsDSQxCPFM-4XD4HKZrhw9uLQXFXk,3157
+langfun/core/structured/__init__.py,sha256=3Wb4ks14D5E5z1nQ5trXSXiLqFN5E_3HvcEJQAfFRl0,3220
 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
@@ -127,11 +127,11 @@ langfun/core/structured/function_generation.py,sha256=g7AOR_e8HxFU6n6Df750aGkgMg
 langfun/core/structured/function_generation_test.py,sha256=LaXYDXf9GlqUrR6v_gtmK_H4kxzonmU7SYbn7XXMgjU,12128
 langfun/core/structured/mapping.py,sha256=vLKH79UT-j0qkQdvqlQBO7SkXXuM-yr2Idm8_HH8qwM,13649
 langfun/core/structured/mapping_test.py,sha256=bHm2ZCXBITq_G8Lvw_olFHeUUc4s_lGXZm9v9JhoPB4,9630
-langfun/core/structured/parsing.py,sha256=lhEkdnvxKzkYwHsTvBdE2j6uLWl-J8uQu6c-3xcsBXM,11770
-langfun/core/structured/parsing_test.py,sha256=-uPiLi0cRBkf0ZycZsgLPIfRLLdwYhRbm2LHHp_pVGE,21475
-langfun/core/structured/querying.py,sha256=3RvK20G39TFrW6yV4lR9s7SbxlIpKRZ43Hey0obtJ3M,17539
-langfun/core/structured/querying_test.py,sha256=CuoClQLtEwW3Pd2otGwAoVzR0KBjT-Keppu9BOpz4mA,29705
-langfun/core/structured/schema.py,sha256=XnIKBDpPU6YmWJ0ncmoW23ha6VLyLzXf7hdY2-Urn-w,28133
+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=sXGhYtiEBac8iOkYOErGXyX8SAHSB1gg69WePhOyGxE,22759
+langfun/core/structured/querying_test.py,sha256=M9Apg83KjQUjT42K9LheBEr74DX3Inwd0YmCanA71kc,31738
+langfun/core/structured/schema.py,sha256=0VUPSfX1JEQ0xu8WvEymCKK_WSGwBNA-rQD2hATErmU,27912
 langfun/core/structured/schema_generation.py,sha256=U3nRQsqmMZg_qIVDh2fiY3K4JLfsAL1LcKzIFP1iXFg,5316
 langfun/core/structured/schema_generation_test.py,sha256=RM9s71kMNg2jTePwInkiW9fK1ACN37eyPeF8OII-0zw,2950
 langfun/core/structured/schema_test.py,sha256=RjYhwTgktQgyqAjzLvo967nTiIK9KWgP-aNGg4e7ihE,25258
@@ -148,8 +148,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.dev202412170805.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202412170805.dist-info/METADATA,sha256=b9Ans5czArE6iUCD8pENhPZzGEgwjLK5oPSF31NPj0s,8281
-langfun-0.1.2.dev202412170805.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-langfun-0.1.2.dev202412170805.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202412170805.dist-info/RECORD,,
+langfun-0.1.2.dev202412190804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202412190804.dist-info/METADATA,sha256=Zr8TfOnhdo83h3aGRNRWXTrJ54h7Sh7E-7Lj95iJVDw,8281
+langfun-0.1.2.dev202412190804.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+langfun-0.1.2.dev202412190804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202412190804.dist-info/RECORD,,