langfun 0.1.2.dev202411140804__py3-none-any.whl → 0.1.2.dev202411160804__py3-none-any.whl

langfun/__init__.py

@@ -53,6 +53,10 @@ PythonFunction = coding.PythonFunction
 from langfun.core import llms
 lm_cache = llms.cache.lm_cache
 
+from langfun.core import agentic
+Action = agentic.Action
+Session = agentic.Session
+
 from langfun.core import memories
 
 from langfun.core import modalities

langfun/core/agentic/__init__.py

@@ -0,0 +1,30 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Langfun agentic framework.."""
+
+# pylint: disable=g-bad-import-order
+# pylint: disable=g-importing-member
+# pylint: disable=g-import-not-at-top
+
+from langfun.core.agentic.action import Action
+from langfun.core.agentic.action import ActionInvocation
+from langfun.core.agentic.action import Session
+
+from langfun.core.agentic.action_eval import ActionEval
+from langfun.core.agentic.action_eval import ActionEvalV1
+
+
+# pylint: enable=g-bad-import-order
+# pylint: enable=g-importing-member
+# pylint: enable=g-import-not-at-top

langfun/core/agentic/action.py

@@ -0,0 +1,250 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Base classes for agentic actions."""
+
+import abc
+from typing import Annotated, Any, Optional, Union
+import langfun.core as lf
+import pyglove as pg
+
+
+class Action(pg.Object):
+  """Base class for agent actions."""
+
+  def _on_bound(self):
+    super()._on_bound()
+    self._result = None
+
+  @property
+  def result(self) -> Any:
+    """Returns the result of the action."""
+    return self._result
+
+  def __call__(
+      self, session: Optional['Session'] = None, **kwargs) -> Any:
+    """Executes the action."""
+    session = session or Session()
+    try:
+      session.begin(self)
+      self._result = self.call(session=session, **kwargs)
+      return self._result
+    finally:
+      session.end(self)
+
+  @abc.abstractmethod
+  def call(self, session: 'Session', **kwargs) -> Any:
+    """Subclasses to implement."""
+
+
+class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
+  """A class for capturing the invocation of an action."""
+  action: Action
+  result: Any = None
+  execution: Annotated[
+      list[Union['ActionInvocation', lf.logging.LogEntry]],
+      'Execution execution.'
+  ] = []
+
+  # Allow symbolic assignment without `rebind`.
+  allow_symbolic_assignment = True
+
+  @property
+  def logs(self) -> list[lf.logging.LogEntry]:
+    """Returns logs from execution sequence."""
+    return [v for v in self.execution if isinstance(v, lf.logging.LogEntry)]
+
+  @property
+  def child_invocations(self) -> list['ActionInvocation']:
+    """Returns child action invocations."""
+    return [v for v in self.execution if isinstance(v, ActionInvocation)]
+
+  def _html_tree_view_summary(
+      self, *, view: pg.views.html.HtmlTreeView, **kwargs
+  ):
+    if isinstance(self.action, RootAction):
+      return None
+    kwargs.pop('title')
+    return view.summary(
+        self,
+        title=view.render(
+            self.action, name='action', collapse_level=0,
+            css_classes='invocation-title',
+        ),
+        **kwargs
+    )
+
+  def _html_tree_view_content(
+      self,
+      *,
+      root_path: pg.KeyPath | None = None,
+      collapse_level: int | None = None,
+      view: pg.views.html.HtmlTreeView,
+      **kwargs
+  ):
+    prepare_phase = []
+    current_phase = prepare_phase
+    action_phases = []
+    for item in self.execution:
+      if isinstance(item, ActionInvocation):
+        current_phase = []
+        action_phases.append(current_phase)
+      current_phase.append(item)
+
+    def _render_phase(
+        phase: list[ActionInvocation | lf.logging.LogEntry]
+    ) -> pg.Html.WritableTypes:
+      return pg.Html.element(
+          'div',
+          [
+              view.render(item) for item in phase
+          ]
+      )
+
+    def _render_action_phases(
+        phases: list[list[ActionInvocation | lf.logging.LogEntry]]
+    ) -> pg.Html.WritableTypes:
+      if len(phases) == 1:
+        return _render_phase(phases[0])
+      return pg.views.html.controls.TabControl(
+          [
+              pg.views.html.controls.Tab(
+                  label=f'Step {i + 1}',
+                  content=_render_phase(phase),
+              )
+              for i, phase in enumerate(phases)
+          ],
+      )
+
+    result_name = 'final_result' if isinstance(
+        self.action, RootAction) else 'result'
+    return pg.Html.element(
+        'div',
+        [
+            view.render(
+                self.result,
+                name=result_name,
+                css_classes=[
+                    f'invocation-{result_name}'.replace('_', '-')
+                ]
+            ),
+            _render_phase(prepare_phase) if prepare_phase else None,
+            _render_action_phases(action_phases)
+        ]
+    )
+
+  @classmethod
+  def _html_tree_view_css_styles(cls) -> list[str]:
+    return super()._html_tree_view_css_styles() + [
+        """
+        details.invocation-title {
+          display: inline-block;
+          background-color: #b1f0ff;
+          border: 1px solid white;
+        }
+        details.invocation-result {
+          border: 1px solid #eee;
+        }
+        details.invocation-final-result {
+          border: 1px solid #eee;
+          background-color: #fef78f;
+        }
+        """
+    ]
+
+
+class RootAction(Action):
+  """A placeholder action for the root of the action tree."""
+
+  def call(self, session: 'Session', **kwargs) -> Any:
+    raise NotImplementedError('Shall not be called.')
+
+
+class Session(pg.Object):
+  """Session for performing an agentic task."""
+
+  root_invocation: ActionInvocation = ActionInvocation(RootAction())
+
+  def _on_bound(self):
+    super()._on_bound()
+    self._invocation_stack = [self.root_invocation]
+
+  @property
+  def final_result(self) -> Any:
+    """Returns the final result of the session."""
+    return self.root_invocation.result
+
+  @property
+  def current_invocation(self) -> ActionInvocation:
+    """Returns the current invocation."""
+    assert self._invocation_stack
+    return self._invocation_stack[-1]
+
+  def begin(self, action: Action):
+    """Signal the beginning of the execution of an action."""
+    new_invocation = ActionInvocation(pg.maybe_ref(action))
+    with pg.notify_on_change(False):
+      self.current_invocation.execution.append(new_invocation)
+    self._invocation_stack.append(new_invocation)
+
+  def end(self, action: Action):
+    """Signal the end of the execution of an action."""
+    assert self._invocation_stack
+    invocation = self._invocation_stack.pop(-1)
+    invocation.rebind(
+        result=action.result, skip_notification=True, raise_on_no_change=False
+    )
+    assert invocation.action is action, (invocation.action, action)
+    assert self._invocation_stack, self._invocation_stack
+
+    if len(self._invocation_stack) == 1:
+      self.root_invocation.rebind(
+          result=invocation.result,
+          skip_notification=True,
+          raise_on_no_change=False
+      )
+
+  def _log(self, level: lf.logging.LogLevel, message: str, **kwargs):
+    with pg.notify_on_change(False):
+      self.current_invocation.execution.append(
+          lf.logging.log(
+              level, message, indent=len(self._invocation_stack) - 1, **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)
+
+  def as_message(self) -> lf.AIMessage:
+    """Returns the session as a message."""
+    return lf.AIMessage(
+        'Agentic task session.',
+        result=self.root_invocation
+    )

langfun/core/agentic/action_eval.py

@@ -0,0 +1,150 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Evaluation (v1) for Langfun agentic actions."""
+
+import io
+import os
+from typing import Annotated, Any
+
+import langfun.core as lf
+from langfun.core import eval as lf_eval
+from langfun.core.agentic import action as action_lib
+import pyglove as pg
+
+
+class ActionEval(lf.eval.v2.Evaluation):
+  """Agent evaluation."""
+
+  action_args: Annotated[
+      dict[str, Any],
+      'Arguments to call the action.'
+  ] = {}
+
+  def process(self, example: pg.Dict) -> tuple[str, dict[str, Any]]:
+    action = example.action
+    session = action_lib.Session()
+    with lf.logging.use_log_level('fatal'):
+      action(session=session, **self.action_args)
+    return session.final_result, dict(session=session)
+
+
+#
+# TODO(daiyip): Remove V1 once V2 is fully launched.
+#
+
+
+@pg.functor()
+def _dummy_schema():
+  return int
+
+
+class ExampleView(pg.Object):
+  id: int
+  input: Any
+  output: Any
+  error: str | None = None
+
+
+class ActionEvalV1(lf_eval.Matching):
+  """Base class for action evaluations.
+
+  The input function should returns a list of pg.Dict, with `action` and
+  `groundtruth` fields.
+  """
+  # We override the schema and prompt to dummy values since they are not used.
+  schema_fn = _dummy_schema()
+  prompt = '<unused>'
+
+  def process(self, example: pg.Dict, **kwargs):
+    action = example.action
+    session = action_lib.Session()
+    action(session=session, lm=self.lm, **kwargs)
+    return session.as_message()
+
+  def answer(self, output: Any, example: pg.Dict) -> Any:
+    return output
+
+  def groundtruth(self, example: Any) -> Any:
+    return example.groundtruth
+
+  def audit(
+      self,
+      example_idx: int,
+      example: Any,
+      message: lf.Message | None,
+      error: Exception | None = None,
+      dryrun: bool = False,
+  ):
+    super().audit(example_idx, example, message, error, dryrun)
+    # Write each example to HTML.
+    if not dryrun and self.dir:
+      def _save_html():
+        ExampleView(
+            example_idx,
+            example,
+            None if message is None else message.result,
+            error
+        ).to_html(
+            collapse_level=None,
+            enable_summary_tooltip=False,
+        ).save(
+            os.path.join(self.dir, f'example_{example_idx}.html')
+        )
+      # Write HTML in a separate thread to avoid blocking the main thread.
+      lf.concurrent.get_executor(
+          'background_eval_io', max_workers=16
+      ).submit(_save_html)
+
+  def _render_mismatches(self, s: io.StringIO) -> None:
+    s.write('<h2> Mismatches (Incorrect) </h2>')
+    first_url = None
+    mismatched_ids = sorted([
+        example_idx for example_idx, *_ in self.mismatches
+    ])
+    for example_idx in mismatched_ids:
+      url = os.path.join(self.dir, f'example_{example_idx}.html')
+      if first_url is None:
+        first_url = url
+      s.write(
+          f'<a href="{url}" style="margin-right: 10px" target="example_view">'
+          f'{example_idx}</a> '
+      )
+    if first_url:
+      s.write(
+          '<iframe style="border:0;width:100%;height:100%" name="example_view"'
+          f'src="{first_url}" title="Example View"></iframe>'
+      )
+    else:
+      s.write('No mismatches found.')
+
+  def _render_matches(self, s: io.StringIO) -> None:
+    s.write('<h2> Matches (correct) </h2>')
+    first_url = None
+    matched_ids = sorted([
+        example_idx for example_idx, *_ in self.matches
+    ])
+    for example_idx in matched_ids:
+      url = os.path.join(self.dir, f'example_{example_idx}.html')
+      if first_url is None:
+        first_url = url
+      s.write(
+          f'<a href="{url}" style="margin-right: 10px">{example_idx}</a> '
+      )
+    if first_url:
+      s.write(
+          '<iframe style="border:0;width:100%;height:100%" name="example_view"'
+          f'src="{first_url}" title="Example View"></iframe>'
+      )
+    else:
+      s.write('No matches found.')

langfun/core/agentic/action_eval_test.py

@@ -0,0 +1,109 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for action evaluation."""
+
+import os
+import tempfile
+import unittest
+
+from langfun.core import eval as lf_eval
+from langfun.core import llms as lf_llms
+from langfun.core.agentic import action as action_lib
+from langfun.core.agentic import action_eval
+import pyglove as pg
+
+
+class Foo(action_lib.Action):
+  x: int
+
+  def call(self, session, **kwargs):
+    del session, kwargs
+    return self.x
+
+
+@pg.functor()
+def foo_inputs():
+  return [
+      pg.Dict(action=Foo(1), groundtruth=1),
+      pg.Dict(action=Foo(2), groundtruth=1),
+  ]
+
+
+class ActionEvalTest(unittest.TestCase):
+
+  def test_basics(self):
+
+    class FooEval(action_eval.ActionEval):
+      inputs = foo_inputs()
+      metrics = [lf_eval.v2.metrics.Match()]
+      action_args = dict(
+          lm=lf_llms.Echo()
+      )
+
+    s = FooEval()
+    root_dir = os.path.join(tempfile.gettempdir(), 'foo_eval')
+    s.run(root_dir, plugins=[])
+    self.assertEqual(s.metrics[0].matches, 0.5)
+    self.assertEqual(s.metrics[0].mismatches, 0.5)
+
+
+class ActionEvalV1Test(unittest.TestCase):
+
+  def test_basics(self):
+
+    class FooEval(action_eval.ActionEvalV1):
+      lm = lf_llms.Echo()
+      inputs = foo_inputs()
+
+    s = FooEval()
+    result = s.run(summary=False)
+    pg.print(result)
+    self.assertEqual(
+        result,
+        dict(
+            experiment_setup=dict(
+                id=s.id,
+                dir=None,
+                model='Echo',
+                prompt_template='<unused>',
+                method='query',
+                schema_fn='_dummy_schema()'
+            ),
+            cache_stats=dict(
+                use_cache=True,
+                num_queries=0,
+                num_hits=0,
+                num_updates=0,
+            ),
+            metrics=dict(
+                total=2,
+                failures=0,
+                failure_rate=0.0,
+                oop_failures=0,
+                oop_failure_rate=0.0,
+                non_oop_failures=0,
+                non_oop_failure_rate=0.0,
+                failure_breakdown={},
+                num_matches=0,
+                match_rate=0.0,
+                num_mismatches=2,
+                mismatch_rate=1.0
+            ),
+            usage=None
+        )
+    )
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/agentic/action_test.py

@@ -0,0 +1,84 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for base action."""
+
+import unittest
+
+import langfun.core as lf
+from langfun.core.agentic import action as action_lib
+
+
+class SessionTest(unittest.TestCase):
+
+  def test_basics(self):
+    test = self
+
+    class Bar(action_lib.Action):
+
+      def call(self, session, **kwargs):
+        test.assertIs(session.current_invocation.action, self)
+        session.info('Begin Bar')
+        return 2
+
+    class Foo(action_lib.Action):
+      x: int
+
+      def call(self, session, **kwargs):
+        test.assertIs(session.current_invocation.action, self)
+        session.info('Begin Foo', x=1)
+        return self.x + Bar()(session)
+
+    session = action_lib.Session()
+    root = session.root_invocation
+    self.assertIsInstance(root.action, action_lib.RootAction)
+    self.assertIs(session.current_invocation, session.root_invocation)
+    self.assertEqual(Foo(1)(session), 3)
+    self.assertEqual(len(session.root_invocation.child_invocations), 1)
+    self.assertEqual(len(session.root_invocation.child_invocations[0].logs), 1)
+    self.assertEqual(
+        len(session.root_invocation.child_invocations[0].child_invocations),
+        1
+    )
+    self.assertEqual(
+        len(session.root_invocation
+            .child_invocations[0].child_invocations[0].logs),
+        1
+    )
+    self.assertEqual(
+        len(session.root_invocation
+            .child_invocations[0].child_invocations[0].child_invocations),
+        0
+    )
+    self.assertIs(session.current_invocation, session.root_invocation)
+    self.assertIs(session.final_result, 3)
+    self.assertIn(
+        'invocation-final-result',
+        session.to_html().content,
+    )
+
+  def test_log(self):
+    session = action_lib.Session()
+    session.debug('hi', x=1, y=2)
+    session.info('hi', x=1, y=2)
+    session.warning('hi', x=1, y=2)
+    session.error('hi', x=1, y=2)
+    session.fatal('hi', x=1, y=2)
+
+  def test_as_message(self):
+    session = action_lib.Session()
+    self.assertIsInstance(session.as_message(), lf.AIMessage)
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/eval/v2/runners.py

@@ -15,6 +15,9 @@
 import abc
 import collections
 import concurrent.futures
+import random
+import threading
+import time
 from typing import Any, Annotated, Callable, Iterator
 
 from langfun import core as lf
@@ -373,6 +376,15 @@ class ParallelRunner(RunnerBase):
       'Timeout for each evaluation example.'
   ] = None
 
+  concurrent_startup_delay: Annotated[
+      tuple[int, int] | None,
+      (
+          'A range of seconds to delay the initial evaluation of each thread '
+          'in the thread pool, helping to prevent a burst in LLM QPS at '
+          'startup. If set to None, no delay will be applied.'
+      )
+  ] = None
+
   def _run(self, evaluations: list[Evaluation]) -> None:
     """Runs the evaluations in parallel."""
     def _run_group(evaluation_group: list[Evaluation]):
@@ -405,8 +417,20 @@ class ParallelRunner(RunnerBase):
       self, evaluation: Evaluation, items: Iterator[Example]
   ) -> None:
     """Override run items to run in parallel."""
+    if self.concurrent_startup_delay is not None:
+      thread_delayed = {}
+      def _evaluate_item(item: Example):
+        thread_id = threading.current_thread().ident
+        if thread_id not in thread_delayed:
+          thread_delayed[thread_id] = True
+          time.sleep(random.randint(*self.concurrent_startup_delay))
+        return self.evaluate_item(evaluation, item)
+    else:
+      def _evaluate_item(item: Example):
+        return self.evaluate_item(evaluation, item)
+
     for _, _, _ in lf.concurrent_map(
-        lambda item: self.evaluate_item(evaluation, item),
+        _evaluate_item,
         items,
         max_workers=evaluation.max_workers,
         timeout=self.timeout,

langfun/core/eval/v2/runners_test.py

@@ -198,7 +198,9 @@ class RunnerTest(unittest.TestCase):
     )
     # Global cache.
     root_dir = os.path.join(tempfile.gettempdir(), 'global_cache')
-    run = exp.run(root_dir, runner='sequential', use_cache='global', plugins=[])
+    run = exp.run(
+        root_dir, 'new', runner='sequential', use_cache='global', plugins=[]
+    )
     self.assertTrue(pg.io.path_exists(run.output_path_for(exp, 'cache.json')))
     self.assertEqual(exp.usage_summary.cached.total.num_requests, 4)
     self.assertEqual(exp.usage_summary.uncached.total.num_requests, 2)
@@ -206,7 +208,8 @@ class RunnerTest(unittest.TestCase):
     # Per-dataset cache.
     root_dir = os.path.join(tempfile.gettempdir(), 'per_dataset')
     run = exp.run(
-        root_dir, runner='sequential', use_cache='per_dataset', plugins=[]
+        root_dir, 'new', runner='sequential',
+        use_cache='per_dataset', plugins=[]
     )
     for leaf in exp.leaf_nodes:
       self.assertTrue(
@@ -226,6 +229,9 @@ class RunnerTest(unittest.TestCase):
     self.assertEqual(exp.usage_summary.cached.total.num_requests, 0)
     self.assertEqual(exp.usage_summary.uncached.total.num_requests, 6)
 
+
+class ParallelRunnerTest(RunnerTest):
+
   def test_parallel_runner(self):
     plugin = TestPlugin()
     exp = test_helper.test_experiment()
@@ -266,6 +272,22 @@ class RunnerTest(unittest.TestCase):
         self.assertEqual(node.progress.num_failed, 0)
         self.assertEqual(node.progress.num_processed, node.progress.num_total)
 
+  def test_concurrent_startup_delay(self):
+    plugin = TestPlugin()
+    exp = test_helper.test_experiment()
+    root_dir = os.path.join(
+        tempfile.gettempdir(), 'test_concurrent_startup_delay'
+    )
+    _ = exp.run(
+        root_dir,
+        runner='parallel',
+        plugins=[plugin],
+        concurrent_startup_delay=(0, 5),
+    )
+
+
+class DebugRunnerTest(RunnerTest):
+
   def test_debug_runner(self):
     plugin = TestPlugin()
     exp = test_helper.test_experiment()

langfun/core/language_model.py

@@ -124,17 +124,18 @@ class LMSamplingUsage(pg.Object):
   def __add__(self, other: Optional['LMSamplingUsage']) -> 'LMSamplingUsage':
     if other is None:
       return self
+    if self.estimated_cost is None:
+      estimated_cost = other.estimated_cost
+    elif other.estimated_cost is None:
+      estimated_cost = self.estimated_cost
+    else:
+      estimated_cost = self.estimated_cost + other.estimated_cost
     return LMSamplingUsage(
         prompt_tokens=self.prompt_tokens + other.prompt_tokens,
         completion_tokens=self.completion_tokens + other.completion_tokens,
         total_tokens=self.total_tokens + other.total_tokens,
         num_requests=self.num_requests + other.num_requests,
-        estimated_cost=(
-            self.estimated_cost + other.estimated_cost    # pylint: disable=g-long-ternary
-            if (self.estimated_cost is not None
-                and other.estimated_cost is not None)
-            else None
-        )
+        estimated_cost=estimated_cost,
     )
 
   def __radd__(self, other: Optional['LMSamplingUsage']) -> 'LMSamplingUsage':
@@ -956,7 +957,9 @@ class UsageSummary(pg.Object, pg.views.HtmlTreeView.Extension):
     if self._usage_badge is not None:
       self._usage_badge.update(
           self._badge_text(),
-          tooltip=pg.format(self.total, verbose=False),
+          tooltip=pg.format(
+              self, verbose=False, custom_format=self._tooltip_format
+          ),
           styles=dict(color=self._badge_color()),
       )
 
@@ -978,6 +981,14 @@ class UsageSummary(pg.Object, pg.views.HtmlTreeView.Extension):
     green = int(255 * (1 - normalized_value))
     return f'rgb({red}, {green}, 0)'
 
+  def _tooltip_format(self, v, root_indent):
+    del root_indent
+    if isinstance(v, int):
+      return f'{v:,}'
+    if isinstance(v, float):
+      return f'{v:,.3f}'
+    return None
+
   def _html_tree_view(
       self,
       *,
@@ -993,7 +1004,9 @@ class UsageSummary(pg.Object, pg.views.HtmlTreeView.Extension):
       if usage_badge is None:
         usage_badge = pg.views.html.controls.Badge(
             self._badge_text(),
-            tooltip=pg.format(self.total, verbose=False),
+            tooltip=pg.format(
+                self, custom_format=self._tooltip_format, verbose=False
+            ),
             css_classes=['usage-summary'],
             styles=dict(color=self._badge_color()),
             interactive=True,

langfun/core/language_model_test.py

@@ -744,6 +744,13 @@ class LMSamplingUsageTest(unittest.TestCase):
     self.assertEqual(usage1 + usage2, usage1 + usage2)
     self.assertIs(usage1 + None, usage1)
     self.assertIs(None + usage1, usage1)
+    usage3 = lm_lib.LMSamplingUsage(100, 200, 300, 4, None)
+    self.assertEqual(
+        usage1 + usage3, lm_lib.LMSamplingUsage(200, 400, 600, 8, 5.0)
+    )
+    self.assertEqual(
+        usage3 + usage1, lm_lib.LMSamplingUsage(200, 400, 600, 8, 5.0)
+    )
 
   def test_usage_not_available(self):
     usage_not_available = lm_lib.UsageNotAvailable()

langfun/core/llms/__init__.py

@@ -100,6 +100,7 @@ from langfun.core.llms.anthropic import Claude3Sonnet
 from langfun.core.llms.anthropic import Claude3Haiku
 from langfun.core.llms.anthropic import VertexAIAnthropic
 from langfun.core.llms.anthropic import VertexAIClaude3_5_Sonnet_20241022
+from langfun.core.llms.anthropic import VertexAIClaude3_5_Sonnet_20240620
 from langfun.core.llms.anthropic import VertexAIClaude3_5_Haiku_20241022
 
 from langfun.core.llms.groq import Groq

langfun/core/llms/anthropic.py

@@ -53,6 +53,13 @@ SUPPORTED_MODELS_AND_SETTINGS = {
         cost_per_1k_input_tokens=0.003,
         cost_per_1k_output_tokens=0.015,
     ),
+    'claude-3-5-sonnet@20240620': pg.Dict(
+        max_tokens=8192,
+        rpm=1000,
+        tpm=100000,
+        cost_per_1k_input_tokens=0.003,
+        cost_per_1k_output_tokens=0.015,
+    ),
     'claude-3-5-haiku@20241022': pg.Dict(
         max_tokens=8192,
         rpm=1000,
@@ -459,6 +466,11 @@ class VertexAIClaude3_5_Sonnet_20241022(VertexAIAnthropic):  # pylint: disable=i
   model = 'claude-3-5-sonnet-v2@20241022'
 
 
+class VertexAIClaude3_5_Sonnet_20240620(VertexAIAnthropic):  # pylint: disable=invalid-name
+  """Anthropic's Claude 3.5 Sonnet model on VertexAI."""
+  model = 'claude-3-5-sonnet@20240620'
+
+
 class VertexAIClaude3_5_Haiku_20241022(VertexAIAnthropic):  # pylint: disable=invalid-name
   """Anthropic's Claude 3.5 Haiku model on VertexAI."""
   model = 'claude-3-5-haiku@20241022'

langfun/core/llms/vertexai_test.py

@@ -199,6 +199,12 @@ class VertexAITest(unittest.TestCase):
     # There is a discrepancy between the `property_ordering` in the
     # Google-internal version and the open-source version.
     actual['response_schema'].pop('property_ordering', None)
+    if pg.KeyPath.parse('response_schema.type_').get(actual):
+      actual['response_schema']['type'] = actual['response_schema'].pop('type_')
+    if pg.KeyPath.parse('response_schema.properties.name.type_').get(actual):
+      actual['response_schema']['properties']['name']['type'] = actual[
+          'response_schema']['properties']['name'].pop('type_')
+
     self.assertEqual(
         actual,
         dict(
@@ -209,9 +215,9 @@ class VertexAITest(unittest.TestCase):
             stop_sequences=['\n'],
             response_mime_type='application/json',
             response_schema={
-                'type_': 'OBJECT',
+                'type': 'OBJECT',
                 'properties': {
-                    'name': {'type_': 'STRING'}
+                    'name': {'type': 'STRING'}
                 },
                 'required': ['name'],
                 'title': 'Person',

{langfun-0.1.2.dev202411140804.dist-info → langfun-0.1.2.dev202411160804.dist-info}/METADATA

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

{langfun-0.1.2.dev202411140804.dist-info → langfun-0.1.2.dev202411160804.dist-info}/RECORD

@@ -1,4 +1,4 @@
-langfun/__init__.py,sha256=mCES7t3R7Z-ZQYvG38-yrVqZubrXNfGCa8tI5HGB7mE,2274
+langfun/__init__.py,sha256=o_HvoQggla5uqNA7uF1126aZhayHnVNP__nd_t5ElEQ,2358
 langfun/core/__init__.py,sha256=xlvFTXc7IKUTs8aCFRFhzOLTmmeuhXgk9yx2InBLNiA,4937
 langfun/core/component.py,sha256=HVrEoTL1Y01iqOHC3FYdbAOnffqfHHtGJXoK1vkdEwo,11583
 langfun/core/component_test.py,sha256=sG-T2wpvBfHqWGZE7sc4NayJj2aj5QFBzSwFiwrGEIc,10376
@@ -8,8 +8,8 @@ langfun/core/console.py,sha256=Fra2_MSWZbFh6rY8HZoYgpGLsrNvhaGuL03znOwQbhM,2529
 langfun/core/console_test.py,sha256=pBOcuNMJdVELywvroptfcRtJMsegMm3wSlHAL2TdxVk,1679
 langfun/core/langfunc.py,sha256=G50YgoVZ0y1GFw2ev41MlOqr6qa8YakbvNC0h_E0PiA,11140
 langfun/core/langfunc_test.py,sha256=fKIAqcSNI_7M6nwoZW77HEam8Oa6vcWhsCNgVJanzb4,8822
-langfun/core/language_model.py,sha256=UtLvclKx55_SAKZ-ajaquudLxKorTARAeZFme5IaPi8,33499
-langfun/core/language_model_test.py,sha256=td81wm4zFPeMb16nmIuIZ6eHtpYhH0k6IeiYLfGgR0o,31525
+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/memory.py,sha256=f-asN1F7Vehgdn_fK84v73GrEUOxRtaW934keutTKjk,2416
@@ -29,6 +29,11 @@ langfun/core/template.py,sha256=_Sae_WsRo_yvwul0nqAPTOa0NOjW1zNYbW0CQpvg7l0,2538
 langfun/core/template_test.py,sha256=Qokz1hQFhRYaTZWBWGqvPJ0NXC9B9ennUpnRYHEf0hE,20542
 langfun/core/text_formatting.py,sha256=d7t9vaY6aCn1dkfkikpNYnBy5E_i93vHbfyDWFclGZU,5284
 langfun/core/text_formatting_test.py,sha256=ck0Xzdd4YF4CtCUj7VE0GybfbAyKQ8p3xkM1FBGrqIk,2096
+langfun/core/agentic/__init__.py,sha256=ndoDX0sAYsa3eVdXuu6nB-a-BH5TaK3urW6zAaFiyVs,1110
+langfun/core/agentic/action.py,sha256=Am5E1EH1ZBAhzagbnDVRnR4vBzI4H6MEtQ58laSPfTg,7515
+langfun/core/agentic/action_eval.py,sha256=ZtjTh34S7XPIUqandQ0YwAtzw-S7ofuZ7rRXnRbUMdQ,4424
+langfun/core/agentic/action_eval_test.py,sha256=tRUkWmOE9p0rpNOq19xAY2oDEnYsEEykjg6sUpAwJk0,2832
+langfun/core/agentic/action_test.py,sha256=CBsUQICD8yPCDUBBFouSkZuyLAcK_C-AWYc28Zts10E,2624
 langfun/core/coding/__init__.py,sha256=5utju_fwEsImaiftx4oXKl9FAM8p281k8-Esdh_-m1w,835
 langfun/core/coding/python/__init__.py,sha256=MJ-vubliz-ebrZH3OBRKBwMi0S9-FrhGCp8YQLR6_I4,1776
 langfun/core/coding/python/correction.py,sha256=WiBdoScL-6C___iA3Tg3vizuYtJWI-_4wy9zcMfVpj8,7020
@@ -71,11 +76,11 @@ langfun/core/eval/v2/progress_tracking.py,sha256=1imwSbllxHWG3zYrzo2NvytBZsVtjqu
 langfun/core/eval/v2/progress_tracking_test.py,sha256=eY2HvZeEXDA5Zyfi2m5NDWO_9kSfQsaAOEcIhkSbWCY,1874
 langfun/core/eval/v2/reporting.py,sha256=TGkli1IDwqfqsCJ_WslOMGk_24JDg7oRRTGXlAJlWpc,4361
 langfun/core/eval/v2/reporting_test.py,sha256=JxffbUPWInUyLjo-AQVFrllga884Mdfm05R86FtxSss,1482
-langfun/core/eval/v2/runners.py,sha256=2OHAVTbqq9hZ3qZpUEvQ--9X-Cr_z8Ghc3MRXCfclpk,13442
-langfun/core/eval/v2/runners_test.py,sha256=s3GgWA-H9x0JyPhPZq2s9-5GXGHo5dSbDD-4faX0h_E,11164
+langfun/core/eval/v2/runners.py,sha256=zJmu-amUiYv1g0Ek4c3mXkBgp-AFvSF7WpXVZCCf7Y4,14245
+langfun/core/eval/v2/runners_test.py,sha256=UeiUNygux_U6iGVG18rhp68ZE4hoWeoT6XsXvSjxNQg,11620
 langfun/core/eval/v2/test_helper.py,sha256=pDpZTBnWRR5xjJv3Uy3NWEzArqlL8FTMOgeR4C53F5M,2348
-langfun/core/llms/__init__.py,sha256=i0m-fVpwuIN_Jno1M-5O9ikzbVbvXWJKFQZO22MFPq8,6272
-langfun/core/llms/anthropic.py,sha256=XPQxjfe9O4b-CygCgqvQU0MPSfe1rU7uErNbo8zth7Q,13606
+langfun/core/llms/__init__.py,sha256=uR2vLghsnZqY6OjZKAs9Lo-YFNxZNunf3A0q6-1GYlc,6346
+langfun/core/llms/anthropic.py,sha256=uJXVgaFONL8okOSVQ4VGMGht_VZ30m1hoLzmDbIjmks,13990
 langfun/core/llms/anthropic_test.py,sha256=-2U4kc_pgBM7wqxu8RuxzyHPGww1EAWqKUvN4PW8Btw,8058
 langfun/core/llms/compositional.py,sha256=csW_FLlgL-tpeyCOTVvfUQkMa_zCN5Y2I-YbSNuK27U,2872
 langfun/core/llms/compositional_test.py,sha256=4eTnOer-DncRKGaIJW2ZQQMLnt5r2R0UIx_DYOvGAQo,2027
@@ -92,7 +97,7 @@ langfun/core/llms/openai_test.py,sha256=_8cd3VRNEUfE0-Ko1RiM6MlC5hjalRj7nYTJNhG1
 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=-KB880Ovab6CQqI-Y5Y6V7RlEA0tAIazmnnG74Ebp4A,18866
-langfun/core/llms/vertexai_test.py,sha256=7uBVOF5VF86xQ9HFAbSTh4J-0NjYLnuotBS1YRm-vgw,10529
+langfun/core/llms/vertexai_test.py,sha256=I8gEHLRXZZGq_d2VDtJAkAIzf-lNSCoB8y2lwFckY-w,10885
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=rt3zwmyw0y9jsSGW-ZbV1vAfLxQ7_3AVk0l2EySlse4,3918
 langfun/core/llms/cache/in_memory.py,sha256=l6b-iU9OTfTRo9Zmg4VrQIuArs4cCJDOpXiEpvNocjo,5004
@@ -143,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.dev202411140804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202411140804.dist-info/METADATA,sha256=LNyNk_qsiVz-CAbbtkN4jzdSwDWGhlR5RkefV6lclFA,8890
-langfun-0.1.2.dev202411140804.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
-langfun-0.1.2.dev202411140804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202411140804.dist-info/RECORD,,
+langfun-0.1.2.dev202411160804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202411160804.dist-info/METADATA,sha256=iHQLFR3kun6zscZZLnzVl_mxeTVMB7-C4bCXS_dLNm8,8890
+langfun-0.1.2.dev202411160804.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
+langfun-0.1.2.dev202411160804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202411160804.dist-info/RECORD,,