diagram_scribe

 17class DiagramScribe:
 18    """Orchestrates diagram generation and refinement.
 19
 20    Routes LLM output to the appropriate backend: ``ExcalidrawAdapter`` for
 21    ``DiagramIR`` (simple spatial diagrams) or ``MermaidAdapter`` for
 22    ``MermaidIR`` (flowcharts, sequence diagrams, ER diagrams, class diagrams).
 23
 24    Args:
 25        llm: An LLM adapter. Defaults to ``OpenRouterAdapter`` using env vars.
 26        backend: A backend adapter for the graph path. Defaults to
 27            ``ExcalidrawAdapter``. Used directly in tests to inject mocks.
 28        output_path: Output ``.excalidraw`` file path. Passed to both
 29            ``ExcalidrawAdapter`` and ``MermaidAdapter``.
 30
 31    Example::
 32
 33        from diagram_scribe import DiagramScribe
 34
 35        ds = DiagramScribe()
 36        ds.draw("Two services: API gateway routes to user service.")
 37        ds.refine("add a database behind the user service")
 38    """
 39
 40    def __init__(
 41        self,
 42        llm: LLMAdapter | None = None,
 43        backend: BackendAdapter | None = None,
 44        output_path: str | None = None,
 45    ):
 46        self._llm = llm or self._default_llm()
 47        self._output_path = output_path
 48        self._excalidraw_backend = backend or self._default_backend(output_path)
 49        self._mermaid_backend: object | None = None
 50        self._current_ir: DiagramIR | MermaidIR | None = None
 51
 52    @staticmethod
 53    def _default_llm() -> LLMAdapter:
 54        import os
 55        from .adapters.llm.openrouter import OpenRouterAdapter
 56        return OpenRouterAdapter(
 57            api_key=os.environ.get("OPENROUTER_API_KEY", ""),
 58            model=os.environ.get("OPENROUTER_MODEL", "nvidia/nemotron-super-49b-v1:free"),
 59        )
 60
 61    @staticmethod
 62    def _default_backend(output_path: str | None = None) -> BackendAdapter:
 63        from .adapters.backend.excalidraw import ExcalidrawAdapter
 64        return ExcalidrawAdapter(output_path=output_path)
 65
 66    def _get_mermaid_backend(self) -> object:
 67        if self._mermaid_backend is None:
 68            from .adapters.backend.mermaid import MermaidAdapter
 69            self._mermaid_backend = MermaidAdapter(output_path=self._output_path)
 70        return self._mermaid_backend
 71
 72    def _render(self, ir: DiagramIR | MermaidIR) -> None:
 73        if isinstance(ir, MermaidIR):
 74            self._get_mermaid_backend().render(ir)
 75        else:
 76            self._excalidraw_backend.render(ir)
 77
 78    def draw(self, description: str) -> DiagramIR | MermaidIR:
 79        """Generate a new diagram from a natural language description.
 80
 81        The LLM decides the output format (Mermaid or graph). The appropriate
 82        backend renders and saves the result as a ``.excalidraw`` file.
 83
 84        Args:
 85            description: Plain English description of the diagram to create.
 86
 87        Returns:
 88            The generated ``DiagramIR`` or ``MermaidIR``.
 89        """
 90        self._current_ir = self._llm.generate(description)
 91        self._render(self._current_ir)
 92        return self._current_ir
 93
 94    def refine(self, feedback: str) -> DiagramIR | MermaidIR:
 95        """Update the current diagram based on feedback.
 96
 97        Must be called after :meth:`draw`. The LLM receives the current
 98        diagram state and the feedback, returns an updated IR of the same type,
 99        and the backend re-renders it.
100
101        Args:
102            feedback: Plain English instruction describing what to change.
103
104        Returns:
105            The updated ``DiagramIR`` or ``MermaidIR``.
106
107        Raises:
108            RuntimeError: If called before :meth:`draw`.
109        """
110        if self._current_ir is None:
111            raise RuntimeError("Call draw() before refine()")
112        self._current_ir = self._llm.refine(feedback, self._current_ir)
113        self._render(self._current_ir)
114        return self._current_ir

 6@dataclass
 7class Node:
 8    """A single node in a diagram.
 9
10    Attributes:
11        id: Unique identifier used to reference this node in edges.
12        label: Display text rendered inside the shape.
13        shape: Visual shape. One of ``"box"``, ``"diamond"``, ``"circle"``,
14            ``"cylinder"``, or ``"text"`` (floating label, no border).
15    """
16    id: str
17    label: str
18    shape: str

