langfun 0.1.2.dev202508250805__py3-none-any.whl → 0.1.2.dev202511110805__py3-none-any.whl

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]):
@@ -62,6 +93,13 @@ class Echo(Fake):
     return lf.AIMessage(prompt.text)
 
 
+class Pseudo(Fake):
+  """A pseudo language model that should never be called."""
+
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    raise ValueError('Pseudo language model should never be called.')
+
+
 @lf.use_init_args(['response'])
 class StaticResponse(Fake):
   """Language model that always gives the same canned response."""

langfun/core/llms/fake_test.py

@@ -20,6 +20,15 @@ import langfun.core as lf
 from langfun.core.llms import fake as fakelm
 
 
+class PseudoTest(unittest.TestCase):
+
+  def test_sample(self):
+    lm = fakelm.Pseudo()
+    self.assertEqual(lm.model_id, 'Pseudo')
+    with self.assertRaises(ValueError):
+      _ = lm.sample(['hi'])
+
+
 class EchoTest(unittest.TestCase):
 
   def test_sample(self):

langfun/core/llms/gemini.py

@@ -195,7 +195,7 @@ SUPPORTED_MODELS = [
         rate_limits=lf.ModelInfo.RateLimits(
             max_requests_per_minute=2000,
             max_tokens_per_minute=4_000_000,
-        )
+        ),
     ),
     # Gemini 2.5 Pro 0605
     GeminiModelInfo(
@@ -218,7 +218,7 @@ SUPPORTED_MODELS = [
         rate_limits=lf.ModelInfo.RateLimits(
             max_requests_per_minute=2000,
             max_tokens_per_minute=4_000_000,
-        )
+        ),
     ),
     # Gemini 2.5 Flash Preview 0520
     GeminiModelInfo(
@@ -264,7 +264,7 @@ SUPPORTED_MODELS = [
         rate_limits=lf.ModelInfo.RateLimits(
             max_requests_per_minute=2000,
             max_tokens_per_minute=4_000_000,
-        )
+        ),
     ),
     # Gemini 2.5 Flash Preview
     GeminiModelInfo(
@@ -614,6 +614,21 @@ SUPPORTED_MODELS = [
     #
     # Experimental models.
     #
+    GeminiModelInfo(
+        model_id='gemini-2.5-flash-image-preview',
+        in_service=True,
+        experimental=True,
+        provider=pg.oneof(['Google GenAI', 'VertexAI']),
+        model_type='instruction-tuned',
+        description='Gemini 2.5 Flash Image Preview model.',
+        release_date=datetime.datetime(2025, 8, 17),
+        input_modalities=GeminiModelInfo.INPUT_IMAGE_TYPES
+        + GeminiModelInfo.INPUT_DOC_TYPES,
+        context_length=lf.ModelInfo.ContextLength(
+            max_input_tokens=32_768,
+            max_output_tokens=32_768,
+        ),
+    ),
     GeminiModelInfo(
         model_id='gemini-2.0-pro-exp-02-05',
         in_service=True,
@@ -681,7 +696,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(
@@ -690,6 +713,12 @@ class Gemini(rest.REST):
       'The name of the model to use.',
   ]
 
+  response_modalities: pg.typing.Annotated[
+      list[str] | None,
+      'Response modalities. It is needed for models whose response modalities '
+      + 'are more than plain text.',
+  ] = None
+
   @functools.cached_property
   def model_info(self) -> GeminiModelInfo:
     return _SUPPORTED_MODELS_BY_ID[self.model]
@@ -731,6 +760,8 @@ class Gemini(rest.REST):
         prompt.as_format('gemini', chunk_preprocessor=modality_conversion)
     )
     request['contents'] = contents
+    if sampling_options.extras:
+      request.update(sampling_options.extras)
     return request
 
   def _generation_config(
@@ -768,6 +799,11 @@ class Gemini(rest.REST):
           'thinkingBudget': options.max_thinking_tokens,
       }
 
+    if self.response_modalities:
+      config['responseModalities'] = self.response_modalities
+      if 'IMAGE' in self.response_modalities:
+        config.pop('responseLogprobs', None)
+        config.pop('logprobs', None)
     return config
 
   def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
@@ -802,9 +838,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)

langfun/core/llms/google_genai.py

@@ -25,7 +25,35 @@ import pyglove as pg
 @lf.use_init_args(['model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
 class GenAI(gemini.Gemini):
-  """Language models provided by Google GenAI."""
+  """Google GenAI models.
+
+  **Quick Start:**
+
+  ```python
+  import langfun as lf
+
+  # Call Gemini 1.5 Flash using API key from environment variable
+  # 'GOOGLE_API_KEY'.
+  lm = lf.llms.Gemini15Flash()
+  r = lm('Who are you?')
+  print(r)
+  ```
+
+  **Setting up API key:**
+
+  The Google API key can be specified in following ways:
+
+  1. At model instantiation:
+
+     ```python
+     lm = lf.llms.Gemini15Flash(api_key='MY_API_KEY')
+     ```
+  2. via environment variable `GOOGLE_API_KEY`.
+
+  **References:**
+
+  *   https://ai.google.dev/docs
+  """
 
   model: pg.typing.Annotated[
       pg.typing.Enum(
@@ -92,6 +120,11 @@ class GenAI(gemini.Gemini):
 #
 
 
+class Gemini25FlashImagePreview(GenAI):
+  """Gemini 2.5 Flash Image Preview model."""
+  model = 'gemini-2.5-flash-image-preview'
+
+
 class Gemini25Pro(GenAI):
   """Gemini 2.5 Pro GA model."""
 

langfun/core/llms/groq.py

@@ -259,10 +259,35 @@ _SUPPORTED_MODELS_BY_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 
 
 @lf.use_init_args(['model'])
-class Groq(openai_compatible.OpenAICompatible):
-  """Groq LLMs through REST APIs (OpenAI compatible).
+class Groq(openai_compatible.OpenAIChatCompletionAPI):
+  """Groq models.
 
-  See https://platform.openai.com/docs/api-reference/chat
+  **Quick Start:**
+
+  ```python
+  import langfun as lf
+
+  # Call Llama 3.3 70B on Groq using API key from environment variable
+  # 'GROQ_API_KEY'.
+  lm = lf.llms.GroqLlama33_70B_Versatile()
+  r = lm('Who are you?')
+  print(r)
+  ```
+
+  **Setting up API key:**
+
+  The Groq API key can be specified in following ways:
+
+  1. At model instantiation:
+
+     ```python
+     lm = lf.llms.GroqLlama33_70B_Versatile(api_key='MY_API_KEY')
+     ```
+  2. via environment variable `GROQ_API_KEY`.
+
+  **References:**
+
+  *   https://console.groq.com/docs
   """
 
   model: pg.typing.Annotated[

langfun/core/llms/llama_cpp.py

@@ -20,11 +20,30 @@ import pyglove as pg
 
 @pg.use_init_args(['url', 'model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
-class LlamaCppRemote(openai_compatible.OpenAICompatible):
-  """The remote LLaMA C++ model.
+class LlamaCppRemote(openai_compatible.OpenAIChatCompletionAPI):
+  """LLaMA C++ models served via a remote server.
 
-  The Remote LLaMA C++ models can be launched via
-  https://github.com/ggerganov/llama.cpp/tree/master/examples/server
+  This class provides an interface to interact with language models
+  hosted on a LLaMA C++ server, which is compatible with the OpenAI
+  Chat Completions API format.
+
+  **Quick Start:**
+
+  Assuming a LLaMA C++ server is running at `http://localhost:8080`,
+  you can interact with it as follows:
+
+  ```python
+  import langfun as lf
+
+  # If model name is not specified, it will use server's default.
+  lm = lf.llms.LlamaCppRemote(url='http://localhost:8080')
+  r = lm('Who are you?')
+  print(r)
+  ```
+
+  **References:**
+
+  *   https://github.com/ggerganov/llama.cpp/tree/master/examples/server
   """
   url: Annotated[
       str,

langfun/core/llms/openai.py

@@ -49,6 +49,53 @@ class OpenAIModelInfo(lf.ModelInfo):
 #
 
 SUPPORTED_MODELS = [
+    # GPT-5 models
+    OpenAIModelInfo(
+        model_id='gpt-5',
+        alias_for='gpt-5-2025-08-07',
+        in_service=True,
+        model_type='instruction-tuned',
+        description='GPT 5 model (latest stable).',
+        url='https://platform.openai.com/docs/models/gpt-5',
+        input_modalities=OpenAIModelInfo.INPUT_IMAGE_TYPES,
+        context_length=lf.ModelInfo.ContextLength(
+            max_input_tokens=400_000,
+            max_output_tokens=128_000,
+        ),
+        pricing=lf.ModelInfo.Pricing(
+            cost_per_1m_cached_input_tokens=0.125,
+            cost_per_1m_input_tokens=1.25,
+            cost_per_1m_output_tokens=10.0,
+        ),
+        # Tier 5 rate limits.
+        rate_limits=lf.ModelInfo.RateLimits(
+            max_requests_per_minute=15_000,
+            max_tokens_per_minute=40_000_000,
+        ),
+    ),
+    OpenAIModelInfo(
+        model_id='gpt-5-mini',
+        alias_for='gpt-5-mini-2025-08-07',
+        in_service=True,
+        model_type='instruction-tuned',
+        description='GPT 5 mini model (latest stable).',
+        url='https://platform.openai.com/docs/models/gpt-5-mini',
+        input_modalities=OpenAIModelInfo.INPUT_IMAGE_TYPES,
+        context_length=lf.ModelInfo.ContextLength(
+            max_input_tokens=400_000,
+            max_output_tokens=128_000,
+        ),
+        pricing=lf.ModelInfo.Pricing(
+            cost_per_1m_cached_input_tokens=0.025,
+            cost_per_1m_input_tokens=0.25,
+            cost_per_1m_output_tokens=2.0,
+        ),
+        # Tier 5 rate limits.
+        rate_limits=lf.ModelInfo.RateLimits(
+            max_requests_per_minute=180_000_000,
+            max_tokens_per_minute=30_000_000,
+        ),
+    ),
     # GPT-4.1 models
     OpenAIModelInfo(
         model_id='gpt-4.1',
@@ -984,8 +1031,36 @@ _SUPPORTED_MODELS_BY_MODEL_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 
 
 @lf.use_init_args(['model'])
-class OpenAI(openai_compatible.OpenAICompatible):
-  """OpenAI model."""
+class OpenAI(openai_compatible.OpenAIResponsesAPI):
+  """OpenAI models.
+
+  **Quick Start:**
+
+  ```python
+  import langfun as lf
+
+  # Call GPT-4o using API key from environment variable 'OPENAI_API_KEY'.
+  lm = lf.llms.Gpt4o()
+  r = lm('Who are you?')
+  print(r)
+  ```
+
+  **Setting up API key:**
+
+  The OpenAI API key can be specified in following ways:
+
+  1. At model instantiation:
+
+     ```python
+     lm = lf.llms.Gpt4o(api_key='MY_API_KEY')
+     ```
+  2. via environment variable `OPENAI_API_KEY`.
+
+  **References:**
+
+  *   https://platform.openai.com/docs/models
+  *   https://platform.openai.com/docs/api-reference
+  """
 
   model: pg.typing.Annotated[
       pg.typing.Enum(
@@ -994,7 +1069,12 @@ class OpenAI(openai_compatible.OpenAICompatible):
       'The name of the model to use.',
   ]
 
-  api_endpoint: str = 'https://api.openai.com/v1/chat/completions'
+  # Disable message storage by default.
+  sampling_options = lf.LMSamplingOptions(
+      extras={'store': False}
+  )
+
+  api_endpoint: str = 'https://api.openai.com/v1/responses'
 
   api_key: Annotated[
       str | None,
@@ -1069,6 +1149,16 @@ class OpenAI(openai_compatible.OpenAICompatible):
     return super()._request_args(options)
 
 
+class Gpt5(OpenAI):
+  """GPT-5."""
+  model = 'gpt-5'
+
+
+class Gpt5Mini(OpenAI):
+  """GPT-5 mini."""
+  model = 'gpt-5-mini'
+
+
 class Gpt41(OpenAI):
   """GPT-4.1."""
   model = 'gpt-4.1'