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

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'
@@ -180,6 +213,33 @@ class VertexAIGemini(VertexAI, gemini.Gemini):
 #
 # Production models.
 #
+class VertexAIGemini3ProPreview(VertexAIGemini):  # pylint: disable=invalid-name
+  """Gemini 3 Pro Preview model launched on 11/18/2025."""
+
+  model = 'gemini-3-pro-preview'
+  location = 'global'
+
+
+class VertexAIGemini3ProImagePreview(VertexAIGemini):  # pylint: disable=invalid-name
+  """Gemini 3 Pro Image Preview model for high-fidelity image generation.
+
+  This model supports:
+  - Text-to-image generation
+  - Image editing (multimodal input)
+  - Visual reasoning
+
+  Key Requirements:
+  - Uses v1beta1 API endpoint
+  - responseModalities must include 'IMAGE'
+  - Supported aspect ratios: 1:1, 16:9, 9:16, 4:3, 3:4
+  - Image sizes: 1K (default), 2K, 4K
+  """
+
+  model = 'gemini-3-pro-image-preview'
+  location = 'global'
+  response_modalities = ['TEXT', 'IMAGE']
+
+
 class VertexAIGemini25Pro(VertexAIGemini):  # pylint: disable=invalid-name
   """Gemini 2.5 Pro GA model launched on 06/17/2025."""
 
@@ -369,6 +429,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 +557,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 +670,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/__init__.py

@@ -0,0 +1,10 @@
+"""Langfun MCP support."""
+
+# pylint: disable=g-importing-member
+
+from langfun.core.mcp.client import McpClient
+from langfun.core.mcp.session import McpSession
+from langfun.core.mcp.tool import McpTool
+from langfun.core.mcp.tool import McpToolInput
+
+# pylint: enable=g-importing-member

langfun/core/mcp/client.py

