Update entity_graph.py

entity_graph.py

@@ -1,17 +1,19 @@
 import networkx as nx
 from typing import List, Dict
-import json
 
 def build_entity_graph(docs: List[Dict]) -> Dict:
     G = nx.Graph()
+
     for d in docs:
-        G.add_node(d["agency"])
-        for word in d["title"].split():
-            if word.isupper():
-                G.add_edge(d["agency"], word)
+        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)
 
-    data = {
-        "nodes": [{"id": n} for n in G.nodes],
+    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]
-    }
-    return data
+    }