orchestrata_ultimate.agents.uml_generator

Reference › orchestrata_ultimate › agents › uml_generator

uml_generator renders structured design data as Mermaid diagram text. It is standalone: it imports nothing from orchestrata and no other reference module imports it. The inputs are ClassDesign (name, methods, attributes, responsibilities, base classes), ModuleDesign (name, description, classes, dependencies) and InteractionStep (sender, receiver, message and optional reply). generate_class_diagram emits a classDiagram with members and inheritance, plus association edges inferred from method names within a module. generate_module_diagram emits a flowchart of module dependencies, and generate_sequence_diagram emits a sequenceDiagram for a list of interaction steps. save_diagrams writes diagram files into an output directory and returns their paths, and generate_all_diagrams produces and saves the whole set at once. _sanitize_id makes names safe for use as Mermaid node ids.

At a glance

Class diagram

Dependency diagram

No internal imports and no importers.

API

Classes

ClassDesign (dataclass)

Decorators: dataclass · Line 16

A designed class with methods and attributes.

ModuleDesign (dataclass)

Decorators: dataclass · Line 36

A designed module containing classes.

InteractionStep (dataclass)

Decorators: dataclass · Line 54

A step in a sequence diagram.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/agents/uml_generator.py

"""UML diagram generation via Mermaid syntax.

Generates class diagrams, module diagrams, and sequence diagrams from
structured design data. Output is Mermaid markdown (.mmd files).

Zero orchestrata imports — standalone utility for spike pipeline.
"""

from dataclasses import dataclass, field
from pathlib import Path

# -- Design Data Types ---------------------------------------------


@dataclass
class ClassDesign:
    """A designed class with methods and attributes."""

    name: str
    methods: list[str] = field(default_factory=list)
    attributes: list[str] = field(default_factory=list)
    responsibilities: str = ""
    base_classes: list[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "methods": self.methods,
            "attributes": self.attributes,
            "responsibilities": self.responsibilities,
            "base_classes": self.base_classes,
        }


@dataclass
class ModuleDesign:
    """A designed module containing classes."""

    name: str
    description: str = ""
    classes: list[ClassDesign] = field(default_factory=list)
    dependencies: list[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "description": self.description,
            "classes": [c.to_dict() for c in self.classes],
            "dependencies": self.dependencies,
        }


@dataclass
class InteractionStep:
    """A step in a sequence diagram."""

    from_actor: str
    to_actor: str
    message: str
    reply: str | None = None


# -- Mermaid Generators --------------------------------------------


def _sanitize_id(name: str) -> str:
    """Sanitize a name for use as a Mermaid node ID."""
    return name.replace("-", "_").replace(".", "_").replace(" ", "_")


def _class_block_lines(cls: ClassDesign) -> list[str]:
    """Render the Mermaid block (members + inheritance) for one class."""
    sid = _sanitize_id(cls.name)
    lines = [f"    class {sid} {{"]

    for attr in cls.attributes:
        attr_clean = attr.replace(":", " :")
        lines.append(f"        +{attr_clean}")

    for method in cls.methods:
        if "(" not in method:
            method = f"{method}()"
        lines.append(f"        +{method}")

    lines.append("    }")

    for base in cls.base_classes:
        base_id = _sanitize_id(base)
        lines.append(f"    {base_id} <|-- {sid}")

    return lines


def _association_lines(module: ModuleDesign) -> list[str]:
    """Render association edges inferred from method names within a module."""
    class_names = [c.name for c in module.classes]
    lines = []
    for cls in module.classes:
        for method in cls.methods:
            for other_cls in class_names:
                if other_cls != cls.name and other_cls in method:
                    sid = _sanitize_id(cls.name)
                    oid = _sanitize_id(other_cls)
                    lines.append(f"    {sid} --> {oid}")
    return lines


def generate_class_diagram(modules: list[ModuleDesign]) -> str:
    """Generate Mermaid classDiagram syntax from module designs.

    Args:
        modules: List of ModuleDesign objects.

    Returns:
        Mermaid classDiagram string.
    """
    lines = ["classDiagram"]

    for module in modules:
        for cls in module.classes:
            lines.extend(_class_block_lines(cls))

    for module in modules:
        lines.extend(_association_lines(module))

    return "\n".join(lines) + "\n"


def generate_module_diagram(modules: list[ModuleDesign]) -> str:
    """Generate Mermaid flowchart showing module dependencies.

    Args:
        modules: List of ModuleDesign objects.

    Returns:
        Mermaid flowchart string.
    """
    lines = ["flowchart TD"]

    module_names = {m.name for m in modules}

    for module in modules:
        sid = _sanitize_id(module.name)
        class_list = ", ".join(c.name for c in module.classes)
        label = module.name
        if class_list:
            label = f"{module.name}\\n({class_list})"
        lines.append(f'    {sid}["{label}"]')

    for module in modules:
        sid = _sanitize_id(module.name)
        for dep in module.dependencies:
            if dep in module_names:
                did = _sanitize_id(dep)
                lines.append(f"    {sid} --> {did}")

    return "\n".join(lines) + "\n"


def generate_sequence_diagram(
    flow: list[InteractionStep],
    title: str = "",
) -> str:
    """Generate Mermaid sequenceDiagram for workflows.

    Args:
        flow: List of InteractionStep objects.
        title: Optional diagram title.

    Returns:
        Mermaid sequenceDiagram string.
    """
    lines = ["sequenceDiagram"]

    if title:
        lines.append(f"    Note over {flow[0].from_actor}: {title}")

    for step in flow:
        from_id = _sanitize_id(step.from_actor)
        to_id = _sanitize_id(step.to_actor)
        lines.append(f"    {from_id}->>+{to_id}: {step.message}")
        if step.reply:
            lines.append(f"    {to_id}-->>-{from_id}: {step.reply}")

    return "\n".join(lines) + "\n"


# -- File Output ---------------------------------------------------


def save_diagrams(diagrams: dict[str, str], output_dir: str) -> list[str]:
    """Write Mermaid diagram files to disk.

    Args:
        diagrams: Dict of {filename: mermaid_content}.
        output_dir: Directory to write .mmd files to.

    Returns:
        List of file paths written.
    """
    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)

    paths = []
    for filename, content in diagrams.items():
        if not filename.endswith(".mmd"):
            filename = f"{filename}.mmd"
        file_path = out / filename
        file_path.write_text(content, encoding="utf-8")
        paths.append(str(file_path))

    return sorted(paths)


def generate_all_diagrams(
    modules: list[ModuleDesign],
    output_dir: str,
    sequence_flow: list[InteractionStep] | None = None,
) -> dict[str, str]:
    """Generate and save all diagram types.

    Args:
        modules: List of ModuleDesign objects.
        output_dir: Directory for .mmd files.
        sequence_flow: Optional interaction steps for sequence diagram.

    Returns:
        Dict of {diagram_type: file_path}.
    """
    diagrams = {
        "class_diagram": generate_class_diagram(modules),
        "module_diagram": generate_module_diagram(modules),
    }

    if sequence_flow:
        diagrams["sequence_diagram"] = generate_sequence_diagram(sequence_flow)

    paths = save_diagrams(diagrams, output_dir)

    result = {}
    for path in paths:
        name = Path(path).stem
        result[name] = path

    return result