kimi联网搜索tool的流式输出教程

仙人球小熊

已于 2024-09-14 18:40:19 修改

阅读量465

点赞数 7

文章标签：人工智能 python 语言模型

于 2024-09-14 18:33:20 首次发布

本文链接：https://blog.csdn.net/weixin_63482830/article/details/142253980

版权

最近，moonshot开放了kimi的联网搜索接口，非常滴好用（直接干掉了博主写的多agent联网搜索框架）。但是，官方所提供的使用代码并不支持流式输出，这导致该接口无法与博主的现有工程兼容，为此，博主自研了一套流式输出方案。

先看kimi开发文档中的代码（不支持流式输出）：

from typing import *
 
import os
import json
 
from openai import OpenAI
from openai.types.chat.chat_completion import Choice
 
client = OpenAI(
    base_url="https://api.moonshot.cn/v1",
    api_key=os.environ.get("MOONSHOT_API_KEY"),
)
 
 
# search 工具的具体实现，这里我们只需要返回参数即可
def search_impl(arguments: Dict[str, Any]) -> Any:
    """
    在使用 Moonshot AI 提供的 search 工具的场合，只需要原封不动返回 arguments 即可，
    不需要额外的处理逻辑。
 
    但如果你想使用其他模型，并保留联网搜索的功能，那你只需要修改这里的实现（例如调用搜索
    和获取网页内容等），函数签名不变，依然是 work 的。
 
    这最大程度保证了兼容性，允许你在不同的模型间切换，并且不需要对代码有破坏性的修改。
    """
    return arguments
 
 
def chat(messages) -> Choice:
    completion = client.chat.completions.create(
        model="moonshot-v1-128k",
        messages=messages,
        temperature=0.3,
        tools=[
            {
                "type": "builtin_function",  # <-- 使用 builtin_function 声明 $web_search 函数，请在每次请求都完整地带上 tools 声明
                "function": {
                    "name": "$web_search",
                },
            }
        ]
    )
    return completion.choices[0]
 
 
def main():
    messages = [
        {"role": "system", "content": "你是 Kimi。"},
    ]
 
    # 初始提问
    messages.append({
        "role": "user",
        "content": "请搜索 Moonshot AI Context Caching 技术，并告诉我它是什么。"
    })
 
    finish_reason = None
    while finish_reason is None or finish_reason == "tool_calls":
        choice = chat(messages)
        finish_reason = choice.finish_reason
        if finish_reason == "tool_calls":  # <-- 判断当前返回内容是否包含 tool_calls
            messages.append(choice.message)  # <-- 我们将 Kimi 大模型返回给我们的 assistant 消息也添加到上下文中，以便于下次请求时 Kimi 大模型能理解我们的诉求
            for tool_call in choice.message.tool_calls:  # <-- tool_calls 可能是多个，因此我们使用循环逐个执行
                tool_call_name = tool_call.function.name
                tool_call_arguments = json.loads(tool_call.function.arguments)  # <-- arguments 是序列化后的 JSON Object，我们需要使用 json.loads 反序列化一下
                if tool_call_name == "$web_search":
                    tool_result = search_impl(tool_call_arguments)
                else:
                    tool_result = f"Error: unable to find tool by name '{tool_call_name}'"
 
                # 使用函数执行结果构造一个 role=tool 的 message，以此来向模型展示工具调用的结果；
                # 注意，我们需要在 message 中提供 tool_call_id 和 name 字段，以便 Kimi 大模型
                # 能正确匹配到对应的 tool_call。
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "name": tool_call_name,
                    "content": json.dumps(tool_result),  # <-- 我们约定使用字符串格式向 Kimi 大模型提交工具调用结果，因此在这里使用 json.dumps 将执行结果序列化成字符串
                })
 
    print(choice.message.content)  # <-- 在这里，我们才将模型生成的回复返回给用户
 
 
if __name__ == '__main__':
    main()

其流程为：

1.当 Kimi 大模型生成了 finish_reason=tool_calls 的响应时，表明 Kimi 大模型已经意识到当前需要执行 $web_search 函数，并且也已经做好执行 $web_search 的一切准备工作；

2.Kimi 大模型会将执行函数所必须得参数以 tool_call.function.arguments 的形式返回给调用方，但这些参数并不由调用方执行，调用方只需要将 tool_call.function.arguments 原封不动地提交给 Kimi 大模型，即可由 Kimi 大模型执行对应的联网搜索流程；

3.当用户将 tool_call.function.arguments 使用 role=tool 的 message 提交时，Kimi 大模型随即开始执行联网搜索流程，并根据搜索和阅读结果生成可供用户阅读的消息，即 finish_reason=stop 的 message。

经过博主试验发现，上述代码的主while循环只会执行两次：

第一次循环中，finish_reason = choice.finish_reason 产生了 finish_reason=tool_calls 的响应。使得 if finish_reason=tool_calls 的判断为True，执行了联网tool的调用，并将联网搜索的结果记录到message中。

第二次循环中，在choice = chat(messages) 一步就产生了最终回答，通过print(choice.message.content) 将模型生成的回复打印到工作台。

可以看到，联网搜索的结果都是在第一次循环时产生的，而第二次循环只是让kimi生成给用户的最终回答。因此，我们只需要在第二次循环中设定kimi的回答方式是流式的即可。

基于这个思路，博主将代码进行了改良：

请将文件命名为：kimi.py

from typing import *
import os
import json
from openai import OpenAI
from openai.types.chat.chat_completion import Choice

# 初始化 Moonshot AI 客户端，设置 API 的基础 URL 和 API 密钥。
client = OpenAI(
    base_url="https://api.moonshot.cn/v1",
    api_key="填写你申请的api-key",
)


# search_impl 函数是一个占位符，用于模拟联网搜索工具的调用。
# 在实际使用中，这里可以调用外部搜索服务，并返回搜索结果。
def search_impl(arguments: Dict[str, Any]) -> Any:
    """
    在使用 Moonshot AI 提供的 search 工具的场合，只需要原封不动返回 arguments 即可，
    不需要额外的处理逻辑。
    但如果你想使用其他模型，并保留联网搜索的功能，那你只需要修改这里的实现（例如调用搜索
    和获取网页内容等），函数签名不变，依然是 work 的。
    """
    return arguments


# chat 函数是主要的聊天处理函数，它发送消息给 Moonshot AI 的 API 并接收回复。
def chat(messages) -> Choice:  # 返回的是Choice类型
    # 使用 Moonshot AI 的 chat.completions.create 方法发送消息并获取回复。
    completion = client.chat.completions.create(
        model="moonshot-v1-128k",  # 使用的模型名称
        messages=messages,  # 发送的消息列表
        temperature=0.3,  # 控制回复的随机性，值越高回复越多样化
        tools=[  # 定义可以使用的工具，这里使用了内置的 web_search 工具
            {
                "type": "builtin_function",
                "function": {
                    "name": "$web_search",
                },
            }
        ]
    )
    usage = completion.usage  # 获取 API 使用情况
    choice = completion.choices[0]  # 获取第一条回复

    # 打印出聊天过程中消耗的 Tokens 数量
    if choice.finish_reason == "stop":
        print(f"chat_total_tokens:           {usage.total_tokens}")

    return choice


def stream(messages: list) -> Choice:
    completion = client.chat.completions.create(
        model="moonshot-v1-128k",
        messages=messages,
        temperature=0.3,
        tools=[  # 定义可以使用的工具
            {
                "type": "builtin_function",
                "function": {
                    "name": "$web_search",
                },
            }
        ],
        stream=True  # 开启流式处理
    )
    return completion


# main 函数是程序的入口点，它处理用户的提问并调用 chat 函数。
def main(question, chat_history: list = []):
    Question = question + "如果联网搜索结果有参考意义，则你可在回答中引用链接"

    messages = [
        {"role": "system", "content": "你是 Kimi。"},  # 系统角色的消息，定义了聊天的上下文
    ]

    # 初始提问
    messages.append({
        "role": "user",
        "content": Question
    })
    flag = 0
    finish_reason = None

    chat_history.append((Question, "联网检索中..."))

    while finish_reason is None or finish_reason == "tool_calls":  # 如果 finish_reason 是 "stop"，则表示聊天已经结束，循环将终止。
        #
        print("循环")
        if flag == 1:
            print("流式")
            choice = stream(messages)
            ans = ""
            for chunk in choice:

                delta = chunk.choices[0].delta  # <-- message 字段被替换成了 delta 字段

                if delta.content:
                    # 我们在打印内容时，由于是流式输出，为了保证句子的连贯性，我们不人为地添加
                    # 换行符，因此通过设置 end="" 来取消 print 自带的换行符。
                    print(delta.content, end="")
                    ans = ans + delta.content
                chat_history[-1] = (question, ans)
                yield "", chat_history
            break

        choice = chat(messages)  # 调用 chat 函数处理消息
        #
        finish_reason = choice.finish_reason  # 获取聊天完成的原因

        if finish_reason == "tool_calls":  # 说明此次循环kimi意图识别为调用搜索工具
            flag = 1
            messages.append(choice.message)  # 将回复添加到消息列表中
            for tool_call in choice.message.tool_calls:
                tool_call_name = tool_call.function.name  # 获取工具调用的名称
                print(tool_call.function)
                tool_call_arguments = json.loads(  # 解析工具调用的参数
                    tool_call.function.arguments)
                print(tool_call_arguments)
                if tool_call_name == "$web_search":
                    # 如果工具调用是 web_search，则调用 search_impl 函数
                    search_content_total_tokens = tool_call_arguments.get("usage", {}).get("total_tokens")
                    print(f"search_content_total_tokens: {search_content_total_tokens}")
                    tool_result = search_impl(tool_call_arguments)
                else:
                    tool_result = f"Error: unable to find tool by name '{tool_call_name}'"

                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "name": tool_call_name,
                    "content": json.dumps(tool_result),  # 将工具调用的结果添加到消息列表中
                })

