langfun 0.1.2.dev202503180804__tar.gz → 0.1.2.dev202503200804__tar.gz

{langfun-0.1.2.dev202503180804 → langfun-0.1.2.dev202503200804}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: langfun
-Version: 0.1.2.dev202503180804
+Version: 0.1.2.dev202503200804
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors
@@ -53,6 +53,7 @@ Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: keywords
 Dynamic: license
+Dynamic: license-file
 Dynamic: provides-extra
 Dynamic: requires-dist
 Dynamic: summary

{langfun-0.1.2.dev202503180804 → langfun-0.1.2.dev202503200804}/langfun/core/agentic/__init__.py

@@ -19,6 +19,8 @@
 
 from langfun.core.agentic.action import Action
 from langfun.core.agentic.action import ActionInvocation
+from langfun.core.agentic.action import ExecutionTrace
+from langfun.core.agentic.action import ParallelExecutions
 from langfun.core.agentic.action import Session
 
 from langfun.core.agentic.action_eval import ActionEval

{langfun-0.1.2.dev202503180804 → langfun-0.1.2.dev202503200804}/langfun/core/agentic/action.py

@@ -16,10 +16,10 @@
 import abc
 import contextlib
 import datetime
+import threading
 import time
-
 import typing
-from typing import Annotated, Any, ContextManager, Iterable, Iterator, Optional, Type, Union
+from typing import Annotated, Any, Callable, Iterable, Iterator, Optional, Type, Union
 import langfun.core as lf
 from langfun.core import structured as lf_structured
 import pyglove as pg
@@ -90,6 +90,7 @@ TracedItem = Union[
     lf_structured.QueryInvocation,
     'ActionInvocation',
     'ExecutionTrace',
+    'ParallelExecutions',
     # NOTE(daiyip): Consider remove log entry once we migrate existing agents.
     lf.logging.LogEntry,
 ]
@@ -206,6 +207,9 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
         child_items.append(item)
       elif isinstance(item, ExecutionTrace):
         child_items.extend(item._child_items(item_cls))  # pylint: disable=protected-access
+      elif isinstance(item, ParallelExecutions):
+        for branch in item.branches:
+          child_items.extend(branch._child_items(item_cls))  # pylint: disable=protected-access
     return child_items
 
   def _all_child_items(self, item_cls: Type[Any]) -> list[Any]:
@@ -217,6 +221,9 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
         child_items.extend(item.execution._all_child_items(item_cls))  # pylint: disable=protected-access
       elif isinstance(item, ExecutionTrace):
         child_items.extend(item._all_child_items(item_cls))  # pylint: disable=protected-access
+      elif isinstance(item, ParallelExecutions):
+        for branch in item.branches:
+          child_items.extend(branch._all_child_items(item_cls))  # pylint: disable=protected-access
     return child_items
 
   def append(self, item: TracedItem) -> None:
@@ -334,6 +341,8 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
       css_class = 'log'
     elif isinstance(item, ExecutionTrace):
       css_class = 'phase'
+    elif isinstance(item, ParallelExecutions):
+      css_class = 'parallel'
     else:
       raise ValueError(f'Unsupported item type: {type(item)}')
 
@@ -383,6 +392,11 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
           item.name or 'Phase',
           tooltip=f'Execution phase {item.name!r}.'
       )
+    elif isinstance(item, ParallelExecutions):
+      return pg.views.html.controls.Label(
+          item.name or 'Parallel',
+          tooltip='Parallel executions.'
+      )
     else:
       raise ValueError(f'Unsupported item type: {type(item)}')
 
@@ -401,6 +415,12 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
           color: purple;
           padding: 10px;
         }
+        .tab-button.parallel > ::before {
+          content: "||";
+          font-weight: bold;
+          color: blue;
+          padding: 10px;
+        }
         .tab-button.query > ::before {
           content: "Q";
           font-weight: bold;
@@ -439,6 +459,79 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
     ]
 
 
