|
4 | 4 | # This source code is licensed under the terms described in the LICENSE file in |
5 | 5 | # the root directory of this source tree. |
6 | 6 |
|
7 | | -from typing import Any, Iterator, Optional, Tuple |
| 7 | +from typing import Any, Iterator, Optional |
8 | 8 |
|
9 | 9 | from termcolor import cprint |
10 | 10 |
|
@@ -55,17 +55,10 @@ def print(self, flush: bool = True) -> None: |
55 | 55 |
|
56 | 56 |
|
57 | 57 | class TurnStreamEventPrinter: |
58 | | - def __init__(self) -> None: |
59 | | - self.previous_event_type: Optional[str] = None |
60 | | - self.previous_step_type: Optional[str] = None |
61 | | - |
62 | 58 | def yield_printable_events(self, chunk: Any) -> Iterator[TurnStreamPrintableEvent]: |
63 | | - for printable_event in self._yield_printable_events(chunk, self.previous_event_type, self.previous_step_type): |
| 59 | + for printable_event in self._yield_printable_events(chunk): |
64 | 60 | yield printable_event |
65 | 61 |
|
66 | | - if not hasattr(chunk, "error"): |
67 | | - self.previous_event_type, self.previous_step_type = self._get_event_type_step_type(chunk) |
68 | | - |
69 | 62 | def _yield_printable_events(self, chunk: Any) -> Iterator[TurnStreamPrintableEvent]: |
70 | 63 | if hasattr(chunk, "error"): |
71 | 64 | yield TurnStreamPrintableEvent(role=None, content=chunk.error["message"], color="red") |
@@ -144,17 +137,6 @@ def _yield_printable_events(self, chunk: Any) -> Iterator[TurnStreamPrintableEve |
144 | 137 | color="green", |
145 | 138 | ) |
146 | 139 |
|
147 | | - def _get_event_type_step_type(self, chunk: Any) -> Tuple[Optional[str], Optional[str]]: |
148 | | - if hasattr(chunk, "event"): |
149 | | - previous_event_type = chunk.event.payload.event_type if hasattr(chunk, "event") else None |
150 | | - previous_step_type = ( |
151 | | - chunk.event.payload.step_type |
152 | | - if previous_event_type not in {"turn_start", "turn_complete", "turn_awaiting_input"} |
153 | | - else None |
154 | | - ) |
155 | | - return previous_event_type, previous_step_type |
156 | | - return None, None |
157 | | - |
158 | 140 |
|
159 | 141 | class EventLogger: |
160 | 142 | def log(self, event_generator: Iterator[Any]) -> Iterator[TurnStreamPrintableEvent]: |
|
0 commit comments