langfun 0.1.2.dev202501050804__py3-none-any.whl → 0.1.2.dev202501090804__py3-none-any.whl

langfun/core/__init__.py

@@ -77,11 +77,6 @@ from langfun.core.concurrent import concurrent_map
 from langfun.core.concurrent import with_context_access
 from langfun.core.concurrent import with_retry
 
-# Utility libraries for text formatting.
-from langfun.core.text_formatting import colored
-from langfun.core.text_formatting import colored_print as print  # pylint: disable=redefined-builtin
-from langfun.core.text_formatting import colored_template
-
 # Interface for natural language formattable.
 from langfun.core.natural_language import NaturalLanguageFormattable
 

langfun/core/coding/python/correction.py

@@ -40,6 +40,7 @@ def run_with_correction(
     sandbox: bool | None = None,
     timeout: int | None = 5,
     returns_code: bool = False,
+    returns_stdout: bool = False,
     outputs_intermediate: bool = False,
 ) -> Any | tuple[Any, str]:
   """Correct code with a language model via self-play.
@@ -62,6 +63,7 @@ def run_with_correction(
       timeout. Applicable only when sandbox is set to True.
     returns_code: If True, the return value is a tuple of (result, final code).
       Otherwise the return value is the result only.
+    returns_stdout: If True, the stdout (a str) will be returned.
     outputs_intermediate: If True, intermediate output will be outputted as a
       dict, with the last line's value accessible by key '__result__'. Otherwise
       the value of the last line will be returned.
@@ -87,6 +89,7 @@ def run_with_correction(
             global_vars=global_vars,
             sandbox=sandbox,
             timeout=timeout,
+            returns_stdout=returns_stdout,
             outputs_intermediate=outputs_intermediate,
         )
     )
@@ -189,9 +192,7 @@ def correct(
 def _error_feedback_str(error: Exception) -> str:
   """Returns the error str for feedback."""
   if isinstance(error, errors.CodeError):
-    return lf.text_formatting.decolored(
-        error.format(include_complete_code=False)
-    )
+    return pg.decolor(error.format(include_complete_code=False))
   else:
     return f"Encountered {error.__class__.__name__}: {error}"
 

langfun/core/coding/python/errors.py

@@ -17,7 +17,8 @@ import io
 import sys
 import textwrap
 import traceback
-import langfun.core as lf
+
+import pyglove as pg
 
 
 class CodeError(RuntimeError):
@@ -62,13 +63,13 @@ class CodeError(RuntimeError):
     if 'line' not in error_message and self.lineno is not None:
       error_message += f' (<unknown>, line {self.lineno})'
     r.write(
-        lf.colored(
+        pg.colored(
             f'{self.cause.__class__.__name__}: {error_message}', 'magenta'))
 
     if self.lineno is not None:
       r.write('\n\n')
       r.write(textwrap.indent(
-          lf.colored(
+          pg.colored(
               self.code_lines(self.lineno - 1, self.end_lineno), 'magenta'),
           ' ' * 2
       ))
@@ -76,14 +77,14 @@ class CodeError(RuntimeError):
 
     if include_complete_code:
       r.write('\n')
-      r.write(lf.colored('[Generated Code]', 'green', styles=['bold']))
+      r.write(pg.colored('[Generated Code]', 'green', styles=['bold']))
       r.write('\n\n')
-      r.write(lf.colored('  ```python\n', 'green'))
+      r.write(pg.colored('  ```python\n', 'green'))
       r.write(textwrap.indent(
-          lf.colored(self.code, 'green'),
+          pg.colored(self.code, 'green'),
           ' ' * 2
       ))
-      r.write(lf.colored('\n  ```\n', 'green'))
+      r.write(pg.colored('\n  ```\n', 'green'))
     return r.getvalue()
 
 
@@ -98,10 +99,10 @@ class SerializationError(RuntimeError):
     r = io.StringIO()
     cause_message = str(self.cause).rstrip()
     if self.message:
-      r.write(lf.colored(self.message, 'magenta'))
+      r.write(pg.colored(self.message, 'magenta'))
       r.write('\n\n')
     r.write(
-        lf.colored(
+        pg.colored(
             f'{self.cause.__class__.__name__}: {cause_message}', 'magenta'
         )
     )

langfun/core/coding/python/execution.py

@@ -57,6 +57,7 @@ def evaluate(
     *,
     global_vars: dict[str, Any] | None = None,
     permission: permissions.CodePermission | None = None,
+    returns_stdout: bool = False,
     outputs_intermediate: bool = False,
 ) -> Any | dict[str, Any]:
   """Executes Python code.
@@ -71,14 +72,17 @@ def evaluate(
     global_vars: An optional dict as the globals that could be referenced by the
       code.
     permission: Permission for the Python code to run.
-    outputs_intermediate: If True, intermediate output will be outputted as a
-      dict, with the last line's value accessible by key '__result__'. Otherwise
-      the value of the last line will be returned.
+    returns_stdout: If True, the stdout (a str) will be returned.
+    outputs_intermediate: Applicable when returns_stdout is False. If True,
+      intermediate output will be outputted as a dict, with the last line's
+      value accessible by key '__result__' and the std output accessible by
+      key '__stdout__'. Otherwise the value of the last line will be returned.
 
   Returns:
-    The value of the last line of the code. Or a dict of variable name to
-    their values if `outputs_intermediate` is set to True, with the final result
-    accessible by key '__result__'.
+    The value of the last line of the code block. Or a dict of variable
+    names of all locals to their evaluated values as the output of the code to
+    run. The value for the last line can be accessed by key '__result__'. Or the
+    stdout as a str.
   """
   # Set up the permission and context.
   permission = permission or permissions.get_permission()
@@ -136,6 +140,8 @@ def evaluate(
         raise errors.CodeError(code, e) from e
       global_vars[RESULT_KEY] = list(global_vars.values())[-1]
 
+  if returns_stdout:
+    return stdout.getvalue()
   if outputs_intermediate:
     outputs = {}
     for k, v in global_vars.items():
@@ -258,6 +264,7 @@ def run(
     *,
     global_vars: dict[str, Any] | None = None,
     permission: permissions.CodePermission | None = None,
+    returns_stdout: bool = False,
     outputs_intermediate: bool = False,
     sandbox: bool | None = None,
     timeout: float | None = None,
@@ -273,9 +280,11 @@ def run(
     code: Python code to run.
     global_vars: An optional dict of
     permission: Permission for the Python code to run.
-    outputs_intermediate: If True, all variables created as locals will be
-      returned, with the final result accessible by key '__result__'. Otherwise
-      only the final result will be returned.
+    returns_stdout: If True, the stdout (a str) will be returned.
+    outputs_intermediate: Applicable when returns_stdout is False. If True,
+      intermediate output will be outputted as a dict, with the last line's
+      value accessible by key '__result__' and the std output accessible by
+      key '__stdout__'. Otherwise the value of the last line will be returned.
     sandbox: If True, run code in sandbox; If False, run code in current
       process. If None, run in sandbox first, if the output could not be
       serialized and pass to current process, run the code again in current
@@ -285,7 +294,8 @@ def run(
   Returns:
     The value of the last line of the code block. Or a dict of variable
     names of all locals to their evaluated values as the output of the code to
-    run. The value for the last line can be accessed by key '__result__'.
+    run. The value for the last line can be accessed by key '__result__'. Or the
+    stdout as a str.
 
   Raises:
     TimeoutError: If the execution time exceeds the timeout.
@@ -293,5 +303,6 @@ def run(
   """
   return call(
       evaluate, code=code, global_vars=global_vars, permission=permission,
-      outputs_intermediate=outputs_intermediate,
-      sandbox=sandbox, timeout=timeout)
+      returns_stdout=returns_stdout, outputs_intermediate=outputs_intermediate,
+      sandbox=sandbox, timeout=timeout
+  )

langfun/core/coding/python/execution_test.py

@@ -63,6 +63,15 @@ class EvaluateTest(unittest.TestCase):
         ),
         3,
     )
