使用 Yi-34B 和 langchain 去实现 OpenAI 的多工具调用

大家好，我是雨飞。最近在研究 Agent 相关的内容，但是现在很多的大模型都是没用 OpenAI 的工具调用（tools）功能，因此很多时候需要对代码进行改造才可以使用。今天就和大家分享下，如何使用 langchain 对没有工具调用的大模型增加工具调用的功能。

Yi-34B 介绍

我们这次采用的大模型是 Yi-34B，官方网址如下：

零一万物大模型开放平台

Yi-34B 的大模型本身没有工具调用的能力，但它的 API 的好处是可以兼容 OpenAI 的接口，节省了我们很多实现上的时间。下面是 API 的简单使用方式，直接使用 OpenAI 的库，替换掉 API_KEY 和 URL 就可以使用，因此也非常方便的就可以和 langchain 进行结合。

import openai
from openai import OpenAI

API_BASE = "https://api.lingyiwanwu.com/v1"
API_KEY = "your key"

client = OpenAI(
    # defaults to os.environ.get("OPENAI_API_KEY")
    api_key=API_KEY,
    base_url=API_BASE
)
completion = client.chat.completions.create(
    model="yi-34b-chat-0205",
    messages=[{"role": "user", "content": "Hi, who are you?"}]
)
print(completion)

实现思路

我们使用 langchain 中的 tool 注解，去声明工具，然后利用提示词工程，去赋予大模型调用工具的能力。本质上可以理解，首先让大模型根据指定的提示词去选择调用的工具名字，并提供调用这个工具所需的参数。然后，根据工具的名字和参数去调用对应的工具，拿到返回结果。最后，将返回结果再输入给大模型，生成最终的答案。

注意：如果一个大模型，第一步无法正确解析出需要调用的工具和参数，那么这个方法是没有办法使用的。

核心的提示词如下：

system_prompt = f"""You are an assistant that has access to the following set of tools. Here are the names and descriptions for each tool:

{rendered_tools}

Given the user input, return the name and input of the tool to use. 
最后的输出以 JSON 格式提供，并包含 'name' 和 'arguments' 两个键."""

sample = """
参考示例
add: add(first_int: int, second_int: int) -> int - Add two integers.
输出示例
{{"name"':"add","arguments":{{"first_int":...,"second_int":...}}}}
"""

下面讲解一些这段提示词的书写技巧，其中 {rendered_tools}就是封装了各个工具的字符串变量，为了方便我们后续的调用，需要将输出的格式转换为 JSON 的格式。

由于Yi-34B 大模型的能力有限，我们采取了「少样本提示」的书写技巧，增加了一个参考示例。针对 add 这个方法，定义它的输出格式，这样会便于大模型理解和输出 JSON 格式的数据。

工具编写

在langchain 中可以很方便的使用 tool 这个注解去声明一个工具，然后进行调用，核心的代码如下：

from langchain_core.tools import tool
@tool
def multiply(first_int: int, second_int: int) -> int:
    """Multiply two integers together."""
    return first_int * second_int

可以直接使用 invoke 方法去测试工具的使用是否正确。

result =multiply.invoke({"first_int": 4, "second_int": 5})
print(result)

完整代码

下面是完整的代码，请将 api_key 换成自己的。

import json,re
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from operator import itemgetter

from langchain.tools.render import render_text_description
from langchain_core.tools import tool
from langchain_core.output_parsers import JsonOutputParser
from  langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI
url_base = "https://api.lingyiwanwu.com/v1"
api_key=""
yi_llm = ChatOpenAI(openai_api_base=url_base, openai_api_key=api_key, model="yi-34b-chat")


@tool
def add(first_int: int, second_int: int) -> int:
    "Add two integers."
    return first_int + second_int

@tool
def exponentiate(base: int, exponent: int) -> int:
    "Exponentiate the base to the exponent power."
    return base**exponent

@tool
def multiply(first_int: int, second_int: int) -> int:
    """Multiply two integers together."""
    return first_int * second_int

def extract_format_json(message: AIMessage):
    """Extracts JSON content from a string where JSON is embedded between ```json and ``` tags.
    """

    text = message.content
    try:
        return json.loads(text)
    except Exception:
        raise ValueError(f"Failed to parse: {message}")

tools = [add, multiply, exponentiate]
def tool_chain(model_output):
    tool_map = {tool.name: tool for tool in tools}
    chosen_tool = tool_map[model_output["name"]]
    return itemgetter("arguments") | chosen_tool

rendered_tools = render_text_description(tools)

system_prompt = f"""You are an assistant that has access to the following set of tools. Here are the names and descriptions for each tool:

{rendered_tools}

Given the user input, return the name and input of the tool to use. 
最后的输出以 JSON 格式提供，并包含 'name' 和 'arguments' 两个键."""

sample = """
参考示例
add: add(first_int: int, second_int: int) -> int - Add two integers.
输出示例
{{"name"':"add","arguments":{{"first_int":...,"second_int":...}}}}
"""
system_prompt  = system_prompt
prompt = ChatPromptTemplate.from_messages(
    [("system", system_prompt),("system",sample), ("user", "{input}")]
)

print(prompt)
chain = prompt | yi_llm |extract_format_json|RunnablePassthrough.assign(output=tool_chain)

result = chain.invoke({"input": "What is 231 * 25?"})
#
print(result)

雨飞同行

• 雨飞
• 主业是推荐算法
• 希望通过自媒体开启自己不上班只工作的美好愿景
• 微信：1060687688
• 欢迎和我交朋友🫰

好了，我写完了，有启发的欢迎点赞🫰。新的一天，愿阳光洒在你的脸上。