Zenn
Closed4

cohereのMulti-step Tool Useを試す

LLM
functioncalling
cohere
kun432kun432

以前にcohereのTool Use（いわゆるFunction Calling）を試した際に、マルチステップでのツールのプランニング・実行ができなかった。

https://zenn.dev/link/comments/fa475361e8dd70

  • 前提
    • ツールAとBがある
    • ツールBはツールAの結果を引数とする（マルチステップで推論する必要がある）
  • 流れ
    • モデルにツールA/B＋クエリを渡す
    • モデルはツールAを推奨、そのパラメータを返す
    • ツールAを実行して結果をモデルに返すのだけど・・・・
      • tool_resultsで返すと、次のツール推論は行われず最終回答が返る
      • SDKやAPIを見る限り、ツールの実行結果を会話履歴に含めるのは非推奨な模様
    • マルチステップにならずに、ハルシネーションされた回答で終わる

以前のcohereのドキュメントで紹介されていたのはLangChainのReActエージェントを使う方法で、モデル・APIの外側でツールの取捨・プランニングを行うというもの。つまり、モデルの機能として実現していたわけではなかった。

issueをあげていたのだけどどうやらできるようになったらしい

https://github.com/cohere-ai/notebooks/issues/152

ドキュメントも更新されていた。

https://docs.cohere.com/docs/multi-step-tool-use

Notice that the model learned information from the first search, which it then used to perform a second web search. This behavior is called "multi-step" because the model tackles the task step by step.

Multi-step tool use with Cohere can also be implemented using the Langchain framework which conveniently comes with many pre-defined tools

ということで、実際に確認してみる。

kun432kun432

以下は前回と同じ

!pip install cohere

!pip freeze | grep cohere

cohere==5.5.7

import cohere
from google.colab import userdata
import os

os.environ["CO_API_KEY"] = userdata.get('CO_API_KEY')
co = cohere.Client()

でツールも同じ。前回同様、「タブレットの値段」を知りたければ、最初に商品名から商品IDを取得、その取得した商品IDで商品情報を取得するというマルチステップを踏まないといけないという流れ。

product_list = {
    'スマートフォン': 'E1001',
    'ノートパソコン': 'E1002',
    'タブレット': 'E1003',
    'Tシャツ': 'C1001',
    'ジーンズ':'C1002',
    'ジャケット': 'C1003',
}

product_catalog = {
    'E1001': {'price': 500, 'stock_level': 20},
    'E1002': { 'price': 1000, 'stock_level': 15},
    'E1003': {'price': 300, 'stock_level': 25},
    'C1001': {'price': 20, 'stock_level': 100},
    'C1002': {'price': 50, 'stock_level': 80},
    'C1003': {'price': 100, 'stock_level': 40},
}

def get_product_id_from_product_name(product_name: str) -> dict:
    return {"product_name": product_name, "product_id": product_list[product_name]}

def get_product_info_from_product_id(product_id: str) -> dict:
    return {"product_id": product_id, "product_info": product_catalog[product_id]}

tools = [
    {
        "name": "get_product_id_from_product_name",
        "description": "「商品名」（product_name）から「商品ID」(product_id)を取得する。",
        "parameter_definitions": {
            "product_name": {
                "description": "「商品ID」(product_id)を取得するための「商品名」を指定する。「商品名」は一般名詞で指定する必要がある。ex. タブレット、ジャケット、等",
                "type": "str",
                "required": True
            }
        }
    },
    {
        "name": "get_product_info_from_product_id",
        "description": "「商品ID」(product_id)から「商品情報（価格、在庫）」（product_info）を取得する。",
        "parameter_definitions": {
            "product_id": {
                "description": "「商品情報（価格、在庫）」（product_info）を取得するための「商品ID」を指定する。「商品ID」は [A-Z]{1}[0-9]{3} で指定すること。ex. X0002、等。",
                "type": "str",
                "required": True
            },
        }
    }
]

# ツール名と関数名のマッピング。実行時に使用。
tool_to_function_map = {
    "get_product_id_from_product_name": get_product_id_from_product_name,
    "get_product_info_from_product_id": get_product_info_from_product_id,
}

ではクエリを投げてみる。説明のため、まずはイテレーションさせずに個別にリクエストを投げて確認する。

preamble = "あなたは親切な日本語のアシスタントです。"
message = "タブレットの在庫を調べて。"
model = "command-r-plus"

first_res = co.chat(
    model=model,
    message=message,
    preamble=preamble,
    force_single_step=False,
    tools=tools
)

レスポンス（first_res）の中身。

