langfun 0.1.2.dev202412030804__tar.gz → 0.1.2.dev202412070804__tar.gz

{langfun-0.1.2.dev202412030804 → langfun-0.1.2.dev202412070804}/PKG-INFO

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

{langfun-0.1.2.dev202412030804 → langfun-0.1.2.dev202412070804}/langfun/__init__.py

@@ -35,6 +35,8 @@ complete = structured.complete
 score = structured.score
 generate_class = structured.generate_class
 
+track_queries = structured.track_queries
+
 # Helper functions for input/output transformations based on
 # `lf.query` (e.g. jax-on-beam could use these for batch processing)
 query_prompt = structured.query_prompt

{langfun-0.1.2.dev202412030804 → langfun-0.1.2.dev202412070804}/langfun/core/agentic/action.py

@@ -14,8 +14,10 @@
 """Base classes for agentic actions."""
 
 import abc
-from typing import Annotated, Any, Optional, Union
+import contextlib
+from typing import Annotated, Any, Iterable, Iterator, Optional, Type, Union
 import langfun.core as lf
+from langfun.core import structured as lf_structured
 import pyglove as pg
 
 
@@ -35,12 +37,9 @@ class Action(pg.Object):
       self, session: Optional['Session'] = None, **kwargs) -> Any:
     """Executes the action."""
     session = session or Session()
-    try:
-      session.begin(self)
+    with session.track(self):
       self._result = self.call(session=session, **kwargs)
       return self._result
-    finally:
-      session.end(self)
 
   @abc.abstractmethod
   def call(self, session: 'Session', **kwargs) -> Any:
@@ -50,9 +49,20 @@ class Action(pg.Object):
 class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
   """A class for capturing the invocation of an action."""
   action: Action
-  result: Any = None
+
+  result: Annotated[
+      Any,
+      'The result of the action.'
+  ] = None
+
   execution: Annotated[
-      list[Union['ActionInvocation', lf.logging.LogEntry]],
+      list[
+          Union[
+              lf_structured.QueryInvocation,
+              'ActionInvocation',
+              lf.logging.LogEntry
+          ]
+      ],
       'Execution execution.'
   ] = []
 
@@ -69,6 +79,18 @@ class ActionInvocation(pg.Object, pg.views.html.HtmlTreeView.Extension):
     """Returns child action invocations."""
     return [v for v in self.execution if isinstance(v, ActionInvocation)]
 
+  def queries(
+      self,
+      include_children: bool = False
+  ) -> Iterable[lf_structured.QueryInvocation]:
+    """Iterates over queries from the current invocation."""
+    for v in self.execution:
+      if isinstance(v, lf_structured.QueryInvocation):
+        yield v
+      elif isinstance(v, ActionInvocation):
+        if include_children:
+          yield from v.queries(include_children=True)
+
   def _html_tree_view_summary(
       self, *, view: pg.views.html.HtmlTreeView, **kwargs
   ):
@@ -190,29 +212,57 @@ class Session(pg.Object):
     assert self._invocation_stack
     return self._invocation_stack[-1]
 
