Spaces:
Sleeping
Sleeping
File size: 1,182 Bytes
da21448 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | from collections.abc import Iterable, Sequence
from techdocs_graphrag.models import GraphFact, canonical_key
SEED_COLOR = "#C7D2FE"
NODE_COLOR = "#F1F5F9"
def facts_to_dot(facts: Sequence[GraphFact], highlighted_keys: Iterable[str] = ()) -> str:
highlighted = set(highlighted_keys)
names = sorted({name for fact in facts for name in (fact.source, fact.target)})
nodes = [
f" {quote(name)} [fillcolor={quote(SEED_COLOR if canonical_key(name) in highlighted else NODE_COLOR)}];"
for name in names
]
edges = [
f" {quote(fact.source)} -> {quote(fact.target)} [label={quote(fact.relation)}];"
for fact in facts
]
return "\n".join(
(
"digraph KnowledgeGraph {",
" rankdir=LR;",
' node [shape=box, style="rounded,filled", fontname="Helvetica", fontcolor="#0F172A"];',
' edge [fontname="Helvetica", fontsize=10, color="#64748B", fontcolor="#64748B"];',
*nodes,
*edges,
"}",
)
)
def quote(value: str) -> str:
escaped = value.replace("\\", "\\\\").replace('"', '\\"').replace("\n", " ")
return f'"{escaped}"'
|