NonStreamedChatResponse(text='商品名から商品IDを検索し、商品情報を取得する。',
generation_id='be80e17c-0cdf-4cb9-8147-ee21723789ae',
citations=None,
documents=None,
is_search_required=None,
search_queries=None,
search_results=None,
finish_reason='COMPLETE',
tool_calls=[
    ToolCall(name='get_product_id_from_product_name',
    parameters={
        'product_name': 'タブレット'
    }
    )
],
chat_history=[
    Message_User(message='タブレットの在庫を調べて。',
    tool_calls=None,
    role='USER'),
    Message_Chatbot(message='商品名から商品IDを検索し、商品情報を取得する。',
    tool_calls=[
        ToolCall(name='get_product_id_from_product_name',
        parameters={
            'product_name': 'タブレット'
        }
        )
    ],
    role='CHATBOT')
],
prompt=None,
meta=ApiMeta(api_version=ApiMetaApiVersion(version='1',
is_deprecated=None,
is_experimental=None),
billed_units=ApiMetaBilledUnits(input_tokens=183,
output_tokens=34,
search_units=None,
classifications=None),
tokens=ApiMetaTokens(input_tokens=958,
output_tokens=34),
warnings=None))

きちんと1つ目のツールだけが提案されている。以前はここでむりやり複数のツールを提案して、当然マルチステップが必要なので2番目のツールの提示でハルシネーションが起きていたりした。

ではこれを元にツールを実行して回答を得る。

tool_results = []
for call in first_res.tool_calls:
    tool_name = call.name
    tool_params = call.parameters
    function_to_call = tool_to_function_map[tool_name]
    # ツール実行の出力は**オブジェクトの配列**にして、callとoutputsからなる結果オブジェクトを作る
    tool_result = {"call": call, "outputs": [function_to_call(**tool_params)]}
    # ツールの結果は配列     
    tool_results.append(tool_result)

second_res = co.chat(
    model="command-r-plus",
    chat_history=first_res.chat_history,
    message="",
    force_single_step=False,
    tools=tools,
    tool_results=tool_results
)

レスポンス内のtool_callsにツール実行に関する情報が含まれているので、これを元にツールの実行結果を生成する。複数のツールがある場合もあるので結果を配列で渡すのはそらそう、なんだけども、各ツールの実行結果の出力もオブジェクトの配列にしないといけないってのがちょっとハマった。

あと、以前は仮に1回目のツール提案が1つだけで正しかったとしても、ここでtool_resultsを返してしまうとそれ以降のツール提案が行われなくなってしまっていたのだけど、レスポンスに含まれているchat_historyを引き継いでいくことで、ツールの実行結果が引き継がれるようになったように見える。

レスポンス（second_res）の中身。

NonStreamedChatResponse(text='タブレットの商品IDはE1003。次に、商品IDから商品情報を取得する。',
generation_id='0ea7046b-8c2f-48de-a74a-48e4c35c0b76',
citations=None,
documents=None,
is_search_required=None,
search_queries=None,
search_results=None,
finish_reason='COMPLETE',
tool_calls=[
    ToolCall(name='get_product_info_from_product_id',
    parameters={
        'product_id': 'E1003'
    }
    )
],
chat_history=[
    Message_User(message='タブレットの在庫を調べて。',
    tool_calls=None,
    role='USER'),
    Message_Chatbot(message='商品名から商品IDを検索し、商品情報を取得する。',
    tool_calls=[
        ToolCall(name='get_product_id_from_product_name',
        parameters={
            'product_name': 'タブレット'
        }
        )
    ],
    role='CHATBOT'),
    Message_Tool(tool_results=[
        ToolResult(call=ToolCall(name='get_product_id_from_product_name',
        parameters={
            'product_name': 'タブレット'
        }
        ),
        outputs=[
            {
                'product_id': 'E1003',
                'product_name': 'タブレット'
            }
        ]
        )
    ],
    role='TOOL'),
    Message_Chatbot(message='タブレットの商品IDはE1003。次に、商品IDから商品情報を取得する。',
    tool_calls=[
        ToolCall(name='get_product_info_from_product_id',
        parameters={
            'product_id': 'E1003'
        }
        )
    ],
    role='CHATBOT')
],
prompt=None,
meta=ApiMeta(api_version=ApiMetaApiVersion(version='1',
is_deprecated=None,
is_experimental=None),
billed_units=ApiMetaBilledUnits(input_tokens=220,
output_tokens=45,
search_units=None,
classifications=None),
tokens=ApiMetaTokens(input_tokens=1222,
output_tokens=45),
warnings=None))

1つ目のツールの結果を踏まえて、2つ目のツールの提案が行われている。

同様にして、2つ目のツールを実行して結果をモデルに渡す。

for call in second_res.tool_calls:
    tool_name = call.name
    tool_params = call.parameters
    function_to_call = tool_to_function_map[tool_name]
    tool_result = {"call": call, "outputs": [function_to_call(**tool_params)]}
    # tool_resultsは前回までのものを引き継いで追加
    tool_results.append(tool_result)

final_res = co.chat(
    model="command-r-plus",
    chat_history=second_res.chat_history,
    message="",
    force_single_step=False,
    tools=tools,
    tool_results=tool_results
)

