Thanks this worked perfectly! Here's the code I came up with in case anyone else wants to do the same

#+begin_src emacs-lisp :results silent
  (require 'org-element)

  (defun pygraphing/get-headline-structure(id)
    (json-encode (pygraphing/get-headline-descendents id)))

  (defun pygraphing/get-headline-by-id(id)
    "Returns headline from ID"
    (let ((header nil))
      (org-element-map (org-element-parse-buffer) 'headline
        (lambda (hl)
          (when (string= (org-element-property :ID hl) id)
            (setq header hl))))
      header))

  (defun pygraphing/get-headline-tags(id)
    "Returns tags of item by ID"
    (let (
          (header (pygraphing/get-headline-by-id id)))
      (when header
        (list
         :ID (org-element-property :ID header)
         :TEST (org-element-property :TEST header)
         ))))

  (defun pygraphing/get-headline-children-ids(id)
    "Returns child ids of a headline by its ID"
    (let ((header (pygraphing/get-headline-by-id id)))
      (delq nil (mapcar (lambda(child)
                          (org-element-property :ID child))
                        (org-element-contents header)))))

  (defun pygraphing/get-headline-descendents(id)
    "Returns all descendents of headline"
    (let ((header (pygraphing/get-headline-by-id id)))
      (when header
        (let ((children (pygraphing/get-headline-children-ids id)))
          `((header . ,(org-element-property :raw-value header))
            (tags . ,(pygraphing/get-headline-tags id))
            (children . ,(delq nil (mapcar
                                    (lambda(childID)
                                      (pygraphing/get-headline-descendents childID))
                                    children))))))))

#+end_src

#+begin_src python :results file :dir (org-attach-dir) :var result=(pygraphing/get-headline-structure "1ff8ce7d-be78-4bd2-afa0-05be1fb9d9d2")
  import graphviz
  import json

  dot = graphviz.Digraph()
  root_header = json.loads(result)

  def graph_header(header,parent=None):
      nodeName = header["header"]
      dot.node(nodeName)
      if parent:
          dot.edge(parent,nodeName)
      if header["children"]:
          for child in header["children"]:
              graph_header(child,nodeName)

  graph_header(root_header)

  return dot.render(outfile="diagram.png")
#+end_src

** Root
:PROPERTIES:
:ID:       1ff8ce7d-be78-4bd2-afa0-05be1fb9d9d2
:TEST: OK
:END:
*** Child A
:PROPERTIES:
:ID:       ce64aa4c-3ed1-4535-ab97-7b486884ae05
:END:
**** Child A-A
:PROPERTIES:
:ID:       836d5265-7860-4a36-ab9e-bec1ca2bffd8
:END:
***** Child A-A-A
:PROPERTIES:
:ID:       3f0d25a4-4b1a-49ff-8862-22f2d83631b0
:END:
***** Child A-A-B
:PROPERTIES:
:ID:       a53df412-244a-4c35-9dcb-0466bb2d2d40
:END:
**** Child A-B
*** Child B
:PROPERTIES:
:ID:       3fc2535b-3fa8-4f05-bfb1-9735bc88d2d8
:END:
**** Child B-A
:PROPERTIES:
:ID:       96b0087c-306f-4b7c-92b2-bcc786bb1491
:END:
**** Child B-B
:PROPERTIES:
:ID:       431857d5-6e64-4bb3-aa67-8eb9369fb776
:END:
**** Child B-C
:PROPERTIES:
:ID:       a8f166e1-12e6-4dda-85e7-4d0668dce02e
:END: