import queue
from langchain.callbacks.base import BaseCallbackHandler

# Handler personnalisé pour capturer les tokens en streaming
class GradioStreamingHandler(BaseCallbackHandler):
    def __init__(self):
        self.tokens_queue = queue.Queue()
        self.full_text = ""
        
    def on_llm_new_token(self, token, **kwargs):
        self.tokens_queue.put(token)
        self.full_text += token