+    with self.assertRaisesRegex(errors.CodeError, 'ValueError'):
+      execution.evaluate(
+          """
+          def foo():
+            raise ValueError("intentional error")
+          foo()
+          """,
+          permission=permissions.CodePermission.ALL
+      )
 
   def test_class_def(self):
     ret = execution.evaluate(
@@ -102,16 +111,20 @@ class EvaluateTest(unittest.TestCase):
     self.assertIs(ret['__result__'], ret['bar'])
 
   def test_function_def_and_call(self):
-    ret = execution.evaluate(
+    code = (
         """
         def foo(x, y):
           return x + y
 
         def bar(z):
+          print(f'z is {z}')
           return z + foo(z, z)
 
         bar(1)
-        """,
+        """
+    )
+    ret = execution.evaluate(
+        code,
         permission=permissions.CodePermission.ALL,
         outputs_intermediate=True,
     )
@@ -119,6 +132,12 @@ class EvaluateTest(unittest.TestCase):
         list(ret.keys()), ['foo', 'bar', '__result__', '__stdout__']
     )
     self.assertEqual(ret['__result__'], 3)
+    ret = execution.evaluate(
+        code,
+        permission=permissions.CodePermission.ALL,
+        returns_stdout=True,
+    )
+    self.assertEqual(ret, 'z is 1\n')
 
   def test_complex(self):
     ret = execution.evaluate(

langfun/core/coding/python/generation.py

@@ -88,6 +88,8 @@ class PythonCode(pg.Object):
       sandbox: bool | None = None,
       timeout: int | None = 5,
       global_vars: dict[str, Any] | None = None,
+      returns_stdout: bool = False,
+      outputs_intermediate: bool = False,
       autofix: int = 3,
       autofix_lm: lf.LanguageModel | None = None,
   ) -> Any:
@@ -101,13 +103,22 @@ class PythonCode(pg.Object):
       timeout: Timeout in seconds. If None, there is no timeout. Applicable when
         sandbox is set to True.
       global_vars: Global variables that could be accessed from the source code.
+      returns_stdout: If True, the stdout (a str) will be returned.
+      outputs_intermediate: Applicable when returns_stdout is False. If True,
+        intermediate output will be outputted as a dict, with the last line's
+        value accessible by key '__result__' and the std output accessible by 
+        key '__stdout__'. Otherwise the value of the last line will be returned.
       autofix: Number of attempts to auto fix the generated code. If 0, autofix
         is disabled.
       autofix_lm: Language model to be used. If not specified, it will try to
         use the `lm` under `lf.context`.
 
     Returns:
-      The value of the last expression in the source code.
+      The value of the last expression in the source code. Or a dict of local
+      variable names defined in the source code to their values if
+      `outputs_intermediate` is set to True. The value for the last line can be
+      accessed by key '__result__'. Or the stdout as a str if `returns_stdout`
+      is set to True.
 
     Raises:
       TimeoutError: If `sandbox` is True and timeout has reached.
@@ -121,6 +132,8 @@ class PythonCode(pg.Object):
         max_attempts=autofix,
         lm=autofix_lm,
         returns_code=True,
+        returns_stdout=returns_stdout,
+        outputs_intermediate=outputs_intermediate,
     )
     self.rebind(source=updated_code)
     return result
@@ -158,18 +171,14 @@ class PythonCode(pg.Object):
       TimeoutError: If `sandbox` is True and timeout has reached.
       Exception: Any errors that the source code has raised.
     """
-    result, updated_code = correction.run_with_correction(
-        self.source,
-        global_vars=global_vars,
+    return self(
         sandbox=sandbox,
         timeout=timeout,
+        global_vars=global_vars,
+        autofix=autofix,
+        autofix_lm=autofix_lm,
         outputs_intermediate=True,
-        max_attempts=autofix,
-        lm=autofix_lm,
-        returns_code=True,
     )
-    self.rebind(source=updated_code)
-    return result
 
 
 class PythonFunction(pg.Object):

langfun/core/concurrent.py

@@ -25,7 +25,6 @@ import time
 from typing import Any, Callable, Iterable, Iterator, Literal, Sequence, Tuple, Type, Union
 
 from langfun.core import component
-from langfun.core import text_formatting
 import pyglove as pg
 
 
@@ -844,10 +843,10 @@ class _ConsoleProgressControl(_ProgressControl):
   def refresh(self):
     s = io.StringIO()
     if self.label is not None:
-      s.write(text_formatting.colored(self.label, 'red', styles=['bold']))
+      s.write(pg.colored(self.label, 'red', styles=['bold']))
       s.write(': ')
     s.write(
-        text_formatting.colored(
+        pg.colored(
             '%d%% (%d/%d)' %
             (
                 self._progress * 100 // self.total,

langfun/core/console.py

@@ -15,7 +15,7 @@
 
 import sys
 from typing import Any
-from langfun.core.text_formatting import colored
+import pyglove as pg
 
 
 def write(
@@ -42,10 +42,15 @@ def write(
   """
   # Print title if present.
   if title is not None:
-    print(colored(title, styles=['bold']))
+    print(pg.colored(title, styles=['bold']))
 
   # Print body.
-  print(colored(str(value), color=color, background=background, styles=styles))
+  print(dir(pg.utils))
+  print(
+      pg.colored(
+          str(value), color=color, background=background, styles=styles
+      )
+  )
 
 
 try:

langfun/core/eval/base.py

@@ -1298,7 +1298,7 @@ class Evaluation(Evaluable):
             id=self.id,
             dir=self.dir,
             model=self.lm.model_id,
-            prompt_template=lf.text_formatting.decolored(str(self.prompt)),
+            prompt_template=pg.decolor(str(self.prompt)),
             method=self.method,
             schema_fn=str(self.schema_fn),
         ),
@@ -2110,8 +2110,7 @@ class Summary(pg.Object):
 
 def _format_error(error: Exception):
   """Formats an error into a string."""
-  return (f'({error.__class__.__name__}) '
-          + lf.text_formatting.decolored(str(error)))
+  return (f'({error.__class__.__name__}) ' + pg.decolor(str(error)))
 
 
 def _error_key(error: Exception) -> str:

langfun/core/eval/v2/reporting.py

@@ -51,6 +51,8 @@ class HtmlReporter(experiment_lib.Plugin):
     self._update_thread = None
     self._stop_update = False
     self._stop_update_experiment_ids = set()
+    self._summary_lock = None
+    self._experiment_index_lock = None
 
   def on_run_start(
       self,
@@ -61,6 +63,10 @@ class HtmlReporter(experiment_lib.Plugin):
     self._last_experiment_report_time = {leaf.id: 0 for leaf in root.leaf_nodes}
     self._stop_update = False
     self._stop_update_experiment_ids = set()
+    self._summary_lock = threading.Lock()
+    self._experiment_index_lock = {
+        leaf.id: threading.Lock() for leaf in root.leaf_nodes
+    }
     self._update_thread = threading.Thread(
         target=self._update_thread_func, args=(runner,)
     )
@@ -133,21 +139,23 @@ class HtmlReporter(experiment_lib.Plugin):
     """Maybe update the summary of current run."""
     run = runner.current_run
     def _summary():
-      run.experiment.to_html(
+      html = run.experiment.to_html(
           collapse_level=None,
           extra_flags=dict(
               current_run=run, interactive=False, card_view=True,
           )
-      ).save(
-          run.output_path_for(run.experiment, _SUMMARY_FILE)
       )
+      with self._summary_lock:
+        html.save(
+            run.output_path_for(run.experiment, _SUMMARY_FILE)
+        )
 
     if force or (time.time() - self._last_summary_time > self.summary_interval):
+      self._last_summary_time = time.time()
       if background:
         runner.background_run(_summary)
       else:
         _summary()
-      self._last_summary_time = time.time()
 
   def _maybe_update_experiment_html(
       self,
@@ -170,7 +178,8 @@ class HtmlReporter(experiment_lib.Plugin):
                   card_view=False,
               ),
           )
-          html.save(index_html_path)
+          with self._experiment_index_lock[experiment.id]:
+            html.save(index_html_path)
           experiment.info(
               f'Updated {index_html_path!r} in {t.elapse:.2f} seconds.',
           )
@@ -185,11 +194,11 @@ class HtmlReporter(experiment_lib.Plugin):
         time.time() - self._last_experiment_report_time[experiment.id]
         > self.experiment_report_interval
     ):
+      self._last_experiment_report_time[experiment.id] = time.time()
       if background:
         runner.background_run(_save)
       else:
         _save()
-      self._last_experiment_report_time[experiment.id] = time.time()
 
   def _save_example_html(
       self, runner: Runner, experiment: Experiment, example: Example

langfun/core/language_model.py

@@ -434,7 +434,10 @@ class LanguageModel(component.Component):
   def __init__(self, *args, **kwargs) -> None:
     """Overrides __init__ to pass through **kwargs to sampling options."""
 
-    sampling_options = kwargs.pop('sampling_options', LMSamplingOptions())
+    sampling_options = kwargs.pop(
+        'sampling_options',
+        pg.clone(self.__schema__.fields['sampling_options'].default_value)
+    )
     sampling_options_delta = {}
 
     for k, v in kwargs.items():
@@ -650,7 +653,7 @@ class LanguageModel(component.Component):
     """Outputs debugging information about the model."""
     title_suffix = ''
     if usage.total_tokens != 0:
-      title_suffix = console.colored(
+      title_suffix = pg.colored(
           f' (total {usage.total_tokens} tokens)', 'red'
       )
 
@@ -669,7 +672,7 @@ class LanguageModel(component.Component):
     """Outputs debugging information about the prompt."""
     title_suffix = ''
     if usage.prompt_tokens != 0:
-      title_suffix = console.colored(f' ({usage.prompt_tokens} tokens)', 'red')
+      title_suffix = pg.colored(f' ({usage.prompt_tokens} tokens)', 'red')
 
     console.write(
         # We use metadata 'formatted_text' for scenarios where the prompt text
@@ -700,7 +703,7 @@ class LanguageModel(component.Component):
     if usage.completion_tokens != 0:
       title_suffix += f'{usage.completion_tokens} tokens '
     title_suffix += f'in {elapse:.2f} seconds)'
-    title_suffix = console.colored(title_suffix, 'red')
+    title_suffix = pg.colored(title_suffix, 'red')
 
     console.write(
         str(response) + '\n',

langfun/core/language_model_test.py

@@ -117,6 +117,21 @@ class LanguageModelTest(unittest.TestCase):
     self.assertEqual(lm.sampling_options.top_k, 2)
     self.assertEqual(lm.max_attempts, 2)
 
+  def test_subclassing(self):
+
+    class ChildModel(lm_lib.LanguageModel):
+
+      sampling_options = lm_lib.LMSamplingOptions(
+          temperature=0.5, top_k=20
+      )
+
+      def _sample(self, *args, **kwargs):
+        pass
+
+    lm = ChildModel(top_k=10)
+    self.assertEqual(lm.sampling_options.temperature, 0.5)
+    self.assertEqual(lm.sampling_options.top_k, 10)
+
   def test_sample(self):
     lm = MockModel(top_k=1)
     self.assertEqual(

langfun/core/llms/__init__.py

@@ -32,16 +32,30 @@ from langfun.core.llms.rest import REST
 
 # Gemini models.
 from langfun.core.llms.google_genai import GenAI
-from langfun.core.llms.google_genai import GeminiFlash2_0ThinkingExp
+from langfun.core.llms.google_genai import GeminiFlash2_0ThinkingExp_20241219
 from langfun.core.llms.google_genai import GeminiFlash2_0Exp
-from langfun.core.llms.google_genai import GeminiExp_20241114
 from langfun.core.llms.google_genai import GeminiExp_20241206
-from langfun.core.llms.google_genai import GeminiFlash1_5
+from langfun.core.llms.google_genai import GeminiExp_20241114
 from langfun.core.llms.google_genai import GeminiPro1_5
-from langfun.core.llms.google_genai import GeminiPro
-from langfun.core.llms.google_genai import GeminiProVision
-from langfun.core.llms.google_genai import Palm2
-from langfun.core.llms.google_genai import Palm2_IT
+from langfun.core.llms.google_genai import GeminiPro1_5_002
+from langfun.core.llms.google_genai import GeminiPro1_5_001
+from langfun.core.llms.google_genai import GeminiFlash1_5
+from langfun.core.llms.google_genai import GeminiFlash1_5_002
+from langfun.core.llms.google_genai import GeminiFlash1_5_001
+from langfun.core.llms.google_genai import GeminiPro1
+
+from langfun.core.llms.vertexai import VertexAI
+from langfun.core.llms.vertexai import VertexAIGeminiFlash2_0ThinkingExp_20241219
+from langfun.core.llms.vertexai import VertexAIGeminiFlash2_0Exp
+from langfun.core.llms.vertexai import VertexAIGeminiExp_20241206
+from langfun.core.llms.vertexai import VertexAIGeminiExp_20241114
+from langfun.core.llms.vertexai import VertexAIGeminiPro1_5
+from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_002
+from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_001
+from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5
+from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_002
+from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_001
+from langfun.core.llms.vertexai import VertexAIGeminiPro1
 
 # OpenAI models.
 from langfun.core.llms.openai import OpenAI
@@ -124,28 +138,13 @@ from langfun.core.llms.groq import GroqGemma_7B_IT
 from langfun.core.llms.groq import GroqWhisper_Large_v3
 from langfun.core.llms.groq import GroqWhisper_Large_v3Turbo
 
-from langfun.core.llms.vertexai import VertexAI
-from langfun.core.llms.vertexai import VertexAIGemini2_0
-from langfun.core.llms.vertexai import VertexAIGeminiFlash2_0Exp
-from langfun.core.llms.vertexai import VertexAIGeminiFlash2_0ThinkingExp
-from langfun.core.llms.vertexai import VertexAIGemini1_5
-from langfun.core.llms.vertexai import VertexAIGeminiPro1_5
-from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_001
-from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_002
-from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_0514
-from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_0409
-from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5
-from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_001
-from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_002
-from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_0514
-from langfun.core.llms.vertexai import VertexAIGeminiPro1
-from langfun.core.llms.vertexai import VertexAIGeminiPro1Vision
-from langfun.core.llms.vertexai import VertexAIEndpoint
-
-
 # LLaMA C++ models.
 from langfun.core.llms.llama_cpp import LlamaCppRemote
 
+# DeepSeek models.
+from langfun.core.llms.deepseek import DeepSeek
+from langfun.core.llms.deepseek import DeepSeekChat
+
 # Placeholder for Google-internal imports.
 
 # Include cache as sub-module.

langfun/core/llms/cache/in_memory.py

@@ -15,6 +15,7 @@
 
 import collections
 import contextlib
+import json
 from typing import Annotated, Any, Iterator
 import langfun.core as lf
 from langfun.core.llms.cache import base
@@ -49,6 +50,11 @@ class InMemory(base.LMCacheBase):
             "Creating a new cache as cache file '%s' does not exist.",
             self.filename,
         )
+      except json.JSONDecodeError:
+        pg.logging.warning(
+            "Creating a new cache as cache file '%s' is corrupted.",
+            self.filename,
+        )
 
   def model_ids(self) -> list[str]:
     """Returns the model ids of cached queires."""

langfun/core/llms/cache/in_memory_test.py

@@ -295,6 +295,11 @@ class InMemoryLMCacheTest(unittest.TestCase):
     self.assertEqual(cache2.stats.num_updates, 2)
     cache2.save()
 
+    # Corrupted file.
+    pg.io.writefile(path, 'bad_content')
+    cache3 = in_memory.InMemory(path)
+    self.assertEqual(len(cache3), 0)
+
 
 class LmCacheTest(unittest.TestCase):