+class ParallelExecutions(pg.Object, pg.views.html.HtmlTreeView.Extension):
+  """A class for encapsulating parallel execution traces."""
+
+  name: Annotated[
+      str | None,
+      'The name of the parallel execution.'
+  ] = None
+
+  branches: Annotated[
+      list[ExecutionTrace],
+      'The branches of the parallel execution.'
+  ] = []
+
+  def __getitem__(self, key: int) -> ExecutionTrace:
+    return self.branches[key]
+
+  def __len__(self) -> int:
+    return len(self.branches)
+
+  def _on_bound(self):
+    super()._on_bound()
+    self._tab_control = None
+    self._lock = threading.Lock()
+
+  def add(self) -> ExecutionTrace:
+    """Appends a branch to the parallel execution."""
+    with self._lock, pg.notify_on_change(False):
+      branch = ExecutionTrace(name=f'[{len(self)}]')
+      self.branches.append(branch)
+      if self._tab_control is not None:
+        self._tab_control.append(self._branch_tab(branch))
+      return branch
+
+  #
+  # HTML views.
+  #
+
+  def _html_tree_view_summary(
+      self,
+      *,
+      name: str | None = None,
+      extra_flags: dict[str, Any] | None = None,
+      view: pg.views.html.HtmlTreeView, **kwargs
+  ):
+    return None
+
+  def _html_tree_view_content(
+      self,
+      *,
+      extra_flags: dict[str, Any] | None = None,
+      **kwargs
+  ):
+    del kwargs
+    extra_flags = extra_flags or {}
+    interactive = extra_flags.get('interactive', True)
+    if interactive or self.branches:
+      self._tab_control = pg.views.html.controls.TabControl(
+          [self._branch_tab(branch) for branch in self.branches],
+          tab_position='left'
+      )
+      return self._tab_control.to_html()
+    return '(no tracked parallel executions)'
+
+  def _branch_tab(self, branch: ExecutionTrace) -> pg.views.html.controls.Tab:
+    return pg.views.html.controls.Tab(
+        label=pg.views.html.controls.Label(
+            branch.name,
+            tooltip=f'Execution thread {branch.name!r}.'
+        ),
+        content=pg.view(branch),
+    )
+
+
 class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
   """A class for capturing the invocation of an action."""
   action: Action
@@ -463,28 +556,8 @@ class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
 
   def _on_bound(self):
     super()._on_bound()
-    self._current_phase = self.execution
     self._tab_control = None
 
-  @property
-  def current_phase(self) -> ExecutionTrace:
-    """Returns the current execution phase."""
-    return self._current_phase
-
-  @contextlib.contextmanager
-  def phase(self, name: str) -> Iterator[ExecutionTrace]:
-    """Context manager for starting a new execution phase."""
-    phase = ExecutionTrace(name=name)
-    phase.start()
-    parent_phase = self._current_phase
-    self._current_phase.append(phase)
-    self._current_phase = phase
-    try:
-      yield phase
-    finally:
-      phase.stop()
-      self._current_phase = parent_phase
-
   @property
   def logs(self) -> list[lf.logging.LogEntry]:
     """Returns immediate child logs from execution sequence."""
@@ -679,52 +752,80 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       'An optional identifier for the sessin, which will be used for logging.'
   ] = None
 
-  def _on_bound(self):
-    super()._on_bound()
-    self._current_action = self.root
+  # NOTE(daiyip): Action execution may involve multi-threading, hence current
+  # action and execution are thread-local.
 
   @property
-  def final_result(self) -> Any:
-    """Returns the final result of the session."""
-    return self.root.result
+  def _current_action(self) -> ActionInvocation:
+    """Returns the current invocation."""
+    return getattr(self._tls, '__current_action__')
+
+  @_current_action.setter
+  def _current_action(self, value: ActionInvocation):
+    setattr(self._tls, '__current_action__', value)
 
   @property
-  def current_action(self) -> ActionInvocation:
-    """Returns the current invocation."""
-    return self._current_action
+  def _current_execution(self) -> ExecutionTrace:
+    """Returns the current execution."""
+    return getattr(self._tls, '__current_execution__')
 
-  def add_metadata(self, **kwargs: Any) -> None:
-    """Adds metadata to the current invocation."""
-    with pg.notify_on_change(False):
-      self._current_action.metadata.update(kwargs)
+  @_current_execution.setter
+  def _current_execution(self, value: ExecutionTrace):
+    setattr(self._tls, '__current_execution__', value)
 
-  def phase(self, name: str) -> ContextManager[ExecutionTrace]:
-    """Context manager for starting a new execution phase."""
-    return self.current_action.phase(name)
+  def _on_bound(self):
+    super()._on_bound()
+    self._tls = threading.local()
+    self._current_action = self.root
+    self._current_execution = self.root.execution
+
+  #
+  # Context-manager for information tracking.
+  #
 
   @contextlib.contextmanager
   def track_action(self, action: Action) -> Iterator[ActionInvocation]:
     """Track the execution of an action."""
-    if not self.root.execution.has_started:
-      self.root.start()
+    if not self._current_execution.has_started:
+      self._current_execution.start()
 
     invocation = ActionInvocation(pg.maybe_ref(action))
     parent_action = self._current_action
