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

langfun/core/mcp/client_test.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 """Tests for MCP client."""
 
-import inspect
 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)
@@ -42,38 +42,11 @@ class McpTest(unittest.TestCase):
     client = lf_mcp.McpClient.from_fastmcp(mcp)
     tools = client.list_tools()
     self.assertEqual(len(tools), 1)
-    tool_cls = tools['add']
-    print(tool_cls.python_definition())
-    self.assertEqual(
-        tool_cls.python_definition(),
-        inspect.cleandoc(
-            '''
-            Add
-
-            ```python
-            class Add:
-              """Adds two integers and returns their sum.
-
-              Args:
-                a: The first integer.
-                b: The second integer.
-
-              Returns:
-                The sum of the two integers.
-              """
-              a: int
-              b: int
-            ```
-            '''
-        )
-    )
-    self.assertEqual(repr(tool_cls), '<tool-class \'Add\'>')
-    self.assertEqual(tool_cls.__name__, 'Add')
-    self.assertEqual(tool_cls.TOOL_NAME, 'add')
-    self.assertEqual(tool_cls(a=1, b=2).input_parameters(), {'a': 1, 'b': 2})
     with client.session() as session:
       self.assertEqual(
-          tool_cls(a=1, b=2)(session), 3
+          # 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):
@@ -86,10 +59,10 @@ class McpTest(unittest.TestCase):
       self.assertEqual(tool_cls.TOOL_NAME, 'add')
       async with client.session() as session:
         self.assertEqual(
-            (await tool_cls(a=1, b=2).acall(
-                session, returns_call_result=True))
-            .structuredContent['result'],
-            3
+            # 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)
 

langfun/core/mcp/session.py

@@ -26,10 +26,37 @@ from mcp.shared import memory
 
 
 class McpSession:
-  """Langfun's MCP session.
+  """Represents a session for interacting with an MCP server.
 
-  Compared to the standard mcp.ClientSession, Langfun's MCP session could be
-  used both synchronously and asynchronously.
+  `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:
@@ -74,11 +101,19 @@ class McpSession:
     self._session = None
 
   def list_tools(self) -> dict[str, Type[mcp_tool.McpTool]]:
-    """Lists all MCP tools synchronously."""
+    """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 MCP tools asynchronously."""
+    """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)
@@ -89,34 +124,37 @@ class McpSession:
       self,
       tool: mcp_tool.McpTool,
       *,
-      returns_call_result: bool = False
+      returns_message: bool = False
   ) -> Any:
-    """Calls a MCP tool synchronously."""
-    return async_support.invoke_sync(
-        self.acall_tool,
-        tool,
-        returns_call_result=returns_call_result
-    )
+    """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_call_result: bool = False
+      returns_message: bool = False
   ) -> Any:
-    """Calls a MCP tool asynchronously."""
-    assert self._session is not None, 'MCP session is not entered.'
-    tool_call_result = await self._session.call_tool(
-        tool.TOOL_NAME, tool.input_parameters()
-    )
-    if returns_call_result:
-      return tool_call_result
-    if (
-        tool_call_result.structuredContent
-        and 'result' in tool_call_result.structuredContent
-    ):
-      return tool_call_result.structuredContent['result']
-    return tool_call_result.content
+    """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(
@@ -124,7 +162,15 @@ class McpSession:
       command: str,
       args: list[str] | None = None
   ) -> 'McpSession':
-    """Creates a MCP session from a command."""
+    """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 [])
@@ -137,7 +183,18 @@ class McpSession:
       url: str,
       headers: dict[str, str] | None = None
   ) -> 'McpSession':
-    """Creates a MCP session from a URL."""
+    """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 {}))
@@ -151,12 +208,20 @@ class McpSession:
       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 a MCP server."""
+  """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):

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()

langfun/core/mcp/tool.py

@@ -13,8 +13,12 @@
 # limitations under the License.
 """MCP tool."""
 
+import base64
 from typing import Annotated, Any, ClassVar
 
+from langfun.core import async_support
+from langfun.core import message as lf_message
+from langfun.core import modalities as lf_modalities
 from langfun.core.structured import schema as lf_schema
 import mcp
 import pyglove as pg
@@ -27,7 +31,31 @@ class _McpToolMeta(pg.symbolic.ObjectMeta):
 
 
 class McpTool(pg.Object, metaclass=_McpToolMeta):
-  """Base class for MCP tools."""
+  """Represents a tool available on an MCP server.
+
+  `McpTool` is the base class for all tool proxies generated from an MCP
+  server's tool definitions. Users do not typically subclass `McpTool` directly.
+  Instead, tool classes are obtained by calling `lf.mcp.McpClient.list_tools()`
+  or `lf.mcp.McpSession.list_tools()`.
+
+  Once a tool class is obtained, it can be instantiated with input parameters
+  and called via an `McpSession` to execute the tool on the server.
+
+  **Example Usage:**
+
+  ```python
+  import langfun as lf
+
+  client = lf.mcp.McpClient.from_command(...)
+  with client.session() as session:
+    # List tools and get the 'math' tool class.
+    math_tool_cls = session.list_tools()['math']
+
+    # Instantiate the tool with parameters and call it.
+    result = math_tool_cls(x=1, y=2, op='+')(session)
+    print(result)
+  ```
+  """
 
   TOOL_NAME: Annotated[
       ClassVar[str],
@@ -36,42 +64,122 @@ class McpTool(pg.Object, metaclass=_McpToolMeta):
 
   @classmethod
   def python_definition(cls, markdown: bool = True) -> str:
-    """Returns the Python definition of the tool."""
-    return lf_schema.Schema.from_value(cls).schema_str(
-        protocol='python', markdown=markdown
-    )
+    """Returns the Python definition of this tool's input schema.
+
+    This is useful for generating prompts that instruct a language model
+    on how to use the tool.
+
+    Args:
+      markdown: If True, formats the output as a Markdown code block.
+
+    Returns:
+      A string containing the Python definition of the tool's input schema.
+    """
+    return lf_schema.Schema.from_value(cls).schema_repr(markdown=markdown)
+
+  @classmethod
+  def result_to_message(
+      cls, result: mcp.types.CallToolResult
+  ) -> lf_message.ToolMessage:
+    """Converts an `mcp.types.CallToolResult` to an `lf.ToolMessage`.
+
+    This method translates results from the MCP protocol, including text,
+    image, and audio content, into a Langfun `ToolMessage`, making it easy
+    to integrate tool results into Langfun workflows.
+
+    Args:
+      result: The `mcp.types.CallToolResult` object to convert.
+
+    Returns:
+      An `lf.ToolMessage` instance representing the tool result.
+    """
+    chunks = []
+    for item in result.content:
+      if isinstance(item, mcp.types.TextContent):
+        chunk = item.text
+      elif isinstance(item, mcp.types.ImageContent):
+        chunk = lf_modalities.Image.from_bytes(_base64_decode(item.data))
+      elif isinstance(item, mcp.types.AudioContent):
+        chunk = lf_modalities.Audio.from_bytes(_base64_decode(item.data))
+      else:
+        raise ValueError(f'Unsupported item type: {type(item)}')
+      chunks.append(chunk)
+    message = lf_message.ToolMessage.from_chunks(chunks)
+    if result.structuredContent:
+      message.metadata.update(result.structuredContent)
+    return message
 
   def __call__(
       self,
       session,
       *,
-      returns_call_result: bool = False) -> Any:
-    """Calls a MCP tool synchronously.
+      returns_message: bool = False) -> Any:
+    """Calls the MCP tool synchronously within a given session.
 
     Args:
-      session: A MCP session.
-      returns_call_result: If True, returns the call result. Otherwise returns
-        the result from structured content, or return content.
+      session: An `McpSession` object.
+      returns_message: If True, the raw `lf.ToolMessage` is returned.
+        If False(default), the method attempts to return a more specific result:
+          - The `result` field of the `ToolMessage` if it's populated.
+          - The `ToolMessage` itself if it contains multi-modal content.
+          - The `text` field of the `ToolMessage` otherwise.
 
     Returns:
-      The call result, or the result from structured content, or content.
+      The result of the tool call, processed according to `returns_message`.
     """
-    return session.call_tool(
-        self,
-        returns_call_result=returns_call_result
+    return async_support.invoke_sync(
+        self.acall,
+        session,
+        returns_message=returns_message
     )
 
   async def acall(
       self,
       session,
       *,
-      returns_call_result: bool = False) -> Any:
-    """Calls a MCP tool asynchronously."""
-    return await session.acall_tool(self, returns_call_result=returns_call_result)
+      returns_message: bool = False
+  ) -> Any:
+    """Calls the MCP tool asynchronously within a given session.
+
+    Args:
+      session: An `McpSession` object or an `mcp.ClientSession`.
+      returns_message: If True, the raw `lf.ToolMessage` is returned.
+        If False(default), the method attempts to return a more specific result:
+          - The `result` field of the `ToolMessage` if it's populated.
+          - The `ToolMessage` itself if it contains multi-modal content.
+          - The `text` field of the `ToolMessage` otherwise.
+
+    Returns:
+      The result of the tool call, processed according to `returns_message`.
+    """
+    if not isinstance(session, mcp.ClientSession):
+      session = getattr(session, '_session', None)
+    assert session is not None, 'MCP session is not entered.'
+    tool_call_result = await session.call_tool(
+        self.TOOL_NAME, self.input_parameters()
+    )
+    message = self.result_to_message(tool_call_result)
+    if returns_message:
+      return message
+    if message.result:
+      return message.result
+    if message.referred_modalities:
+      return message
+    return message.text
 
   def input_parameters(self) -> dict[str, Any]:
-    """Returns the input parameters of the tool."""
-    json = self.to_json()
+    """Returns the input parameters for the tool call.
+
+    Returns:
+      A dictionary containing the input parameters, formatted for an
+      MCP `call_tool` request.
+    """
+    # Optional fields are represented as fields with default values. Therefore,
+    # we need to remove the default values from the JSON representation of the
+    # tool.
+    json = self.to_json(hide_default_values=True)
+
+    # Remove the type name key from the JSON representation of the tool.
     def _transform(path: pg.KeyPath, x: Any) -> Any:
       del path
       if isinstance(x, dict):
@@ -81,7 +189,15 @@ class McpTool(pg.Object, metaclass=_McpToolMeta):
 
   @classmethod
   def make_class(cls, tool_definition: mcp.Tool) -> type['McpTool']:
-    """Makes a MCP tool class from tool definition."""
+    """Creates an `McpTool` subclass from an MCP tool definition.
+
+    Args:
+      tool_definition: An `mcp.Tool` object containing the tool's metadata.
+
+    Returns:
+      A dynamically generated class that inherits from `McpTool` and
+      represents the defined tool.
+    """
 
     class _McpTool(cls):
       auto_schema = False
@@ -104,11 +220,19 @@ class _McpToolInputMeta(pg.symbolic.ObjectMeta):
 
 
 class McpToolInput(pg.Object, metaclass=_McpToolInputMeta):
-  """Base class for MCP tool inputs."""
+  """Base class for generated MCP tool input schemas."""
 
   @classmethod
   def make_class(cls, name: str, schema: pg.Schema):
-    """Converts a schema to an input class."""
+    """Creates an `McpToolInput` subclass from a schema.
+
+    Args:
+      name: The name of the input class to generate.
+      schema: A `pg.Schema` object defining the input fields.
+
+    Returns:
+      A dynamically generated class that inherits from `McpToolInput`.
+    """
 
     class _McpToolInput(cls):
       pass
@@ -123,3 +247,8 @@ class McpToolInput(pg.Object, metaclass=_McpToolInputMeta):
 def _snake_to_camel(name: str) -> str:
   """Converts a snake_case name to a CamelCase name."""
   return ''.join(x.capitalize() for x in name.split('_'))
+
+
+def _base64_decode(data: str) -> bytes:
+  """Decodes a base64 string."""
+  return base64.b64decode(data.encode('utf-8'))