0shin0/GraphDB_GraphRAG_Workshop

1from neo4j import GraphDatabase, basic_auth
2from neo4j.time import Date
3
4def get_node_datatype(value):
5    """
6        입력된 노드 Value의 데이터 타입을 반환하는 함수
7    """
8    if isinstance(value, str):
9        return "STRING"
10    elif isinstance(value, int):
11        return "INTEGER"
12    elif isinstance(value, float):
13        return "FLOAT"
14    elif isinstance(value, bool):
15        return "BOOLEAN"
16    elif isinstance(value, list):
17        return f"LIST[{get_node_datatype(value[0])}]" if value else "LIST"
18    elif isinstance(value, Date):
19        return "DATE"
20    else:
21        return "UNKNOWN"
22
23def get_schema(uri, user, password):
24    """
25        Graph DB의 정보를 받아 노드 및 관계의 프로퍼티를 추출하고 스키마 딕셔너리를 반환하는 함수
26    """
27    driver = GraphDatabase.driver(
28        uri,
29        auth=basic_auth(user, password))
30
31    with driver.session() as session:
32        # 노드 프로퍼티 및 타입 추출
33        node_query = """
34        MATCH (n)
35        WITH DISTINCT labels(n) AS node_labels, keys(n) AS property_keys, n
36        UNWIND node_labels AS label
37        UNWIND property_keys AS key
38        RETURN label, key, n[key] AS sample_value
39        """
40        nodes = session.run(node_query)
41
42        # 관계 프로퍼티 및 타입 추출
43        rel_query = """
44        MATCH ()-[r]->()
45        WITH DISTINCT type(r) AS rel_type, keys(r) AS property_keys, r
46        UNWIND property_keys AS key
47        RETURN rel_type, key, r[key] AS sample_value
48        """
49        relationships = session.run(rel_query)
50
51        # 관계 유형 및 방향 추출
52        rel_direction_query = """
53        MATCH (a)-[r]->(b)
54        RETURN DISTINCT labels(a) AS start_label, type(r) AS rel_type, labels(b) AS end_label
55        ORDER BY start_label, rel_type, end_label
56        """
57        rel_directions = session.run(rel_direction_query)
58
59        # 스키마 딕셔너리 생성
60        schema = {"nodes": {}, "relationships": {}, "relations": []}
61
62        for record in nodes:
63            label = record["label"]
64            key = record["key"]
65            sample_value = record["sample_value"] # 데이터 타입을 추론하기 위한 샘플 데이터
66            inferred_type = get_node_datatype(sample_value)
67            if label not in schema["nodes"]:
68                schema["nodes"][label] = {}
69            schema["nodes"][label][key] = inferred_type
70
71        for record in relationships:
72            rel_type = record["rel_type"]
73            key = record["key"]
74            sample_value = record["sample_value"] # 데이터 타입을 추론하기 위한 샘플 데이터
75            inferred_type = get_node_datatype(sample_value)
76            if rel_type not in schema["relationships"]:
77                schema["relationships"][rel_type] = {}
78            schema["relationships"][rel_type][key] = inferred_type
79
80        for record in rel_directions:
81            start_label = record["start_label"][0]
82            rel_type = record["rel_type"]
83            end_label = record["end_label"][0]
84            schema["relations"].append(f"(:{start_label})-[:{rel_type}]->(:{end_label})")
85
86        return schema
87
88def format_schema(schema):
89    """
90        스키마 딕셔너리를 LLM에 제공하기 위해 원하는 형태로 formatting 하는 함수
91    """
92    result = []
93
94    # 노드 프로퍼티 출력
95    result.append("Node properties:")
96    for label, properties in schema["nodes"].items():
97        props = ", ".join(f"{k}: {v}" for k, v in properties.items())
98        result.append(f"{label} {{{props}}}")
99
100    # 관계 프로퍼티 출력
101    result.append("Relationship properties:")
102    for rel_type, properties in schema["relationships"].items():
103        props = ", ".join(f"{k}: {v}" for k, v in properties.items())
104        result.append(f"{rel_type} {{{props}}}")
105
106    # 관계 프로퍼티 출력
107    result.append("The relationships:")
108    for relation in schema["relations"]:
109        result.append(relation)
110
111    return "\n".join(result)