Python中如何让生成器yield另一个类函数的返回结果？

核心思路是让GRPCTokenStreamingHandler把收到的token存入队列，然后在AskQuestion生成器中通过后台线程执行Pipeline，同时从队列中逐个取出token并yield出去。

修改后的GRPCTokenStreamingHandler代码 #

import queue
import threading

class GRPCTokenStreamingHandler(TokenStreamingHandler):
    def __init__(self):
        super().__init__()
        # 存储流式返回的token
        self.token_queue = queue.Queue()
        # 标记流式输出是否完成
        self.done = False

    def __call__(self, token_received, **kwargs) -> str:
        # 将收到的token放入队列
        self.token_queue.put(token_received)
        return token_received

    def mark_done(self):
        # 标记流式输出结束
        self.done = True

修改后的AskQuestion生成器代码 #

class ChatBot(chat_pb2_grpc.ChatBotServicer):

    def AskQuestion(self, request, context):
        query = request.query
        custom_handler = GRPCTokenStreamingHandler()
        prompt_node = PromptNode(
            "gpt-4",
            default_prompt_template=lfqa_prompt,
            api_key=api_key,
            max_length=4096,
            model_kwargs={"stream": True, "stream_handler": custom_handler}
        )

        pipe = Pipeline()
        pipe.add_node(component=retriever, name="retriever", inputs=["Query"])
        pipe.add_node(component=prompt_node, name="prompt_node", inputs=["retriever"])

        # 后台线程执行Pipeline，避免阻塞生成器
        def run_pipeline():
            try:
                pipe.run(query=query)
            finally:
                # 无论成功或异常，都标记流式输出结束
                custom_handler.mark_done()

        # 启动守护线程，随生成器结束自动退出
        threading.Thread(target=run_pipeline, daemon=True).start()

        # 循环获取token并yield
        while not custom_handler.done or not custom_handler.token_queue.empty():
            try:
                # 超时检查，避免一直阻塞在取队列操作
                token = custom_handler.token_queue.get(timeout=0.1)
                yield chat_pb2.Response(token=token, final=False)
            except queue.Empty:
                continue

        # 返回最终结束标记
        yield chat_pb2.Response(token="", final=True)

关键改动说明 #

• 给处理器添加队列：让GRPCTokenStreamingHandler把每个收到的token存入队列，实现和生成器的解耦。
• 后台线程执行Pipeline：pipe.run()是阻塞调用，放在后台线程中执行，生成器才能持续处理队列中的token。
• 循环监听队列：通过循环检查结束标记和队列状态，确保所有token都被yield出去，最后返回结束信号。

内容的提问来源于stack exchange，提问作者serlingpa