Spaces:
Sleeping
Sleeping
| 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}"' | |