@@ -0,0 +1,177 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MCP client."""
+
+import abc
+from typing import Annotated, Type
+
+from langfun.core.mcp import session as mcp_session
+from langfun.core.mcp import tool as mcp_tool
+from mcp.server import fastmcp as fastmcp_lib
+import pyglove as pg
+
+
+class McpClient(pg.Object):
+  """Interface for Model Context Protocol (MCP) client.
+
+  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
+
+  # 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())
+
+  with client.session() as session:
+    result = tool_cls(x=1, y=2)(session)
+    print(result)
+
+  # 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(tool_cls.python_definition())
+
+    async with client.session() as session:
+      result = await tool_cls(x=1, y=2).acall(session)
+      print(result)
+  ```
+  """
+
+  def _on_bound(self):
+    super()._on_bound()
+    self._tools = None
+
+  def list_tools(
+      self, refresh: bool = False
+  ) -> dict[str, Type[mcp_tool.McpTool]]:
+    """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()
+    return self._tools
+
+  @abc.abstractmethod
+  def session(self) -> mcp_session.McpSession:
+    """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 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
+  def from_url(
+      cls,
+      url: str,
+      headers: dict[str, str] | None = None
+  ) -> 'McpClient':
+    """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 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)
+
+
+class _StdioMcpClient(McpClient):
+  """Stdio-based MCP client."""
+
+  command: Annotated[str, 'Command to execute.']
+  args: Annotated[list[str], 'Arguments to pass to the command.']
+
+  def session(self) -> mcp_session.McpSession:
+    """Creates an McpSession from command."""
+    return mcp_session.McpSession.from_command(self.command, self.args)
+
+
+class _HttpMcpClient(McpClient):
+  """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 an McpSession from URL."""
+    return mcp_session.McpSession.from_url(self.url, self.headers)
+
+
+class _InMemoryFastMcpClient(McpClient):
+  """In-memory MCP client."""
+
+  fastmcp: Annotated[fastmcp_lib.FastMCP, 'MCP server to connect to.']
+
+  def session(self) -> mcp_session.McpSession:
+    """Creates an McpSession from an in-memory FastMCP instance."""
+    return mcp_session.McpSession.from_fastmcp(self.fastmcp)

langfun/core/mcp/client_test.py

@@ -0,0 +1,71 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for MCP client."""
+
+import unittest
+from langfun.core import async_support
+from langfun.core import mcp as lf_mcp
+from langfun.core import message as lf_message
+from mcp.server import fastmcp as fastmcp_lib
+
+mcp = fastmcp_lib.FastMCP(host='0.0.0.0', port=1234)
+
+
+@mcp.tool()
+async def add(a: int, b: int) -> int:
+  """Adds two integers and returns their sum.
+
+  Args:
+    a: The first integer.
+    b: The second integer.
+
+  Returns:
+    The sum of the two integers.
+  """
+  return a + b
+
+
+class McpTest(unittest.TestCase):
+
+  def test_sync_usages(self):
+    client = lf_mcp.McpClient.from_fastmcp(mcp)
+    tools = client.list_tools()
+    self.assertEqual(len(tools), 1)
+    with client.session() as session:
+      self.assertEqual(
+          # Test `session.call_tool` method as `tool.__call__` is already tested
+          # in `tool_test.py`.
+          session.call_tool(tools['add'](a=1, b=2)), 3
+      )
+
+  def test_async_usages(self):
+    async def _test():
+      client = lf_mcp.McpClient.from_fastmcp(mcp)
+      tools = client.list_tools()
+      self.assertEqual(len(tools), 1)
+      tool_cls = tools['add']
+      self.assertEqual(tool_cls.__name__, 'Add')
+      self.assertEqual(tool_cls.TOOL_NAME, 'add')
+      async with client.session() as session:
+        self.assertEqual(
+            # Test `session.acall_tool` method as `tool.acall` is already
+            # tested in `tool_test.py`.
+            await session.acall_tool(tool_cls(a=1, b=2), returns_message=True),
+            lf_message.ToolMessage(text='3', result=3)
+        )
+    async_support.invoke_sync(_test)
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/mcp/session.py

@@ -0,0 +1,241 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MCP session."""
+
+import contextlib
+from typing import Any, Type
+import anyio
+from langfun.core import async_support
+from langfun.core.mcp import tool as mcp_tool
+import mcp
+from mcp.client import sse
+from mcp.client import streamable_http
+from mcp.server import fastmcp as fastmcp_lib
+from mcp.shared import memory
+
+
+class McpSession:
+  """Represents a session for interacting with an MCP server.
+
+  `McpSession` provides the context for making calls to tools hosted on an
+  MCP server. It wraps the standard `mcp.ClientSession` to offer both
+  synchronous and asynchronous usage patterns.
+
+  Sessions are created using `lf.mcp.McpClient.session()` and should be used
+  as context managers (either sync or async) to ensure proper initialization
+  and teardown of the connection to the server.
+
+  **Example Sync Usage:**
+
+  ```python
+  import langfun as lf
+
+  client = lf.mcp.McpClient.from_command(...)
+  with client.session() as session:
+    tools = session.list_tools()
+    result = tools['my_tool'](x=1)(session)
+  ```
+
+  **Example Async Usage:**
+
+  ```python
+  import langfun as lf
+
+  client = lf.mcp.McpClient.from_url(...)
+  async with client.session() as session:
+    tools = await session.alist_tools()
+    result = await tools['my_tool'](x=1).acall(session)
+  ```
+  """
+
+  def __init__(self, stream) -> None:
+    self._stream = stream
+    self._session = None
+    self._session_exit_stack = None
+    self._in_session = False
+
+    # For supporting sync context manager.
+    self._sync_context_manager_exit_stack = None
+
+  def __enter__(self) -> 'McpSession':
+    exit_stack = contextlib.ExitStack()
+    exit_stack.enter_context(async_support.sync_context_manager(self))
+    self._sync_context_manager_exit_stack = exit_stack
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    assert self._sync_context_manager_exit_stack is not None
+    self._sync_context_manager_exit_stack.close()
+
+  async def __aenter__(self) -> 'McpSession':
+    assert self._session_exit_stack is None, 'Session cannot be re-entered.'
+
+    self._session_exit_stack = contextlib.AsyncExitStack()
+    stream_output = await self._session_exit_stack.enter_async_context(
+        self._stream
+    )
+    assert isinstance(stream_output, tuple) and len(stream_output) in (2, 3)
+    read, write = stream_output[:2]
+    self._session = mcp.ClientSession(read, write)
+    await self._session_exit_stack.enter_async_context(self._session)
+    await self._session.initialize()
+    return self
+
+  async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+    del exc_type, exc_val, exc_tb
+    if self._session is None:
+      return
+    assert self._session_exit_stack is not None
+    await self._session_exit_stack.aclose()
+    self._session = None
+
+  def list_tools(self) -> dict[str, Type[mcp_tool.McpTool]]:
+    """Lists all available tools on the MCP server synchronously.
+
+    Returns:
+      A dictionary mapping tool names to their corresponding `McpTool` classes.
+    """
+    return async_support.invoke_sync(self.alist_tools)
+
+  async def alist_tools(self) -> dict[str, Type[mcp_tool.McpTool]]:
+    """Lists all available tools on the MCP server asynchronously.
+
+    Returns:
+      A dictionary mapping tool names to their corresponding `McpTool` classes.
+    """
+    assert self._session is not None, 'MCP session is not entered.'
+    return {
+        t.name: mcp_tool.McpTool.make_class(t)
+        for t in (await self._session.list_tools()).tools
+    }
+
+  def call_tool(
+      self,
+      tool: mcp_tool.McpTool,
+      *,
+      returns_message: bool = False
+  ) -> Any:
+    """Calls an MCP tool synchronously.
+
+    Args:
+      tool: The `McpTool` instance to call.
+      returns_message: If True, the tool call will return an `mcp.Message`
+        object; otherwise, it returns the tool's direct result.
+
+    Returns:
+      The result of the tool call.
+    """
+    return tool(self, returns_message=returns_message)
+
+  async def acall_tool(
+      self,
+      tool: mcp_tool.McpTool,
+      *,
+      returns_message: bool = False
+  ) -> Any:
+    """Calls an MCP tool asynchronously.
+
+    Args:
+      tool: The `McpTool` instance to call.
+      returns_message: If True, the tool call will return an `mcp.Message`
+        object; otherwise, it returns the tool's direct result.
+
+    Returns:
+      The result of the tool call.
+    """
+    return await tool.acall(self, returns_message=returns_message)
+
+  @classmethod
+  def from_command(
+      cls,
+      command: str,
+      args: list[str] | None = None
+  ) -> 'McpSession':
+    """Creates an MCP session from a command-line executable.
+
+    Args:
+      command: The command to execute.
+      args: An optional list of arguments to pass to the command.
+
+    Returns:
+      An `McpSession` instance.
+    """
+    return cls(
+        mcp.stdio_client(
+            mcp.StdioServerParameters(command=command, args=args or [])
+        )
+    )
+
+  @classmethod
+  def from_url(
+      cls,
+      url: str,
+      headers: dict[str, str] | None = None
+  ) -> 'McpSession':
+    """Creates an MCP session from an HTTP URL.
+
+    The transport protocol (e.g., 'mcp' or 'sse') is inferred from the
+    last part of the URL path.
+
+    Args:
+      url: The URL of the MCP server.
+      headers: An optional dictionary of HTTP headers to include in requests.
+
+    Returns:
+      An `McpSession` instance.
+    """
+    transport = url.removesuffix('/').split('/')[-1].lower()
+    if transport == 'mcp':
+      return cls(streamable_http.streamablehttp_client(url, headers or {}))
+    elif transport == 'sse':
+      return cls(sse.sse_client(url, headers or {}))
+    else:
+      raise ValueError(f'Unsupported transport: {transport}')
+
+  @classmethod
+  def from_fastmcp(
+      cls,
+      fastmcp: fastmcp_lib.FastMCP
+  ):
+    """Creates an MCP session from an in-memory FastMCP instance.
+
+    Args:
+      fastmcp: An instance of `fastmcp_lib.FastMCP`.
+
+    Returns:
+      An `McpSession` instance.
+    """
+    return cls(_client_streams_from_fastmcp(fastmcp))
+
+
+@contextlib.asynccontextmanager
+async def _client_streams_from_fastmcp(fastmcp: fastmcp_lib.FastMCP):
+  """Creates client streams from an in-memory FastMCP instance."""
+  server = fastmcp._mcp_server  # pylint: disable=protected-access
+  async with memory.create_client_server_memory_streams(
+      ) as (client_streams, server_streams):
+    client_read, client_write = client_streams
+    server_read, server_write = server_streams
+
+    # Create a cancel scope for the server task
+    async with anyio.create_task_group() as tg:
+      tg.start_soon(
+          lambda: server.run(
+              server_read,
+              server_write,
+              server.create_initialization_options(),
+              raise_exceptions=True,
+          )
+      )
+      yield client_read, client_write

langfun/core/mcp/session_test.py

@@ -0,0 +1,54 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for MCP session."""
+
+import unittest
+from unittest import mock
+
+from langfun.core.mcp import session as mcp_session
+import mcp
+from mcp.client import sse
+from mcp.client import streamable_http
+
+
+class McpSessionTest(unittest.TestCase):
+
+  @mock.patch.object(mcp, 'stdio_client', autospec=True)
+  def test_from_command(self, mock_stdio_client):
+    mcp_session.McpSession.from_command('my-command', ['--foo'])
+    mock_stdio_client.assert_called_once_with(
+        mcp.StdioServerParameters(command='my-command', args=['--foo'])
+    )
+
+  @mock.patch.object(streamable_http, 'streamablehttp_client', autospec=True)
+  def test_from_url_mcp(self, mock_streamablehttp_client):
+    mcp_session.McpSession.from_url(
+        'http://localhost/mcp', headers={'k': 'v'}
+    )
+    mock_streamablehttp_client.assert_called_once_with(
+        'http://localhost/mcp', {'k': 'v'}
+    )
+
+  @mock.patch.object(sse, 'sse_client', autospec=True)
+  def test_from_url_sse(self, mock_sse_client):
+    mcp_session.McpSession.from_url('http://localhost/sse', headers={'k': 'v'})
+    mock_sse_client.assert_called_once_with('http://localhost/sse', {'k': 'v'})
+
+  def test_from_url_unsupported(self):
+    with self.assertRaisesRegex(ValueError, 'Unsupported transport: foo'):
+      mcp_session.McpSession.from_url('http://localhost/foo')
+
+
+if __name__ == '__main__':
+  unittest.main()