GodsDevProject/FOIA_Doc_Search
1
1import networkx as nx2from typing import List, Dict3 4def build_entity_graph(docs: List[Dict]) -> Dict:5 G = nx.Graph()6 7 for d in docs:8 agency = d.get("agency", "Unknown")9 G.add_node(agency, group="agency")10 11 for token in d.get("content", "").split():12 if token.isupper() and len(token) > 2:13 G.add_node(token, group="entity")14 G.add_edge(agency, token)15 16 return {17 "nodes": [{"id": n, "group": G.nodes[n]["group"]} for n in G.nodes],18 "links": [{"source": u, "target": v} for u, v in G.edges]19 }