mail-swarms 1.3.2__py3-none-any.whl

mail/utils/logger.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 Addison Kline
+
+import datetime
+import logging
+import os
+
+from rich.logging import RichHandler
+
+
+def get_loggers():
+    """
+    Get the loggers for the application.
+    """
+    return list(logging.root.manager.loggerDict.keys())
+
+
+def init_logger():
+    """
+    Initialize the logger for the application.
+    """
+    # Create logs directory if it doesn't exist
+    os.makedirs("logs", exist_ok=True)
+
+    # File handler
+    file_handler = logging.FileHandler(
+        f"logs/mail_{datetime.datetime.now(datetime.UTC).strftime('%Y-%m-%d')}.log"
+    )
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(
+        logging.Formatter("%(asctime)s [%(levelname)s] [%(name)s] :: %(message)s")
+    )
+
+    # for all loggers that are not mail, clear all handlers
+    # and then add only the file handler above
+    for logger in get_loggers():
+        if not logger.startswith("mail"):
+            logging.getLogger(logger).propagate = False
+            logging.getLogger(logger).handlers.clear()
+            logging.getLogger(logger).addHandler(file_handler)
+
+    # Rich handler for colored console output
+    # All `mail.*` loggers should use this handler
+    # Use `mailquiet.*` for loggers that should not be verbose
+    console_handler = RichHandler(
+        rich_tracebacks=True,
+        show_time=True,
+        show_level=True,
+        show_path=True,
+        markup=True,
+        tracebacks_show_locals=True,
+    )
+    console_handler.setLevel(logging.DEBUG)
+
+    # Configure the mail logger (using the actual module name)
+    mail_logger = logging.getLogger("mail")
+    mail_logger.setLevel(logging.DEBUG)
+    mail_logger.propagate = False  # Prevent double logging
+
+    # Clear any existing handlers
+    mail_logger.handlers.clear()
+
+    # Add our handlers
+    mail_logger.addHandler(console_handler)
+    mail_logger.addHandler(file_handler)
+
+    # Configure root logger to avoid conflicts with Uvicorn
+    root_logger = logging.getLogger()
+    root_logger.setLevel(logging.INFO)
+
+    # Only add file handler to root if it doesn't already have handlers
+    if not root_logger.handlers:
+        root_logger.addHandler(file_handler)

mail/utils/openai.py

