langfun 0.1.2.dev202505120804__tar.gz → 0.1.2.dev202505140804__tar.gz

{langfun-0.1.2.dev202505120804 → langfun-0.1.2.dev202505140804}/PKG-INFO

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

{langfun-0.1.2.dev202505120804 → langfun-0.1.2.dev202505140804}/langfun/core/agentic/action.py

@@ -209,7 +209,7 @@ class Action(pg.Object):
   ) -> Any:
     """Executes the action."""
     if session is None:
-      session = Session()
+      session = Session(verbose=verbose)
       session.start()
 
       if show_progress:
@@ -220,45 +220,13 @@ class Action(pg.Object):
     else:
       self._session = None
 
-    with session.track_action(self) as invocation:
-      if verbose:
-        session.info('Action execution started.', keep=False, action=self)
-
+    with session.track_action(self):
       try:
-        result = self.call(session=session, verbose=verbose, **kwargs)
+        result = self.call(session=session, **kwargs)
         self._invocation.end(result)
-        if verbose:
-          session.info(
-              (
-                  f'Action execution succeeded in '
-                  f'{self._invocation.execution.elapse:.2f} seconds.'
-              ),
-              keep=False,
-              result=result
-          )
       except BaseException as e:
         error = pg.utils.ErrorInfo.from_exception(e)
         self._invocation.end(result=None, error=error)
-        if invocation.parent_action is session.root:
-          session.error(
-              (
-                  f'Top-level action execution failed in '
-                  f'{self._invocation.execution.elapse:.2f} seconds.'
-              ),
-              keep=True,
-              action=self,
-              error=error
-          )
-        else:
-          session.warning(
-              (
-                  f'Action execution failed in '
-                  f'{self._invocation.execution.elapse:.2f} seconds.'
-              ),
-              keep=False,
-              action=self,
-              error=error
-          )
         if self._session is not None:
           self._session.end(result=None, error=error)
         raise
@@ -477,21 +445,26 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
   def __getitem__(self, index: int) -> TracedItem:
     return self.items[index]
 
+  def merge_usage_summary(self, usage_summary: lf.UsageSummary) -> None:
+    if usage_summary.total.num_requests == 0:
+      return
+    current_invocation = self
+    while current_invocation is not None:
+      current_invocation.usage_summary.merge(usage_summary)
+      current_invocation = typing.cast(
+          ExecutionTrace,
+          current_invocation.sym_ancestor(
+              lambda x: isinstance(x, ExecutionTrace)
+          )
+      )
+
   def append(self, item: TracedItem) -> None:
     """Appends an item to the sequence."""
     with pg.notify_on_change(False):
       self.items.append(item)
 
     if isinstance(item, lf_structured.QueryInvocation):
-      current_invocation = self
-      while current_invocation is not None:
-        current_invocation.usage_summary.merge(item.usage_summary)
-        current_invocation = typing.cast(
-            ExecutionTrace,
-            current_invocation.sym_ancestor(
-                lambda x: isinstance(x, ExecutionTrace)
-            )
-        )
+      self.merge_usage_summary(item.usage_summary)
 
     if self._tab_control is not None:
       self._tab_control.append(self._execution_item_tab(item))
@@ -519,15 +492,46 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
   def execution_summary(self) -> dict[str, Any]:
     """Execution summary string."""
     return pg.Dict(
-        num_queries=len(self.queries),
-        execution_breakdown=[
-            dict(
-                action=action.action.__class__.__name__,
-                usage=action.usage_summary.total,
-                execution_time=action.execution.elapse,
-            )
-            for action in self.actions
-        ]
+        subtree=dict(
+            num_actions=len(self.all_actions),
+            num_action_failures=len([
+                a for a in self.all_actions if a.has_error
+            ]),
+            num_queries=len(self.all_queries),
+            num_oop_failures=len([
+                q for q in self.all_queries if q.has_oop_error
+            ]),
+            num_non_oop_failures=len([
+                q for q in self.all_queries
+                if q.has_error and not q.has_oop_error
+            ]),
+            total_query_time=sum(q.elapse for q in self.all_queries),
+        ),
+        current_level=dict(
+            num_actions=len(self.actions),
+            num_action_failures=len([
+                a for a in self.actions if a.has_error
+            ]),
+            num_queries=len(self.queries),
+            num_oop_failures=len([
+                q for q in self.queries if q.has_oop_error
+            ]),
+            num_non_oop_failures=len([
+                q for q in self.queries
+                if q.has_error and not q.has_oop_error
+            ]),
+            execution_breakdown=[
+                dict(
+                    action=action.action.__class__.__name__,
+                    usage=dict(
+                        total_tokens=action.usage_summary.total.total_tokens,
+                        estimated_cost=action.usage_summary.total.estimated_cost,
+                    ),
+                    execution_time=action.execution.elapse,
+                )
+                for action in self.actions
+            ]
+        )
     )
 
   #
