llama-index

1
"""Struct store schema."""
2

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

6
from dataclasses_json import DataClassJsonMixin
7

8
from llama_index.legacy.data_structs.data_structs import IndexStruct
9
from llama_index.legacy.data_structs.struct_type import IndexStructType
10

11

12
@dataclass
13
class StructDatapoint(DataClassJsonMixin):
14
    """Struct outputs."""
15

16
    # map from field name to StructValue
17
    fields: Dict[str, Any]
18

19

20
@dataclass
21
class BaseStructTable(IndexStruct):
22
    """Struct outputs."""
23

24

25
@dataclass
26
class SQLStructTable(BaseStructTable):
27
    """SQL struct outputs."""
28

29
    context_dict: Dict[str, str] = field(default_factory=dict)
30

31
    @classmethod
32
    def get_type(cls) -> IndexStructType:
33
        """Get type."""
34
        # TODO: consolidate with IndexStructType
35
        return IndexStructType.SQL
36

37

38
@dataclass
39
class PandasStructTable(BaseStructTable):
40
    """Pandas struct outputs."""
41

42
    @classmethod
43
    def get_type(cls) -> IndexStructType:
44
        """Get type."""
45
        return IndexStructType.PANDAS
46