Upload app.py

app.py

@@ -310,6 +310,67 @@ def build_entity_graph(entities_by_doc: List[List[Tuple[str, str]]],
         "top_entities": [{"entity": e[0][0], "type": e[0][1], "docs": e[1]} for e in top_entities]
     }
 
+def build_entity_to_entity_graph(entities_by_doc: List[List[Tuple[str, str]]]) -> Dict[str, Any]:
+    """
+    Constrói grafo onde os NÓS são entidades e ARESTAS são co-ocorrências.
+    Entidades que aparecem no mesmo documento são conectadas.
+    """
+    # Contar co-ocorrências
+    cooccurrence = defaultdict(int)
+    entity_doc_count = defaultdict(int)
+    
+    for entities in entities_by_doc:
+        unique_entities = list(set(entities))
+        
+        # Contar docs por entidade
+        for ent in unique_entities:
+            entity_doc_count[ent] += 1
+        
+        # Criar pares de co-ocorrência
+        for i in range(len(unique_entities)):
+            for j in range(i + 1, len(unique_entities)):
+                pair = tuple(sorted([unique_entities[i], unique_entities[j]], key=str))
+                cooccurrence[pair] += 1
+    
+    # Construir nós (entidades com >= 2 docs)
+    nodes = []
+    entity_to_id = {}
+    for idx, (entity, count) in enumerate(sorted(entity_doc_count.items(), key=lambda x: x[1], reverse=True)):
+        if count >= 2:
+            entity_to_id[entity] = idx
+            nodes.append({
+                "id": idx,
+                "entity": entity[0],
+                "type": entity[1],
+                "docs": count,
+                # Posição inicial em círculo
+                "x": 2 * np.cos(2 * np.pi * idx / max(len(entity_doc_count), 1)),
+                "y": 2 * np.sin(2 * np.pi * idx / max(len(entity_doc_count), 1)),
+                "z": 0
+            })
+    
+    # Construir arestas (co-ocorrências)
+    entity_edges = []
+    for (ent1, ent2), weight in cooccurrence.items():
+        if ent1 in entity_to_id and ent2 in entity_to_id and weight >= 1:
+            entity_edges.append({
+                "source": entity_to_id[ent1],
+                "target": entity_to_id[ent2],
+                "weight": weight,
+                "source_entity": ent1[0],
+                "target_entity": ent2[0]
+            })
+    
+    # Ordenar arestas por peso
+    entity_edges.sort(key=lambda x: x["weight"], reverse=True)
+    
+    return {
+        "nodes": nodes,
+        "edges": entity_edges[:200],  # Limitar a 200 arestas mais fortes
+        "node_count": len(nodes),
+        "edge_count": len(entity_edges)
+    }
+
 
 # ==============================================================================
 # PIPELINE DE PROCESSAMENTO DE DADOS
@@ -843,7 +904,11 @@ async def entity_graph_api(job_id: str = Form(...)):
             for ents in entities_by_doc
         ]
         
-        logging.info(f"Grafo construído: {graph_data['edge_count']} arestas, {len(graph_data['top_entities'])} entidades principais")
+        # Adicionar grafo entidade-entidade
+        entity_network = build_entity_to_entity_graph(entities_by_doc)
+        graph_data["entity_network"] = entity_network
+        
+        logging.info(f"Grafo construído: {graph_data['edge_count']} arestas doc-doc, {entity_network['node_count']} nós entidade")
         
         return graph_data