@@ -0,0 +1,212 @@
+import asyncio
+import uuid
+from datetime import datetime
+from typing import Any
+
+import ujson
+from openai.types.responses import (
+    Response,
+    ResponseFunctionToolCall,
+    ResponseOutputMessage,
+    ResponseOutputText,
+)
+from pydantic import BaseModel, ValidationError
+
+from mail.api import MAILAction, MAILSwarm, MAILSwarmTemplate
+from mail.utils.serialize import dump_mail_result
+
+
+async def async_lambda(x: Any) -> Any:
+    return x
+
+
+def build_oai_clients_dict() -> dict[str, "SwarmOAIClient"]:
+    """
+    Build a dictionary of SwarmOAIClient instances for each API key in the environment.
+    """
+    return {}
+
+
+class SwarmOAIClient:
+    def __init__(
+        self,
+        template: MAILSwarmTemplate,
+        instance: MAILSwarm | None = None,
+        validate_responses: bool = True,
+    ):
+        self.responses = self.Responses(self)
+        self.template = template
+        self.result_dumps: dict[str, list[Any]] = {}
+        self.swarm = instance
+        self.validate_responses = validate_responses
+
+    class Responses:
+        def __init__(self, owner: "SwarmOAIClient"):
+            self.owner = owner
+
+        async def create(
+            self,
+            input: list[dict[str, Any]],
+            tools: list[dict[str, Any]],
+            instructions: str | None = None,
+            previous_response_id: str | None = None,
+            tool_choice: str | dict[str, str] = "auto",
+            parallel_tool_calls: bool = True,
+            api_key: str | None = None,
+            **kwargs: Any,
+        ) -> Response:
+            if self.owner.swarm is None:
+                new_swarm = self.owner.template
+                complete_agent = next(
+                    (a for a in new_swarm.agents if a.can_complete_tasks), None
+                )
+                assert complete_agent is not None
+                if instructions is not None:
+                    raw_sys_msg = {"content": instructions}
+                else:
+                    raw_sys_msg = next(
+                        (
+                            input_item  # type: ignore
+                            for input_item in input
+                            if (
+                                "role" in input_item
+                                and input_item["role"] == "system"
+                                or input_item["role"] == "developer"
+                            )
+                        ),
+                        {"content": ""},
+                    )
+                complete_agent.agent_params["system"] = (
+                    complete_agent.agent_params["system"] + raw_sys_msg["content"]
+                )
+                if len(tools) > 0:
+                    new_actions: list[MAILAction] = []
+                    for tool in tools:
+                        name = tool["name"]
+                        description = tool["description"]
+                        parameters = tool["parameters"]
+                        new_actions.append(
+                            MAILAction(
+                                name=name,
+                                description=description,
+                                parameters=parameters,
+                                function=async_lambda,
+                            )
+                        )
+                        complete_agent.actions += new_actions
+                        new_swarm.actions += new_actions
+                        complete_agent.agent_params["system"] = (
+                            complete_agent.agent_params["system"]
+                            + f"\n\nYou can perform actions in the environment by calling one of the following tools: {', '.join([a.name for a in new_actions])}"
+                        )
+                        new_swarm.breakpoint_tools = [a.name for a in new_actions]
+
+                self.owner.swarm = new_swarm.instantiate({"user_token": ""})
+                asyncio.create_task(self.owner.swarm.run_continuous())
+            swarm = self.owner.swarm
+            body = ""
+            if "type" in input[-1] and input[-1]["type"] == "function_call_output":
+                tool_responses: list[dict[str, Any]] = []
+                for input_item in reversed(input):
+                    if (
+                        "type" not in input_item
+                        or input_item["type"] == "function_call"
+                    ):
+                        break
+                    if input_item["type"] == "function_call_output":
+                        tool_responses.append(
+                            {
+                                "call_id": input_item["call_id"],
+                                "content": input_item["output"],
+                            }
+                        )
+                out, events = await swarm.post_message(
+                    body="",
+                    subject="Tool Response",
+                    task_id=previous_response_id,
+                    show_events=True,
+                    resume_from="breakpoint_tool_call",
+                    breakpoint_tool_call_result=tool_responses,
+                )
+            else:
+                for input_item in reversed(input):
+                    if isinstance(input_item, BaseModel):
+                        input_item = input_item.model_dump()
+                    if (
+                        ("role" in input_item and input_item["role"] == "assistant")
+                        or "type" in input_item
+                        and (
+                            input_item["type"]
+                            in ["function_call_output", "function_call"]
+                        )
+                    ):
+                        break
+                    body = f"{input_item['content']}\n\n{body}"
+                out, events = await swarm.post_message(
+                    body=body,
+                    subject="Task Request",
+                    task_id=previous_response_id,
+                    show_events=True,
+                )
+            response_id = out["message"]["task_id"]
+            dump = dump_mail_result(result=out, events=events, verbose=True)
+            if response_id not in self.owner.result_dumps:
+                self.owner.result_dumps[response_id] = []
+            self.owner.result_dumps[response_id].append(dump)
+            has_called_tools = out["message"]["subject"] == "::breakpoint_tool_call::"
+            if not has_called_tools:
+                response = Response(
+                    id=response_id,
+                    created_at=float(datetime.now().timestamp()),
+                    model=f"{swarm.name}",
+                    object="response",
+                    tools=tools,  # type: ignore
+                    output=[
+                        ResponseOutputMessage(
+                            type="message",
+                            id=str(uuid.uuid4()),
+                            status="completed",
+                            role="assistant",
+                            content=[
+                                ResponseOutputText(
+                                    type="output_text",
+                                    text=out["message"]["body"],
+                                    annotations=[],
+                                )
+                            ],
+                        )
+                    ],
+                    parallel_tool_calls=parallel_tool_calls,
+                    tool_choice=tool_choice,  # type: ignore
+                )
+                return response
+
+            tool_calls: list[ResponseFunctionToolCall] = []
+            body = ujson.loads(out["message"]["body"])
+            for tool_call in body:
+                tool_calls.append(
+                    ResponseFunctionToolCall(
+                        call_id=tool_call["call_id"],
+                        name=tool_call["name"],
+                        arguments=tool_call["arguments"],
+                        type="function_call",
+                        id=tool_call["id"],
+                        status=tool_call["status"],
+                    )
+                )
+            try:
+                return Response(
+                    id=response_id,
+                    created_at=float(datetime.now().timestamp()),
+                    model=f"{swarm.name}",
+                    object="response",
+                    tools=tools,  # type: ignore
+                    output=tool_calls,  # type: ignore
+                    parallel_tool_calls=parallel_tool_calls,
+                    tool_choice=tool_choice,  # type: ignore
+                )
+            except ValidationError as e:
+                if self.owner.validate_responses:
+                    raise e
+                else:
+                    return response

mail/utils/parsing.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 Addison Kline
+
+import importlib
+import json
+from typing import Any
+
+import httpx
+
+from mail.core import parse_agent_address
+
+PYTHON_STRING_PREFIX = "python::"
+URL_STRING_PREFIX = "url::"
+
+
+def read_python_string(string: str) -> Any:
+    """
+    Resolve an import string to a Python object.
+
+    Accepts strings in the format ``module:variable`` or with the explicit
+    ``python::`` prefix used in swarm configuration files, e.g.
+    ``python::package.module:object``.
+    """
+
+    if string.startswith(PYTHON_STRING_PREFIX):
+        string = string[len(PYTHON_STRING_PREFIX) :]
+
+    try:
+        module_str, attribute_path = string.split(":", 1)
+    except ValueError as err:  # pragma: no cover - defensive guard
+        raise ValueError(
+            f"Invalid python reference '{string}'. Expected 'module:object' format."
+        ) from err
+
+    module = importlib.import_module(module_str)
+    obj: Any = module
+    for attr in attribute_path.split("."):
+        obj = getattr(obj, attr)
+
+    return obj
+
+
+def resolve_prefixed_string_references(value: Any) -> Any:
+    """
+    Recursively resolve strings prefixed with ``python::`` or ``url::`` to Python objects or strings, respectively.
+    """
+
+    if isinstance(value, dict):
+        return {
+            key: resolve_prefixed_string_references(item) for key, item in value.items()
+        }
+    if isinstance(value, list):
+        return [resolve_prefixed_string_references(item) for item in value]
+    if isinstance(value, str):
+        if value.startswith(PYTHON_STRING_PREFIX):
+            return read_python_string(value)
+        if value.startswith(URL_STRING_PREFIX):
+            return read_url_string(value)
+        return value
+
+    return value
+
+
+def read_url_string(string: str, raise_on_error: bool = False) -> str:
+    """
+    Resolve a URL to a string.
+
+    Accepts strings in the format ``url::`` prefix used in swarm configuration files, e.g.
+    ``url::https://example.com``.
+    """
+
+    if string.startswith(URL_STRING_PREFIX):
+        string = string[len(URL_STRING_PREFIX) :]
+
+    try:
+        response = httpx.get(string)
+        return json.dumps(response.json())
+    except Exception as e:
+        if raise_on_error:
+            raise RuntimeError(f"error reading URL string: '{str(e)}'")
+        return string
+
+
+def target_address_is_interswarm(address: str) -> bool:
+    """
+    Check if a target address is an interswarm address.
+    """
+
+    return parse_agent_address(address)[1] is not None

mail/utils/serialize.py

@@ -0,0 +1,292 @@
+from collections.abc import Mapping
+from dataclasses import asdict, is_dataclass
+from typing import Any
+
+import ujson
+from sse_starlette import ServerSentEvent
+
+from mail.core.message import MAILMessage
+from mail.utils.version import get_version
+
+_REDACT_KEYS = {
+    "id",
+    "task_id",
+    "request_id",
+    "broadcast_id",
+    "message_id",
+    "event_id",
+    "timestamp",
+    "created_at",
+    "updated_at",
+    "sent_at",
+    "received_at",
+}
+
+
+def dump_mail_result(
+    result: MAILMessage,
+    events: list[ServerSentEvent],
+    verbose: bool = False,
+) -> str:
+    """
+    For a completed MAIL task, create an LLM-friendly string dump of the result and events.
+    """
+    serialized_result = serialize_mail_value(result, exclude_keys=_REDACT_KEYS)
+    if not verbose:
+        return str(serialized_result["message"]["body"])
+
+    serialized_events = []
+    for event in events:
+        serialized = _serialize_event(event, exclude_keys=_REDACT_KEYS)
+        if serialized is not None:
+            serialized_events.append(serialized)
+
+    event_sections = _format_event_sections(serialized_events)
+    result_section = str(serialized_result["message"]["body"])
+
+    sections: list[str] = [
+        "=== MAIL EVENTS ===",
+        event_sections if event_sections else "(no events)",
+        "=== FINAL ANSWER ===",
+        result_section,
+    ]
+    return "\n".join(sections)
+
+
+def serialize_mail_value(
+    value: Any, *, exclude_keys: frozenset[str] | set[str] | None = None
+) -> Any:
+    """
+    Convert MAIL runtime objects into JSON-friendly primitives.
+    """
+    exclude_keys = frozenset() if exclude_keys is None else frozenset(exclude_keys)
+
+    if value is None or isinstance(value, str | int | float | bool):
+        return value
+    if isinstance(value, bytes):
+        return value.decode("utf-8", errors="replace")
+    if isinstance(value, Mapping):
+        return {
+            k: serialize_mail_value(v, exclude_keys=exclude_keys)
+            for k, v in value.items()
+            if k not in exclude_keys
+        }
+    if isinstance(value, list | tuple | set | frozenset):
+        return [serialize_mail_value(v, exclude_keys=exclude_keys) for v in value]
+    if is_dataclass(value):
+        return serialize_mail_value(asdict(value), exclude_keys=exclude_keys)  # type: ignore[arg-type]
+    if hasattr(value, "model_dump"):
+        try:
+            dumped = value.model_dump()
+        except TypeError:
+            dumped = value.model_dump(mode="python")
+        return serialize_mail_value(dumped, exclude_keys=exclude_keys)
+    if hasattr(value, "dict"):
+        try:
+            dumped = value.dict()
+        except TypeError:
+            dumped = value.dict()
+        return serialize_mail_value(dumped, exclude_keys=exclude_keys)
+    if hasattr(value, "__dict__"):
+        return {
+            k: serialize_mail_value(v, exclude_keys=exclude_keys)
+            for k, v in vars(value).items()
+            if not k.startswith("_")
+        }
+    return str(value)
+
+
+def _normalise_event_payload(payload: Any) -> Any:
+    if isinstance(payload, str):
+        try:
+            return ujson.loads(payload)
+        except ValueError:
+            return payload
+    return payload
+
+
+def _serialize_event(
+    event: ServerSentEvent, *, exclude_keys: frozenset[str] | set[str] | None
+) -> dict[str, Any] | None:
+    payload = _normalise_event_payload(event.data)
+    if _should_skip_event(payload):
+        return None
+    event_type = _standardise_event_type(getattr(event, "event", None))
+
+    description = ""
+    if isinstance(payload, Mapping):
+        description = str(dict(payload).get("description", ""))
+
+    result = {
+        "event": event_type,
+        "description": description,
+    }
+    return result
+
+
+def _standardise_event_type(event_name: Any) -> str:
+    if not isinstance(event_name, str):
+        return ""
+    normalised = event_name.strip()
+    normalised_key = normalised.casefold().replace(" ", "_")
+    if normalised_key == "action_complete":
+        return "action_output"
+    return normalised
+
+
+def _should_skip_event(payload: Any) -> bool:
+    return _is_action_complete_broadcast(payload)
+
+
+def _is_action_complete_broadcast(payload: Any) -> bool:
+    """
+    Check if the payload is an `action_complete_broadcast` message.
+    """
+    if not isinstance(payload, Mapping):
+        return False
+    description = payload.get("description")
+    if not isinstance(description, str):
+        return False
+    if "<subject>::action_complete_broadcast::</subject>" in description:
+        return True
+    return False
+
+
+def _format_event_sections(events: list[Any]) -> str:
+    sections: list[str] = []
+    for idx, event in enumerate(events, start=1):
+        sections.append(f"--- Event {idx} ---\n{_format_json(event)}")
+    return "\n\n".join(sections)
+
+
+def _format_json(payload: Any) -> str:
+    if isinstance(payload, str | int | float | bool) or payload is None:
+        return str(payload)
+    return ujson.dumps(payload, indent=2)
+
+
+def extract_task_body(raw_result: Any, serialized_result: Any | None = None) -> Any:
+    """
+    Extract the body from a MAIL result.
+    """
+    for candidate in (raw_result, serialized_result):
+        if candidate is None:
+            continue
+        if isinstance(candidate, Mapping):
+            if "body" in candidate:
+                return candidate["body"]
+            message = candidate.get("message")
+            if isinstance(message, Mapping) and "body" in message:
+                return message["body"]
+        if hasattr(candidate, "body"):
+            body_attr = getattr(candidate, "body")
+            if body_attr is not None:
+                return body_attr
+        if hasattr(candidate, "message"):
+            message_obj = getattr(candidate, "message")
+            if isinstance(message_obj, Mapping) and "body" in message_obj:
+                return message_obj["body"]
+            if hasattr(message_obj, "body"):
+                body_attr = getattr(message_obj, "body")
+                if body_attr is not None:
+                    return body_attr
+    return None
+
+
+def export(swarms: list[Any]) -> str:
+    """
+    Export a `MAILSwarm` or `MAILSwarmTemplate` to a JSON string compatible
+    with the client-side `Generation`/`Swarm` interfaces.
+    """
+
+    def _format_python_reference(reference: Any) -> str | None:
+        if reference is None:
+            return None
+        if isinstance(reference, str):
+            return reference
+        if callable(reference):
+            module = getattr(reference, "__module__", "")
+            qualname = getattr(
+                reference, "__qualname__", getattr(reference, "__name__", "")
+            )
+            if module and qualname:
+                return f"python::{module}:{qualname}"
+            if qualname:
+                return qualname
+        return None
+
+    def _serialize_agent(agent: Any) -> dict[str, Any]:
+        agent_dict: dict[str, Any] = {
+            "name": getattr(agent, "name", ""),
+            "comm_targets": list(getattr(agent, "comm_targets", [])),
+            "enable_entrypoint": bool(getattr(agent, "enable_entrypoint", False)),
+            "enable_interswarm": bool(getattr(agent, "enable_interswarm", False)),
+            "can_complete_tasks": bool(getattr(agent, "can_complete_tasks", False)),
+            "tool_format": getattr(agent, "tool_format", "responses"),
+            "agent_params": serialize_mail_value(getattr(agent, "agent_params", {})),
+        }
+
+        factory_ref = _format_python_reference(getattr(agent, "factory", None))
+        if factory_ref:
+            agent_dict["factory"] = factory_ref
+
+        actions = getattr(agent, "actions", [])
+        if actions:
+            agent_dict["actions"] = [
+                getattr(action, "name", serialize_mail_value(action))
+                for action in actions
+            ]
+
+        breakpoint_tools = getattr(agent, "breakpoint_tools", None)
+        if breakpoint_tools:
+            agent_dict["breakpoint_tools"] = list(breakpoint_tools)
+
+        return agent_dict
+
+    def _serialize_action(action: Any) -> dict[str, Any]:
+        action_dict: dict[str, Any] = {
+            "name": getattr(action, "name", ""),
+            "description": getattr(action, "description", ""),
+            "parameters": serialize_mail_value(getattr(action, "parameters", {})),
+        }
+        function_ref = _format_python_reference(getattr(action, "function", None))
+        if function_ref:
+            action_dict["function"] = function_ref
+        return action_dict
+
+    swarm_payloads: list[dict[str, Any]] = []
+    for swarm in swarms:
+        agent_payloads = [
+            _serialize_agent(agent) for agent in getattr(swarm, "agents", [])
+        ]
+        action_payloads = [
+            _serialize_action(action) for action in getattr(swarm, "actions", [])
+        ]
+
+        swarm_payload: dict[str, Any] = {
+            "name": getattr(swarm, "name", ""),
+            "version": get_version(),
+            "entrypoint": getattr(swarm, "entrypoint", ""),
+            "enable_interswarm": bool(getattr(swarm, "enable_interswarm", False)),
+            "agents": agent_payloads,
+        }
+
+        breakpoint_tools = getattr(swarm, "breakpoint_tools", None)
+        if breakpoint_tools is not None:
+            swarm_payload["breakpoint_tools"] = list(breakpoint_tools)
+
+        if action_payloads:
+            swarm_payload["actions"] = action_payloads
+
+        user_id = getattr(swarm, "user_id", None)
+        if user_id is not None:
+            swarm_payload["user_id"] = user_id
+
+        swarm_payloads.append(swarm_payload)
+
+    swarms_payload = {
+        "swarms": swarm_payloads,
+        "n": len(swarm_payloads),
+    }
+
+    return ujson.dumps(swarms_payload)