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

langfun/core/language_model_test.py

@@ -656,11 +656,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
 
@@ -151,6 +152,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 +172,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',
@@ -658,6 +768,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 +791,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/cache/base.py

@@ -121,4 +121,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_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/deepseek.py

@@ -93,7 +93,7 @@ _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):
+class DeepSeek(openai_compatible.OpenAIChatCompletionAPI):
   """DeepSeek model."""
 
   model: pg.typing.Annotated[

langfun/core/llms/gemini.py

@@ -752,11 +752,8 @@ class Gemini(rest.REST):
         prompt.as_format('gemini', chunk_preprocessor=modality_conversion)
     )
     request['contents'] = contents
-    # Users could use `metadata_gemini_tools` to pass Gemini tools. For example,
-    # for enabling Search Grounding, users could pass:
-    # metadata_gemini_tools=[{'google_search': {}}]
-    if tools := prompt.metadata.get('gemini_tools'):
-      request['tools'] = tools
+    if sampling_options.extras:
+      request.update(sampling_options.extras)
     return request
 
   def _generation_config(

langfun/core/llms/groq.py

@@ -259,7 +259,7 @@ _SUPPORTED_MODELS_BY_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 
 
 @lf.use_init_args(['model'])
-class Groq(openai_compatible.OpenAICompatible):
+class Groq(openai_compatible.OpenAIChatCompletionAPI):
   """Groq LLMs through REST APIs (OpenAI compatible).
 
   See https://platform.openai.com/docs/api-reference/chat

langfun/core/llms/llama_cpp.py

@@ -20,7 +20,7 @@ import pyglove as pg
 
 @pg.use_init_args(['url', 'model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
-class LlamaCppRemote(openai_compatible.OpenAICompatible):
+class LlamaCppRemote(openai_compatible.OpenAIChatCompletionAPI):
   """The remote LLaMA C++ model.
 
   The Remote LLaMA C++ models can be launched via

langfun/core/llms/openai.py

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

langfun/core/llms/openai_compatible.py

@@ -23,8 +23,13 @@ import pyglove as pg
 
 
 @lf.use_init_args(['api_endpoint', 'model'])
-class OpenAICompatible(rest.REST):
-  """Base for OpenAI compatible models."""
+class OpenAIChatCompletionAPI(rest.REST):
+  """Base for OpenAI compatible models based on ChatCompletion API.
+
+  See https://platform.openai.com/docs/api-reference/chat
+  As of 2025-10-23, OpenAI is migrating from ChatCompletion API to Responses
+  API.
+  """
 
   model: Annotated[
       str, 'The name of the model to use.',
@@ -42,12 +47,14 @@ class OpenAICompatible(rest.REST):
     # Reference:
     # https://platform.openai.com/docs/api-reference/completions/create
     # NOTE(daiyip): options.top_k is not applicable.
-    args = dict(
-        n=options.n,
-        top_logprobs=options.top_logprobs,
-    )
+    args = {}
+
     if self.model:
       args['model'] = self.model
+    if options.n != 1:
+      args['n'] = options.n
+    if options.top_logprobs is not None:
+      args['top_logprobs'] = options.top_logprobs
     if options.logprobs:
       args['logprobs'] = options.logprobs
     if options.temperature is not None:
@@ -62,6 +69,8 @@ class OpenAICompatible(rest.REST):
       args['seed'] = options.random_seed
     if options.reasoning_effort is not None:
       args['reasoning_effort'] = options.reasoning_effort
+    if options.extras:
+      args.update(options.extras)
     return args
 
   def request(
@@ -72,27 +81,13 @@ class OpenAICompatible(rest.REST):
     """Returns the JSON input for a message."""
     request_args = self._request_args(sampling_options)
 
-    # Users could use `metadata_json_schema` to pass additional
-    # request arguments.
-    json_schema = prompt.metadata.get('json_schema')
-    if json_schema is not None:
-      if not isinstance(json_schema, dict):
-        raise ValueError(
-            f'`json_schema` must be a dict, got {json_schema!r}.'
-        )
-      if 'title' not in json_schema:
-        raise ValueError(
-            f'The root of `json_schema` must have a `title` field, '
-            f'got {json_schema!r}.'
-        )
+    # Handle structured output.
+    output_schema = self._structure_output_schema(prompt)
+    if output_schema is not None:
       request_args.update(
           response_format=dict(
               type='json_schema',
-              json_schema=dict(
-                  schema=json_schema,
-                  name=json_schema['title'],
-                  strict=True,
-              )
+              json_schema=output_schema,
           )
       )
       prompt.metadata.formatted_text = (
@@ -118,17 +113,43 @@ class OpenAICompatible(rest.REST):
       assert isinstance(system_message, lf.SystemMessage), type(system_message)
       messages.append(
           system_message.as_format(
-              'openai', chunk_preprocessor=modality_check
+              'openai_chat_completion_api', chunk_preprocessor=modality_check
           )
       )
     messages.append(
-        prompt.as_format('openai', chunk_preprocessor=modality_check)
+        prompt.as_format(
+            'openai_chat_completion_api',
+            chunk_preprocessor=modality_check
+        )
     )
     request = dict()
     request.update(request_args)
     request['messages'] = messages
     return request
 
+  def _structure_output_schema(
+      self, prompt: lf.Message
+  ) -> dict[str, Any] | None:
+    # Users could use `metadata_json_schema` to pass additional
+    # request arguments.
+    json_schema = prompt.metadata.get('json_schema')
+    if json_schema is not None:
+      if not isinstance(json_schema, dict):
+        raise ValueError(
+            f'`json_schema` must be a dict, got {json_schema!r}.'
+        )
+      if 'title' not in json_schema:
+        raise ValueError(
+            f'The root of `json_schema` must have a `title` field, '
+            f'got {json_schema!r}.'
+        )
+      return dict(
+          schema=json_schema,
+          name=json_schema['title'],
+          strict=True,
+      )
+    return None
+
   def _parse_choice(self, choice: dict[str, Any]) -> lf.LMSample:
     # Reference:
     # https://platform.openai.com/docs/api-reference/chat/object
@@ -144,7 +165,10 @@ class OpenAICompatible(rest.REST):
           for t in choice_logprobs['content']
       ]
     return lf.LMSample(
-        lf.Message.from_value(choice['message'], format='openai'),
+        lf.Message.from_value(
+            choice['message'],
+            format='openai_chat_completion_api'
+        ),
         score=0.0,
         logprobs=logprobs,
     )
@@ -169,3 +193,88 @@ class OpenAICompatible(rest.REST):
         or (status_code == 400 and b'string_above_max_length' in content)):
       return lf.ContextLimitError(f'{status_code}: {content}')
     return super()._error(status_code, content)
+
+
+class OpenAIResponsesAPI(OpenAIChatCompletionAPI):
+  """Base for OpenAI compatible models based on Responses API.
+
+  https://platform.openai.com/docs/api-reference/responses/create
+  """
+
+  def _request_args(
+      self, options: lf.LMSamplingOptions) -> dict[str, Any]:
+    """Returns a dict as request arguments."""
+    if options.logprobs:
+      raise ValueError('logprobs is not supported on Responses API.')
+    if options.n != 1:
+      raise ValueError('n must be 1 for Responses API.')
+    return super()._request_args(options)
+
+  def request(
+      self,
+      prompt: lf.Message,
+      sampling_options: lf.LMSamplingOptions
+  ) -> dict[str, Any]:
+    """Returns the JSON input for a message."""
+    request_args = self._request_args(sampling_options)
+
+    # Handle structured output.
+    output_schema = self._structure_output_schema(prompt)
+    if output_schema is not None:
+      output_schema['type'] = 'json_schema'
+      request_args.update(text=dict(format=output_schema))
+      prompt.metadata.formatted_text = (
+          prompt.text
+          + '\n\n [RESPONSE FORMAT (not part of prompt)]\n'
+          + pg.to_json_str(request_args['text'], json_indent=2)
+      )
+
+    request = dict()
+    request.update(request_args)
+
+    # Users could use `metadata_system_message` to pass system message.
+    system_message = prompt.metadata.get('system_message')
+    if system_message:
+      assert isinstance(system_message, lf.SystemMessage), type(system_message)
+      request['instructions'] = system_message.text
+
+    # Prepare input.
+    def modality_check(chunk: str | lf.Modality) -> Any:
+      if (isinstance(chunk, lf_modalities.Mime)
+          and not self.supports_input(chunk.mime_type)):
+        raise ValueError(
+            f'Unsupported modality: {chunk!r}.'
+        )
+      return chunk
+
+    request['input'] = [
+        prompt.as_format(
+            'openai_responses_api',
+            chunk_preprocessor=modality_check
+        )
+    ]
+    return request
+
+  def _parse_output(self, output: dict[str, Any]) -> lf.LMSample:
+    for item in output:
+      if isinstance(item, dict) and item.get('type') == 'message':
+        return lf.LMSample(
+            lf.Message.from_value(item, format='openai_responses_api'),
+            score=0.0,
+        )
+    raise ValueError('No message found in output.')
+
+  def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
+    """Returns a LMSamplingResult from a JSON response."""
+    usage = json['usage']
+    return lf.LMSamplingResult(
+        samples=[self._parse_output(json['output'])],
+        usage=lf.LMSamplingUsage(
+            prompt_tokens=usage['input_tokens'],
+            completion_tokens=usage['output_tokens'],
+            total_tokens=usage['total_tokens'],
+            completion_tokens_details=usage.get(
+                'output_tokens_details', None
+            ),
+        ),
+    )