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

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

langfun/core/mcp/tool_test.py

@@ -0,0 +1,197 @@
+# 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 tool."""
+
+import base64
+import inspect
+import unittest
+
+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.mcp import client as mcp_client
+from langfun.core.mcp import tool as mcp_tool
+import mcp
+from mcp.server import fastmcp as fastmcp_lib
+import pyglove as pg
+
+
+# MCP server setup for testing.
+_mcp_server = fastmcp_lib.FastMCP(host='0.0.0.0', port=1235)
+
+
+@_mcp_server.tool()
+async def add(a: int, b: int) -> int:
+  """Adds two integers."""
+  return a + b
+
+
+class McpToolTest(unittest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.client = mcp_client.McpClient.from_fastmcp(_mcp_server)
+    self.tools = self.client.list_tools()
+
+  def test_snake_to_camel(self):
+    self.assertEqual(mcp_tool._snake_to_camel('foo_bar'), 'FooBar')
+    self.assertEqual(mcp_tool._snake_to_camel('foo'), 'Foo')
+
+  def test_base64_decode(self):
+    self.assertEqual(
+        mcp_tool._base64_decode(base64.b64encode(b'foo').decode('utf-8')),
+        b'foo'
+    )
+
+  def test_make_input_class(self):
+    schema = pg.Schema(
+        description='Foo input.',
+        fields=[
+            pg.typing.Field('x', pg.typing.Int(), 'Integer x.'),
+            pg.typing.Field('y', pg.typing.Str(), 'String y.'),
+        ],
+    )
+    input_cls = mcp_tool.McpToolInput.make_class('foo_input', schema)
+    self.assertTrue(issubclass(input_cls, mcp_tool.McpToolInput))
+    self.assertEqual(input_cls.__name__, 'FooInput')
+    self.assertEqual(input_cls.__doc__, 'Foo input.')
+    s = input_cls.__schema__
+    self.assertEqual(list(s.fields.keys()), ['x', 'y'])
+    self.assertEqual(repr(input_cls), "<input-class 'FooInput'>")
+    self.assertEqual(
+        repr(input_cls(x=1, y='abc')),
+        "FooInput(x=1, y='abc')",
+    )
+
+  def test_make_tool_class(self):
+    tool_def = mcp.Tool(
+        name='my_tool',
+        inputSchema={
+            'type': 'object',
+            'properties': {
+                'a': {'type': 'integer', 'description': 'Integer a.'},
+                'b': {'type': 'string', 'description': 'String b.'},
+            },
+            'required': ['a'],
+        },
+        description='My tool.',
+    )
+    tool_cls = mcp_tool.McpTool.make_class(tool_def)
+    self.assertTrue(issubclass(tool_cls, mcp_tool.McpTool))
+    self.assertEqual(tool_cls.__name__, 'MyTool')
+    self.assertEqual(tool_cls.TOOL_NAME, 'my_tool')
+    self.assertEqual(tool_cls.__doc__, 'My tool.')
+    s = tool_cls.__schema__
+    self.assertEqual(list(s.fields.keys()), ['a', 'b'])
+    self.assertEqual(repr(tool_cls), "<tool-class 'MyTool'>")
+    self.assertEqual(s.fields['a'].description, 'Integer a.')
+    self.assertEqual(s.fields['b'].description, 'String b.')
+
+    self.assertEqual(
+        tool_cls.python_definition(markdown=True),
+        inspect.cleandoc(
+            """
+            MyTool
+
+            ```python
+            class MyTool:
+              \"\"\"My tool.\"\"\"
+              # Integer a.
+              a: int
+              # String b.
+              b: str | None
+            ```
+            """
+        ),
+    )
+    self.assertEqual(
+        tool_cls.python_definition(markdown=False),
+        inspect.cleandoc(
+            """
+            MyTool
+
+            class MyTool:
+              \"\"\"My tool.\"\"\"
+              # Integer a.
+              a: int
+              # String b.
+              b: str | None
+            """
+        ),
+    )
+
+  def test_input_parameters(self):
+    tool_cls = self.tools['add']
+    self.assertEqual(tool_cls(a=1, b=2).input_parameters(), {'a': 1, 'b': 2})
+
+  def test_result_to_message(self):
+    img_data = base64.b64encode(b'image-data').decode('utf-8')
+    audio_data = base64.b64encode(b'audio-data').decode('utf-8')
+
+    tool_def = self.tools['add']
+    result = mcp.types.CallToolResult(
+        content=[
+            mcp.types.TextContent(type='text', text='hello'),
+            mcp.types.ImageContent(
+                type='image', data=img_data, mimeType='image/png'
+            ),
+            mcp.types.AudioContent(
+                type='audio', data=audio_data, mimeType='audio/wav'
+            ),
+        ],
+        structuredContent={'x': 1},
+    )
+    message = tool_def.result_to_message(result)
+    self.assertIsInstance(message, lf_message.ToolMessage)
+    self.assertIn('hello', message.text)
+    self.assertIn('<<[[image', message.text)
+    self.assertIn('<<[[audio', message.text)
+    self.assertEqual(message.metadata, {'x': 1})
+    modalities = message.modalities()
+    self.assertEqual(len(modalities), 2)
+    self.assertIsInstance(modalities[0], lf_modalities.Image)
+    self.assertEqual(modalities[0].to_bytes(), b'image-data')
+    self.assertIsInstance(modalities[1], lf_modalities.Audio)
+    self.assertEqual(modalities[1].to_bytes(), b'audio-data')
+
+  def test_sync_call(self):
+    add_tool_cls = self.tools['add']
+    with self.client.session() as session:
+      # Test returning structured content.
+      self.assertEqual(add_tool_cls(a=1, b=2)(session), 3)
+
+      # Test returning message.
+      self.assertEqual(
+          add_tool_cls(a=1, b=2)(session, returns_message=True),
+          lf_message.ToolMessage(text='3', result=3),
+      )
+
+  def test_async_call(self):
+    async def _test():
+      add_tool_cls = self.tools['add']
+      async with self.client.session() as session:
+        # Test returning structured content.
+        self.assertEqual(await add_tool_cls(a=1, b=2).acall(session), 3)
+
+        # Test returning message.
+        self.assertEqual(
+            await add_tool_cls(a=1, b=2).acall(session, returns_message=True),
+            lf_message.ToolMessage(text='3', result=3),
+        )
+
+    async_support.invoke_sync(_test)
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/memory.py

@@ -43,6 +43,7 @@ class Memory(NaturalLanguageFormattable, Component):
       value: Any,
       **kwargs
   ) -> None:
+    """Remembers a value."""
     self._remember(value, **kwargs)
 
   def reset(self, **kwargs) -> None: