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

langfun/core/llms/openai_compatible_test.py

@@ -38,7 +38,7 @@ def mock_chat_completion_request(url: str, json: dict[str, Any], **kwargs):
     response_format = ''
 
   choices = []
-  for k in range(json['n']):
+  for k in range(json.get('n', 1)):
     if json.get('logprobs'):
       logprobs = dict(
           content=[
@@ -89,7 +89,7 @@ def mock_chat_completion_request_vision(
       c['image_url']['url']
       for c in json['messages'][0]['content'] if c['type'] == 'image_url'
   ]
-  for k in range(json['n']):
+  for k in range(json.get('n', 1)):
     choices.append(pg.Dict(
         message=pg.Dict(
             content=f'Sample {k} for message: {"".join(urls)}'
@@ -111,12 +111,88 @@ def mock_chat_completion_request_vision(
   return response
 
 
-class OpenAIComptibleTest(unittest.TestCase):
+def mock_responses_request(url: str, json: dict[str, Any], **kwargs):
+  del url, kwargs
+  _ = json['input']
+
+  system_message = ''
+  if 'instructions' in json:
+    system_message = f' system={json["instructions"]}'
+
+  response_format = ''
+  if 'text' in json and 'format' in json['text']:
+    response_format = f' format={json["text"]["format"]["type"]}'
+
+  output = [
+      dict(
+          type='message',
+          content=[
+              dict(
+                  type='output_text',
+                  text=(
+                      f'Sample 0 for message.{system_message}{response_format}'
+                  )
+              )
+          ],
+      )
+  ]
+
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str(
+      dict(
+          output=output,
+          usage=dict(
+              input_tokens=100,
+              output_tokens=100,
+              total_tokens=200,
+          ),
+      )
+  ).encode()
+  return response
+
+
+def mock_responses_request_vision(
+    url: str, json: dict[str, Any], **kwargs
+):
+  del url, kwargs
+  urls = [
+      c['image_url']
+      for c in json['input'][0]['content']
+      if c['type'] == 'input_image'
+  ]
+  output = [
+      pg.Dict(
+          type='message',
+          content=[
+              pg.Dict(
+                  type='output_text',
+                  text=f'Sample 0 for message: {"".join(urls)}',
+              )
+          ],
+      )
+  ]
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str(
+      dict(
+          output=output,
+          usage=dict(
+              input_tokens=100,
+              output_tokens=100,
+              total_tokens=200,
+          ),
+      )
+  ).encode()
+  return response
+
+
+class OpenAIChatCompletionAPITest(unittest.TestCase):
   """Tests for OpenAI compatible language model."""
 
   def test_request_args(self):
     self.assertEqual(
-        openai_compatible.OpenAICompatible(
+        openai_compatible.OpenAIChatCompletionAPI(
             api_endpoint='https://test-server',
             model='test-model'
         )._request_args(
@@ -126,8 +202,6 @@ class OpenAIComptibleTest(unittest.TestCase):
         ),
         dict(
             model='test-model',
-            top_logprobs=None,
-            n=1,
             temperature=1.0,
             stop=['\n'],
             seed=123,
@@ -137,7 +211,7 @@ class OpenAIComptibleTest(unittest.TestCase):
   def test_call_chat_completion(self):
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_chat_completion_request
-      lm = openai_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model',
       )
       self.assertEqual(
@@ -148,7 +222,7 @@ class OpenAIComptibleTest(unittest.TestCase):
   def test_call_chat_completion_with_logprobs(self):
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_chat_completion_request
-      lm = openai_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model',
       )
       results = lm.sample(['hello'], logprobs=True)
@@ -214,13 +288,14 @@ class OpenAIComptibleTest(unittest.TestCase):
       def mime_type(self) -> str:
         return 'image/png'
 
+    image = FakeImage.from_uri('https://fake/image')
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_chat_completion_request_vision
-      lm_1 = openai_compatible.OpenAICompatible(
+      lm_1 = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server',
           model='test-model1',
       )
-      lm_2 = openai_compatible.OpenAICompatible(
+      lm_2 = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server',
           model='test-model2',
       )
@@ -228,15 +303,15 @@ class OpenAIComptibleTest(unittest.TestCase):
         self.assertEqual(
             lm(
                 lf.UserMessage(
-                    'hello <<[[image]]>>',
-                    image=FakeImage.from_uri('https://fake/image')
+                    f'hello <<[[{image.id}]]>>',
+                    referred_modalities=[image],
                 ),
                 sampling_options=lf.LMSamplingOptions(n=2)
             ),
             'Sample 0 for message: https://fake/image',
         )
 
-    class TextOnlyModel(openai_compatible.OpenAICompatible):
+    class TextOnlyModel(openai_compatible.OpenAIChatCompletionAPI):
 
       class ModelInfo(lf.ModelInfo):
         input_modalities: list[str] = lf.ModelInfo.TEXT_INPUT_ONLY
@@ -251,15 +326,15 @@ class OpenAIComptibleTest(unittest.TestCase):
     with self.assertRaisesRegex(ValueError, 'Unsupported modality'):
       lm_3(
           lf.UserMessage(
-              'hello <<[[image]]>>',
-              image=FakeImage.from_uri('https://fake/image')
+              f'hello <<[[{image.id}]]>>',
+              referred_modalities=[image],
           ),
       )
 
   def test_sample_chat_completion(self):
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_chat_completion_request
-      lm = openai_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model'
       )
       results = lm.sample(
@@ -400,7 +475,7 @@ class OpenAIComptibleTest(unittest.TestCase):
   def test_sample_with_contextual_options(self):
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_chat_completion_request
-      lm = openai_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model'
       )
       with lf.use_settings(sampling_options=lf.LMSamplingOptions(n=2)):
@@ -458,7 +533,7 @@ class OpenAIComptibleTest(unittest.TestCase):
   def test_call_with_system_message(self):
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_chat_completion_request
-      lm = openai_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model'
       )
       self.assertEqual(
@@ -475,7 +550,7 @@ class OpenAIComptibleTest(unittest.TestCase):
   def test_call_with_json_schema(self):
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_chat_completion_request
-      lm = openai_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model'
       )
       self.assertEqual(
@@ -515,7 +590,7 @@ class OpenAIComptibleTest(unittest.TestCase):
 
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_context_limit_error
-      lm = openai_compatible.OpenAICompatible(
+      lm = openai_compatible.OpenAIChatCompletionAPI(
           api_endpoint='https://test-server', model='test-model'
       )
       with self.assertRaisesRegex(
@@ -524,5 +599,117 @@ class OpenAIComptibleTest(unittest.TestCase):
         lm(lf.UserMessage('hello'))
 
 
+class OpenAIResponsesAPITest(unittest.TestCase):
+  """Tests for OpenAI compatible language model on Responses API."""
+
+  def test_request_args(self):
+    lm = openai_compatible.OpenAIResponsesAPI(
+        api_endpoint='https://test-server', model='test-model'
+    )
+    # Test valid args.
+    self.assertEqual(
+        lm._request_args(
+            lf.LMSamplingOptions(
+                temperature=1.0, stop=['\n'], n=1, random_seed=123
+            )
+        ),
+        dict(
+            model='test-model',
+            temperature=1.0,
+            stop=['\n'],
+            seed=123,
+        ),
+    )
+    # Test unsupported n.
+    with self.assertRaisesRegex(ValueError, 'n must be 1 for Responses API.'):
+      lm._request_args(lf.LMSamplingOptions(n=2))
+
+    # Test unsupported logprobs.
+    with self.assertRaisesRegex(
+        ValueError, 'logprobs is not supported on Responses API.'
+    ):
+      lm._request_args(lf.LMSamplingOptions(logprobs=True))
+
+  def test_call_responses(self):
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_responses_request
+      lm = openai_compatible.OpenAIResponsesAPI(
+          api_endpoint='https://test-server',
+          model='test-model',
+      )
+      self.assertEqual(lm('hello'), 'Sample 0 for message.')
+
+  def test_call_responses_vision(self):
+    class FakeImage(lf_modalities.Image):
+      @property
+      def mime_type(self) -> str:
+        return 'image/png'
+
+    image = FakeImage.from_uri('https://fake/image')
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_responses_request_vision
+      lm = openai_compatible.OpenAIResponsesAPI(
+          api_endpoint='https://test-server',
+          model='test-model1',
+      )
+      self.assertEqual(
+          lm(
+              lf.UserMessage(
+                  f'hello <<[[{image.id}]]>>',
+                  referred_modalities=[image],
+              )
+          ),
+          'Sample 0 for message: https://fake/image',
+      )
+
+  def test_call_with_system_message(self):
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_responses_request
+      lm = openai_compatible.OpenAIResponsesAPI(
+          api_endpoint='https://test-server', model='test-model'
+      )
+      self.assertEqual(
+          lm(
+              lf.UserMessage(
+                  'hello',
+                  system_message=lf.SystemMessage('hi'),
+              )
+          ),
+          'Sample 0 for message. system=hi',
+      )
+
+  def test_call_with_json_schema(self):
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_responses_request
+      lm = openai_compatible.OpenAIResponsesAPI(
+          api_endpoint='https://test-server', model='test-model'
+      )
+      self.assertEqual(
+          lm(
+              lf.UserMessage(
+                  'hello',
+                  json_schema={
+                      'type': 'object',
+                      'properties': {
+                          'name': {'type': 'string'},
+                      },
+                      'required': ['name'],
+                      'title': 'Person',
+                  },
+              )
+          ),
+          'Sample 0 for message. format=json_schema',
+      )
+
+    # Test bad json schema.
+    with self.assertRaisesRegex(ValueError, '`json_schema` must be a dict'):
+      lm(lf.UserMessage('hello', json_schema='foo'))
+
+    with self.assertRaisesRegex(
+        ValueError, 'The root of `json_schema` must have a `title` field'
+    ):
+      lm(lf.UserMessage('hello', json_schema={}))
+
+
 if __name__ == '__main__':
   unittest.main()

langfun/core/llms/openai_test.py

@@ -61,8 +61,6 @@ class OpenAITest(unittest.TestCase):
         ),
         dict(
             model='gpt-4',
-            top_logprobs=None,
-            n=1,
             temperature=1.0,
             stop=['\n'],
             seed=123,

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,

langfun/core/llms/vertexai.py

@@ -44,13 +44,32 @@ except ImportError:
 
 @pg.use_init_args(['api_endpoint'])
 class VertexAI(rest.REST):
-  """Base class for VertexAI models.
+  """Base class for models served on Vertex AI.
 
-  This class handles the authentication of vertex AI models. Subclasses
-  should implement `request` and `result` methods, as well as the `api_endpoint`
-  property. Or let users to provide them as __init__ arguments.
+  This class handles authentication for Vertex AI models. Subclasses,
+  such as `VertexAIGemini`, `VertexAIAnthropic`, and `VertexAILlama`,
+  provide specific implementations for different model families hosted
+  on Vertex AI.
 
-  Please check out VertexAIGemini in `gemini.py` as an example.
+  **Quick Start:**
+
+  If you are using Langfun from a Google Cloud environment (e.g., GCE, GKE)
+  that has service account credentials, authentication is handled automatically.
+  Otherwise, you might need to set up credentials:
+
+  ```bash
+  gcloud auth application-default login
+  ```
+
+  Then you can use a Vertex AI model:
+
+  ```python
+  import langfun as lf
+
+  lm = lf.llms.VertexAIGemini25Flash(project='my-project', location='global')
+  r = lm('Who are you?')
+  print(r)
+  ```
   """
 
   model: pg.typing.Annotated[
@@ -158,7 +177,21 @@ class VertexAI(rest.REST):
 @pg.use_init_args(['model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
 class VertexAIGemini(VertexAI, gemini.Gemini):
-  """Gemini models served by Vertex AI.."""
+  """Gemini models served on Vertex AI.
+
+  **Quick Start:**
+
+  ```python
+  import langfun as lf
+
+  # Call Gemini 1.5 Flash on Vertex AI.
+  # If project and location are not specified, they will be read from
+  # environment variables 'VERTEXAI_PROJECT' and 'VERTEXAI_LOCATION'.
+  lm = lf.llms.VertexAIGemini25Flash(project='my-project', location='global')
+  r = lm('Who are you?')
+  print(r)
+  ```
+  """
 
   # Set default location to us-central1.
   location = 'us-central1'
@@ -369,6 +402,16 @@ class VertexAIAnthropic(VertexAI, anthropic.Anthropic):
 # pylint: disable=invalid-name
 
 
+class VertexAIClaude45Haiku_20251001(VertexAIAnthropic):
+  """Anthropic's Claude 4.5 Haiku model on VertexAI."""
+  model = 'claude-haiku-4-5@20251001'
+
+
+class VertexAIClaude45Sonnet_20250929(VertexAIAnthropic):
+  """Anthropic's Claude 4.5 Sonnet model on VertexAI."""
+  model = 'claude-sonnet-4-5@20250929'
+
+
 class VertexAIClaude4Opus_20250514(VertexAIAnthropic):
   """Anthropic's Claude 4 Opus model on VertexAI."""
   model = 'claude-opus-4@20250514'
@@ -487,7 +530,7 @@ _LLAMA_MODELS_BY_MODEL_ID = {m.model_id: m for m in LLAMA_MODELS}
 
 @pg.use_init_args(['model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
-class VertexAILlama(VertexAI, openai_compatible.OpenAICompatible):
+class VertexAILlama(VertexAI, openai_compatible.OpenAIChatCompletionAPI):
   """Llama models on VertexAI."""
 
   model: pg.typing.Annotated[
@@ -600,7 +643,7 @@ _MISTRAL_MODELS_BY_MODEL_ID = {m.model_id: m for m in MISTRAL_MODELS}
 
 @pg.use_init_args(['model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
-class VertexAIMistral(VertexAI, openai_compatible.OpenAICompatible):
+class VertexAIMistral(VertexAI, openai_compatible.OpenAIChatCompletionAPI):
   """Mistral AI models on VertexAI."""
 
   model: pg.typing.Annotated[

langfun/core/logging.py

@@ -310,7 +310,7 @@ def warning(
     console: bool = False,
     **kwargs
 ) -> LogEntry:
-  """Logs an info message to the session."""
+  """Logs a warning message to the session."""
   return log('warning', message, indent=indent, console=console, **kwargs)
 
 

langfun/core/mcp/client.py

@@ -23,33 +23,53 @@ import pyglove as pg
 
 
 class McpClient(pg.Object):
-  """Base class for MCP client.
+  """Interface for Model Context Protocol (MCP) client.
 
-  Usage:
+  An MCP client serves as a bridge to an MCP server, enabling users to interact
+  with tools hosted on the server. It provides methods for listing available
+  tools and creating sessions for tool interaction.
+
+  There are three types of MCP clients:
+
+  * **Stdio-based client**: Ideal for interacting with tools exposed as
+    command-line executables through stdin/stdout.
+    Created by `lf.mcp.McpClient.from_command`.
+  * **HTTP-based client**: Designed for tools accessible via HTTP,
+    supporting Server-Sent Events (SSE) for streaming.
+    Created by `lf.mcp.McpClient.from_url`.
+  * **In-memory client**: Useful for testing or embedding MCP servers
+    within the same process.
+    Created by `lf.mcp.McpClient.from_fastmcp`.
+
+  **Example Usage:**
 
   ```python
+  import langfun as lf
 
-  def tool_use():
-    client = lf.mcp.McpClient.from_command('<MCP_CMD>', ['<ARG1>', 'ARG2'])
-    tools = client.list_tools()
-    tool_cls = tools['<TOOL_NAME>']
+  # Example 1: Stdio-based client
+  client = lf.mcp.McpClient.from_command('<MCP_CMD>', ['<ARG1>', 'ARG2'])
+  tools = client.list_tools()
+  tool_cls = tools['<TOOL_NAME>']
 
-    # Print the python definition of the tool.
-    print(tool_cls.python_definition())
+  # Print the Python definition of the tool.
+  print(tool_cls.python_definition())
 
-    with client.session() as session:
-      return tool_cls(x=1, y=2)(session)
+  with client.session() as session:
+    result = tool_cls(x=1, y=2)(session)
+    print(result)
 
-  async def tool_use_async_version():
+  # Example 2: HTTP-based client (async)
+  async def main():
     client = lf.mcp.McpClient.from_url('http://localhost:8000/mcp')
     tools = client.list_tools()
     tool_cls = tools['<TOOL_NAME>']
 
-    # Print the python definition of the tool.
+    # Print the Python definition of the tool.
     print(tool_cls.python_definition())
 
     async with client.session() as session:
-      return await tool_cls(x=1, y=2).acall(session)
+      result = await tool_cls(x=1, y=2).acall(session)
+      print(result)
   ```
   """
 
@@ -60,7 +80,15 @@ class McpClient(pg.Object):
   def list_tools(
       self, refresh: bool = False
   ) -> dict[str, Type[mcp_tool.McpTool]]:
-    """Lists all MCP tools."""
+    """Lists all available tools on the MCP server.
+
+    Args:
+      refresh: If True, forces a refresh of the tool list from the server.
+        Otherwise, a cached list may be returned.
+
+    Returns:
+      A dictionary mapping tool names to their corresponding `McpTool` classes.
+    """
     if self._tools is None or refresh:
       with self.session() as session:
         self._tools = session.list_tools()
@@ -68,11 +96,23 @@ class McpClient(pg.Object):
 
   @abc.abstractmethod
   def session(self) -> mcp_session.McpSession:
-    """Creates a MCP session."""
+    """Creates a new session for interacting with MCP tools.
+
+    Returns:
+      An `McpSession` object.
+    """
 
   @classmethod
   def from_command(cls, command: str, args: list[str]) -> 'McpClient':
-    """Creates a MCP client from a tool."""
+    """Creates an MCP client from a command-line executable.
+
+    Args:
+      command: The command to execute.
+      args: A list of arguments to pass to the command.
+
+    Returns:
+      A `McpClient` instance that communicates via stdin/stdout.
+    """
     return _StdioMcpClient(command=command, args=args)
 
   @classmethod
@@ -81,12 +121,27 @@ class McpClient(pg.Object):
       url: str,
       headers: dict[str, str] | None = None
   ) -> 'McpClient':
-    """Creates a MCP client from a URL."""
+    """Creates an MCP client from an HTTP URL.
+
+    Args:
+      url: The URL of the MCP server.
+      headers: An optional dictionary of HTTP headers to include in requests.
+
+    Returns:
+      A `McpClient` instance that communicates via HTTP.
+    """
     return _HttpMcpClient(url=url, headers=headers or {})
 
   @classmethod
   def from_fastmcp(cls, fastmcp: fastmcp_lib.FastMCP) -> 'McpClient':
-    """Creates a MCP client from a MCP server."""
+    """Creates an MCP client from an in-memory FastMCP instance.
+
+    Args:
+      fastmcp: An instance of `fastmcp_lib.FastMCP`.
+
+    Returns:
+      A `McpClient` instance that communicates with the in-memory server.
+    """
     return _InMemoryFastMcpClient(fastmcp=fastmcp)
 
 
@@ -97,18 +152,18 @@ class _StdioMcpClient(McpClient):
   args: Annotated[list[str], 'Arguments to pass to the command.']
 
   def session(self) -> mcp_session.McpSession:
-    """Creates a MCP session."""
+    """Creates an McpSession from command."""
     return mcp_session.McpSession.from_command(self.command, self.args)
 
 
 class _HttpMcpClient(McpClient):
-  """Server-Sent Events (SSE)/Streamable HTTP-based MCP client."""
+  """HTTP-based MCP client."""
 
   url: Annotated[str, 'URL to connect to.']
   headers: Annotated[dict[str, str], 'Headers to send with the request.'] = {}
 
   def session(self) -> mcp_session.McpSession:
-    """Creates a MCP session."""
+    """Creates an McpSession from URL."""
     return mcp_session.McpSession.from_url(self.url, self.headers)
 
 
@@ -118,5 +173,5 @@ class _InMemoryFastMcpClient(McpClient):
   fastmcp: Annotated[fastmcp_lib.FastMCP, 'MCP server to connect to.']
 
   def session(self) -> mcp_session.McpSession:
-    """Creates a MCP session."""
+    """Creates an McpSession from an in-memory FastMCP instance."""
     return mcp_session.McpSession.from_fastmcp(self.fastmcp)