@@ -1025,9 +1029,6 @@ class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
                         self.usage_summary.to_html(  # pylint: disable=g-long-ternary
                             extra_flags=dict(as_badge=True)
                         )
-                        if (interactive
-                            or self.usage_summary.total.num_requests > 0)
-                        else None
                     ),
                 ],
                 css_classes=['execution-tab-title']
@@ -1069,12 +1070,24 @@ class RootAction(Action):
 class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
   """Session for performing an agentic task."""
 
-  root: ActionInvocation = ActionInvocation(RootAction())
+  root: Annotated[
+      ActionInvocation,
+      'The root action invocation of the session.'
+  ] = ActionInvocation(RootAction())
+
   id: Annotated[
       str | None,
       'An optional identifier for the sessin, which will be used for logging.'
   ] = None
 
+  verbose: Annotated[
+      bool,
+      (
+          'If True, the session will be logged with verbose action and query '
+          'activities.'
+      )
+  ] = False
+
   # NOTE(daiyip): Action execution may involve multi-threading, hence current
   # action and execution are thread-local.
 
@@ -1169,8 +1182,45 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       self._current_execution = invocation.execution
       # Start the execution of the current action.
       self._current_action.start()
+      if self.verbose:
+        self.info(
+            'Action execution started.',
+            action=invocation.action,
+            keep=False,
+        )
       yield invocation
     finally:
+      if invocation.has_error:
+        if invocation.parent_action is self.root:
+          self.error(
+              (
+                  f'Top-level action execution failed in '
+                  f'{invocation.execution.elapse:.2f} seconds.'
+              ),
+              action=invocation.action,
+              error=invocation.error,
+              keep=True,
+          )
+        else:
+          self.warning(
+              (
+                  f'Action execution failed in '
+                  f'{invocation.execution.elapse:.2f} seconds.'
+              ),
+              action=invocation.action,
+              error=invocation.error,
+              keep=True,
+          )
+      elif self.verbose:
+        self.info(
+            (
+                f'Action execution succeeded in '
+                f'{invocation.execution.elapse:.2f} seconds.'
+            ),
+            action=invocation.action,
+            result=invocation.result,
+            keep=False,
+        )
       self._current_execution = parent_execution
       self._current_action = parent_action
 
@@ -1208,18 +1258,63 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       A list of `lf.QueryInvocation` objects, each for a single `lf.query`
       call.
     """
-    with self.track_phase(phase) as execution:
-      with lf_structured.track_queries(include_child_scopes=False) as queries:
-        try:
-          yield queries
-        finally:
-          for i, query in enumerate(queries):
-            query.rebind(
-                id=f'{execution.id}/q{len(execution.queries) + i + 1}',
-                skip_notification=False,
-                raise_on_no_change=False
-            )
-          execution.extend(queries)
+    def _query_start(invocation: lf_structured.QueryInvocation):
+      execution = self._current_execution
+      invocation.rebind(
+          id=f'{execution.id}/q{len(execution.queries) + 1}',
+          skip_notification=False, raise_on_no_change=False
+      )
+      execution.append(invocation)
+      if self.verbose:
+        self.info(
+            'Querying LLM started.',
+            lm=invocation.lm.model_id,
+            output_type=(
+                lf_structured.annotation(invocation.schema.spec)
+                if invocation.schema is not None else None
+            ),
+            keep=False,
+        )
+
+    def _query_end(invocation: lf_structured.QueryInvocation):
+      self._current_execution.merge_usage_summary(invocation.usage_summary)
+      if invocation.has_error:
+        self.warning(
+            (
+                f'Querying LLM failed in '
+                f'{time.time() - invocation.start_time:.2f} seconds.'
+            ),
+            lm=invocation.lm.model_id,
+            output_type=(
+                lf_structured.annotation(invocation.schema.spec)
+                if invocation.schema is not None else None
+            ),
+            error=invocation.error,
+            keep=True,
+        )
+      elif self.verbose:
+        self.info(
+            (
+                f'Querying LLM succeeded in '
+                f'{time.time() - invocation.start_time:.2f} seconds.'
+            ),
+            lm=invocation.lm.model_id,
+            output_type=(
+                lf_structured.annotation(invocation.schema.spec)
+                if invocation.schema is not None else None
+            ),
+            keep=False,
+        )
+
+    with self.track_phase(phase), lf_structured.track_queries(
+        include_child_scopes=False,
+        start_callabck=_query_start,
+        end_callabck=_query_end,
+    ) as queries:
+      try:
+        yield queries
+      finally:
+        pass
 
   #
   # Operations with activity tracking.
@@ -1272,24 +1367,14 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       The result of the query.
     """
     with self.track_queries():
