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

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'