langfun 0.1.2.dev202412170805__tar.gz → 0.1.2.dev202412190804__tar.gz

{langfun-0.1.2.dev202412170805 → langfun-0.1.2.dev202412190804}/PKG-INFO

@@ -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 → langfun-0.1.2.dev202412190804}/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-0.1.2.dev202412170805 → langfun-0.1.2.dev202412190804}/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-0.1.2.dev202412170805 → langfun-0.1.2.dev202412190804}/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-0.1.2.dev202412170805 → langfun-0.1.2.dev202412190804}/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-0.1.2.dev202412170805 → langfun-0.1.2.dev202412190804}/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-0.1.2.dev202412170805 → langfun-0.1.2.dev202412190804}/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-0.1.2.dev202412170805 → langfun-0.1.2.dev202412190804}/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-0.1.2.dev202412170805 → langfun-0.1.2.dev202412190804}/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-0.1.2.dev202412170805 → langfun-0.1.2.dev202412190804}/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-0.1.2.dev202412170805 → langfun-0.1.2.dev202412190804}/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):