langfun 0.1.2.dev202510230805__py3-none-any.whl → 0.1.2.dev202511270805__py3-none-any.whl

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
@@ -1244,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',
@@ -1334,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(
@@ -1411,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(
@@ -1439,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.
@@ -1512,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()

langfun/core/llms/__init__.py

@@ -30,7 +30,8 @@ from langfun.core.llms.compositional import RandomChoice
 
 # Base models by request/response protocol.
 from langfun.core.llms.rest import REST
-from langfun.core.llms.openai_compatible import OpenAICompatible
+from langfun.core.llms.openai_compatible import OpenAIChatCompletionAPI
+from langfun.core.llms.openai_compatible import OpenAIResponsesAPI
 from langfun.core.llms.gemini import Gemini
 from langfun.core.llms.anthropic import Anthropic
 
@@ -41,6 +42,7 @@ from langfun.core.llms.azure_openai import AzureOpenAI
 
 # Gemini models.
 from langfun.core.llms.google_genai import GenAI
+from langfun.core.llms.google_genai import Gemini3ProPreview
 from langfun.core.llms.google_genai import Gemini25Pro
 from langfun.core.llms.google_genai import Gemini25Flash
 from langfun.core.llms.google_genai import Gemini25ProPreview_20250605
@@ -89,6 +91,7 @@ from langfun.core.llms.vertexai import VertexAIGemini25ProPreview_20250605
 from langfun.core.llms.vertexai import VertexAIGemini25Pro
 from langfun.core.llms.vertexai import VertexAIGemini25Flash
 from langfun.core.llms.vertexai import VertexAIGemini25FlashImagePreview
+from langfun.core.llms.vertexai import VertexAIGemini3ProPreview
 
 # For backward compatibility.
 GeminiPro1_5 = Gemini15Pro
@@ -151,6 +154,9 @@ from langfun.core.llms.openai import Gpt35
 
 # Anthropic models.
 
+from langfun.core.llms.anthropic import Claude45
+from langfun.core.llms.anthropic import Claude45Haiku_20251001
+from langfun.core.llms.anthropic import Claude45Sonnet_20250929
 from langfun.core.llms.anthropic import Claude4
 from langfun.core.llms.anthropic import Claude4Sonnet_20250514
 from langfun.core.llms.anthropic import Claude4Opus_20250514
@@ -168,6 +174,8 @@ from langfun.core.llms.anthropic import Claude3Haiku
 from langfun.core.llms.anthropic import Claude3Haiku_20240307
 
 from langfun.core.llms.vertexai import VertexAIAnthropic
+from langfun.core.llms.vertexai import VertexAIClaude45Haiku_20251001
+from langfun.core.llms.vertexai import VertexAIClaude45Sonnet_20250929
 from langfun.core.llms.vertexai import VertexAIClaude4Opus_20250514
 from langfun.core.llms.vertexai import VertexAIClaude4Sonnet_20250514
 from langfun.core.llms.vertexai import VertexAIClaude37Sonnet_20250219

langfun/core/llms/anthropic.py

@@ -59,6 +59,60 @@ class AnthropicModelInfo(lf.ModelInfo):
 
 
 SUPPORTED_MODELS = [
+    AnthropicModelInfo(
+        model_id='claude-haiku-4-5-20251001',
+        provider='Anthropic',
+        in_service=True,
+        description='Claude 4.5 Haiku model (10/15/2025).',
+        release_date=datetime.datetime(2025, 10, 15),
+        input_modalities=(
+            AnthropicModelInfo.INPUT_IMAGE_TYPES
+            + AnthropicModelInfo.INPUT_DOC_TYPES
+        ),
+        context_length=lf.ModelInfo.ContextLength(
+            max_input_tokens=200_000,
+            max_output_tokens=64_000,
+        ),
+        pricing=lf.ModelInfo.Pricing(
+            cost_per_1m_cached_input_tokens=0.1,
+            cost_per_1m_input_tokens=1,
+            cost_per_1m_output_tokens=5,
+        ),
+        rate_limits=AnthropicModelInfo.RateLimits(
+            # Tier 4 rate limits
+            max_requests_per_minute=4000,
+            max_input_tokens_per_minute=4_000_000,
+            max_output_tokens_per_minute=800_000,
+        ),
+    ),
+    AnthropicModelInfo(
+        model_id='claude-sonnet-4-5-20250929',
+        provider='Anthropic',
+        in_service=True,
+        description='Claude 4.5 Sonnet model (9/29/2025).',
+        release_date=datetime.datetime(2025, 9, 29),
+        input_modalities=(
+            AnthropicModelInfo.INPUT_IMAGE_TYPES
+            + AnthropicModelInfo.INPUT_DOC_TYPES
+        ),
+        context_length=lf.ModelInfo.ContextLength(
+            max_input_tokens=200_000,
+            max_output_tokens=64_000,
+        ),
+        pricing=lf.ModelInfo.Pricing(
+            cost_per_1m_cached_input_tokens=0.3,
+            cost_per_1m_input_tokens=3,
+            cost_per_1m_output_tokens=15,
+        ),
+        rate_limits=AnthropicModelInfo.RateLimits(
+            # Tier 4 rate limits
+            # This rate limit is a total limit that applies to combined traffic
+            # across both Sonnet 4 and Sonnet 4.5.
+            max_requests_per_minute=4000,
+            max_input_tokens_per_minute=2_000_000,
+            max_output_tokens_per_minute=400_000,
+        ),
+    ),
     AnthropicModelInfo(
         model_id='claude-4-opus-20250514',
         provider='Anthropic',
@@ -190,6 +244,62 @@ SUPPORTED_MODELS = [
             max_output_tokens_per_minute=80_000,
         ),
     ),
+    AnthropicModelInfo(
+        model_id='claude-haiku-4-5@20251001',
+        alias_for='claude-haiku-4-5-20251001',
+        provider='VertexAI',
+        in_service=True,
+        description='Claude 4.5 Haiku model served on VertexAI (10/15/2025).',
+        release_date=datetime.datetime(2025, 10, 15),
+        input_modalities=(
+            AnthropicModelInfo.INPUT_IMAGE_TYPES
+            + AnthropicModelInfo.INPUT_DOC_TYPES
+        ),
+        context_length=lf.ModelInfo.ContextLength(
+            max_input_tokens=200_000,
+            max_output_tokens=64_000,
+        ),
+        pricing=lf.ModelInfo.Pricing(
+            # For global endpoint
+            cost_per_1m_cached_input_tokens=0.1,
+            cost_per_1m_input_tokens=1,
+            cost_per_1m_output_tokens=5,
+        ),
+        rate_limits=AnthropicModelInfo.RateLimits(
+            # For global endpoint
+            max_requests_per_minute=2500,
+            max_input_tokens_per_minute=200_000,
+            max_output_tokens_per_minute=0,
+        ),
+    ),
+    AnthropicModelInfo(
+        model_id='claude-sonnet-4-5@20250929',
+        alias_for='claude-sonnet-4-5-20250929',
+        provider='VertexAI',
+        in_service=True,
+        description='Claude 4.5 Sonnet model (9/29/2025).',
+        release_date=datetime.datetime(2025, 9, 29),
+        input_modalities=(
+            AnthropicModelInfo.INPUT_IMAGE_TYPES
+            + AnthropicModelInfo.INPUT_DOC_TYPES
+        ),
+        context_length=lf.ModelInfo.ContextLength(
+            max_input_tokens=200_000,
+            max_output_tokens=64_000,
+        ),
+        pricing=lf.ModelInfo.Pricing(
+            # For global endpoint
+            cost_per_1m_cached_input_tokens=0.3,
+            cost_per_1m_input_tokens=3,
+            cost_per_1m_output_tokens=15,
+        ),
+        rate_limits=AnthropicModelInfo.RateLimits(
+            # For global endpoint
+            max_requests_per_minute=1500,
+            max_input_tokens_per_minute=200_000,
+            max_output_tokens_per_minute=0,
+        ),
+    ),
     AnthropicModelInfo(
         model_id='claude-opus-4@20250514',
         alias_for='claude-opus-4-20250514',
@@ -540,9 +650,34 @@ _SUPPORTED_MODELS_BY_MODEL_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 
 @lf.use_init_args(['model'])
 class Anthropic(rest.REST):
-  """Anthropic LLMs (Claude) through REST APIs.
+  """Anthropic Claude models.
+
+  **Quick Start:**
+
+  ```python
+  import langfun as lf
+
+  # Call Claude 3.5 Sonnet using API key from environment variable
+  # 'ANTHROPIC_API_KEY'.
+  lm = lf.llms.Claude35Sonnet()
+  r = lm('Who are you?')
+  print(r)
+  ```
+
+  **Setting up API key:**
+
+  The Anthropic API key can be specified in following ways:
+
+  1. At model instantiation:
+
+     ```python
+     lm = lf.llms.Claude35Sonnet(api_key='MY_API_KEY')
+
+  2. via environment variable `ANTHROPIC_API_KEY`.
 
-  See https://docs.anthropic.com/claude/reference/messages_post
+  **References:**
+
+  *   https://docs.anthropic.com/claude/reference/messages_post
   """
 
   model: pg.typing.Annotated[
@@ -658,6 +793,8 @@ class Anthropic(rest.REST):
       args.pop('temperature', None)
       args.pop('top_k', None)
       args.pop('top_p', None)
+    if options.extras:
+      args.update(options.extras)
     return args
 
   def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
@@ -679,6 +816,24 @@ class Anthropic(rest.REST):
     return super()._error(status_code, content)
 
 
+class Claude45(Anthropic):
+  """Base class for Claude 4.5 models."""
+
+
+# pylint: disable=invalid-name
+class Claude45Haiku_20251001(Claude45):
+  """Claude 4.5 Haiku model 20251001."""
+
+  model = 'claude-haiku-4-5-20251001'
+
+
+# pylint: disable=invalid-name
+class Claude45Sonnet_20250929(Claude45):
+  """Claude 4.5 Sonnet model 20250929."""
+
+  model = 'claude-sonnet-4-5-20250929'
+
+
 class Claude4(Anthropic):
   """Base class for Claude 4 models."""
 

langfun/core/llms/azure_openai.py

@@ -23,23 +23,35 @@ import pyglove as pg
 @lf.use_init_args(['model', 'deployment_name'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
 class AzureOpenAI(openai.OpenAI):
-  """Azure OpenAI model service.
-
-  This service interacts with the Azure OpenAI API to generate chat completions.
-  It uses the deployment_name and API version to construct the endpoint, and
-  authenticates using an API key provided via parameter or the
-  AZURE_OPENAI_API_KEY environment variable.
-
-  Example:
-      lm = AzureOpenAI(
-          model='gpt-4o',
-          deployment_name='gpt-4o',
-          api_version='2024-08-01-preview',
-          azure_endpoint='https://trackname.openai.azure.com/',
-          api_key='token'
-      )
-      response = lf.query(prompt="what the capital of France", lm=lm)
-      print(response)
+  """Azure OpenAI models.
+
+  **Quick Start:**
+
+  ```python
+  import langfun as lf
+
+  # Call GPT-4o on Azure using API key from environment variable
+  # 'AZURE_OPENAI_API_KEY'.
+  lm = lf.llms.AzureOpenAI(
+      model='gpt-4o',
+      deployment_name='my-gpt4o-deployment',
+      api_version='2024-08-01-preview',
+      azure_endpoint='https://my-resource.openai.azure.com/',
+  )
+  r = lm('Who are you?')
+  print(r)
+  ```
+
+  **Setting up API key:**
+
+  The Azure OpenAI API key can be specified in following ways:
+
+  1. At model instantiation:
+
+     ```python
+     lm = lf.llms.AzureOpenAI(..., api_key='MY_API_KEY')
+     ```
+  2. via environment variable `AZURE_OPENAI_API_KEY`.
   """
 
   deployment_name: Annotated[

langfun/core/llms/cache/base.py

@@ -22,13 +22,33 @@ import langfun.core as lf
 
 @dataclasses.dataclass(frozen=True)
 class LMCacheEntry:
-  """LM cache entry."""
+  """Represents a single entry in the language model cache.
+
+  An `LMCacheEntry` stores the result of a language model sampling operation
+  and an optional expiration timestamp.
+  """
   result: lf.LMSamplingResult
   expire: datetime.datetime | None = None
 
 
 class LMCacheBase(lf.LMCache):
-  """The common LMCache base."""
+  """Base class for language model cache implementations.
+
+  `LMCacheBase` provides the core logic for a key-value based cache,
+  handling key generation, expiration (TTL), and statistics tracking.
+  Subclasses must implement the abstract methods `_get`, `_put`, and `_delete`
+  to provide the specific storage mechanism (e.g., in-memory, file-based).
+
+  **Key Features:**
+
+  *   **Customizable Keying**: Allows specifying a custom function to generate
+      cache keys based on the language model, prompt, and seed. If not provided,
+      a default key based on prompt text, sampling options, and seed is used.
+  *   **Time-to-Live (TTL)**: Supports setting an expiration time for cache
+      entries, after which they are considered invalid and removed upon access.
+  *   **Cache Statistics**: Tracks metrics like hits, misses, updates,
+      deletions, and expired hits through the `stats` property.
+  """
 
   key: Annotated[
       Callable[[lf.LanguageModel, lf.Message, int], Any] | None,
@@ -121,4 +141,6 @@ class LMCacheBase(lf.LMCache):
 
 def default_key(lm: lf.LanguageModel, prompt: lf.Message, seed: int) -> Any:
   """Default key for LM cache."""
-  return (prompt.text_with_modality_hash, lm.sampling_options.cache_key(), seed)
+  # prompt text already contains the modality id for referenced modality
+  # objects, so no need to include them in the key.
+  return (prompt.text, lm.sampling_options.cache_key(), seed)