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

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)

langfun/core/llms/cache/in_memory.py

@@ -24,7 +24,32 @@ import pyglove as pg
 
 @pg.use_init_args(['filename', 'ttl', 'key'])
 class InMemory(base.LMCacheBase):
-  """In memory cache."""
+  """An in-memory cache for language model lookups.
+
+  `InMemory` stores LM prompts and their corresponding responses in memory,
+  providing a simple and fast caching mechanism for a single session.
+  Optionally, it can persist the cache to a JSON file on disk, allowing
+  results to be reused across sessions.
+
+  When a filename is provided, the cache will be loaded from the file upon
+  initialization and saved to the file when `save()` is called. This is
+  useful for caching results in interactive environments like Colab or
+  when running batch jobs.
+
+  Example:
+
+  ```python
+  import langfun as lf
+  # Using in-memory cache without persistence
+  lm = lf.llms.GeminiPro(cache=lf.llms.cache.InMemory())
+  r = lm.query('hello')
+
+  # Using in-memory cache with persistence
+  lm = lf.llms.GeminiPro(cache=lf.llms.cache.InMemory('cache.json'))
+  r = lm.query('hello')
+  lm.cache.save()
+  ```
+  """
 
   filename: Annotated[
       str | None,
@@ -144,17 +169,33 @@ class InMemory(base.LMCacheBase):
 
 @contextlib.contextmanager
 def lm_cache(filename: str | None = None) -> Iterator[InMemory]:
-  """Context manager to enable cache for LMs under the context.
+  """Context manager to enable in-memory cache for LMs in the current context.
+
+  This context manager sets an `InMemory` cache as the default cache for
+  any Langfun language model instantiated within its scope, unless a model
+  is explicitly configured with a different cache.
+
+  If a `filename` is provided, the cache will be loaded from the specified
+  file at the beginning of the context and automatically saved back to the
+  file upon exiting the context. This is a convenient way to manage
+  persistent caching for a block of code.
+
+  Example:
 
-  If LMs under the context manager have explicitly specified cache, they will
-  use their own cache. Otherwise they will use the cache created by the context
-  manager.
+  ```python
+  import langfun as lf
+  with lf.lm_cache('my_cache.json'):
+    # LMs created here will use 'my_cache.json' for caching.
+    lm = lf.llms.GeminiPro()
+    print(lm.query('hello'))
+  ```
 
   Args:
-    filename: If not None, JSON file to load and save the cache.
+    filename: If provided, specifies the JSON file for loading and saving
+      the cache.
 
   Yields:
-    A cache object created.
+    The `InMemory` cache instance created for this context.
   """
   cache = InMemory(filename)
   try:

langfun/core/llms/cache/in_memory_test.py

@@ -175,18 +175,28 @@ class InMemoryLMCacheTest(unittest.TestCase):
 
     cache = in_memory.InMemory()
     lm = fake.StaticSequence(['1', '2', '3', '4', '5', '6'], cache=cache)
-    lm(lf.UserMessage('hi <<[[image]]>>', image=CustomModality('foo')))
-    lm(lf.UserMessage('hi <<[[image]]>>', image=CustomModality('bar')))
+    image_foo = CustomModality('foo')
+    image_bar = CustomModality('bar')
+    lm(
+        lf.UserMessage(
+            f'hi <<[[{image_foo.id}]]>>', referred_modalities=[image_foo]
+        )
+    )
+    lm(
+        lf.UserMessage(
+            f'hi <<[[{image_bar.id}]]>>', referred_modalities=[image_bar]
+        )
+    )
     self.assertEqual(
         list(cache.keys()),
         [
             (
-                'hi <<[[image]]>><image>acbd18db</image>',
+                f'hi <<[[{image_foo.id}]]>>',
                 (None, None, 1, 40, None, None),
                 0,
             ),
             (
-                'hi <<[[image]]>><image>37b51d19</image>',
+                f'hi <<[[{image_bar.id}]]>>',
                 (None, None, 1, 40, None, None),
                 0,
             ),

langfun/core/llms/compositional.py

@@ -21,7 +21,31 @@ import pyglove as pg
 
 @pg.use_init_args(['candidates', 'seed'])
 class RandomChoice(lf.LanguageModel):
-  """Random choice of a list of LLM models."""
+  """A composite language model that randomly selects from a list of candidates.
+
+  `RandomChoice` acts as a proxy that forwards each request (`sample`, `score`,
+  `tokenize`, or `__call__`) to one of the `candidates` selected randomly.
+  This can be useful for load balancing across multiple LLM endpoints,
+  for A/B testing different models, or for ensembling model outputs
+  by calling it multiple times.
+
+  The selection is determined by the provided `seed`, ensuring reproducibility
+  if needed.
+
+  Example:
+
+  ```python
+  import langfun as lf
+
+  lm = lf.llms.RandomChoice([
+      lf.llms.GeminiPro(),
+      lf.llms.GPT4(),
+  ])
+
+  # This call will be handled by either GeminiPro or GPT4, chosen randomly.
+  r = lm.sample('hello')
+  ```
+  """
 
   candidates: Annotated[
       list[lf.LanguageModel],

langfun/core/llms/deepseek.py

@@ -93,8 +93,36 @@ _SUPPORTED_MODELS_BY_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 # DeepSeek API uses an API format compatible with OpenAI.
 # Reference: https://api-docs.deepseek.com/
 @lf.use_init_args(['model'])
-class DeepSeek(openai_compatible.OpenAICompatible):
-  """DeepSeek model."""
+class DeepSeek(openai_compatible.OpenAIChatCompletionAPI):
+  """DeepSeek models.
+
+  **Quick Start:**
+
+  ```python
+  import langfun as lf
+
+  # Call DeepSeek-V3 using API key from environment variable
+  # 'DEEPSEEK_API_KEY'.
+  lm = lf.llms.DeepSeekV3()
+  r = lm('Who are you?')
+  print(r)
+  ```
+
+  **Setting up API key:**
+
+  The DeepSeek API key can be specified in following ways:
+
+  1. At model instantiation:
+
+     ```python
+     lm = lf.llms.DeepSeekV3(api_key='MY_API_KEY')
+     ```
+  2. via environment variable `DEEPSEEK_API_KEY`.
+
+  **References:**
+
+  *   https://api-docs.deepseek.com/
+  """
 
   model: pg.typing.Annotated[
       pg.typing.Enum(

langfun/core/llms/fake.py

@@ -20,7 +20,38 @@ import langfun.core as lf
 
 
 class Fake(lf.LanguageModel):
-  """The base class for all fake language models."""
+  """Base class for fake language models, used for testing.
+
+  Fake models simulate the behavior of real language models but return
+  pre-defined responses, making them useful for testing prompts,
+  data processing logic, and agent behavior without incurring API costs
+  or relying on external services.
+
+  Langfun provides several fake models:
+  * `lf.llms.Echo`: Echoes the prompt back as the response.
+  * `lf.llms.StaticResponse`: Returns a fixed, pre-defined response for
+    any prompt.
+  * `lf.llms.StaticMapping`: Returns responses based on a prompt-to-response
+    dictionary.
+  * `lf.llms.StaticSequence`: Returns responses from a pre-defined sequence
+    in order.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+
+  # Use Echo model for testing
+  lm = lf.llms.Echo()
+  response = lm('hello')
+  assert response.text == 'hello'
+
+  # Use StaticResponse model
+  lm = lf.llms.StaticResponse('world')
+  response = lm('hello')
+  assert response.text == 'world'
+  ```
+  """
 
   def _score(self, prompt: lf.Message| list[lf.Message],
              completions: list[lf.Message]):

langfun/core/llms/gemini.py

@@ -151,6 +151,32 @@ SUPPORTED_MODELS = [
     #
     # Production models.
     #
+    # Gemini 3 Pro Preview
+    GeminiModelInfo(
+        model_id='gemini-3-pro-preview',
+        in_service=True,
+        provider=pg.oneof(['Google GenAI', 'VertexAI']),
+        model_type='instruction-tuned',
+        description='Gemini 3 Pro Preview.',
+        release_date=datetime.datetime(2025, 11, 18),
+        input_modalities=GeminiModelInfo.ALL_SUPPORTED_INPUT_TYPES,
+        context_length=lf.ModelInfo.ContextLength(
+            max_input_tokens=1_048_576,
+            max_output_tokens=65_536,
+        ),
+        pricing=GeminiModelInfo.Pricing(
+            cost_per_1m_cached_input_tokens=0.2,
+            cost_per_1m_input_tokens=2.0,
+            cost_per_1m_output_tokens=12.0,
+            cost_per_1m_cached_input_tokens_with_prompt_longer_than_128k=0.4,
+            cost_per_1m_input_tokens_with_prompt_longer_than_128k=4.0,
+            cost_per_1m_output_tokens_with_prompt_longer_than_128k=18.0,
+        ),
+        rate_limits=lf.ModelInfo.RateLimits(
+            max_requests_per_minute=2000,
+            max_tokens_per_minute=4_000_000,
+        ),
+    ),
     # Gemini 2.5 Flash
     GeminiModelInfo(
         model_id='gemini-2.5-flash',
@@ -696,7 +722,15 @@ _SUPPORTED_MODELS_BY_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 
 @pg.use_init_args(['model'])
 class Gemini(rest.REST):
-  """Language models provided by Google GenAI."""
+  """Base class for Gemini models served on Google GenAI and Vertex AI.
+
+  This class implements the Gemini API protocol, shared by
+  `lf.llms.GoogleGenAI` and `lf.llms.VertexAI`, providing common request
+  formatting and response parsing for Gemini models.
+
+  It is not intended to be used directly. Please use `lf.llms.GoogleGenAI` or
+  `lf.llms.VertexAI` instead.
+  """
 
   model: pg.typing.Annotated[
       pg.typing.Enum(
@@ -752,6 +786,13 @@ class Gemini(rest.REST):
         prompt.as_format('gemini', chunk_preprocessor=modality_conversion)
     )
     request['contents'] = contents
+    request['toolConfig'] = {
+        'functionCallingConfig': {
+            'mode': 'NONE',
+        }
+    }
+    if sampling_options.extras:
+      request.update(sampling_options.extras)
     return request
 
   def _generation_config(
@@ -783,11 +824,18 @@ class Gemini(rest.REST):
           + '\n\n [RESPONSE FORMAT (not part of prompt)]\n'
           + pg.to_json_str(json_schema, json_indent=2)
       )
+    thinking_config_data = {}
     if options.max_thinking_tokens is not None:
-      config['thinkingConfig'] = {
-          'includeThoughts': options.max_thinking_tokens > 0,
-          'thinkingBudget': options.max_thinking_tokens,
-      }
+      thinking_config_data['includeThoughts'] = options.max_thinking_tokens > 0
+      thinking_config_data['thinkingBudget'] = options.max_thinking_tokens
+    if options.thinking_level is not None:
+      thinking_config_data['thinkingLevel'] = options.thinking_level
+    if thinking_config_data:
+      config['thinkingConfig'] = thinking_config_data
+
+    # This is the new feature since Gemini 3.
+    if self.model_id.startswith('gemini-3'):
+      config['mediaResolution'] = 'MEDIA_RESOLUTION_HIGH'
 
     if self.response_modalities:
       config['responseModalities'] = self.response_modalities
@@ -803,10 +851,14 @@ class Gemini(rest.REST):
           '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 candidates
-    ]
+
+    messages = []
+    for candidate in candidates:
+      message = lf.Message.from_value(candidate['content'], format='gemini')
+      if finish_reason := candidate.get('finishReason'):
+        message.metadata['finish_reason'] = finish_reason
+      messages.append(message)
+
     usage = json['usageMetadata']
     input_tokens = usage['promptTokenCount']
     # NOTE(daiyip): We saw cases that `candidatesTokenCount` is not present.
@@ -828,9 +880,9 @@ class Gemini(rest.REST):
     )
 
   def _error(self, status_code: int, content: str) -> lf.LMError:
-    if (
-        status_code == 400
-        and b'exceeds the maximum number of tokens' in content
+    if status_code == 400 and (
+        b'exceeds the maximum number of tokens' in content
+        or b'Reduce the input token count and try again.' in content
     ):
       return lf.ContextLimitError(f'{status_code}: {content}')
     return super()._error(status_code, content)