llama-index

1
from llama_index.legacy.constants import DATA_KEY, TYPE_KEY
2
from llama_index.legacy.schema import (
3
    BaseNode,
4
    Document,
5
    ImageDocument,
6
    ImageNode,
7
    IndexNode,
8
    NodeRelationship,
9
    RelatedNodeInfo,
10
    TextNode,
11
)
12

13

14
def doc_to_json(doc: BaseNode) -> dict:
15
    return {
16
        DATA_KEY: doc.dict(),
17
        TYPE_KEY: doc.get_type(),
18
    }
19

20

21
def json_to_doc(doc_dict: dict) -> BaseNode:
22
    doc_type = doc_dict[TYPE_KEY]
23
    data_dict = doc_dict[DATA_KEY]
24
    doc: BaseNode
25

26
    if "extra_info" in data_dict:
27
        return legacy_json_to_doc(doc_dict)
28
    else:
29
        if doc_type == Document.get_type():
30
            doc = Document.parse_obj(data_dict)
31
        elif doc_type == ImageDocument.get_type():
32
            doc = ImageDocument.parse_obj(data_dict)
33
        elif doc_type == TextNode.get_type():
34
            doc = TextNode.parse_obj(data_dict)
35
        elif doc_type == ImageNode.get_type():
36
            doc = ImageNode.parse_obj(data_dict)
37
        elif doc_type == IndexNode.get_type():
38
            doc = IndexNode.parse_obj(data_dict)
39
        else:
40
            raise ValueError(f"Unknown doc type: {doc_type}")
41

42
        return doc
43

44

45
def legacy_json_to_doc(doc_dict: dict) -> BaseNode:
46
    """Todo: Deprecated legacy support for old node versions."""
47
    doc_type = doc_dict[TYPE_KEY]
48
    data_dict = doc_dict[DATA_KEY]
49
    doc: BaseNode
50

51
    text = data_dict.get("text", "")
52
    metadata = data_dict.get("extra_info", {}) or {}
53
    id_ = data_dict.get("doc_id", None)
54

55
    relationships = data_dict.get("relationships", {})
56
    relationships = {
57
        NodeRelationship(k): RelatedNodeInfo(node_id=v)
58
        for k, v in relationships.items()
59
    }
60

61
    if doc_type == Document.get_type():
62
        doc = Document(
63
            text=text, metadata=metadata, id=id_, relationships=relationships
64
        )
65
    elif doc_type == TextNode.get_type():
66
        doc = TextNode(
67
            text=text, metadata=metadata, id=id_, relationships=relationships
68
        )
69
    elif doc_type == ImageNode.get_type():
70
        image = data_dict.get("image", None)
71
        doc = ImageNode(
72
            text=text,
73
            metadata=metadata,
74
            id=id_,
75
            relationships=relationships,
76
            image=image,
77
        )
78
    elif doc_type == IndexNode.get_type():
79
        index_id = data_dict.get("index_id", None)
80
        doc = IndexNode(
81
            text=text,
82
            metadata=metadata,
83
            id=id_,
84
            relationships=relationships,
85
            index_id=index_id,
86
        )
87
    else:
88
        raise ValueError(f"Unknown doc type: {doc_type}")
89

90
    return doc
91