langfun 0.1.2.dev202511030805__py3-none-any.whl → 0.1.2.dev202511050805__py3-none-any.whl

langfun/core/language_model.py

@@ -478,7 +478,7 @@ class UsageNotAvailable(LMSamplingUsage):
 
 
 class LMSamplingResult(pg.Object):
-  """Language model response."""
+  """The result from a language model sampling."""
 
   samples: Annotated[
       list[LMSample],
@@ -681,13 +681,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()
@@ -1448,7 +1526,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.
@@ -1521,7 +1599,7 @@ class _ConcurrencyControl:
 
 
 class UsageSummary(pg.Object, pg.views.HtmlTreeView.Extension):
-  """Usage sumary."""
+  """Usage summary."""
 
   class AggregatedUsage(pg.Object):
     """Aggregated usage."""

langfun/core/llms/anthropic.py

@@ -650,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.
 
-  See https://docs.anthropic.com/claude/reference/messages_post
+  **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`.
+
+  **References:**
+
+  *   https://docs.anthropic.com/claude/reference/messages_post
   """
 
   model: pg.typing.Annotated[

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,

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/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

@@ -94,7 +94,35 @@ _SUPPORTED_MODELS_BY_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 # Reference: https://api-docs.deepseek.com/
 @lf.use_init_args(['model'])
 class DeepSeek(openai_compatible.OpenAIChatCompletionAPI):
-  """DeepSeek model."""
+  """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

@@ -696,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(

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(

langfun/core/llms/groq.py

@@ -260,9 +260,34 @@ _SUPPORTED_MODELS_BY_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 
 @lf.use_init_args(['model'])
 class Groq(openai_compatible.OpenAIChatCompletionAPI):
-  """Groq LLMs through REST APIs (OpenAI compatible).
+  """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

@@ -21,10 +21,29 @@ import pyglove as pg
 @pg.use_init_args(['url', 'model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
 class LlamaCppRemote(openai_compatible.OpenAIChatCompletionAPI):
-  """The remote LLaMA C++ model.
+  """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

@@ -1032,7 +1032,35 @@ _SUPPORTED_MODELS_BY_MODEL_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 
 @lf.use_init_args(['model'])
 class OpenAI(openai_compatible.OpenAIResponsesAPI):
-  """OpenAI model."""
+  """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(

langfun/core/llms/openai_compatible.py

@@ -24,11 +24,16 @@ import pyglove as pg
 
 @lf.use_init_args(['api_endpoint', 'model'])
 class OpenAIChatCompletionAPI(rest.REST):
-  """Base for OpenAI compatible models based on ChatCompletion API.
+  """Base class for models compatible with OpenAI's Chat Completion API.
 
-  See https://platform.openai.com/docs/api-reference/chat
-  As of 2025-10-23, OpenAI is migrating from ChatCompletion API to Responses
-  API.
+  This class provides a common interface for language models that adhere to
+  the OpenAI Chat Completion API format, which is used by providers like
+  Groq, DeepSeek, and others. It standardizes request formatting and
+  response parsing for these models.
+
+  **References:**
+
+  *   https://platform.openai.com/docs/api-reference/chat
   """
 
   model: Annotated[
@@ -196,9 +201,16 @@ class OpenAIChatCompletionAPI(rest.REST):
 
 
 class OpenAIResponsesAPI(OpenAIChatCompletionAPI):
-  """Base for OpenAI compatible models based on Responses API.
+  """Base class for models compatible with OpenAI's Responses API.
+
+  This class provides a common interface for language models that adhere to
+  the new OpenAI Responses API format. It standardizes request formatting
+  and response parsing for these models, including handling instructions
+  (system messages) and structured outputs.
+
+  **References:**
 
-  https://platform.openai.com/docs/api-reference/responses/create
+  *   https://platform.openai.com/docs/api-reference/responses
   """
 
   def _request_args(

langfun/core/llms/rest.py

@@ -22,7 +22,18 @@ import requests
 
 
 class REST(lf.LanguageModel):
-  """REST-based language model."""
+  """Base class for language models accessed via REST APIs.
+
+  The `REST` class provides a foundation for implementing language models
+  that are accessed through RESTful endpoints. It handles the details of
+  making HTTP requests, managing sessions, and handling common errors like
+  timeouts and connection issues.
+
+  Subclasses need to implement the `request` and `result` methods to
+  convert Langfun messages to API-specific request formats and to parse
+  API responses back into `LMSamplingResult` objects. They also need to
+  provide the `api_endpoint` and can override `headers` for authentication.
+  """
 
   api_endpoint: Annotated[
       str,