更新时间:2024-11-21 GMT+08:00
分享

添加Agent流式输出(Python SDK)

Agent用于工具调用场景,与普通的LLM流式输出相比,区分了文本流与工具流。

文本流将输出模型的思考过程和最终结果;工具流将输出工具的调用过程,而工具的调用的执行结果是通过监听获取的。

通过如下接口为Agent添加流式输出的回调:

from pangukitsappdev.callback.StreamCallbackHandler import StreamCallbackHandler
# 以下为两个自定义StreamCallbackHandler示例
class TextStreamCallBack(StreamCallbackHandler):
    def __init__(self):
        # agent文本输出工具类
        super().__init__()
        self.stream = ""

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[BaseMessage]],
        **kwargs: Any,
    ) -> Any:
        self.stream += "[Text stream start]"

    def on_llm_new_token(
        self,
        token: str,
        **kwargs: Any,
    ) -> Any:
        self.stream += token

    def on_llm_end(
        self,
        response: LLMResult,
        **kwargs: Any,
    ) -> Any:
        self.stream += "[Text stream end]"


class ToolStreamCallBack(StreamCallbackHandler):
    def __init__(self):
        # agent工具输出工具流
        super().__init__()
        self.stream = ""

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[BaseMessage]],
        **kwargs: Any,
    ) -> Any:
        self.stream += "[Tool stream start]"

    def on_llm_new_token(
        self,
        token: str,
        **kwargs: Any,
    ) -> Any:
        self.stream += token

    def on_llm_end(
        self,
        response: LLMResult,
        **kwargs: Any,
    ) -> Any:
        self.stream += "[Tool stream end]"
text_stream_callback = TextStreamCallBack()
tool_stream_callback = ToolStreamCallBack()
agent.set_stream_callback(text_stream_callback, tool_stream_callback)

StreamCallBack的实现与定义与LLM的回调完全相同。

相关文档