langfun 0.1.2.dev202509120804__py3-none-any.whl → 0.1.2.dev202512150805__py3-none-any.whl

langfun/core/langfunc.py

@@ -32,146 +32,43 @@ _TLS_LFUN_CALL_STACK = '_langfunc_callstack'
 # NOTE(daiyip): Only the template string belongs to the positional arguments,
 # all others are keyword-only for clarity.
 @pg.use_init_args(['template_str'])
-class LangFunc(
-    template_lib.Template,
-):
-  r"""Base class for natural-language driven component.
-
-  ``LangFunc`` is a language-driven component that enables users to
-  seamlessly interact with Language Models (LLMs) using a blend of natural
-  language and code. It empowers users to easily modularize prompt/execution
-  logics, compose them, and simplify the creation of Language Model (LLM)-based
-  components and applications.
-
-  LangFunc can be conceptualized as a string template with embeddable code,
-  but it distinguishes itself from traditional template systems in four key
-  ways.
-
-  Firstly, it enables easy modularization of templates along with the required
-  values with OO principles, providing a reusable way for LLM-based content
-  generation. For example:
-
-    ```
-    class FewshotExamples(lf.LangFunc):
-      '''Base for fewshot prompt.
-
-      {% for example in examples %}
-      {{ example }}
-      {% endfor %}
-      '''
-
-    # Usage 1: __init__ time binding.
-    assert FewshotPrompt(examples=['foo', 'bar'])() == 'foo\nbar'
-
-    # Usage 2: __call__ time binding.
-    assert FewshotPrompt()(examples=['foo', 'bar']) == 'foo\nbar'
-
-    class ToolDescription(lf.LangFunc):
-      '''Tool descriptions.
-
-      {% for tool in tools %}
-      {{ tool.description }}
-      {% endfor %}
-      '''
-      # We want to constrain tools to be a list of `Tool` objects.
-      tools: list[Tool]
-
-    # Raises: runtime type checking will fail on [1, 2, 3].
-    ToolDescription(tools=[1, 2, 3])
-    ```
-
-  Secondly, it has the capability to compose multiple LangFuncs together,
-  enabling the accomplishment of complex language tasks with maximum reuse.
-  It allows users to provide program inputs to all the LangFuncs within a
-  composition at the top level, significantly simplifying the process of
-  providing context for users. For example:
-
-    ```
-    class ReAct(lf.LangFunc):
-      '''ReAct prompt for tool-use.
-
-      {{ preamble }}
-      {{ tool_description }}
-      {{ tool_examples }}
-      {{ user_input }}
-      '''
-      # Default preamble, which could be overriden from subclass
-      # or parsed from the `__init__` argument.
-      preamble = 'Please help me on my task based on the following tools.',
-
-    react = ReAct(
-        tool_description=ToolDescription()
-        tool_examples=FewshotExamples(),
-        # Partially bind `tools` and `examples`.
-        tools=my_tools,
-        examples=[t.examples for t in my_tools]
-        )
-
-    # Late bind `user_input` at __call__ time.
-    react(user_input='Help me get a lunch to go, veggie please.' )
-    ```
-
-  Thirdly, it allows the flexibility to encapsulate complex compositions to
-  reusable classes and modify them. For example:
-
-    ```
-    # The compound decorator converts a function into a LangFunc.
-    @lf.compound
-    def react_with_tools(preamble, tools: list[Tool]):
-      return ReAct(
-          preamble=preamble,
-          tool_description=ToolDescription()
-          tool_examples=FewshotExamples(),
-          # Partially bind `tools` and `examples`.
-          tools=my_tools,
-          examples=[t.examples for t in my_tools]
-      )
+class LangFunc(template_lib.Template):
+  r"""Base class for Language-based functions.
 
-    # Actually, the entire chat application is a LangFunc.
-    class Chat(lt.LangFunc):
-      '''LLM-based Chat application.
+  LangFunc represents a function powered by a language model. It is a subclass
+  of `lf.Template` and can be thought of as a `lf.Template` augmented with an LM
+  and an output transformation. Calling a `lf.LangFunc` is equivalent to calling
+  the LM with the rendered prompt and transforming the output.
 
-      llm({{ prompt }})
-      '''
+  LangFunc can be directly constructed and used.
 