最後のレスポンス

NonStreamedChatResponse(text='タブレットは在庫が25個あり、価格は300円です。',
generation_id='aad365c0-69de-4bdb-b0e5-edf33e9c0e58',
citations=[
    ChatCitation(start=9,
    end=12,
    text='25個',
    document_ids=[
        'get_product_info_from_product_id:1:4:0'
    ]
    ),
    ChatCitation(start=18,
    end=22,
    text='300円',
    document_ids=[
        'get_product_info_from_product_id:1:4:0'
    ]
    )
],
documents=[
    {
        'id': 'get_product_info_from_product_id:1:4:0',
        'product_id': 'E1003',
        'product_info': '{"price":300,"stock_level":25}',
        'tool_name': 'get_product_info_from_product_id'
    }
],
is_search_required=None,
search_queries=None,
search_results=None,
finish_reason='COMPLETE',
tool_calls=None,
chat_history=[
    Message_User(message='タブレットの在庫を調べて。',
    tool_calls=None,
    role='USER'),
    Message_Chatbot(message='商品名から商品IDを検索し、商品情報を取得する。',
    tool_calls=[
        ToolCall(name='get_product_id_from_product_name',
        parameters={
            'product_name': 'タブレット'
        }
        )
    ],
    role='CHATBOT'),
    Message_Tool(tool_results=[
        ToolResult(call=ToolCall(name='get_product_id_from_product_name',
        parameters={
            'product_name': 'タブレット'
        }
        ),
        outputs=[
            {
                'product_id': 'E1003',
                'product_name': 'タブレット'
            }
        ]
        )
    ],
    role='TOOL'),
    Message_Chatbot(message='タブレットの商品IDはE1003。次に、商品IDから商品情報を取得する。',
    tool_calls=[
        ToolCall(name='get_product_info_from_product_id',
        parameters={
            'product_id': 'E1003'
        }
        )
    ],
    role='CHATBOT'),
    Message_Tool(tool_results=[
        ToolResult(call=ToolCall(name='get_product_id_from_product_name',
        parameters={
            'product_name': 'タブレット'
        }
        ),
        outputs=[
            {
                'product_id': 'E1003',
                'product_name': 'タブレット'
            }
        ]
        ),
        ToolResult(call=ToolCall(name='get_product_info_from_product_id',
        parameters={
            'product_id': 'E1003'
        }
        ),
        outputs=[
            {
                'product_id': 'E1003',
                'product_info': {
                    'price': 300,
                    'stock_level': 25
                }
            }
        ]
        )
    ],
    role='TOOL'),
    Message_Chatbot(message='タブレットは在庫が25個あり、価格は300円です。',
    tool_calls=None,
    role='CHATBOT')
],
prompt=None,
meta=ApiMeta(api_version=ApiMetaApiVersion(version='1',
is_deprecated=None,
is_experimental=None),
billed_units=ApiMetaBilledUnits(input_tokens=321,
output_tokens=20,
search_units=None,
classifications=None),
tokens=ApiMetaTokens(input_tokens=1445,
output_tokens=20),
warnings=None))

マルチステップで最終回答を生成できているのがわかる。

上記をまとめると以下となる。

preamble = "あなたは親切な日本語のアシスタントです。"
message = "タブレットの在庫を調べて。"
model = "command-r-plus"

res = co.chat(
    model=model,
    message=message,
    preamble=preamble,
    force_single_step=False,
    tools=tools
)

while res.tool_calls:
    print(res.text)
    tool_results = []
    for call in res.tool_calls:
        tool_name = call.name
        tool_params = call.parameters
        function_to_call = tool_to_function_map[tool_name]
        tool_result = {"call": call, "outputs": [function_to_call(**tool_params)]}
        print(tool_result)
        tool_results.append(tool_result)
    
    res = co.chat(
        model="command-r-plus",
        chat_history=res.chat_history,
        message="",
        force_single_step=False,
        tools=tools,
        tool_results=tool_results
    )

print(res.text)

商品名から商品IDを検索し、在庫を調べる。
{'call': ToolCall(name='get_product_id_from_product_name', parameters={'product_name': 'タブレット'}), 'outputs': [{'product_name': 'タブレット', 'product_id': 'E1003'}]}
タブレットの商品IDはE1003。
{'call': ToolCall(name='get_product_info_from_product_id', parameters={'product_id': 'E1003'}), 'outputs': [{'product_id': 'E1003', 'product_info': {'price': 300, 'stock_level': 25}}]}
タブレット（商品ID: E1003）の在庫は25です。

kun432kun432

まとめ

やっと本来のMulti-step Tool Useができるようになった。これでOpenAIやAnthropicと同じような感覚で使える。

kun432kun432

つまり、LiteLLMでプロキシさせる場合に互換性が出来たともいえるのではないか。LiteLLMの対応が待たれる。

このスクラップは5ヶ月前にクローズされました