-    parent_action.current_phase.append(invocation)
+    parent_execution = self._current_execution
+    parent_execution.append(invocation)
 
     try:
       self._current_action = invocation
+      self._current_execution = invocation.execution
       # Start the execution of the current action.
       self._current_action.start()
       yield invocation
     finally:
       # Stop the execution of the current action.
-      self._current_action.end(action.result, action.metadata)
+      invocation.end(action.result, action.metadata)
+      self._current_execution = parent_execution
       self._current_action = parent_action
       if parent_action is self.root:
-        parent_action.end(
-            result=action.result, metadata=action.metadata,
-        )
+        parent_action.end(result=action.result, metadata=action.metadata)
+
+  @contextlib.contextmanager
+  def track_phase(self, name: str | None) -> Iterator[ExecutionTrace]:
+    """Context manager for starting a new execution phase."""
+    parent_execution = self._current_execution
+    if name is None:
+      phase = parent_execution
+    else:
+      phase = ExecutionTrace(name=name)
+      phase.start()
+      parent_execution.append(phase)
+
+    try:
+      self._current_execution = phase
+      yield phase
+    finally:
+      if phase is not parent_execution:
+        phase.stop()
+        self._current_execution = parent_execution
 
   @contextlib.contextmanager
   def track_queries(
@@ -741,12 +842,21 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       A list of `lf.QueryInvocation` objects, each for a single `lf.query`
       call.
     """
-    with self.phase(phase) if phase else contextlib.nullcontext():
+    with self.track_phase(phase) as execution:
       with lf_structured.track_queries(include_child_scopes=False) as queries:
         try:
           yield queries
         finally:
-          self._current_action.current_phase.extend(queries)
+          execution.extend(queries)
+
+  #
+  # Operations with activity tracking.
+  #
+
+  def add_metadata(self, **kwargs: Any) -> None:
+    """Adds metadata to the current invocation."""
+    with pg.notify_on_change(False):
+      self._current_action.metadata.update(kwargs)
 
   def query(
       self,
@@ -785,7 +895,7 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       lm: The language model to use for the query.
       examples: The examples to use for the query.
       **kwargs: Additional keyword arguments to pass to `lf.query`.
-    
+
     Returns:
       The result of the query.
     """
@@ -799,6 +909,49 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
           **kwargs
       )
 
