entity_graph.py

import networkx as nx
from typing import List, Dict

def build_entity_graph(docs: List[Dict]) -> Dict:
    G = nx.Graph()

    for d in docs:
        agency = d.get("agency", "Unknown")
        G.add_node(agency, group="agency")

        for token in d.get("content", "").split():
            if token.isupper() and len(token) > 2:
                G.add_node(token, group="entity")
                G.add_edge(agency, token)

    return {
        "nodes": [{"id": n, "group": G.nodes[n]["group"]} for n in G.nodes],
        "links": [{"source": u, "target": v} for u, v in G.edges]
    }