llama-index

1
"""Data struct for document summary index."""
2

3
from dataclasses import dataclass, field
4
from typing import Dict, List
5

6
from llama_index.legacy.data_structs.data_structs import IndexStruct
7
from llama_index.legacy.data_structs.struct_type import IndexStructType
8
from llama_index.legacy.schema import BaseNode
9

10

11
@dataclass
12
class IndexDocumentSummary(IndexStruct):
13
    """A simple struct containing a mapping from summary node_id to doc node_ids.
14

15
    Also mapping vice versa.
16

17
    """
18

19
    summary_id_to_node_ids: Dict[str, List[str]] = field(default_factory=dict)
20
    node_id_to_summary_id: Dict[str, str] = field(default_factory=dict)
21

22
    # track mapping from doc id to node summary id
23
    doc_id_to_summary_id: Dict[str, str] = field(default_factory=dict)
24

25
    def add_summary_and_nodes(
26
        self,
27
        summary_node: BaseNode,
28
        nodes: List[BaseNode],
29
    ) -> str:
30
        """Add node and summary."""
31
        summary_id = summary_node.node_id
32
        ref_doc_id = summary_node.ref_doc_id
33
        if ref_doc_id is None:
34
            raise ValueError(
35
                "ref_doc_id of node cannot be None when building a document "
36
                "summary index"
37
            )
38
        self.doc_id_to_summary_id[ref_doc_id] = summary_id
39

40
        for node in nodes:
41
            node_id = node.node_id
42
            if summary_id not in self.summary_id_to_node_ids:
43
                self.summary_id_to_node_ids[summary_id] = []
44
            self.summary_id_to_node_ids[summary_id].append(node_id)
45

46
            self.node_id_to_summary_id[node_id] = summary_id
47

48
        return summary_id
49

50
    @property
51
    def summary_ids(self) -> List[str]:
52
        """Get summary ids."""
53
        return list(self.summary_id_to_node_ids.keys())
54

55
    def delete(self, doc_id: str) -> None:
56
        """Delete a document and its nodes."""
57
        summary_id = self.doc_id_to_summary_id[doc_id]
58
        del self.doc_id_to_summary_id[doc_id]
59
        node_ids = self.summary_id_to_node_ids[summary_id]
60
        for node_id in node_ids:
61
            del self.node_id_to_summary_id[node_id]
62
        del self.summary_id_to_node_ids[summary_id]
63

64
    def delete_nodes(self, node_ids: List[str]) -> None:
65
        for node_id in node_ids:
66
            summary_id = self.node_id_to_summary_id[node_id]
67
            self.summary_id_to_node_ids[summary_id].remove(node_id)
68
            del self.node_id_to_summary_id[node_id]
69

70
    @classmethod
71
    def get_type(cls) -> IndexStructType:
72
        """Get type."""
73
        return IndexStructType.DOCUMENT_SUMMARY
74