21@dataclass
22class Edge:
23    """A directed connection between two nodes.
24
25    Attributes:
26        from_id: ``id`` of the source node.
27        to_id: ``id`` of the target node.
28        label: Optional text rendered along the arrow.
29    """
30    from_id: str
31    to_id: str
32    label: str | None = None

35@dataclass
36class DiagramIR:
37    """Intermediate representation of a diagram.
38
39    This is the contract between LLM adapters and backend adapters.
40    LLM adapters produce a ``DiagramIR``; backend adapters consume one.
41    Neither side needs to know anything about the other.
42
43    Attributes:
44        nodes: All nodes in the diagram, in no particular order.
45        edges: All directed edges, referencing nodes by ``id``.
46    """
47    nodes: list[Node] = field(default_factory=list)
48    edges: list[Edge] = field(default_factory=list)

11class OpenRouterAdapter:
12    """LLM adapter that calls models via OpenRouter.
13
14    OpenRouter provides access to hundreds of models — free and paid —
15    under a single API key at https://openrouter.ai. The adapter uses the
16    OpenAI-compatible chat completions endpoint.
17
18    Free models have a ``:free`` suffix, e.g.
19    ``meta-llama/llama-3.1-8b-instruct:free``. Browse models at
20    https://openrouter.ai/models.
21
22    Args:
23        api_key: OpenRouter API key. Get one at https://openrouter.ai.
24        model: Model ID to use. Defaults to
25            ``"meta-llama/llama-3.1-8b-instruct:free"`` (free, no billing required).
26
27    Example::
28
29        from diagram_scribe.adapters.llm.openrouter import OpenRouterAdapter
30        adapter = OpenRouterAdapter(api_key="sk-or-...", model="anthropic/claude-sonnet-4-6")
31        ir = adapter.generate("CI/CD pipeline")
32    """
33
34    def __init__(
35        self,
36        api_key: str,
37        model: str = "meta-llama/llama-3.1-8b-instruct:free",
38    ):
39        self._client = OpenAI(
40            base_url="https://openrouter.ai/api/v1",
41            api_key=api_key,
42        )
43        self._model = model
44
45    def _call(self, messages: list[dict]) -> str:
46        response = self._client.chat.completions.create(
47            model=self._model,
48            messages=[{"role": "system", "content": SYSTEM_PROMPT}] + messages,
49        )
50        return response.choices[0].message.content
51
52    def generate(self, description: str) -> DiagramIR | MermaidIR:
53        return parse_response(self._call(build_generate_messages(description)))
54
55    def refine(self, feedback: str, current: DiagramIR | MermaidIR) -> DiagramIR | MermaidIR:
56        if isinstance(current, MermaidIR):
57            return parse_response(self._call(build_mermaid_refine_messages(feedback, current)))
58        return parse_response(self._call(build_refine_messages(feedback, current)))

263class ExcalidrawAdapter:
264    """Backend adapter that renders diagrams as Excalidraw files.
265
266    Writes the diagram to ``~/Documents/diagram-scribe.excalidraw`` by
267    default (or a custom path if provided). Prints the file path after
268    each render so the user knows where to find it.
269
270    Args:
271        output_path: Path to write the ``.excalidraw`` file. Defaults to
272            ``~/Documents/diagram-scribe.excalidraw``.
273
274    Example::
275
276        from diagram_scribe.adapters.backend.excalidraw import ExcalidrawAdapter
277        adapter = ExcalidrawAdapter(output_path="/tmp/my-diagram.excalidraw")
278        adapter.render(ir)
279    """
280
281    def __init__(self, output_path: str | None = None):
282        self._output_path = output_path or _DEFAULT_PATH
283
284    def render(self, ir: DiagramIR) -> None:
285        data = _to_excalidraw(ir)
286        os.makedirs(os.path.dirname(self._output_path), exist_ok=True)
287        with open(self._output_path, "w", encoding="utf-8") as f:
288            json.dump(data, f, indent=2)
289        print(f"[diagram saved to {self._output_path}]")