-      start_time = time.time()
-      try:
-        return lf_structured.query(
-            prompt,
-            schema=schema,
-            default=default,
-            lm=lm,
-            examples=examples,
-            **kwargs
-        )
-      except BaseException as e:
-        elapse = time.time() - start_time
-        self.warning(
-            f'Failed to query LLM ({lm.model_id}) in {elapse:.2f} seconds.',
-            error=pg.utils.ErrorInfo.from_exception(e),
-            keep=False,
-        )
-        raise
+      return lf_structured.query(
+          prompt,
+          schema=schema,
+          default=default,
+          lm=lm,
+          examples=examples,
+          **kwargs
+      )
 
   def concurrent_map(
       self,

{langfun-0.1.2.dev202505120804 → langfun-0.1.2.dev202505140804}/langfun/core/agentic/action_eval.py

@@ -36,17 +36,15 @@ class ActionEval(lf.eval.v2.Evaluation):
     action = example_input.action
 
     # We explicitly create a session here to use a custom session ID.
-    with action_lib.Session(id=f'{self.id}#example-{example.id}') as session:
+    with action_lib.Session(
+        id=f'{self.id}#example-{example.id}', verbose=True
+    ) as session:
 
       # NOTE(daiyip): Setting session as metadata before action execution, so we
       # could use `Evaluation.state.in_progress_examples` to access the session
       # for status reporting from other threads.
       example.metadata['session'] = session
-
-      with lf.logging.use_log_level('fatal'):
-        kwargs = self.action_args.copy()
-        kwargs.update(verbose=True)
-        action(session=session, **kwargs)
+      action(session=session, **self.action_args)
 
     return session.final_result, dict(session=session)
 

{langfun-0.1.2.dev202505120804 → langfun-0.1.2.dev202505140804}/langfun/core/coding/python/correction.py

@@ -37,6 +37,7 @@ def run_with_correction(
     lm: lf.LanguageModel | None = None,
     max_attempts: int = 5,
     sandbox: bool | None = None,
+    permission: pg.coding.CodePermission = pg.coding.CodePermission.ALL,
     timeout: int | None = 5,
     returns_code: bool = False,
     returns_stdout: bool = False,
@@ -58,6 +59,7 @@ def run_with_correction(
       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
       process.
+    permission: The permission to run the code.
     timeout: The timeout for running the corrected code. If None, there is no
       timeout. Applicable only when sandbox is set to True.
     returns_code: If True, the return value is a tuple of (result, final code).
@@ -88,6 +90,7 @@ def run_with_correction(
             global_vars=global_vars,
             sandbox=sandbox,
             timeout=timeout,
+            permission=permission,
             returns_stdout=returns_stdout,
             outputs_intermediate=outputs_intermediate,
         )
@@ -102,6 +105,7 @@ def run_with_correction(
               global_vars=global_vars,
               sandbox=sandbox,
               timeout=timeout,
+              permission=permission,
               outputs_intermediate=outputs_intermediate,
           )
       )

{langfun-0.1.2.dev202505120804 → langfun-0.1.2.dev202505140804}/langfun/core/console.py

@@ -52,10 +52,13 @@ def write(
   )
 
 
+_notebook = None
 try:
-  _notebook = sys.modules['IPython'].display
-except Exception:  # pylint: disable=broad-except
-  _notebook = None
+  ipython_module = sys.modules['IPython']
+  if 'IPKernelApp' in ipython_module.get_ipython().config:
+    _notebook = ipython_module.display
+except (KeyError, AttributeError):  # pylint: disable=broad-except
+  pass
 
 
 def under_notebook() -> bool:

{langfun-0.1.2.dev202505120804 → langfun-0.1.2.dev202505140804}/langfun/core/language_model.py

@@ -1019,11 +1019,21 @@ class LanguageModel(component.Component):
       ] = RetryableLMError,
   ) -> list[Any]:
     """Helper method for subclasses for implementing _sample."""
