|
| 1 | +"""YAML output formatter for Journey tools. |
| 2 | +
|
| 3 | +Converts dict/list tool responses into a compact, LLM-friendly YAML format. |
| 4 | +
|
| 5 | +Design decisions: |
| 6 | +- Lists use flow-style (inline) to save tokens: [foo, bar, baz] |
| 7 | +- Multiline strings use literal block scalar (|) for readability |
| 8 | +- None, empty string, and empty list/dict values are stripped |
| 9 | +- 2-space indentation (YAML default) |
| 10 | +- No line-wrapping (width=1000) to prevent arbitrary breaks |
| 11 | +- Values containing colons are forced-quoted for parser safety |
| 12 | +- Values containing commas or spaces are auto-quoted by PyYAML |
| 13 | +
|
| 14 | +Missing-field semantics: |
| 15 | + Fields with None, "", [], or {} values are omitted from output. |
| 16 | + Consumers should treat a missing field as empty/unset. |
| 17 | + Falsy-but-meaningful values (False, 0) are preserved. |
| 18 | +
|
| 19 | +Behavioral contracts: |
| 20 | +- Input: any dict, list, or primitive returned by a Journey tool |
| 21 | +- Output: YAML string for dict/list inputs; passthrough for primitives |
| 22 | +- Failure mode: on any exception, returns original data unchanged (no data loss) |
| 23 | +""" |
| 24 | + |
| 25 | +import logging |
| 26 | +from collections.abc import Callable |
| 27 | +from functools import wraps |
| 28 | +from typing import Any, TypeVar, cast |
| 29 | + |
| 30 | +import yaml |
| 31 | + |
| 32 | +logger = logging.getLogger(__name__) |
| 33 | + |
| 34 | +T = TypeVar("T", bound=Callable[..., Any]) |
| 35 | + |
| 36 | + |
| 37 | +class FlowList(list): |
| 38 | + """A list subclass that YAML serializes in flow style: [a, b, c]. |
| 39 | +
|
| 40 | + PyYAML's SafeDumper will render normal lists as block-style (- item). |
| 41 | + By registering a custom representer for this subclass, we get inline |
| 42 | + arrays that are more token-efficient for LLM consumption. |
| 43 | + """ |
| 44 | + |
| 45 | + pass |
| 46 | + |
| 47 | + |
| 48 | +def _flow_list_representer(dumper: yaml.SafeDumper, data: FlowList) -> yaml.Node: |
| 49 | + return dumper.represent_sequence("tag:yaml.org,2002:seq", data, flow_style=True) |
| 50 | + |
| 51 | + |
| 52 | +def _str_presenter(dumper: yaml.SafeDumper, data: str) -> yaml.Node: |
| 53 | + """Use literal block scalar for multiline strings. |
| 54 | +
|
| 55 | + Forces single-quoting for values containing colons to prevent |
| 56 | + YAML parser ambiguity (e.g. cursor values like '3:-1:1'). |
| 57 | + """ |
| 58 | + if len(data.splitlines()) > 1: |
| 59 | + return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|") |
| 60 | + if ":" in data: |
| 61 | + return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="'") |
| 62 | + return dumper.represent_scalar("tag:yaml.org,2002:str", data) |
| 63 | + |
| 64 | + |
| 65 | +yaml.SafeDumper.add_representer(FlowList, _flow_list_representer) |
| 66 | +yaml.SafeDumper.add_representer(str, _str_presenter) |
| 67 | + |
| 68 | + |
| 69 | +def _clean_and_convert(obj: Any) -> Any: |
| 70 | + """Recursively strip empty values and convert lists to FlowList. |
| 71 | +
|
| 72 | + Removes: None, "", [], {} |
| 73 | + Preserves: False, 0, and other falsy-but-meaningful values. |
| 74 | + """ |
| 75 | + if isinstance(obj, dict): |
| 76 | + cleaned = {} |
| 77 | + for k, v in obj.items(): |
| 78 | + val = _clean_and_convert(v) |
| 79 | + if val not in (None, "", [], {}): |
| 80 | + cleaned[k] = val |
| 81 | + return cleaned |
| 82 | + elif isinstance(obj, list): |
| 83 | + cleaned_list = [_clean_and_convert(v) for v in obj] |
| 84 | + is_primitive = all(not isinstance(item, (dict, list)) for item in cleaned_list) |
| 85 | + if is_primitive: |
| 86 | + return FlowList(cleaned_list) |
| 87 | + return cleaned_list |
| 88 | + return obj |
| 89 | + |
| 90 | + |
| 91 | +def format_as_yaml(raw_data: Any) -> Any: |
| 92 | + """Convert a dict or list to compact YAML string. |
| 93 | +
|
| 94 | + Args: |
| 95 | + raw_data: Tool output (typically dict or list). |
| 96 | +
|
| 97 | + Returns: |
| 98 | + YAML string if input is dict/list, original value otherwise. |
| 99 | + On formatting failure, returns original data unchanged. |
| 100 | + """ |
| 101 | + if not isinstance(raw_data, (dict, list)): |
| 102 | + return raw_data |
| 103 | + |
| 104 | + try: |
| 105 | + cleaned_data = _clean_and_convert(raw_data) |
| 106 | + yaml_str = yaml.safe_dump( |
| 107 | + cleaned_data, |
| 108 | + allow_unicode=True, |
| 109 | + sort_keys=False, |
| 110 | + default_flow_style=False, |
| 111 | + width=1000, |
| 112 | + ) |
| 113 | + return yaml_str.strip() |
| 114 | + except Exception as e: |
| 115 | + logger.debug("YAML formatting failed: %s", e) |
| 116 | + return raw_data |
| 117 | + |
| 118 | + |
| 119 | +def with_yaml(func: T) -> T: |
| 120 | + """Decorator that converts a tool's dict/list return value to YAML. |
| 121 | +
|
| 122 | + Must be applied OUTSIDE (above) mcp_error_boundary so that error |
| 123 | + dicts also get YAML-formatted. Stack order: |
| 124 | +
|
| 125 | + @mcp.tool() |
| 126 | + @with_yaml |
| 127 | + @mcp_error_boundary |
| 128 | + def my_tool(...) -> dict: ... |
| 129 | + """ |
| 130 | + |
| 131 | + @wraps(func) |
| 132 | + def wrapper(*args: Any, **kwargs: Any) -> Any: |
| 133 | + result = func(*args, **kwargs) |
| 134 | + return format_as_yaml(result) |
| 135 | + |
| 136 | + if 'return' in wrapper.__annotations__: |
| 137 | + wrapper.__annotations__['return'] = Any |
| 138 | + |
| 139 | + return cast(T, wrapper) |
0 commit comments