-    chat = Chat(
-        llm=Bard24B(),
-        prompt=react_with_tools(
-            preamble=(
-                f'Please help me solve my problem using tools. '
-                f'Current time is {{datetime.datetime.now()}}'),
-            tools=my_tools))
+  ```python
+  import langfun as lf
 
-    chat(user_input='Help me get a lunch to go, veggie please.')
-    ```
+  func = lf.LangFunc("Hello, {{name}}!")
+  print(func(name="Gemini", lm=lf.llms.Gemini25Flash()))
+  # Output: Hello, how are you today?
+  ```
 
-  Fourthly, LangFunc is built on top of PyGlove symbolic programming power,
-  it could be manipulated programmatically, turned into a space for data
-  sampling, or even tuned by AutoML. For example:
+  Or it can be subclassed:
 
-    ```
-    import pyglove as pg
+  ```python
+  import langfun as lf
 
-    prompt_space = react_with_tools(
-        preamble=pg.oneof([
-            'Help me solve my problem using the following tools:',
-            'Help me with the tools below:',
-            ...
-        ])
-        # Choose any two of the tools for generating data.
-        tools=pg.manyof(2, [
-            google_search(...),
-            doordash(...),
-            ...
-        ])
+  class Compute(lf.LangFunc):
+    '''Compute a simple arithmetic expression.
 
-    for prompt in pg.random_sample(prompt_space):
-      print(prompt(user_input='Help me book a conf room please.'))
+    {{expression}} = ?
+    '''
+    expression: str
 
-    ```
+    def transform_output(self, lm_output: lf.Message) -> lf.Message:
+      lm_output.metadata.result = float(lm_output.text)
+      return lm_output
 
-  For more capabilities on symbolic programming with PyGlove, please checkout
-  https://pyglove.readthedocs.io/en/latest/.
+  r = Compute(expression="1 + 1")(lm=lf.llms.Gemini25Flash())
+  print(r.result)
+  # Output: 2.0
 
   Final note: always include these capitalized words if you don't want to treat
   the docstr as the template str: THIS IS NOT A TEMPLATE. So as a result, this
@@ -305,6 +202,24 @@ class LangFunc(
       message_cls: Type[message_lib.Message] = message_lib.UserMessage,
       **kwargs,
   ) -> message_lib.Message:
+    """Renders the template and transforms it as LM input message.
+
+    Args:
+      allow_partial: If True, allows partial rendering, which leaves unresolved
+        variables in place in the output text. Otherwise, raises error when
+        there are unresolved variables.
+      implicit: If True, reuse the rendering output if a parent `lf.Template`
+        is rendering current `lf.Template` multiple times. This is important
+        for making sure all references to the same `lf.Template` within a single
+        top-level rendering would return the same result. If False, every call
+        to `render` will trigger the actual rendering process.
+      message_cls: The message class used for creating the return value.
+      **kwargs: Values for template variables, which override values from
+        member attributes or context.
+
+    Returns:
+      A Message object containing the rendered result.
+    """
     lm_input = super().render(
         allow_partial=allow_partial,
         implicit=implicit,

langfun/core/langfunc_test.py

@@ -82,7 +82,7 @@ class LangFuncCallTest(unittest.TestCase):
 
     i = l.render()
     self.assertEqual(i, 'Hello')
-    self.assertEqual(i, message.UserMessage('Hello'))
+    self.assertEqual(i, message.UserMessage('Hello', __template_input__={}))
     self.assertEqual(i.tags, ['rendered'])
 
     r = l()
@@ -96,7 +96,9 @@ class LangFuncCallTest(unittest.TestCase):
     self.assertEqual(r.tags, ['lm-response', 'lm-output'])
     self.assertEqual(
         r.source,
-        message.UserMessage('Hello', metadata=dict(cache_seed=0))
+        message.UserMessage(
+            'Hello', metadata=dict(cache_seed=0, __template_input__={})
+        )
     )
     self.assertEqual(r.source.tags, ['rendered', 'lm-input'])
 
@@ -107,9 +109,9 @@ class LangFuncCallTest(unittest.TestCase):
         ' lm=ExcitedEchoer(sampling_options=LMSamplingOptions(temperature=None,'
         ' max_tokens=None, n=1, top_k=40, top_p=None, stop=None,'
         ' random_seed=None, logprobs=False, top_logprobs=None,'
-        ' max_thinking_tokens=None, reasoning_effort=None), cache=None,'
-        ' max_concurrency=None, timeout=120.0, max_attempts=5,'
-        ' retry_interval=(5, 60), exponential_backoff=True,'
+        ' max_thinking_tokens=None, thinking_level=None, reasoning_effort=None,'
+        ' extras={}), cache=None, max_concurrency=None, timeout=120.0,'
+        ' max_attempts=5, retry_interval=(5, 60), exponential_backoff=True,'
         ' max_retry_interval=300, debug=False))',
     )
 