+    if self.max_concurrency is None:
+      execute = action
+      executor = None
+      max_workers = len(inputs)
+    else:
+      execute = lambda x: _ConcurrencyControl.get(
+          self.resource_id, self.max_concurrency)(action, x)
+      executor = self.resource_id if len(inputs) > 1 else None
+      max_workers = self.max_concurrency
+
     executed_jobs = concurrent.concurrent_execute(
-        action,
+        execute,
         inputs,
-        executor=self.resource_id if self.max_concurrency else None,
-        max_workers=self.max_concurrency or len(inputs),
+        executor=executor,
+        max_workers=max_workers,
         retry_on_errors=retry_on_errors,
         max_attempts=self.max_attempts,
         retry_interval=self.retry_interval,
@@ -1323,6 +1333,46 @@ class LanguageModel(component.Component):
     return None
 
 
+class _ConcurrencyControl:
+  """Controls the max concurrent LLM calls for a given model."""
+
+  _MODEL_CONCURRENCY: ClassVar[dict[str, '_ConcurrencyControl']] = {}
+
+  def __init__(self, max_concurrency: int):
+    self.max_concurrency = max_concurrency
+    self._concurrency = 0
+
+  @property
+  def concurrency(self) -> int:
+    """Returns the current concurrency."""
+    return self._concurrency
+
+  def __call__(self, fn: Callable[..., Any], *args, **kwargs):
+    """Calls the function with concurrency control."""
+    while self._concurrency >= self.max_concurrency:
+      time.sleep(0.01)
+
+    try:
+      # Increment/decrement is atomic in Python, so we don't need to protect it
+      # with a lock.
+      self._concurrency += 1
+      return fn(*args, **kwargs)
+    finally:
+      self._concurrency -= 1
+
+  @classmethod
+  def get(
+      cls, model_id: str, max_concurrency: int | None = None
+  ) -> '_ConcurrencyControl':
+    """Returns the concurrency control for the given model ID."""
+    control = cls._MODEL_CONCURRENCY.get(model_id, None)
+    if control is None:
+      assert max_concurrency is not None
+      control = cls(max_concurrency)
+      cls._MODEL_CONCURRENCY[model_id] = control
+    return control
+
+
 class UsageSummary(pg.Object, pg.views.HtmlTreeView.Extension):
   """Usage sumary."""
 
@@ -1403,7 +1453,8 @@ class UsageSummary(pg.Object, pg.views.HtmlTreeView.Extension):
       self._usage_badge.update(
           self._badge_text(),
           tooltip=pg.format(
-              self, verbose=False, custom_format=self._tooltip_format
+              self, verbose=False, custom_format=self._tooltip_format,
+              hide_default_values=True,
           ),
           styles=dict(color=self._badge_color()),
       )
@@ -1450,7 +1501,8 @@ class UsageSummary(pg.Object, pg.views.HtmlTreeView.Extension):
         usage_badge = pg.views.html.controls.Badge(
             self._badge_text(),
             tooltip=pg.format(
-                self, custom_format=self._tooltip_format, verbose=False
+                self, custom_format=self._tooltip_format, verbose=False,
+                hide_default_values=True,
             ),
             css_classes=['usage-summary'],
             styles=dict(color=self._badge_color()),

{langfun-0.1.2.dev202505120804 → langfun-0.1.2.dev202505140804}/langfun/core/llms/gemini.py

@@ -678,9 +678,15 @@ class Gemini(rest.REST):
     return config
 
   def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
+    candidates = json.get('candidates')
+    if candidates is None:
+      raise lf.TemporaryLMError(
+          'No candidates found in response. This is a Gemini API issue that '
+          'happens occasionally, and retrying should fix it. '
+      )
     messages = [
         lf.Message.from_value(candidate['content'], format='gemini')
-        for candidate in json.get('candidates', [])
+        for candidate in candidates
     ]
     usage = json['usageMetadata']
     input_tokens = usage['promptTokenCount']

{langfun-0.1.2.dev202505120804 → langfun-0.1.2.dev202505140804}/langfun/core/logging.py

@@ -268,14 +268,13 @@ def log(level: LogLevel,
       metadata=kwargs,
   )
 
-  if entry.should_output(get_log_level()):
+  if console and entry.should_output(get_log_level()):
     if console_lib.under_notebook():
       console_lib.display(entry)
-    elif console:
+    else:
       # TODO(daiyip): Improve the console output formatting.
       console_lib.write(entry)
-
-  if not console:
+  elif not console:
     if kwargs:
       message = f'{message} (metadata: {pg.format(kwargs)})'
     system_log_func(level)(message)

{langfun-0.1.2.dev202505120804 → langfun-0.1.2.dev202505140804}/langfun/core/structured/mapping.py

@@ -278,6 +278,11 @@ class Mapping(lf.LangFunc):
       'A `lf.structured.Schema` object that constrains mapping output ',
   ] = None
 
+  permission: Annotated[
+      pg.coding.CodePermission,
+      'The permission to run the LLM generated code.'
+  ] = pg.coding.CodePermission.ASSIGN | pg.coding.CodePermission.CALL
+
   @property
   def mapping_request(self) -> MappingExample:
     """Returns a MappingExample as the mapping request."""
@@ -434,6 +439,7 @@ class Mapping(lf.LangFunc):
         additional_context=self.globals(),
         autofix=self.autofix,
         autofix_lm=self.autofix_lm or self.lm,
+        permission=self.permission,
     )
 
   def postprocess_response(self, response: lf.Message) -> lf.Message: