llama-index

1
from queue import Queue
2
from threading import Event
3
from typing import Any, Generator, List, Optional
4
from uuid import UUID
5

6
from llama_index.legacy.bridge.langchain import BaseCallbackHandler, LLMResult
7

8

9
class StreamingGeneratorCallbackHandler(BaseCallbackHandler):
10
    """Streaming callback handler."""
11

12
    def __init__(self) -> None:
13
        self._token_queue: Queue = Queue()
14
        self._done = Event()
15

16
    def __deepcopy__(self, memo: Any) -> "StreamingGeneratorCallbackHandler":
17
        # NOTE: hack to bypass deepcopy in langchain
18
        return self
19

20
    def on_llm_new_token(self, token: str, **kwargs: Any) -> Any:
21
        """Run on new LLM token. Only available when streaming is enabled."""
22
        self._token_queue.put_nowait(token)
23

24
    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
25
        self._done.set()
26

27
    def on_llm_error(
28
        self,
29
        error: BaseException,
30
        *,
31
        run_id: UUID,
32
        parent_run_id: Optional[UUID] = None,
33
        tags: Optional[List[str]] = None,
34
        **kwargs: Any,
35
    ) -> None:
36
        self._done.set()
37

38
    def get_response_gen(self) -> Generator:
39
        while True:
40
            if not self._token_queue.empty():
41
                token = self._token_queue.get_nowait()
42
                yield token
43
            elif self._done.is_set():
44
                break
45