138 lines
4.6 KiB
Python
138 lines
4.6 KiB
Python
"""
|
|
Script test ZepEntityReader — đọc graph đã build, xem logic chọn entity làm agent
|
|
Chạy từ thư mục backend/
|
|
|
|
Usage:
|
|
python scripts/test_entity_reader.py <graph_id>
|
|
python scripts/test_entity_reader.py <graph_id> --type Student
|
|
python scripts/test_entity_reader.py <graph_id> --uuid <entity_uuid>
|
|
python scripts/test_entity_reader.py --from-log case_01_academic_scandal
|
|
"""
|
|
import sys
|
|
import json
|
|
import os
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from app.services.zep_entity_reader import ZepEntityReader
|
|
|
|
OUTPUT_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "logs")
|
|
|
|
|
|
def load_graph_id_from_log(case_id: str) -> str:
|
|
log_path = os.path.join(OUTPUT_DIR, f"test_graph_output_{case_id}.json")
|
|
if not os.path.exists(log_path):
|
|
print(f"Không tìm thấy file log: {log_path}")
|
|
sys.exit(1)
|
|
with open(log_path, encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
return data["graph_id"]
|
|
|
|
|
|
def print_entity(entity, show_edges: bool = True):
|
|
entity_type = entity.get_entity_type() or "(unknown)"
|
|
print(f"\n ── [{entity_type}] {entity.name}")
|
|
print(f" uuid : {entity.uuid}")
|
|
if entity.summary:
|
|
print(f" summary : {entity.summary[:120]}")
|
|
if entity.attributes:
|
|
for k, v in entity.attributes.items():
|
|
if k not in ("name",) and v and v != "null":
|
|
print(f" attr : {k} = {v}")
|
|
if show_edges and entity.related_edges:
|
|
print(f" edges ({len(entity.related_edges)}):")
|
|
for e in entity.related_edges:
|
|
arrow = "→" if e["direction"] == "outgoing" else "←"
|
|
print(f" {arrow} [{e['edge_name']}] {e.get('fact', '')[:80]}")
|
|
if entity.related_nodes:
|
|
names = [n["name"] for n in entity.related_nodes]
|
|
print(f" related : {', '.join(names)}")
|
|
|
|
|
|
def cmd_all(reader: ZepEntityReader, graph_id: str):
|
|
"""Lấy tất cả entity hợp lệ (có custom label) → danh sách agent tiềm năng"""
|
|
print(f"\nGraph: {graph_id}")
|
|
print("Đang lọc entity...")
|
|
result = reader.filter_defined_entities(graph_id, enrich_with_edges=True)
|
|
|
|
print(f"\n{'=' * 60}")
|
|
print(f"Tổng nodes : {result.total_count}")
|
|
print(f"Entity hợp lệ : {result.filtered_count} ← đây là các agent tiềm năng")
|
|
print(f"Entity types : {sorted(result.entity_types)}")
|
|
print(f"{'=' * 60}")
|
|
|
|
# Nhóm theo type để dễ đọc
|
|
by_type: dict = {}
|
|
for e in result.entities:
|
|
t = e.get_entity_type() or "unknown"
|
|
by_type.setdefault(t, []).append(e)
|
|
|
|
for entity_type, entities in sorted(by_type.items()):
|
|
print(f"\n[{entity_type}] — {len(entities)} entity")
|
|
for e in entities:
|
|
print_entity(e)
|
|
|
|
# Lưu output
|
|
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
out_path = os.path.join(OUTPUT_DIR, f"test_entity_reader_{graph_id}.json")
|
|
with open(out_path, "w", encoding="utf-8") as f:
|
|
json.dump(result.to_dict(), f, ensure_ascii=False, indent=2)
|
|
print(f"\nĐã lưu: {out_path}")
|
|
|
|
|
|
def cmd_by_type(reader: ZepEntityReader, graph_id: str, entity_type: str):
|
|
print(f"\nLấy entity type [{entity_type}] từ graph {graph_id}...")
|
|
entities = reader.get_entities_by_type(graph_id, entity_type, enrich_with_edges=True)
|
|
print(f"Tìm thấy {len(entities)} entity loại [{entity_type}]:")
|
|
for e in entities:
|
|
print_entity(e)
|
|
|
|
|
|
def cmd_by_uuid(reader: ZepEntityReader, graph_id: str, uuid: str):
|
|
print(f"\nLấy chi tiết entity {uuid}...")
|
|
entity = reader.get_entity_with_context(graph_id, uuid)
|
|
if not entity:
|
|
print("Không tìm thấy entity.")
|
|
return
|
|
print_entity(entity, show_edges=True)
|
|
|
|
|
|
def main():
|
|
args = sys.argv[1:]
|
|
graph_id = None
|
|
entity_type = None
|
|
entity_uuid = None
|
|
|
|
if not args:
|
|
print(__doc__)
|
|
sys.exit(1)
|
|
|
|
if args[0] == "--from-log":
|
|
if len(args) < 2:
|
|
print("Usage: --from-log <case_id>")
|
|
sys.exit(1)
|
|
graph_id = load_graph_id_from_log(args[1])
|
|
print(f"Graph ID từ log [{args[1]}]: {graph_id}")
|
|
args = args[2:]
|
|
else:
|
|
graph_id = args[0]
|
|
args = args[1:]
|
|
|
|
if "--type" in args:
|
|
entity_type = args[args.index("--type") + 1]
|
|
if "--uuid" in args:
|
|
entity_uuid = args[args.index("--uuid") + 1]
|
|
|
|
reader = ZepEntityReader()
|
|
|
|
if entity_uuid:
|
|
cmd_by_uuid(reader, graph_id, entity_uuid)
|
|
elif entity_type:
|
|
cmd_by_type(reader, graph_id, entity_type)
|
|
else:
|
|
cmd_all(reader, graph_id)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|