langfun/core/language_model.py

@@ -53,6 +53,10 @@ class RetryableLMError(LMError):
   """Base class for LLM errors that can be solved by retrying."""
 
 
+class EmptyGenerationError(RetryableLMError):
+  """Error for empty generaition."""
+
+
 class RateLimitError(RetryableLMError):
   """Error for rate limit reached."""
 
@@ -478,7 +482,7 @@ class UsageNotAvailable(LMSamplingUsage):
 
 
 class LMSamplingResult(pg.Object):
-  """Language model response."""
+  """The result from a language model sampling."""
 
   samples: Annotated[
       list[LMSample],
@@ -575,6 +579,14 @@ class LMSamplingOptions(component.Component):
       int | None, 'Number of max thinking tokens.'
   ] = None
 
+  thinking_level: Annotated[
+      Literal['low', 'high'] | None,
+      (
+          'Thinking level for Gemini models. High is for complex tasks, '
+          'while low is for faster responses.'
+      ),
+  ] = None
+
   reasoning_effort: Annotated[
       Literal['low', 'medium', 'high'] | None,
       (
@@ -584,6 +596,15 @@ class LMSamplingOptions(component.Component):
       ),
   ] = None
 
+  extras: Annotated[
+      dict[str, Any],
+      (
+          'Extra arguments (e.g. configuration for tool calls) to pass to '
+          'the model. This is model-specific, please check model '
+          'implementation to see how to use this.'
+      ),
+  ] = {}
+
   def cache_key(self) -> tuple[Any, ...]:
     """Returns a tuple of current values as cache key."""
     return (
@@ -672,13 +693,91 @@ class LMDebugMode(enum.IntFlag):
 
 
 class LanguageModel(component.Component):
-  """Interface of a language model.
-
-  Language models are at the center of LLM-based agents. ``LanguageModel``
-  is the interface to interact with different language modles.
-
-  In langfun, users can use different language models with the same agents,
-  allowing fast prototype, as well as side-by-side comparisons.
+  """Interface for language model.
+
+  `lf.LanguageModel` is the cornerstone of Langfun, providing a consistent
+  interface for interacting with various language models, such as those from
+  Google, OpenAI, Anthropic, and more. It abstracts away provider-specific
+  details, allowing users to switch between models seamlessly.
+
+  All language models in Langfun can be accessed via `lf.llms`. For example,
+  `lf.llms.Gpt4()` creates an instance for OpenAI's GPT-4, and
+  `lf.llms.GeminiPro()` creates an instance for Google's Gemini Pro.
+
+  **Key Features:**
+
+  *   **Unified API**: Provides `sample`, `score`, and `tokenize` methods
+      across all supported models.
+  *   **Sampling**: The `__call__` method and `sample` method allow generating
+      text completions or chat responses.
+  *   **Scoring**: The `score` method computes the likelihood of completions
+      given a prompt.
+  *   **Tokenization**: The `tokenize` method breaks text into tokens
+      according to the model's tokenizer.
+  *   **Caching**: Built-in support for caching LLM requests to save cost and
+      time via the `cache` attribute.
+  *   **Concurrency**: Manages concurrency to respect API rate limits via
+      `max_concurrency`.
+  *   **Retries**: Automatic retries with exponential backoff for transient
+      errors via `max_attempts` and `retry_interval`.
+
+  **1. Creating a Language Model:**
+  You can create a language model by instantiating its class or by using
+  `lf.LanguageModel.get`:
+
+  ```python
+  # Direct instantiation
+  gpt4 = lf.llms.Gpt4()
+  gemini = lf.llms.GeminiPro()
+
+  # Creation via lf.LanguageModel.get()
+  gpt4 = lf.LanguageModel.get('gpt-4')
+  ```
+
+  **2. Customizing Sampling Options:**
+  Sampling options like `temperature`, `max_tokens`, etc., can be customized
+  at model creation, or overridden at call time or via `lf.context`.
+
+  ```python
+  # Set temperature to 0 at model creation
+  lm = lf.llms.Gpt4(temperature=0.0)
+
+  # Override temperature to 0.5 for a single call
+  response = lm('1 + 1 =', temperature=0.5)
+
+  # Override temperature to 1.0 using lf.context
+  with lf.context(temperature=1.0):
+    response = lm('1 + 1 =')
+  ```
+
+  **3. Sampling:**
+  Use `lm()`, `lm.sample()`, or `lf.query()` to generate text:
+
+  ```python
+  lm = lf.llms.Gpt4()
+  response = lm('1 + 1 =')
+  print(response.text)
+  # Output: 2
+  ```
+
+  **4. Scoring:**
+  Use `lm.score()` to score completions:
+
+  ```python
+  lm = lf.llms.Gpt4()
+  results = lm.score('Weather in SF is', completions=['sunny', 'cloudy'])
+  print(results[0].score)
+  # Output: -1.0
+  ```
+
+  **5. Tokenization:**
+  Use `lm.tokenize()` to get tokens:
+  ```python
+  lm = lf.llms.Gpt4()
+  tokens = lm.tokenize('hello world')
+  print(tokens)
+  # Output: [('hello', 15339), (' world', 1917)]
+  ```
   """
 
   sampling_options: LMSamplingOptions = LMSamplingOptions()
@@ -989,10 +1088,32 @@ class LanguageModel(component.Component):
     prompts = [message_lib.UserMessage.from_value(p) for p in prompts]
 
     with component.context(override_attrs=True, **kwargs):
-      if self.cache is None:
-        results = self._sample(prompts)
-      else:
-        results = self._sample_with_cache_lookup(prompts, cache_seed)
+
+      def _sample_with_retry():
+        if self.cache is None:
+          results = self._sample(prompts)
+        else:
+          results = self._sample_with_cache_lookup(prompts, cache_seed)
+
+        for i, result in enumerate(results):
+          for sample in result.samples:
+            if not sample.response.text:
+              if self.cache is not None:
+                self.cache.delete(self, prompts[i], seed=cache_seed)
+              raise EmptyGenerationError(
+                  f'Empty generation encountered from model {self.model_id}.'
+              )
+        return results
+
+      retry_fn = concurrent.with_retry(
+          _sample_with_retry,
+          retry_on_errors=EmptyGenerationError,
+          max_attempts=self.max_attempts,
+          retry_interval=self.retry_interval,
+          exponential_backoff=self.exponential_backoff,
+          max_retry_interval=self.max_retry_interval,
+      )
+      results = retry_fn()
 
       for prompt, result in zip(prompts, results):
 
@@ -1001,7 +1122,6 @@ class LanguageModel(component.Component):
 
         for sample in result.samples:
           # Update metadata for response message.
-
           response = sample.response
           response.metadata.score = sample.score
           response.metadata.logprobs = sample.logprobs
@@ -1159,21 +1279,35 @@ class LanguageModel(component.Component):
   ) -> message_lib.Message:
     """Returns the first candidate."""
     prompt = message_lib.UserMessage.from_value(prompt)
-    with component.context(override_attrs=True, **kwargs):
-      sampling_options = self.sampling_options
-      if sampling_options.n != 1:
-        sampling_options = sampling_options.clone(override=dict(n=1))
-
-      call_counter = self._call_counter
-      self._call_counter += 1
-      request_start = time.time()
-      result = self.sample(
-          [prompt], sampling_options=sampling_options, cache_seed=cache_seed
-      )[0]
-      elapse = time.time() - request_start
-      response = result.samples[0].response
-      self._debug(prompt, response, call_counter, result.usage, elapse)
-      return response
+    start_time = time.time()
+    error_tag = ''
+    try:
+      with component.context(override_attrs=True, **kwargs):
+        sampling_options = self.sampling_options
+        if sampling_options.n != 1:
+          sampling_options = sampling_options.clone(override=dict(n=1))
+
+        call_counter = self._call_counter
+        self._call_counter += 1
+        request_start = time.time()
+        result = self.sample(
+            [prompt], sampling_options=sampling_options, cache_seed=cache_seed
+        )[0]
+        elapse = time.time() - request_start
+        response = result.samples[0].response
+        self._debug(prompt, response, call_counter, result.usage, elapse)
+        return response
+    except BaseException as e:
+      error_tag = pg.ErrorInfo.from_exception(e).tag
+      raise e
+    finally:
+      _METRICS.language_model_calls.increment(
+          model=self.model_id, error=error_tag
+      )
+      _METRICS.language_model_call_duration_ms.record(
+          int((time.time() - start_time) * 1000),
+          model=self.model_id, error=error_tag,
+      )
 
   def _debug(
       self,
@@ -1230,11 +1364,11 @@ class LanguageModel(component.Component):
         title=f'\n[{call_counter}] PROMPT SENT TO LM{title_suffix}:',
         color='green',
     )
-    referred_modalities = prompt.referred_modalities()
-    if referred_modalities:
+    if prompt.referred_modalities:
       console.write(
           pg.object_utils.kvlist_str(
-              [(k, repr(v), None) for k, v in referred_modalities.items()]
+              [(k, repr(v), None)
+               for k, v in prompt.referred_modalities.items()]
           ),
           title=f'\n[{call_counter}] MODALITY OBJECTS SENT TO LM:',
           color='green',
@@ -1320,9 +1454,9 @@ class LanguageModel(component.Component):
           color='green',
       )
       if isinstance(prompt, list):
-        referred_modalities_lst = [p.referred_modalities() for p in prompt]
+        referred_modalities_lst = [p.referred_modalities for p in prompt]
       else:
-        referred_modalities_lst = [prompt.referred_modalities(),]
+        referred_modalities_lst = [prompt.referred_modalities,]
       if referred_modalities_lst:
         for referred_modalities in referred_modalities_lst:
           console.write(
@@ -1397,7 +1531,7 @@ class LanguageModel(component.Component):
           title=f'\n[{call_counter}] PROMPT TO TOKENIZE:',
           color='green',
       )
-      referred_modalities_lst = [prompt.referred_modalities(),]
+      referred_modalities_lst = [prompt.referred_modalities,]
       if referred_modalities_lst:
         for referred_modalities in referred_modalities_lst:
           console.write(
@@ -1425,7 +1559,7 @@ class LanguageModel(component.Component):
       max_requests_per_minute: int | None,
       average_tokens_per_request: int = 250
   ) -> int | None:
-    """Estimates max concurrency concurrency based on the rate limits."""
+    """Estimates max concurrency based on the rate limits."""
     # NOTE(daiyip): max concurrency is estimated based on the rate limit.
     # We assume each request has approximately 250 tokens, and each request
     # takes 1 second to complete. This might not be accurate for all models.
@@ -1438,6 +1572,25 @@ class LanguageModel(component.Component):
     return None
 
 
+class _Metrics:
+  """Metrics for Langfun."""
+
+  def __init__(self):
+    self._metrics = pg.monitoring.metric_collection('/third_party/langfun')
+    self.language_model_calls = self._metrics.get_counter(
+        'language_model_calls',
+        'Number of calls to the language model.',
+        parameters={'model': str, 'error': str},
+    )
+    self.language_model_call_duration_ms = self._metrics.get_distribution(
+        'language_model_call_duration_ms',
+        'Duration of calls to the language model in milliseconds.',
+        parameters={'model': str, 'error': str},
+    )
+
+_METRICS = _Metrics()
+
+
 class _ConcurrencyControl:
   """Controls the max concurrent LLM calls for a given model."""
 
@@ -1479,7 +1632,7 @@ class _ConcurrencyControl:
 
 
 class UsageSummary(pg.Object, pg.views.HtmlTreeView.Extension):
-  """Usage sumary."""
+  """Usage summary."""
 
   class AggregatedUsage(pg.Object):
     """Aggregated usage."""

langfun/core/language_model_test.py

@@ -591,6 +591,51 @@ class LanguageModelTest(unittest.TestCase):
     lm = MockModel(cache=cache, top_k=1)
     self.assertEqual(lm('a'), 'a')
 
+  def test_empty_generation_error(self):
+    class MockModelWithEmptyResponse(MockModel):
+      def _sample(self,
+                  prompts: list[message_lib.Message]
+                  ) -> list[lm_lib.LMSamplingResult]:
+        return [lm_lib.LMSamplingResult(
+            [lm_lib.LMSample(response='')],
+            usage=lm_lib.LMSamplingUsage(100, 0, 100, 1, 1.0)
+        )]
+    lm = MockModelWithEmptyResponse(max_attempts=1, retry_interval=0)
+    with self.assertRaisesRegex(
+        concurrent.RetryError, 'Empty generation encountered'
+    ):
+      lm('a')
+
+  def test_empty_generation_retry(self):
+    class MockModelWithEmptyThenValid(MockModel):
+      attempt_count: int = 0
+
+      def _sample(
+          self, prompts: list[message_lib.Message]
+      ) -> list[lm_lib.LMSamplingResult]:
+        self.rebind(attempt_count=self.attempt_count + 1)
+        if self.attempt_count == 1:
+          # First attempt returns empty
+          return [
+              lm_lib.LMSamplingResult(
+                  [lm_lib.LMSample(response='')],
+                  usage=lm_lib.LMSamplingUsage(100, 0, 100, 1, 1.0),
+              )
+          ]
+        else:
+          # Subsequent attempts return valid response
+          return [
+              lm_lib.LMSamplingResult(
+                  [lm_lib.LMSample(response='valid response')],
+                  usage=lm_lib.LMSamplingUsage(100, 100, 200, 1, 1.0),
+              )
+          ]
+
+    lm = MockModelWithEmptyThenValid(max_attempts=3, retry_interval=0)
+    result = lm('a')
+    self.assertEqual(result.text, 'valid response')
+    self.assertEqual(lm.attempt_count, 2)
+
   def test_estimate_max_concurrency(self):
     self.assertIsNone(lm_lib.LanguageModel.estimate_max_concurrency(None, None))
     self.assertEqual(
@@ -656,11 +701,17 @@ class LanguageModelTest(unittest.TestCase):
 
     string_io = io.StringIO()
     lm = MockModel(sampling_options=lm_lib.LMSamplingOptions(top_k=1))
+    image = Image()
     with contextlib.redirect_stdout(string_io):
       self.assertEqual(
-          lm(message_lib.UserMessage(
-              'hi <<[[image]]>>', image=Image()), debug=True),
-          'hi <<[[image]]>>'
+          lm(
+              message_lib.UserMessage(
+                  f'hi <<[[{image.id}]]>>',
+                  referred_modalities=[image],
+              ),
+              debug=True
+          ),
+          f'hi <<[[{image.id}]]>>'
       )
 
     debug_info = string_io.getvalue()