REFACTOR: rename message_parsing.py to message_proc.py; move standalone message util functions into message_proc.py

Author: mdw771

src/eaa/agents/base.py

@@ -97,13 +97,18 @@
 import asyncio
 
 import numpy as np
-from openai.types.chat import ChatCompletionMessage
 
 from eaa.api.memory import MemoryManagerConfig
+from eaa.message_proc import (
+    generate_openai_message, 
+    has_tool_call, 
+    get_tool_call_info, 
+    print_message
+)
 from eaa.tools.base import BaseTool, ToolReturnType, ExposedToolSpec, generate_openai_tool_schema
 from eaa.tools.mcp import MCPTool
 from eaa.comms import get_api_key
-from eaa.util import encode_image_base64, get_image_path_from_text
+from eaa.util import get_image_path_from_text
 from eaa.api.llm_config import LLMConfig
 from eaa.agents.memory import (
     MemoryManager,
@@ -919,256 +924,3 @@ def register_message_hook(self, hook: Callable) -> None:
             The hook function.
         """
         self.message_hooks.append(hook)
-
-
-def to_dict(message: ChatCompletionMessage | Dict[str, Any]) -> dict:
-    """Convert a ChatCompletionMessage to a dictionary.
-    """
-    if isinstance(message, ChatCompletionMessage):
-        return message.to_dict()
-    else:
-        return message
-
-
-def generate_openai_message(
-    content: str,
-    role: Literal["user", "system", "tool"] = "user",
-    tool_call_id: str = None,
-    image: np.ndarray = None,
-    image_path: str = None,
-    encoded_image: str = None
-) -> Dict[str, Any]:
-    """Generate a dictionary in OpenAI-compatible format 
-    containing the message to be sent to the agent.
-
-    Parameters
-    ----------
-    content : str
-        The content of the message.
-    role : Literal["user", "system", "tool"], optional
-        The role of the sender.
-    image : np.ndarray, optional
-        The image to be sent to the agent. Exclusive with `encoded_image` and `image_path`.
-    image_path : str, optional
-        The path to the image to be sent to the agent. Exclusive with `image` and `encoded_image`.
-    encoded_image : str, optional
-        The base-64 encoded image to be sent to the agent. Exclusive with `image` and `image_path`.
-    """
-    if sum([image is not None, encoded_image is not None, image_path is not None]) > 1:
-        raise ValueError("Only one of `image`, `encoded_image`, or `image_path` should be provided.")
-    if role not in ["user", "system", "tool"]:
-        raise ValueError("Invalid role. Must be one of `user`, `system`, or `tool`.")
-
-    if image is not None or image_path is not None:
-        encoded_image = encode_image_base64(image=image, image_path=image_path)
-
-    if role == "user":
-        message = {
-            "role": "user",
-            "content": content
-        }
-    elif role == "system":
-        message = {
-            "role": "system",
-            "content": content
-        }
-    elif role == "tool":
-        message = {
-            "role": "tool",
-            "content": content,
-            "tool_call_id": tool_call_id
-        }
-
-    if encoded_image is not None:
-        message["content"] = [
-            {
-                "type": "text",
-                "text": content
-            },
-            {
-                "type": "image_url",
-                "image_url": {
-                    "url": f"data:image/png;base64,{encoded_image}"
-                }
-            }
-        ]
-    return message
-
-
-def has_tool_call(message: dict | ChatCompletionMessage) -> bool:
-    """Check if the message has a tool call.
-    
-    Parameters
-    ----------
-    message : dict | ChatCompletionMessage
-        A message in OpenAI-compatible format.
-        
-    Returns
-    -------
-    """
-    message = to_dict(message)
-    if "tool_calls" in message.keys():
-        return True
-    else:
-        return False
-
-
-def get_tool_call_info(
-    message: dict | ChatCompletionMessage,
-    index: Optional[int] = 0
-) -> str | List[str]:
-    """Get the tool call ID from the message.
-
-    Parameters
-    ----------
-    message : dict | ChatCompletionMessage
-        The message to get the tool call ID from. The message
-        should be in OpenAI-compatible format.
-    index : int, optional
-        The index of the tool call to get the ID from. If None,
-        all tool calls are returned as a list.
-
-    Returns
-    -------
-    str | List[str]
-        The tool call(s).
-    """
-    message = to_dict(message)
-    if index is None:
-        return message["tool_calls"]
-    else:
-        return message["tool_calls"][index]
-
-
-def get_message_elements_as_text(message: Dict[str, Any]) -> Dict[str, Any]:
-    """Get the elements of the message as human readable text.
-
-    Parameters
-    ----------
-    message : Dict[str, Any]
-        The message to get the elements from. The message
-        should be in OpenAI-compatible format.
-
-    Returns
-    -------
-    Dict[str, Any]
-        The elements of the message.
-    """
-    role = message["role"]
-
-    image = None
-    content = ""
-    if "content" in message.keys():
-        if isinstance(message["content"], str):
-            content += message["content"] + "\n"
-        elif isinstance(message["content"], list):
-            for item in message["content"]:
-                if item["type"] == "text":
-                    content += item["text"] + "\n"
-                elif item["type"] == "image_url":
-                    content += "<image> \n"
-                    image = item["image_url"]["url"]
-
-    tool_calls = None
-    if "tool_calls" in message.keys():
-        tool_calls = ""
-        for tool_call in message["tool_calls"]:
-            tool_calls += f"{tool_call['id']}: {tool_call['function']['name']}\n"
-            tool_calls += f"Arguments: {tool_call['function']['arguments']}\n"
-
-    return {
-        "role": role,
-        "content": content,
-        "tool_calls": tool_calls,
-        "image": image
-    }
-    
-
-def get_message_elements(message: Dict[str, Any]) -> Dict[str, Any]:
-    """Get the elements of the message as a structured dictionary.
-
-    Parameters
-    ----------
-    message : Dict[str, Any]
-        The message to get the elements from. The message
-        should be in OpenAI-compatible format.
-
-    Returns
-    -------
-    Dict[str, Any]
-        The elements of the message.
-    """
-    role = message["role"]
-
-    image = []
-    content = []
-    if "content" in message.keys():
-        if isinstance(message["content"], str):
-            content.append(message["content"])
-        elif isinstance(message["content"], list):
-            content = message["content"]
-            for item in content:
-                if item["type"] == "image_url":
-                    image.append(item["image_url"]["url"])
-
-    tool_calls = None
-    if "tool_calls" in message.keys():
-        tool_calls = message["tool_calls"]
-        
-    tool_response_id = None
-    if "tool_call_id" in message.keys():
-        tool_response_id = message["tool_call_id"]
-
-    return {
-        "role": role,
-        "content": content,
-        "tool_calls": tool_calls,
-        "image": image,
-        "tool_response_id": tool_response_id
-    }
-
-
-def print_message(
-    message: Dict[str, Any],
-    response_requested: Optional[bool] = None,
-    return_string: bool = False
-) -> None:
-    """Print the message.
-
-    Parameters
-    ----------
-    message : Dict[str, Any]
-        The message to be printed. The message should be in
-        OpenAI-compatible format.
-    response_requested : bool, optional
-        Whether a response is requested for the message.
-    return_string : bool, optional
-        If True, the message is returned as a string instead of printed.
-    """
-    color_dict = {
-        "user": "\033[94m",
-        "system": "\033[92m",
-        "tool": "\033[93m",
-        "assistant": "\033[91m"
-    }
-    color = color_dict[message["role"]]
-
-    text = f"[Role] {message['role']}\n"
-    if response_requested is not None:
-        text += f"[Response requested] {response_requested}\n"
-
-    elements = get_message_elements_as_text(message)
-
-    text += "[Content]\n"
-    text += elements["content"] + "\n"
-
-    if elements["tool_calls"] is not None:
-        text += "[Tool call]\n"
-        text += elements["tool_calls"] + "\n"
-
-    text += "\n ========================================= \n"
-
-    if return_string:
-        return text
-    else:
-        print(f"{color}{text}\033[0m")