100行代码搞定多智能体？这个极简AI框架PocketFlow有点东西 原创

class AsyncHinter(AsyncNode):
    asyncdef prep_async(self, shared):
        # 1. 从队列中等待猜词者的消息
        guess = await shared["hinter_queue"].get()
        if guess == "GAME_OVER":
            returnNone
        # 2. 准备LLM的输入
        return shared["target_word"], shared["forbidden_words"], shared.get("past_guesses", [])

    asyncdef exec_async(self, inputs):
        # ... (调用LLM生成提示)
        target, forbidden, past_guesses = inputs
        prompt = f"Generate hint for '{target}'\nForbidden words: {forbidden}"
        if past_guesses:
            prompt += f"\nPrevious wrong guesses: {past_guesses}\nMake hint more specific."
        hint = call_llm(prompt)
        print(f"\nHinter: Here's your hint - {hint}")
        return hint

    asyncdef post_async(self, shared, prep_res, exec_res):
        if exec_res isNone:
            return"end"
        # 3. 将生成的提示放入猜词者的队列
        await shared["guesser_queue"].put(exec_res)
        return"continue"# 返回Action，驱动流程继续

class AsyncGuesser(AsyncNode):
    asyncdef prep_async(self, shared):
        # 1. 从队列中等待提示者的提示
        hint = await shared["guesser_queue"].get()
        return hint, shared.get("past_guesses", [])

    asyncdef exec_async(self, inputs):
        # ... (调用LLM生成猜测)
        hint, past_guesses = inputs
        prompt = f"Given hint: {hint}, past wrong guesses: {past_guesses}, make a new guess."
        guess = call_llm(prompt)
        print(f"Guesser: I guess it's - {guess}")
        return guess

    asyncdef post_async(self, shared, prep_res, exec_res):
        # 2. 检查答案
        if exec_res.lower() == shared["target_word"].lower():
            print("Game Over - Correct guess!")
            await shared["hinter_queue"].put("GAME_OVER")
            return"end"
        
        # 3. 如果猜错，将错误答案发回给提示者，以便其给出更好的提示
        shared.setdefault("past_guesses", []).append(exec_res)
        await shared["hinter_queue"].put(exec_res)
        return"continue"

async def main():
    # ... (初始化shared, 包括target_word, forbidden_words, 和两个queue)

    # 创建节点和流
    hinter = AsyncHinter()
    guesser = AsyncGuesser()

    hinter_flow = AsyncFlow(start=hinter)
    guesser_flow = AsyncFlow(start=guesser)

    # 定义循环：当post返回"continue"时，节点会再次执行自己
    hinter - "continue" >> hinter
    guesser - "continue" >> guesser

    # 使用asyncio.gather并发运行两个智能体流
    await asyncio.gather(
        hinter_flow.run_async(shared),
        guesser_flow.run_async(shared)
    )

+-----------+
  Shared   |           |
  Store <--|   Node    |<-- Params（仅 Batch 用）
           +-----------+
                |
             Action
                v
           +-----------+
           |   Node    |
           +-----------+

class BaseNode:
    def __init__(self): self.params,self.successors={},{}
    def set_params(self,params): self.params=params
    def next(self,node,actinotallow="default"):
        if action in self.successors: warnings.warn(f"Overwriting successor for action '{action}'")
        self.successors[action]=node; return node
    def prep(self,shared): pass
    def exec(self,prep_res): pass
    def post(self,shared,prep_res,exec_res): pass
    def _exec(self,prep_res): return self.exec(prep_res)
    def _run(self,shared): p=self.prep(shared); e=self._exec(p); return self.post(shared,p,e)
    def run(self,shared): 
        if self.successors: warnings.warn("Node won't run successors. Use Flow.")  
        return self._run(shared)

class BaseNode:
    ......
    def __rshift__(self,other): return self.next(other)
    def __sub__(self,action):
        if isinstance(action,str): return _ConditionalTransition(self,action)
        raise TypeError("Action must be a string")

class _ConditionalTransition:
    def __init__(self,src,action): self.src,self.actinotallow=src,action
    def __rshift__(self,tgt): return self.src.next(tgt,self.action)

class Node(BaseNode):
    def __init__(self,max_retries=1,wait=0): super().__init__(); self.max_retries,self.wait=max_retries,wait
    def exec_fallback(self,prep_res,exc): raise exc
    def _exec(self,prep_res):
        for self.cur_retry in range(self.max_retries):
            try: return self.exec(prep_res)
            except Exception as e:
                if self.cur_retry==self.max_retries-1: return self.exec_fallback(prep_res,e)
                if self.wait>0: time.sleep(self.wait)

class Flow(BaseNode):
    def __init__(self,start=None): super().__init__(); self.start_node=start
    # ...
    def get_next_node(self,curr,action):
        nxt=curr.successors.get(action or"default")
        ifnot nxt and curr.successors: warnings.warn(f"Flow ends: '{action}' not found in {list(curr.successors)}")
        return nxt
    def _orch(self,shared,params=None):
        curr,p,last_action =copy.copy(self.start_node),(params or {**self.params}),None
        while curr:
            curr.set_params(p)
            last_actinotallow=curr._run(shared)
            curr=copy.copy(self.get_next_node(curr,last_action))
        return last_action