可以看到，在第二次循环时，执行的是如下的代码：

     if flag == 1:
            print("流式")
            choice = stream(messages)
            ans = ""
            for chunk in choice:

                delta = chunk.choices[0].delta  # <-- message 字段被替换成了 delta 字段

                if delta.content:
                    # 我们在打印内容时，由于是流式输出，为了保证句子的连贯性，我们不人为地添加
                    # 换行符，因此通过设置 end="" 来取消 print 自带的换行符。
                    print(delta.content, end="")
                    ans = ans + delta.content
                chat_history[-1] = (question, ans)
                yield "", chat_history
            break

为了实现流式输出，main函数的返回方式改为for+yield的组合形式，而非return。

同时，为了方便演示，博主将输出由工作台动态打印的同时，也加入了基于Gradio库的界面展示，代码如下：

详细原理请参考博主的Gradio相关教程。

import gradio as gr
# 检错
from kimi import main
import time

# init
print("init")
start_time = time.time()

# 加主题风格theme
google_font_primary = gr.themes.GoogleFont("Roboto")  # 例如使用 "Roboto" 作为主要字体
google_font_mono = gr.themes.GoogleFont("Roboto Mono")  # 使用 "Roboto Mono" 作为等宽字体
theme = gr.themes.Glass(

    font=[google_font_primary],
    font_mono=[google_font_mono],
    primary_hue="gray",
    secondary_hue="blue"

).set(

    background_fill_secondary="*primary_300",
    input_background_fill="linear-gradient(180deg, *primary_100 0%, white 100%)",  # 输入框

)

# 创建一个 Web 界面
with gr.Blocks(theme=theme) as demo:
    demo.title = "kimi"
    demo.description = "kimi for research"
    demo.css = """
            .title {
                text-align: center; /* 文字居中 */
                background:  #0F0000 ; /* 背景色 */
                border-radius: 10px; /* 圆角边框 */
                padding: 18px; /* 内边距 */     
                font-size: 2em; /* 字体大小 */
                text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.5); /* 文字阴影 */
                margin: 10px 0; /* 外边距 */
                box-shadow: 0 2px 4px rgba(0, 0, 0, 0.5); /* 边框阴影效果 */

            }

        """

    with gr.Row(equal_height=True):
        with gr.Column(scale=15):
            gr.Markdown("""<h1 class="title" style="color: #FFFFFF;" >kimi for research</h1>""")

    with gr.Row():
        with gr.Column(scale=4):
            # 创建一个聊天机器人对象
            chatbot = gr.Chatbot(height=450, show_copy_button=True)
            # 创建一个文本框组件，用于输入 prompt。
            query = gr.Textbox(label="Prompt/问题")
            with gr.Row():
                # 创建提交按钮。
                db_wo_his_btn = gr.Button("对话")
            with gr.Row():
                # 创建一个清除按钮，用于清除聊天机器人组件的内容。
                clear = gr.ClearButton(
                    components=[chatbot], value="清空记录")

        # 设置按钮的点击事件。当点击时，调用上面定义的 qa_chain_self_answer 函数，并传入用户的消息和聊天历史记录，然后更新文本框和聊天机器人组件。
        # 上传
        db_wo_his_btn.click(main, inputs=[query, chatbot], outputs=[query, chatbot])

    gr.Markdown("""提醒：<br>
    1. 联网搜索的时间可能比表长，请耐心等待

      """)
    gr.close_all()
    demo.queue()
    demo.launch()

运行之，效果如下图所示：

输入的搜索内容：广东省数字孪生人重点实验室

1.结果动态打印在工作台