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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140 | class Event(BaseModel):
"""
Base class for event types that mimics dict interface.
PrivateAttr:
_data (Dict[str, Any]): Underlying Python dict.
Examples:
Basic example usage
```python
from llama_index.core.workflows.events import Event
evt = Event(a=1, b=2)
# can use dot access to get values of `a` and `b`
print((evt.a, evt.b))
# can also set the attrs
evt.a = 2
```
Custom event with additional Fields/PrivateAttr
```python
from llama_index.core.workflows.events import Event
from llama_index.core.bridge.pydantic import Field, PrivateAttr
class CustomEvent(Event):
field_1: int = Field(description="my custom field")
_private_attr_1: int = PrivateAttr()
evt = CustomEvent(a=1, b=2, field_1=3, _private_attr_1=4)
# `field_1` and `_private_attr_1` get set as they do with Pydantic BaseModel
print(evt.field_1)
print(evt._private_attr_1)
# `a` and `b` get set in the underlying dict, namely `evt._data`
print((evt.a, evt.b))
```
"""
model_config = ConfigDict(arbitrary_types_allowed=True)
_data: Dict[str, Any] = PrivateAttr(default_factory=dict)
def __init__(self, **params: Any):
"""
__init__.
NOTE: fields and private_attrs are pulled from params by name.
"""
# extract and set fields, private attrs and remaining shove in _data
fields = {}
private_attrs = {}
data = {}
for k, v in params.items():
if k in self.model_fields:
fields[k] = v
elif k in self.__private_attributes__:
private_attrs[k] = v
else:
data[k] = v
super().__init__(**fields)
for private_attr, value in private_attrs.items():
super().__setattr__(private_attr, value)
if data:
self._data.update(data)
def __getattr__(self, __name: str) -> Any:
if __name in self.__private_attributes__ or __name in self.model_fields:
return super().__getattr__(__name) # type: ignore
else:
try:
return self._data[__name]
except KeyError:
raise AttributeError(
f"'{self.__class__.__name__}' object has no attribute '{__name}'"
)
def __setattr__(self, name: str, value: Any) -> None:
if name in self.__private_attributes__ or name in self.model_fields:
super().__setattr__(name, value)
else:
self._data.__setitem__(name, value)
def __getitem__(self, key: str) -> Any:
return self._data[key]
def __setitem__(self, key: str, value: Any) -> None:
self._data[key] = value
def get(self, key: str, default: Any = None) -> Any:
return self._data.get(key, default)
def __contains__(self, key: str) -> bool:
return key in self._data
def keys(self) -> "dict_keys[str, Any]":
return self._data.keys()
def values(self) -> "dict_values[str, Any]":
return self._data.values()
def items(self) -> "dict_items[str, Any]":
return self._data.items()
def __len__(self) -> int:
return len(self._data)
def __iter__(self) -> Any:
return iter(self._data)
def dict(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
return self._data
def __bool__(self) -> bool:
"""Make test `if event:` pass on Event instances."""
return True
@model_serializer(mode="wrap")
def custom_model_dump(self, handler: Any) -> Dict[str, Any]:
data = handler(self)
# include _data in serialization
if self._data:
data["_data"] = self._data
return data
|
