from langchain.callbacks.base import BaseCallbackHandler


class QueueCallbackHandler(BaseCallbackHandler):
    def __init__(self, queue, target_chain_name):
        super(QueueCallbackHandler, self).__init__()
        self.queue = queue
        self.target_chain_name = target_chain_name
        self.target_parent_run_id = None

    def on_chain_start(self, serialized, inputs, **kwargs) -> None:
        if serialized['id'][-1] == self.target_chain_name:
            self.target_parent_run_id = kwargs.get('run_id')

    def on_llm_new_token(self, token, **kwargs) -> None:
        if kwargs.get('parent_run_id') == self.target_parent_run_id:
            self.queue.put_nowait(token)

    def on_llm_error(self, error, **kwargs) -> None:
        self.queue.put_nowait(error)

    def on_chain_error(self, error, **kwargs) -> None:
        self.queue.put_nowait(error)
        
    def on_chain_end(self, outputs, **kwargs) -> None:
        if not kwargs.get('parent_run_id'):
            self.queue.put_nowait(None)