-  def begin(self, action: Action):
-    """Signal the beginning of the execution of an action."""
+  @contextlib.contextmanager
+  def track(self, action: Action) -> Iterator[ActionInvocation]:
+    """Track 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
+    try:
+      yield new_invocation
+    finally:
+      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 query(
+      self,
+      prompt: Union[str, lf.Template, Any],
+      schema: Union[
+          lf_structured.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
+      ] = None,
+      default: Any = lf.RAISE_IF_HAS_ERROR,
+      *,
+      lm: lf.LanguageModel | None = None,
+      examples: list[lf_structured.MappingExample] | None = None,
+      **kwargs
+      ) -> Any:
+    """Calls `lf.query` and associates it with the current invocation."""
+    with lf_structured.track_queries() as queries:
+      output = lf_structured.query(
+          prompt,
+          schema=schema,
+          default=default,
+          lm=lm,
+          examples=examples,
+          **kwargs
+      )
+    with pg.notify_on_change(False):
+      self.current_invocation.execution.extend(queries)
+    return output
 
   def _log(self, level: lf.logging.LogLevel, message: str, **kwargs):
     with pg.notify_on_change(False):

{langfun-0.1.2.dev202412030804 → langfun-0.1.2.dev202412070804}/langfun/core/agentic/action_test.py

@@ -17,6 +17,7 @@ import unittest
 
 import langfun.core as lf
 from langfun.core.agentic import action as action_lib
+from langfun.core.llms import fake
 
 
 class SessionTest(unittest.TestCase):
@@ -26,25 +27,35 @@ class SessionTest(unittest.TestCase):
 
     class Bar(action_lib.Action):
 
-      def call(self, session, **kwargs):
+      def call(self, session, *, lm, **kwargs):
         test.assertIs(session.current_invocation.action, self)
         session.info('Begin Bar')
+        session.query('bar', lm=lm)
         return 2
 
     class Foo(action_lib.Action):
       x: int
 
-      def call(self, session, **kwargs):
+      def call(self, session, *, lm, **kwargs):
         test.assertIs(session.current_invocation.action, self)
         session.info('Begin Foo', x=1)
-        return self.x + Bar()(session)
+        session.query('foo', lm=lm)
+        return self.x + Bar()(session, lm=lm)
 
+    lm = fake.StaticResponse('lm response')
     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(Foo(1)(session, lm=lm), 3)
     self.assertEqual(len(session.root_invocation.child_invocations), 1)
+    self.assertEqual(len(list(session.root_invocation.queries())), 0)
+    self.assertEqual(
+        len(list(session.root_invocation.queries(include_children=True))), 2
+    )
+    self.assertEqual(
+        len(list(session.root_invocation.child_invocations[0].queries())), 1
+    )
     self.assertEqual(len(session.root_invocation.child_invocations[0].logs), 1)
     self.assertEqual(
         len(session.root_invocation.child_invocations[0].child_invocations),
@@ -55,6 +66,11 @@ class SessionTest(unittest.TestCase):
             .child_invocations[0].child_invocations[0].logs),
         1
     )
+    self.assertEqual(
+        len(list(session.root_invocation
+                 .child_invocations[0].child_invocations[0].queries())),
+        1
+    )
     self.assertEqual(
         len(session.root_invocation
             .child_invocations[0].child_invocations[0].child_invocations),

{langfun-0.1.2.dev202412030804 → langfun-0.1.2.dev202412070804}/langfun/core/eval/v2/runners.py

@@ -261,6 +261,9 @@ class RunnerBase(Runner):
       if cache is not None:
         self.background_run(cache.save)
 
+      # Wait for the background tasks to finish.
+      self._io_pool.shutdown(wait=True)
+
   @abc.abstractmethod
   def _run(self, evaluations: list[Evaluation]) -> None:
     """Runs multiple evaluations."""

{langfun-0.1.2.dev202412030804 → langfun-0.1.2.dev202412070804}/langfun/core/structured/__init__.py

@@ -69,6 +69,8 @@ from langfun.core.structured.prompting import query
 from langfun.core.structured.prompting import query_prompt
 from langfun.core.structured.prompting import query_output
 from langfun.core.structured.prompting import query_reward
+from langfun.core.structured.prompting import QueryInvocation
+from langfun.core.structured.prompting import track_queries
 
 from langfun.core.structured.description import DescribeStructure
 from langfun.core.structured.description import describe

{langfun-0.1.2.dev202412030804 → langfun-0.1.2.dev202412070804}/langfun/core/structured/function_generation.py

@@ -16,7 +16,7 @@
 import functools
 import inspect
 import re
-from typing import Any, Callable, Optional, Tuple
+from typing import Any, Callable, Literal, Optional, Tuple
 
 from langfun.core import language_model
 from langfun.core import template
@@ -25,7 +25,7 @@ from langfun.core.structured import prompting
 import pyglove as pg
 
 
-def unittest_gen(signature, lm, num_retries=10):
+def unittest_gen(signature, lm, num_retries=1):
   """Generates unit tests for a python function signature."""
 
   class UnitTest(pg.Object):
@@ -78,10 +78,13 @@ def _function_gen(
     func: Callable[..., Any],
     signature: str,
     lm: language_model.LanguageModel,
-    num_retries: int = 10,
+    num_retries: int = 1,
     unittest: Optional[
-        Callable[[Callable[..., Any]], None] | list[Tuple[Any, Any]]
+        Callable[[Callable[..., Any]], None]
+        | list[Tuple[Any, Any]]
+        | Literal["auto"]
     ] = None,
+    unittest_num_retries: int = 1,
 ):
   """Generates a python function with LLM and verify its quality with unit testing."""
 
@@ -131,9 +134,11 @@ def _function_gen(
     """
 
   unittest_examples = None
