Skip to content

Neo4j vector

Neo4jVector

Source code in Docs2KG/rag/neo4j_vector.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class Neo4jVector:
    def __init__(self, uri, user, password):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))

    def close(self):
        self.driver.close()

    def add_embedding(self):
        """
        Loop the whole graph
        For all nodes without embedding field, get the nodes out, and grab the embedding.

        Returns:

        """
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH (n)
                WHERE n.content_embedding IS NULL
                RETURN n
                """
            )
            records = list(result)
            total_records = len(records)
            with tqdm(total=total_records, desc="Adding embedding") as pbar:
                for record in records:
                    node = record["n"]
                    logger.debug(node)
                    node_properties = dict(node.items())
                    logger.debug(node_properties)
                    content_embedding = get_openai_embedding(
                        node_properties.get("content", "")
                    )
                    meta_embedding = get_openai_embedding(json.dumps(node_properties))
                    logger.debug(content_embedding)
                    logger.debug(meta_embedding)
                    session.run(
                        """
                        MATCH (n)
                        WHERE id(n) = $id
                        SET n.content_embedding = $content_embedding
                        SET n.meta_embedding = $meta_embedding
                        """,
                        id=node.id,
                        content_embedding=content_embedding,
                        meta_embedding=meta_embedding,
                    )
                    pbar.update(1)

add_embedding()

Loop the whole graph For all nodes without embedding field, get the nodes out, and grab the embedding.

Returns:

Source code in Docs2KG/rag/neo4j_vector.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def add_embedding(self):
    """
    Loop the whole graph
    For all nodes without embedding field, get the nodes out, and grab the embedding.

    Returns:

    """
    with self.driver.session() as session:
        result = session.run(
            """
            MATCH (n)
            WHERE n.content_embedding IS NULL
            RETURN n
            """
        )
        records = list(result)
        total_records = len(records)
        with tqdm(total=total_records, desc="Adding embedding") as pbar:
            for record in records:
                node = record["n"]
                logger.debug(node)
                node_properties = dict(node.items())
                logger.debug(node_properties)
                content_embedding = get_openai_embedding(
                    node_properties.get("content", "")
                )
                meta_embedding = get_openai_embedding(json.dumps(node_properties))
                logger.debug(content_embedding)
                logger.debug(meta_embedding)
                session.run(
                    """
                    MATCH (n)
                    WHERE id(n) = $id
                    SET n.content_embedding = $content_embedding
                    SET n.meta_embedding = $meta_embedding
                    """,
                    id=node.id,
                    content_embedding=content_embedding,
                    meta_embedding=meta_embedding,
                )
                pbar.update(1)