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

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:

langfun/core/message.py

@@ -20,7 +20,7 @@ import contextlib
 import functools
 import inspect
 import io
-from typing import Annotated, Any, ClassVar, Optional, Type, Union
+from typing import Annotated, Any, Callable, ClassVar, Optional, Type, Union
 
 from langfun.core import modality
 from langfun.core import natural_language
@@ -32,15 +32,49 @@ class Message(
     pg.Object,
     pg.views.HtmlTreeView.Extension
 ):
-  """Message.
+  """Message between users, LLMs and tools.
 
-  ``Message`` is the protocol for users and the system to interact with
-  LLMs. It consists of a text in the form of natural language, 
-  an identifier of the sender, and a dictionary of Python values as structured
-  meta-data.
+  `lf.Message` is the fundamental unit of communication in Langfun. It
+  standardizes interactions with LLMs by encapsulating not only text but also
+  multi-modal content, as well as the sender's role and structured metadata.
 
-  The subclasses of ``Message`` represent messages sent from different roles.
-  Agents may use the roles to decide the orchastration logic.
+  **Key Components:**
+
+  *   **`text`**: The natural language content of the message.
+  *   **`sender`**: An identifier for the message originator (e.g., 'User',
+    'AI', 'System').
+  *   **`metadata`**: A dictionary for structured data, such as tool inputs/
+    outputs, scores, or other contextual information.
+  *   **`referred_modalities`**: A dictionary of modality objects (e.g.,
+    `lf.Image`, `lf.Audio`) referenced within the message text via placeholders
+    like `<<[[image_id]]>>`.
+
+  Subclasses like `lf.UserMessage`, `lf.AIMessage`, and `lf.ToolMessage`
+  represent messages from specific roles, enabling more complex conversational
+  flows and agentic behaviors.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+
+  # Creating a user message with an image
+  image = lf.Image.from_path('/path/to/image.png')
+  user_message = lf.UserMessage(
+      f'What is in this image <<[[{image.id}]]>>?',
+      referred_modalities=[image])
+
+  # Creating an AI message with structured results
+  ai_message = lf.AIMessage(
+      'It is a cat.',
+      metadata=dict(result=dict(label='cat', confidence=0.9)))
+
+  print(user_message.chunk())
+  # Output: ['What is in this image', <lf.Image object>, '?']
+
+  print(ai_message.result)
+  # Output: {'label': 'cat', 'confidence': 0.9}
+  ```
   """
 
   #
@@ -86,6 +120,11 @@ class Message(
 
   sender: Annotated[str, 'The sender of the message.']
 
+  referred_modalities: Annotated[
+      dict[str, pg.Ref[modality.Modality]],
+      'The modality objects referred in the message.'
+  ] = pg.Dict()
+
   metadata: Annotated[
       dict[str, Any],
       (
@@ -111,6 +150,11 @@ class Message(
       *,
       # Default sender is specified in subclasses.
       sender: str | pg.object_utils.MissingValue = pg.MISSING_VALUE,
+      referred_modalities: (
+          list[modality.Modality]
+          | dict[str, modality.Modality]
+          | None
+      ) = None,
       metadata: dict[str, Any] | None = None,
       tags: list[str] | None = None,
       source: Optional['Message'] = None,
@@ -125,6 +169,7 @@ class Message(
     Args:
       text: The text in the message.
       sender: The sender name of the message.
+      referred_modalities: The modality objects referred in the message.
       metadata: Structured meta-data associated with this message.
       tags: Tags for the message.
       source: The source message of the current message.
@@ -138,9 +183,13 @@ class Message(
     """
     metadata = metadata or {}
     metadata.update(kwargs)
+    if isinstance(referred_modalities, list):
+      referred_modalities = {m.id: pg.Ref(m) for m in referred_modalities}
+
     super().__init__(
         text=text,
         metadata=metadata,
+        referred_modalities=referred_modalities or {},
         tags=tags or [],
         sender=sender,
         allow_partial=allow_partial,
@@ -186,7 +235,7 @@ class Message(
       A message created from the value.
     """
     if isinstance(value, modality.Modality):
-      return cls('<<[[object]]>>', object=value)
+      return cls(f'<<[[{value.id}]]>>', referred_modalities=[value])
     if isinstance(value, Message):
       return value
     if isinstance(value, str):
@@ -224,6 +273,11 @@ class Message(
     """
     return MessageConverter.get(format_or_type, **kwargs).to_value(self)
 
+  @classmethod
+  def is_convertible(cls, format_or_type: str | Type[Any]) -> bool:
+    """Returns True if the value can be converted to a message."""
+    return MessageConverter.is_convertible(format_or_type)
+
   @classmethod
   def convertible_formats(cls) -> list[str]:
     """Returns supported format for message conversion."""
@@ -280,8 +334,7 @@ class Message(
     if key_path == Message.PATH_TEXT:
       return self.text
     else:
-      v = self.metadata.sym_get(key_path, default, use_inferred=True)
-      return v.value if isinstance(v, pg.Ref) else v
+      return self.metadata.sym_get(key_path, default, use_inferred=True)
 
   #
   # API for accessing the structured result and error.
@@ -361,46 +414,63 @@ class Message(
   # API for supporting modalities.
   #
 
+  def modalities(
+      self,
+      filter: (  # pylint: disable=redefined-builtin
+          Type[modality.Modality]
+          | Callable[[modality.Modality], bool]
+          | None
+      ) = None  # pylint: disable=bad-whitespace
+  ) -> list[modality.Modality]:
+    """Returns the modality objects referred in the message."""
+    if inspect.isclass(filter) and issubclass(filter, modality.Modality):
+      filter_fn = lambda v: isinstance(v, filter)  # pytype: disable=wrong-arg-types
+    elif filter is None:
+      filter_fn = lambda v: True
+    else:
+      filter_fn = filter
+    return [v for v in self.referred_modalities.values() if filter_fn(v)]
+
   @property
-  def text_with_modality_hash(self) -> str:
-    """Returns text with modality object placeheld by their 8-byte MD5 hash."""
-    parts = [self.text]
-    for name, modality_obj in self.referred_modalities().items():
-      parts.append(
-          f'<{name}>{modality_obj.hash}</{name}>'
-      )
-    return ''.join(parts)
+  def images(self) -> list[modality.Modality]:
+    """Returns the image objects referred in the message."""
+    assert False, 'Overridden in core/modalities/__init__.py'
+
+  @property
+  def videos(self) -> list[modality.Modality]:
+    """Returns the video objects referred in the message."""
+    assert False, 'Overridden in core/modalities/__init__.py'
+
+  @property
+  def audios(self) -> list[modality.Modality]:
+    """Returns the audio objects referred in the message."""
+    assert False, 'Overridden in core/modalities/__init__.py'
 
   def get_modality(
-      self, var_name: str, default: Any = None, from_message_chain: bool = True
+      self,
+      var_name: str,
+      default: Any = None
   ) -> modality.Modality | None:
-    """Gets the modality object referred in the message.
+    """Returns modality object referred in the message by its variable name.
 
     Args:
       var_name: The referred variable name for the modality object.
       default: default value.
-      from_message_chain: If True, the look up will be performed from the
-        message chain. Otherwise it will be performed in current message.
 
     Returns:
       A modality object if found, otherwise None.
     """
-    obj = self.get(var_name, None)
-    if isinstance(obj, modality.Modality):
-      return obj
-    elif obj is None and self.source is not None:
-      return self.source.get_modality(var_name, default, from_message_chain)
-    return default
-
-  def referred_modalities(self) -> dict[str, modality.Modality]:
-    """Returns modality objects attached on this message."""
-    chunks = self.chunk()
-    return {
-        m.referred_name: m for m in chunks if isinstance(m, modality.Modality)
-    }
+    return self.referred_modalities.get(var_name, default)
 
   def chunk(self, text: str | None = None) -> list[str | modality.Modality]:
-    """Chunk a message into a list of str or modality objects."""
+    """Chunks message into a list of text and modality chunks.
+
+    Args:
+      text: The text to chunk. If None, use `self.text`.
+
+    Returns:
+      A list of text and modality chunks.
+    """
     chunks = []
 
     def add_text_chunk(text_piece: str) -> None:
@@ -425,20 +495,25 @@ class Message(
 
       var_name = text[var_start:ref_end].strip()
       var_value = self.get_modality(var_name)
-      if var_value is not None:
-        add_text_chunk(text[chunk_start:ref_start].strip(' '))
-        chunks.append(var_value)
-        chunk_start = ref_end + len(modality.Modality.REF_END)
+      if var_value is None:
+        raise ValueError(
+            f'Unknown modality reference: {var_name!r}. '
+            'Please make sure the modality object is present in '
+            f'`referred_modalities` when creating {self.__class__.__name__}.'
+        )
+      add_text_chunk(text[chunk_start:ref_start].strip(' '))
+      chunks.append(var_value)
+      chunk_start = ref_end + len(modality.Modality.REF_END)
     return chunks
 
   @classmethod
   def from_chunks(
       cls, chunks: list[str | modality.Modality], separator: str = ' '
   ) -> 'Message':
-    """Assembly a message from a list of string or modality objects."""
+    """Assembles a message from a list of string or modality objects."""
     fused_text = io.StringIO()
-    ref_index = 0
     metadata = dict()
+    referred_modalities = dict()
     last_char = None
     for i, chunk in enumerate(chunks):
       if i > 0 and last_char not in ('\t', ' ', '\n', None):
@@ -451,14 +526,16 @@ class Message(
           last_char = None
       else:
         assert isinstance(chunk, modality.Modality), chunk
-        var_name = f'obj{ref_index}'
-        fused_text.write(modality.Modality.text_marker(var_name))
+        fused_text.write(modality.Modality.text_marker(chunk.id))
         last_char = modality.Modality.REF_END[-1]
         # Make a reference if the chunk is already owned by another object
         # to avoid copy.
-        metadata[var_name] = pg.maybe_ref(chunk)
-        ref_index += 1
-    return cls(fused_text.getvalue().strip(), metadata=metadata)
+        referred_modalities[chunk.id] = pg.Ref(chunk)
+    return cls(
+        fused_text.getvalue().strip(),
+        referred_modalities=referred_modalities,
+        metadata=metadata,
+    )
 
   #
   # Tagging
@@ -523,7 +600,7 @@ class Message(
     return self.trace(Message.TAG_LM_OUTPUT)
 
   def last(self, tag: str) -> Optional['Message']:
-    """Return the last message wih certain tag."""
+    """Returns the last message with a given tag."""
     current = self
     while current is not None:
       if tag in current.tags:
@@ -551,6 +628,11 @@ class Message(
   #
 
   def natural_language_format(self) -> str:
+    """Returns the natural language format representation."""
+    # Propagate the modality references to parent context if any.
+    if capture_context := modality.get_modality_capture_context():
+      for v in self.referred_modalities.values():
+        capture_context.capture(v)
     return self.text
 
   def __eq__(self, other: Any) -> bool:
@@ -568,8 +650,7 @@ class Message(
   def __getattr__(self, key: str) -> Any:
     if key not in self.metadata:
       raise AttributeError(key)
-    v = self.metadata[key]
-    return v.value if isinstance(v, pg.Ref) else v
+    return self.metadata[key]
 
   def _html_tree_view_content(
       self,
@@ -646,15 +727,14 @@ class Message(
           s.write(s.escape(chunk))
         else:
           assert isinstance(chunk, modality.Modality), chunk
-          child_path = pg.KeyPath(['metadata', chunk.referred_name], root_path)
           s.write(
               pg.Html.element(
                   'div',
                   [
                       view.render(
                           chunk,
-                          name=chunk.referred_name,
-                          root_path=child_path,
+                          name=chunk.id,
+                          root_path=chunk.sym_path,
                           collapse_level=(
                               0 if collapse_modalities_in_text else 1
                           ),
@@ -667,7 +747,7 @@ class Message(
                   css_classes=['modality-in-text'],
               )
           )
-          referred_chunks[chunk.referred_name] = chunk
+          referred_chunks[chunk.id] = chunk
       s.write('</div>')
       return s
 
@@ -874,6 +954,12 @@ class _MessageConverterRegistry:
     if converter.OUTPUT_TYPE is not None:
       self._type_to_converters[converter.OUTPUT_TYPE].append(converter)
 
+  def unregister(self, converter: Type['MessageConverter']) -> None:
+    """Unregisters a message converter."""
+    self._name_to_converter.pop(converter.FORMAT_ID, None)
+    if converter.OUTPUT_TYPE is not None:
+      self._type_to_converters[converter.OUTPUT_TYPE].remove(converter)
+
   def get_by_type(self, t: Type[Any], **kwargs) -> 'MessageConverter':
     """Returns a message converter for the given type."""
     t = self._type_to_converters[t]
@@ -904,6 +990,13 @@ class _MessageConverterRegistry:
     assert isinstance(format_or_type, type), format_or_type
     return self.get_by_type(format_or_type, **kwargs)
 
+  def is_convertible(self, format_or_type: str | Type[Any]) -> bool:
+    """Returns whether the message is convertible to the given format or type."""
+    if isinstance(format_or_type, str):
+      return format_or_type in self._name_to_converter
+    assert isinstance(format_or_type, type), format_or_type
+    return bool(self._type_to_converters.get(format_or_type))
+
   def convertible_formats(self) -> list[str]:
     """Returns a list of converter names."""
     return sorted(list(self._name_to_converter.keys()))
@@ -995,6 +1088,11 @@ class MessageConverter(pg.Object):
     """Returns a message converter for the given type."""
     return cls._REGISTRY.get_by_type(t, **kwargs)
 
+  @classmethod
+  def is_convertible(cls, format_or_type: str | Type[Any]) -> bool:
+    """Returns whether the message is convertible to the given format or type."""
+    return cls._REGISTRY.is_convertible(format_or_type)
+
   @classmethod
   def convertible_formats(cls) -> list[str]:
     """Returns a list of converter names."""
@@ -1036,3 +1134,10 @@ class MemoryRecord(Message):
   """Message used as a memory record."""
 
   sender = 'Memory'
+
+
+@pg.use_init_args(['text', 'sender', 'metadata'])
+class ToolMessage(Message):
+  """Message used as a tool call."""
+
+  sender = 'Tool'