-  if unittest is None:
-    unittest_examples = unittest_gen(signature, lm=lm)
-  elif not callable(unittest):
+  if unittest == "auto":
+    unittest_examples = unittest_gen(
+        signature, lm=lm, num_retries=unittest_num_retries
+    )
+  elif isinstance(unittest, list):
     unittest_examples = unittest
 
   for _ in range(num_retries):
@@ -145,11 +150,16 @@ def _function_gen(
 
       # Check whether the sigantures are the same.
       if inspect.signature(f) != inspect.signature(func):
+        pg.logging.warning(
+            "Signature mismatch. Expected: %s, Actual: %s",
+            inspect.signature(func),
+            inspect.signature(f),
+        )
         continue
 
       if callable(unittest):
         unittest(f)
-      else:
+      elif unittest_examples:
         unittest_with_test_cases(f, unittest_examples)
 
       return f, source_code
@@ -172,10 +182,13 @@ def _process_signature(signature):
 def function_gen(
     lm: language_model.LanguageModel,
     cache_filename: str | None = None,
-    num_retries: int = 10,
+    num_retries: int = 1,
     unittest: Optional[
-        Callable[[Callable[..., Any]], None] | list[Tuple[Any, Any]]
+        Callable[[Callable[..., Any]], None]
+        | list[Tuple[Any, Any]]
+        | Literal["auto"]
     ] = None,
+    unittest_num_retries: int = 1,
 ):
   """A decorator for automating function generation using a language model.
 
@@ -192,9 +205,12 @@ def function_gen(
         make to generate a suitable function implementation.
       unittest: This optional parameter enables the definition of custom unit
         tests. You can either provide a list of test cases as tuples of inputs
-        and outputs, or a function that throws an error if a test fails. If left
-        as None (the default setting), the LLM will automatically create the
-        unit test cases.
+        and outputs, or a function that throws an error if a test fails, or let
+        LLM automatically create the unit test cases. If a generated function is
+        and returned, it should pass all the unittests.
+      unittest_num_retries: If unittest is set to "auto", this parameter
+        specifies the number of times the LLM's attempts to generate unit test
+        cases.
 
   Returns:
       The implemented function object.
@@ -226,7 +242,12 @@ def function_gen(
           return func.__function__(*args, **kwargs)
 
       func.__function__, func.__source_code__ = _function_gen(
-          func, signature, lm, num_retries=num_retries, unittest=unittest
+          func,
+          signature,
+          lm,
+          num_retries=num_retries,
+          unittest=unittest,
+          unittest_num_retries=unittest_num_retries,
       )
       if func.__function__ is None:
         raise ValueError(f"Function generation failed. Signature:\n{signature}")

{langfun-0.1.2.dev202412030804 → langfun-0.1.2.dev202412070804}/langfun/core/structured/function_generation_test.py

@@ -63,6 +63,42 @@ class FunctionGenerationTest(unittest.TestCase):
 
     lm = fake.StaticSequence([unittest_lm_response, function_gen_lm_response])
 
+    @function_generation.function_gen(lm=lm, unittest='auto')
+    def linear_search(items, target):  # pylint: disable=unused-argument
+      """Performs a linear search on a list to find a target value.
+
+      Args:
+          items (list): The list to search within.
+          target: The value to search for.
+
+      Returns:
+          int: The index of the target value if found, otherwise -1.
+      """
+
+    self.assertEqual(linear_search(['a', 'b', 'c'], 'c'), 2)
+    self.assertEqual(linear_search.source(), function_gen_lm_response)
+
+  def test_generate_function_without_unittest(self):
+    function_gen_lm_response = inspect.cleandoc("""
+        def linear_search(items, target):
+            \"\"\"
+            Performs a linear search on a list to find a target value.
+
+            Args:
+                items (list): The list to search within.
+                target: The value to search for.
+
+            Returns:
+                int: The index of the target value if found, otherwise -1.
+            \"\"\"
+            for i, item in enumerate(items):
+                if item == target:
+                    return i
+            return -1
+        """)
+
+    lm = fake.StaticSequence([function_gen_lm_response])
+
     @function_generation.function_gen(lm=lm)
     def linear_search(items, target):  # pylint: disable=unused-argument
       """Performs a linear search on a list to find a target value.
@@ -258,7 +294,9 @@ class FunctionGenerationTest(unittest.TestCase):
     cache_file = os.path.join(cache_file_dir, 'cache_file.json')
 
     @function_generation.function_gen(
-        lm=lm, unittest=_unittest_fn, cache_filename=cache_file
+        lm=lm,
+        unittest=_unittest_fn,
+        cache_filename=cache_file,
     )
     def linear_search(items, target):  # pylint: disable=unused-argument
       """Performs a linear search on a list to find a target value.
@@ -310,7 +348,9 @@ class FunctionGenerationTest(unittest.TestCase):
 
     custom_unittest = _unittest_fn
 
-    @function_generation.function_gen(lm=lm, unittest=custom_unittest)
+    @function_generation.function_gen(
+        lm=lm, unittest=custom_unittest, num_retries=2
+    )
     def linear_search(items, target):  # pylint: disable=unused-argument
       """Performs a linear search on a list to find a target value.
 

{langfun-0.1.2.dev202412030804 → langfun-0.1.2.dev202412070804}/langfun/core/structured/prompting.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 """Symbolic query."""
 
+import contextlib
 import functools
-from typing import Any, Callable, Type, Union
+from typing import Annotated, Any, Callable, Iterator, Type, Union
 
 import langfun.core as lf
 from langfun.core.llms import fake
@@ -102,7 +103,7 @@ def _query_structure_cls(
 
 
 def query(
-    prompt: Union[str, pg.Symbolic],
+    prompt: Union[str, lf.Template, Any],
     schema: Union[
         schema_lib.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
     ] = None,
@@ -119,7 +120,7 @@ def query(
     skip_lm: bool = False,
     **kwargs,
 ) -> Any:
-  """Parse a natural langugage message based on schema.
+  """Queries an language model for a (maybe) structured output.
 
   Examples:
 
@@ -189,59 +190,93 @@ def query(
   """
     # Internal usage logging.
 
+  # Normalize query schema.
   # When `lf.query` is used for symbolic completion, schema is automatically
   # inferred when it is None.
   if isinstance(prompt, pg.Symbolic) and prompt.sym_partial and schema is None:
     schema = prompt.__class__
 
-  # Create a copy of the prompt if it has a parent object, so all child modality
-  # objects could be referred by path relative to the prompt.
-  if isinstance(prompt, lf.Template) and prompt.sym_parent:
-    prompt = prompt.clone()
-
-  if schema in (None, str):
-    # Query with natural language output.
-    output = lf.LangFunc.from_value(prompt, **kwargs)(
-        lm=lm, cache_seed=cache_seed, skip_lm=skip_lm
+  # Normalize query input.
+  if isinstance(prompt, (lf.Message, str)):
+    # Query with structured output.
+    prompt_kwargs = kwargs.copy()
+    prompt_kwargs.pop('template_str', None)
+    query_input = lf.Template.from_value(prompt, **prompt_kwargs)
+  elif isinstance(prompt, lf.Template):
+    # Create a copy of the prompt if it has a parent object, so all child
+    # modality objects could be referred by path relative to the prompt.
+    query_input = prompt.clone() if prompt.sym_parent is not None else prompt
+
+    # Attach template metadata from kwargs. This is used to pass through fields
+    # from kwargs to the rendered message.
+    template_metadata = {
+        k: v for k, v in kwargs.items() if k.startswith('metadata_')
+    }
+    query_input.rebind(
+        template_metadata, skip_notification=True, raise_on_no_change=False
     )
-    if response_postprocess:
-      processed_text = response_postprocess(output.text)
-      if processed_text != output.text:
-        output = lf.AIMessage(processed_text, source=output)
-    return output if returns_message else output.text
-
-  # Query with structured output.
-  prompt_kwargs = kwargs.copy()
-
-  # NOTE(daiyip): when `template_str` is passed in, it's intended to modify the
-  # QueryStructure template string. Therefore, we pop out the argument for
-  # prompt rendering.
-  prompt_kwargs.pop('template_str', None)
-
-  if isinstance(prompt, (str, lf.Message, lf.Template)):
-    prompt = lf.Template.from_value(prompt, **prompt_kwargs).render(lm=lm)
+  elif pg.MISSING_VALUE == prompt:
+    query_input = lf.UserMessage('')
   else:
-    prompt = schema_lib.mark_missing(prompt)
-
-  output = _query_structure_cls(protocol)(
-      input=prompt,
-      schema=schema,
-      default=default,
-      examples=examples,
-      response_postprocess=response_postprocess,
-      autofix=autofix if protocol == 'python' else 0,
-      **kwargs,
-  )(
-      lm=lm,
-      autofix_lm=autofix_lm or lm,
-      cache_seed=cache_seed,
-      skip_lm=skip_lm,
-  )
-  return output if returns_message else output.result
+    query_input = schema_lib.mark_missing(prompt)
+
+  with lf.track_usages() as usage_summary:
+    if schema in (None, str):
+      # Query with natural language output.
+      output_message = lf.LangFunc.from_value(query_input, **kwargs)(
+          lm=lm, cache_seed=cache_seed, skip_lm=skip_lm
+      )
+      if response_postprocess:
+        processed_text = response_postprocess(output_message.text)
+        if processed_text != output_message.text:
+          output_message = lf.AIMessage(processed_text, source=output_message)
+    else:
+      # Query with structured output.
+      output_message = _query_structure_cls(protocol)(
+          input=(
+              query_input.render(lm=lm)
+              if isinstance(query_input, lf.Template)
+              else query_input
+          ),
+          schema=schema,
+          default=default,
+          examples=examples,
+          response_postprocess=response_postprocess,
+          autofix=autofix if protocol == 'python' else 0,
+          **kwargs,
+      )(
+          lm=lm,
+          autofix_lm=autofix_lm or lm,
+          cache_seed=cache_seed,
+          skip_lm=skip_lm,
+      )
+
+  def _result(message: lf.Message):
+    return message.text if schema in (None, str) else message.result
+
+  # Track the query invocations.
+  if pg.MISSING_VALUE != prompt and not skip_lm:
+    trackers = lf.context_value('__query_trackers__', [])
+    if trackers:
+      invocation = QueryInvocation(
+          input=pg.Ref(query_input),
+          schema=(
+              schema_lib.Schema.from_value(schema)
+              if schema not in (None, str) else None
+          ),
+          output=pg.Ref(_result(output_message)),
+          lm=pg.Ref(lm),
+          examples=pg.Ref(examples) if examples else [],
+          usage_summary=usage_summary,
+      )
+      for i, (tracker, include_child_scopes) in enumerate(trackers):
+        if i == 0 or include_child_scopes:
+          tracker.append(invocation)
+  return output_message if returns_message else _result(output_message)
 
 
 def query_prompt(
-    prompt: Union[str, pg.Symbolic],
+    prompt: Union[str, lf.Template, Any],
     schema: Union[
         schema_lib.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
     ] = None,
@@ -264,7 +299,7 @@ def query_output(
   kwargs.pop('prompt', None)
   kwargs.pop('lm', None)
   return query(
-      'Unused prompt', schema, lm=fake.StaticResponse(response), **kwargs
+      pg.MISSING_VALUE, schema, lm=fake.StaticResponse(response), **kwargs
   )
 
 
@@ -320,3 +355,69 @@ def _reward_fn(cls) -> Callable[
     args = [self, input, expected_output, metadata]
     return cls.__reward__(*args[:num_args])
   return _reward
+
+
+class QueryInvocation(pg.Object):
+  """A class to represent the invocation of `lf.query`."""
+
+  input: Annotated[
+      Union[lf.Template, pg.Symbolic],
+      'Mapping input of `lf.query`.'
+  ]
+  schema: pg.typing.Annotated[
+      schema_lib.schema_spec(noneable=True),
+      'Schema of `lf.query`.'
+  ]
+  output: Annotated[
+      Any,
+      'Mapping output of `lf.query`.'
+  ]
+  lm: Annotated[
+      lf.LanguageModel,
+      'Language model used for `lf.query`.'
+  ]
+  examples: Annotated[
+      list[mapping.MappingExample],
+      'Fewshot exemplars for `lf.query`.'
+  ]
+  usage_summary: Annotated[
+      lf.UsageSummary,
+      'Usage summary for `lf.query`.'
+  ]
+
+
+@contextlib.contextmanager
+def track_queries(
+    include_child_scopes: bool = True
+) -> Iterator[list[QueryInvocation]]:
+  """Track all queries made during the context.
+
+  Example:
+
+    ```
+    with lf.track_queries() as queries:
+      lf.query('hi', lm=lm)
+      lf.query('What is this {{image}}?', lm=lm, image=image)
+
+    print(queries)
+    ```
+
+  Args:
+    include_child_scopes: If True, the queries made in child scopes will be
+      included in the returned list. Otherwise, only the queries made in the
+      current scope will be included.
+
+  Yields:
+    A list of `QueryInvocation` objects representing the queries made during
+    the context.
+  """
+  trackers = lf.context_value('__query_trackers__', [])
+  tracker = []
+
+  with lf.context(
+      __query_trackers__=[(tracker, include_child_scopes)] + trackers
+  ):
+    try:
+      yield tracker
+    finally:
+      pass