+  def concurrent_map(
+      self,
+      func: Callable[[Any], Any],
+      parallel_inputs: Iterable[Any],
+      *,
+      phase: str | None = None,
+      max_workers: int = 32,
+      timeout: int | None = None,
+      silence_on_errors: Union[
+          Type[BaseException], tuple[Type[BaseException], ...], None
+      ] = Exception
+  ) -> Iterator[Any]:
+    """Starts and tracks parallel execution with `lf.concurrent_map`."""
+    parallel_inputs = list(parallel_inputs)
+    parallel_execution = ParallelExecutions(name=phase)
+    self._current_execution.append(parallel_execution)
+    parent_action = self._current_action
+
+    def _map_single(input_value):
+      execution = parallel_execution.add()
+
+      # This happens on a new thread. Therefore, we update the thread-local
+      # states from the parent thread.
+      self._current_execution = execution
+      self._current_action = parent_action
+      execution.start()
+      try:
+        with self.track_queries():
+          return func(input_value)
+      finally:
+        execution.stop()
+
+    for input_value, result, error in lf.concurrent_map(
+        _map_single,
+        parallel_inputs,
+        max_workers=max_workers,
+        timeout=timeout,
+        silence_on_errors=silence_on_errors
+    ):
+      yield input_value, result, error
+
+  # NOTE(daiyip): Clean up `query_prompt` and `query_output` once TS
+  # code migration is done.
   def query_prompt(
       self,
       prompt: Union[str, lf.Template, Any],
@@ -878,7 +1031,7 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       return lf_structured.query_output(response, schema=schema, **kwargs)
 
   def _log(self, level: lf.logging.LogLevel, message: str, **kwargs):
-    self._current_action.current_phase.append(
+    self._current_execution.append(
         lf.logging.LogEntry(
             level=level,
             time=datetime.datetime.now(),
@@ -914,6 +1067,16 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
         result=self.root
     )
 
+  @property
+  def final_result(self) -> Any:
+    """Returns the final result of the session."""
+    return self.root.result
+
+  @property
+  def current_action(self) -> ActionInvocation:
+    """Returns the current invocation."""
+    return self._current_action
+
   #
   # HTML views.
   #

{langfun-0.1.2.dev202503180804 → langfun-0.1.2.dev202503200804}/langfun/core/agentic/action_test.py

@@ -41,12 +41,23 @@ class SessionTest(unittest.TestCase):
 
       def call(self, session, *, lm, **kwargs):
         test.assertIs(session.current_action.action, self)
-        with session.phase('prepare'):
+        with session.track_phase('prepare'):
           session.info('Begin Foo', x=1)
           session.query('foo', lm=lm)
         with session.track_queries():
           self.make_additional_query(lm)
         session.add_metadata(note='foo')
+
+        def _sub_task(i):
+          session.add_metadata(**{f'subtask_{i}': i})
+          return lf_structured.query(f'subtask_{i}', lm=lm)
+
+        for i, output, error in session.concurrent_map(
+            _sub_task, range(3), max_workers=2, silence_on_errors=None,
+        ):
+          assert isinstance(i, int), i
+          assert isinstance(output, str), output
+          assert error is None, error
         return self.x + Bar()(session, lm=lm)
 
       def make_additional_query(self, lm):
@@ -73,21 +84,24 @@ class SessionTest(unittest.TestCase):
     self.assertTrue(root.execution.has_stopped)
     self.assertGreater(root.execution.elapse, 0)
     self.assertEqual(root.result, 3)
-    self.assertEqual(root.metadata, dict(note='foo'))
+    self.assertEqual(
+        root.metadata,
+        dict(note='foo', subtask_0=0, subtask_1=1, subtask_2=2)
+    )
 
     # The root space should have one action (foo), no queries, and no logs.
     self.assertEqual(len(list(root.actions)), 1)
     self.assertEqual(len(list(root.queries)), 0)
     self.assertEqual(len(list(root.logs)), 0)
-    # 1 query from Bar and 2 from Foo.
-    self.assertEqual(len(list(root.all_queries)), 3)
+    # 1 query from Bar, 2 from Foo and 3 from parallel executions.
+    self.assertEqual(len(list(root.all_queries)), 6)
     # 1 log from Bar and 1 from Foo.
     self.assertEqual(len(list(root.all_logs)), 2)
-    self.assertEqual(root.usage_summary.total.num_requests, 3)
+    self.assertEqual(root.usage_summary.total.num_requests, 6)
 
     # Inspecting the top-level action (Foo)
     foo_invocation = root.execution.items[0]
-    self.assertEqual(len(foo_invocation.execution.items), 3)
+    self.assertEqual(len(foo_invocation.execution.items), 4)
 
     # Prepare phase.
     prepare_phase = foo_invocation.execution.items[0]
@@ -104,8 +118,16 @@ class SessionTest(unittest.TestCase):
     self.assertIsInstance(query_invocation, lf_structured.QueryInvocation)
     self.assertIs(query_invocation.lm, lm)
 
+    # Tracked parallel executions.
+    parallel_executions = foo_invocation.execution.items[2]
+    self.assertIsInstance(parallel_executions, action_lib.ParallelExecutions)
+    self.assertEqual(len(parallel_executions), 3)
+    self.assertEqual(len(parallel_executions[0].queries), 1)
+    self.assertEqual(len(parallel_executions[1].queries), 1)
+    self.assertEqual(len(parallel_executions[2].queries), 1)
+
     # Invocation to Bar.
-    bar_invocation = foo_invocation.execution.items[2]
+    bar_invocation = foo_invocation.execution.items[3]
     self.assertIsInstance(bar_invocation, action_lib.ActionInvocation)
     self.assertIsInstance(bar_invocation.action, Bar)
     self.assertEqual(bar_invocation.result, 2)

{langfun-0.1.2.dev202503180804 → langfun-0.1.2.dev202503200804}/langfun/core/llms/gemini.py

@@ -640,7 +640,10 @@ class Gemini(rest.REST):
     ]
     usage = json['usageMetadata']
     input_tokens = usage['promptTokenCount']
-    output_tokens = usage['candidatesTokenCount']
+    # NOTE(daiyip): We saw cases that `candidatesTokenCount` is not present.
+    # Therefore, we use 0 as the default value.
+    output_tokens = usage.get('candidatesTokenCount', 0)
+
     return lf.LMSamplingResult(
         [lf.LMSample(message) for message in messages],
         usage=lf.LMSamplingUsage(

{langfun-0.1.2.dev202503180804 → langfun-0.1.2.dev202503200804}/langfun.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: langfun
-Version: 0.1.2.dev202503180804
+Version: 0.1.2.dev202503200804
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors
@@ -53,6 +53,7 @@ Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: keywords
 Dynamic: license
+Dynamic: license-file
 Dynamic: provides-extra
 Dynamic: requires-dist
 Dynamic: summary