| 92 | |
| 93 | |
| 94 | class EventCallbackHandler(BaseCallbackHandler): |
| 95 | _aqueue: asyncio.Queue |
| 96 | is_done: bool = False |
| 97 | |
| 98 | def __init__( |
| 99 | self, |
| 100 | ): |
| 101 | """Initialize the base callback handler.""" |
| 102 | ignored_events = [ |
| 103 | CBEventType.CHUNKING, |
| 104 | CBEventType.NODE_PARSING, |
| 105 | CBEventType.EMBEDDING, |
| 106 | CBEventType.LLM, |
| 107 | CBEventType.TEMPLATING, |
| 108 | ] |
| 109 | super().__init__(ignored_events, ignored_events) |
| 110 | self._aqueue = asyncio.Queue() |
| 111 | |
| 112 | def on_event_start( |
| 113 | self, |
| 114 | event_type: CBEventType, |
| 115 | payload: Optional[Dict[str, Any]] = None, |
| 116 | event_id: str = "", |
| 117 | **kwargs: Any, |
| 118 | ) -> str: |
| 119 | event = CallbackEvent(event_id=event_id, event_type=event_type, payload=payload) |
| 120 | if event.to_response() is not None: |
| 121 | self._aqueue.put_nowait(event) |
| 122 | |
| 123 | def on_event_end( |
| 124 | self, |
| 125 | event_type: CBEventType, |
| 126 | payload: Optional[Dict[str, Any]] = None, |
| 127 | event_id: str = "", |
| 128 | **kwargs: Any, |
| 129 | ) -> None: |
| 130 | event = CallbackEvent(event_id=event_id, event_type=event_type, payload=payload) |
| 131 | if event.to_response() is not None: |
| 132 | self._aqueue.put_nowait(event) |
| 133 | |
| 134 | def start_trace(self, trace_id: Optional[str] = None) -> None: |
| 135 | """No-op.""" |
| 136 | |
| 137 | def end_trace( |
| 138 | self, |
| 139 | trace_id: Optional[str] = None, |
| 140 | trace_map: Optional[Dict[str, List[str]]] = None, |
| 141 | ) -> None: |
| 142 | """No-op.""" |
| 143 | |
| 144 | async def async_event_gen(self) -> AsyncGenerator[CallbackEvent, None]: |
| 145 | while not self._aqueue.empty() or not self.is_done: |
| 146 | try: |
| 147 | yield await asyncio.wait_for(self._aqueue.get(), timeout=0.1) |
| 148 | except asyncio.TimeoutError: |
| 149 | pass |