Closed13

LlamaIndexのエージェントを改めて試す

kun432kun432

LlamaIndexといえばRAG、みたいな印象があるが、最近は"Agentic RAG"というキーワードをよく使っているように思える。

https://www.llamaindex.ai/blog/agentic-rag-with-llamaindex-2721b8a49ff6

以前にLlamaIndexのモジュールガイドを全部ためしたときには"Data Agent"とかという言い方であくまでもRAGがメインだったように思えるし、自分もRAG向けの機能を中心に確認していたのであまり細かくはみていなかった。

https://zenn.dev/link/comments/6480954595a810

いろいろエージェント関連のモジュールや実装も増えているように思えるので、改めてLlamaIndexのエージェント機能について確認していく。

kun432kun432

こういうのが出てきた

https://www.llamaindex.ai/blog/introducing-llama-agents-a-powerful-framework-for-building-production-multi-agent-ai-systems

ざっと見た感じ、

  • 複数のエージェントをLlamaIndexで作る
  • llama-agentsは、これらのエージェントのオーケストレーション、メッセージキューを提供するコントロールプレーンとして動作する
  • それぞれのエージェントはマイクロサービスとして独立して動作する

という感じに見える。

なので、LlamaIndexのエージェントを使うことになるので、やっぱり一通りやっておこう。

kun432kun432

ということで、以下に従って進めていく。

https://docs.llamaindex.ai/en/stable/module_guides/deploying/agents/

コンセプト

データエージェントは、LlamaIndexにおいてはLLMで強化されたナレッジワーカーで、データに対して様々なタスクを "読み込み "と "書き込み "の両方の機能において、インテリジェントに実行することができる。以下のようなことが可能:

  • 非構造化データ、半構造化データ、構造化データなど、さまざまなタイプのデータに対して、自動検索を実行。
  • 構造化された方法で外部サービスAPIを呼び出し、レスポンスを処理し、あとの処理のために保存する。

これらの意味でエージェントは、静的なデータソースから "読み取る "だけでなく、様々な異なるツールから動的にデータを取り込み、修正することができるという点で、我々のクエリエンジンを一歩超えている。

データエージェントの構築には、以下のコアコンポーネントが必要である:

  • 推論ループ
  • ツールの抽象化

データエージェントは、相互にやりとりするためのAPI(ツール)のセットで初期化される。入力タスクが与えられると、データエージェントは推論ループを使用して、どのツールをどの順序で使用するか、各ツールを呼び出すパラメータを決定する。

推論ループ

推論ループはエージェントのタイプに依存する。以下のエージェントをサポートしている:

ツール抽象化

ツールの抽象化については、ツールセクションを参照。

ブログ記事

詳細については、詳細なブログ記事をご覧いただきたい。

低レベルAPI: ステップワイズ実行

デフォルトでは、エージェントは、エンドツーエンドでユーザクエリを実行する、querychatを公開している。

また、エージェントのステップ実行を可能にするより低レベルなAPIも提供している。これにより、タスクを作成し、タスク内の各ステップの入出力を分析し、実行することができる。

ガイドをチェックしてほしい。

使用パターン

データエージェントは、以下の方法で使用できる(例ではOpenAI Function APIを使用している)。

from llama_index.agent.openai import OpenAIAgent
from llama_index.llms.openai import OpenAI

# ツールのインポートして定義する
...

# llm を初期化する
llm = OpenAI(model="gpt-3.5-turbo-0613")

# openaiエージェントを初期化する
agent = OpenAIAgent.from_tools(tools, llm=llm, verbose=True)

詳しくは使用パターンガイドを参照のこと。

モジュール

モジュールガイドで、さまざまなエージェントの種類と使用例について詳しく学ぶことができる。

また、エージェントのランナーとワーカーのための低レベルAPI ガイドもある。

ツールのセクションも見てみよう!

kun432kun432

Usage Pattern

ということでUsage Patternを実際に試していく。Colaboratoryで。

https://docs.llamaindex.ai/en/stable/module_guides/deploying/agents/usage_pattern/

Getting Started

事前準備

パッケージインストール。トレーシングのためにArize Phoenixのインテグレーションもあわせてインストールしておく。

!pip install llama-index llama-index-callbacks-arize-phoenix
!pip freeze | egrep "llama-|arize"
arize-phoenix==4.5.0
llama-cloud==0.0.6
llama-index==0.10.50
llama-index-agent-openai==0.2.7
llama-index-callbacks-arize-phoenix==0.1.5
llama-index-cli==0.1.12
llama-index-core==0.10.50
llama-index-embeddings-openai==0.1.10
llama-index-indices-managed-llama-cloud==0.2.1
llama-index-legacy==0.9.48
llama-index-llms-openai==0.1.23
llama-index-multi-modal-llms-openai==0.1.6
llama-index-program-openai==0.1.6
llama-index-question-gen-openai==0.1.3
llama-index-readers-file==0.1.25
llama-index-readers-llama-parse==0.1.4
llama-parse==0.4.4
openinference-instrumentation-llama-index==2.0.0

OpenAI APIキーの読み込み

from google.colab import userdata
import os

os.environ["OPENAI_API_KEY"] = userdata.get('OPENAI_API_KEY')

Arize Phoenixのトレーシングを有効化。

import phoenix as px
import llama_index.core

px.launch_app()
llama_index.core.set_global_handler("arize_phoenix")

表示されたURLを別タブで開いておく。

🌍 To view the Phoenix app in your browser, visit https://XXXXXXXXXXXXXXXX-colab.googleusercontent.com/
📖 For more information on how to use Phoenix, check out https://docs.arize.com/phoenix

notebookなのでイベントループのネストを有効化。

import nest_asyncio

nest_asyncio.apply()

ではエージェントをやっていく。エージェントは複数のツールから初期化する。以下はReActエージェントの初期化例。

from llama_index.core.tools import FunctionTool
from llama_index.llms.openai import OpenAI
from llama_index.core.agent import ReActAgent


# ツールを定義
def multiply(a: int, b: int) -> int:
    """2つの整数を乗算し、結果を整数で返す。"""
    return a * b


multiply_tool = FunctionTool.from_defaults(fn=multiply)

# llmを初期化
llm = OpenAI(model="gpt-3.5-turbo-0613")

# ReActエージェントを初期化
agent = ReActAgent.from_tools([multiply_tool], llm=llm, verbose=True)

エージェントは、ChatEngineQueryEngineから継承された、chatqueryの2つのエンドポイントを持っている。

agent.chat("2123 かける 215123 は?")
AgentChatResponse(response='2123 かける 215123 は 456706129 です。', sources=[ToolOutput(content='456706129', tool_name='multiply', raw_input={'args': (), 'kwargs': {'a': 2123, 'b': 215123}}, raw_output=456706129, is_error=False)], source_nodes=[], is_dummy_stream=False, metadata=None)

ReActエージェントの途中の思考過程も表示される。

queryだとこんな感じ。

agent.query("2123 かける 215123 は?")
Response(response='2123 かける 215123 は 456706129 です。', source_nodes=[], metadata=None)

LLMによって最も最適なエージェントモジュールを自動的に選択させることもできる。AgentRunnerエージェントをオーケストレーションする最上位のクラスらしい。このモジュールのfrom_llmメソッドでエージェントを生成する。

from llama_index.core.agent import AgentRunner

agent = AgentRunner.from_llm([multiply_tool], llm=llm, verbose=True)
print(type(agent))

OpenAIAgentが選択されている。

<class 'llama_index.agent.openai.base.OpenAIAgent'>

このエージェントを使って同じようにクエリを投げてみる。

agent.chat("2123 かける 215123 は?")
Added user message to memory: 2123 かける 215123 は?
=== Calling Function ===
Calling function: multiply with args: {
  "a": 2123,
  "b": 215123
}
Got output: 456706129
========================

AgentChatResponse(response='2123 かける 215123 は 456,706,129 です。', sources=[ToolOutput(content='456706129', tool_name='multiply', raw_input={'args': (), 'kwargs': {'a': 2123, 'b': 215123}}, raw_output=456706129, is_error=False)], source_nodes=[], is_dummy_stream=False, metadata=None)
agent.query("2123 かける 215123 は?")
Added user message to memory: 2123 かける 215123 は?
=== Calling Function ===
Calling function: multiply with args: {
  "a": 2123,
  "b": 215123
}
Got output: 456706129
========================

Response(response='2123 かける 215123 は 456,706,129 です。', source_nodes=[], metadata=None)

どちらもFunction Callingでツールを実行、その結果を元に回答を生成している。

ツールを定義する

最初の例のように関数をツールとして定義することもできるが、LlamaIndexで作成したQuery Engineをツールとしてラップしてエージェントで使うことができる。

以下を元にQuery Engineを作成する。

https://ja.wikipedia.org/wiki/キタサンブラック

https://ja.wikipedia.org/wiki/ドゥラメンテ

from pathlib import Path
import requests
import re

def replace_heading(match):
    level = len(match.group(1))
    return '#' * level + ' ' + match.group(2).strip()

# Wikipediaからのデータ読み込み
wiki_titles = ["キタサンブラック", "ドゥラメンテ"]
for title in wiki_titles:
    response = requests.get(
        "https://ja.wikipedia.org/w/api.php",
        params={
            "action": "query",
            "format": "json",
            "titles": title,
            "prop": "extracts",
            # 'exintro': True,
            "explaintext": True,
        },
    ).json()
    page = next(iter(response["query"]["pages"].values()))
    wiki_text = f"# {title}\n\n## 概要\n\n"
    wiki_text += page["extract"]

    wiki_text = re.sub(r"(=+)([^=]+)\1", replace_heading, wiki_text)
    wiki_text = re.sub(r"\t+", "", wiki_text)
    wiki_text = re.sub(r"\n{3,}", "\n\n", wiki_text)
    data_path = Path("data")
    if not data_path.exists():
        Path.mkdir(data_path)

    with open(data_path / f"{title}.txt", "w") as fp:
        fp.write(wiki_text)
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.embeddings.openai import OpenAIEmbedding

kitasan_docs = SimpleDirectoryReader(input_files=["data/キタサンブラック.txt"]).load_data()
duramente_docs = SimpleDirectoryReader(input_files=["data/ドゥラメンテ.txt"]).load_data()

embed_model = OpenAIEmbedding(model="text-embedding-3-small")

kitasan_index = VectorStoreIndex.from_documents(
    kitasan_docs,
    embed_model=embed_model,
    show_progress=True
)
duramente_index = VectorStoreIndex.from_documents(
    duramente_docs,
    embed_model=embed_model,
    show_progress=True
)

kitasan_query_engine = kitasan_index.as_query_engine(similarity_top_k=5)
duramente_query_engine = duramente_index.as_query_engine(similarity_top_k=5)

以下のようにQueryEngineToolを使って、Query Engineをツールにして、ReActエージェントを作成する。

from llama_index.core.agent import ReActAgent
from llama_index.core.tools import QueryEngineTool, ToolMetadata

query_engine_tools = [
    QueryEngineTool(
        query_engine=kitasan_query_engine,
        metadata=ToolMetadata(
            name="kitasan_black_tool",
            description="日本の競走馬「キタサンブラック」に関する情報を提供する。 "
            "プレーンなテキストで詳細に書かれた質問をツールの入力として使用する。",
        ),
    ),
    QueryEngineTool(
        query_engine=duramente_query_engine,
        metadata=ToolMetadata(
            name="duramente_tool",
            description="日本の競走馬「ドゥラメンテ」に関する情報を提供する。 "
            "プレーンなテキストで詳細に書かれた質問をツールの入力として使用する。",
        ),
    ),
]

agent = ReActAgent.from_tools(query_engine_tools, llm=llm, verbose=True)

クエリを投げてみる。

response = agent.chat("キタサンブラックの代表産駒について教えて。")

print("回答:", response.response)
print("根拠:")
print("===")
print("\n---\n".join([n.text.replace("\n","")[:100] + "..." for n in response.source_nodes]))
print("===")
回答: キタサンブラックの代表産駒の一つはイクイノックスです。
根拠:
===
これまで500万円が最高だった種付け料だったが6年目に跳ね上がり、大台の1000万円に到達した。値上げとなったがむしろさらに繁殖牝馬を集め、6年目は242頭に増加している。これを受けて、2024年度の...
---
武は落鉄と敗戦の因果は不明とし、続けて「力負けとは思っていません」と回顧した。### 有馬記念引退レースとなる有馬記念に出走し勝利。2017年当時の史上最多タイとなるJRAGI7勝目(後述)を挙げ、テ...
---
しかし翌2016年、天皇賞(春)とジャパンカップを優勝しGI2勝を記録し、年度代表馬と最優秀4歳以上牡馬を受賞している。この年は、モーリスとサトノダイヤモンドもGI級競走を複数勝利しており、モーリスは...
---
2017年、春の古馬GI3連勝が懸かった宝塚記念こそ9着大敗するも、それ以外ではすべて馬券圏内で応えていた。ただ勝利で応えたのは引退レース、挑戦3回目の有馬記念だけだった。前々年の有馬記念は3着、前年...
---
## 競走成績以下の内容は、netkeiba並びにJBISサーチの情報に基づく。タイム欄のRはレコード勝ちを示す。## 種牡馬成績### 年度別成績### 重賞優勝産駒一覧#### GI級競走優勝産駒...
===

response = agent.chat("ドゥラメンテの代表産駒について教えて。")

print("回答:", response.response)
print("根拠:")
print("===")
print("\n---\n".join([n.text.replace("\n","")[:100] + "..." for n in response.source_nodes]))
print("===")
回答: ドゥラメンテの代表産駒には、タイトルホルダー、スターズオンアース、リバティアイランドなどがあります。
根拠:
===
2022年11月3日、ヴァレーデラルナがJBCレディスクラシックを制し、産駒のダートGI級競走初勝利を果たした。2022年12月11日、リバティアイランドが阪神ジュベナイルフィリーズを制し、産駒の2歳...
---
長らくディープインパクトが守ってきたJRAリーディングサイアーのタイトルを獲得した。ディープインパクト以外の種牡馬がJRAリーディングサイアーを獲得したのは、自身の父であるキングカメハメハが2011年...
---
母が現役時代にGI3勝を挙げたスイープトウショウだったため、両親のGI勝利数を合計して「五冠ベイビー」と称された。同馬は後にクリーンスイープと命名され、シルクレーシングの所有馬、美浦の国枝栄厩舎の所属...
---
### 4歳(2016年)2月28日に行われた中山記念で9か月ぶりに復帰。M.デムーロ騎乗で単勝2.1倍1番人気に推され勝利。今後のレースは3月28日にメイダン競馬場で行われるドバイ国際競走出走を予定...
---
# ドゥラメンテ## 概要ドゥラメンテ(欧字名:Duramente、2012年3月22日 - 2021年8月31日) は、日本の競走馬・種牡馬。馬名の意味は、イタリア語の "duramente" (荒...
===

インデックスを雑に作っているのでクエリによってはハルシネーションするのだけども、クエリに応じてそれぞれのツールが呼び出され、Query Engineでインデックスから情報を取得しているのがわかる。

また、エージェントはBaseQueryEngineを継承しているので、QueryEngineToolを通じてエージェントそのものをツールとして使うこともできる。(ここはいい感じにサンプルを思いつかなかったので、ドキュメントのママ)

from llama_index.core.tools import QueryEngineTool

query_engine_tools = [
    QueryEngineTool(
        query_engine=sql_agent,
        metadata=ToolMetadata(
            name="sql_agent", description="Agent that can execute SQL queries."
        ),
    ),
    QueryEngineTool(
        query_engine=gmail_agent,
        metadata=ToolMetadata(
            name="gmail_agent",
            description="Tool that can send emails on Gmail.",
        ),
    ),
]

outer_agent = ReActAgent.from_tools(query_engine_tools, llm=llm, verbose=True)
kun432kun432

エージェントとプランニング

大きなタスクを小さなサブタスクに分解させてるというのはエージェントにおける1つのパターンで、LlamaIndexにはこのためのエージェントプランニングモジュールがある。

from llama_index.agent.openai import OpenAIAgentWorker
from llama_index.core.agent import (
    StructuredPlannerAgent,
    FunctionCallingAgentWorker,
)

worker = FunctionCallingAgentWorker.from_tools(tools, llm=llm)
agent = StructuredPlannerAgent(worker)

StructuredPlannerAgentというのがおそらくそれで、FunctionCallingAgentWorkerが実際のサブタスクを実行するものなのだろうと思う。

StructuredPlannerAgentAgentRunnerクラスを継承したものだと思うが、他のAgentRunnerの派生に比べると、プランニングを行う分だけレスポンスに時間が掛かる、そして、プランニングは非常に難しいため、性能が高いLLM(gpt-3.5-turboよりもgpt-4-turbo)が必要になる、というトレードオフがある。ここは納得感がある。

StructuredPlannerAgentについては以下にnotebookがあるので、別途試してみたいと思う。
https://docs.llamaindex.ai/en/stable/examples/agent/structured_planner/

低レベルAPI

ここまでに出てきたLlamaIndexのエージェントモジュールには、OpenAIAgentReActAgentなどがあったが、これらはAgentWorkerとやりとりするAgentRunnerに基づいたシンプルなラッパーとなっている。

「すべて」のエージェントはこれに倣って定義ができる。OpenAIAgentの場合

from llama_index.core.agent import AgentRunner
from llama_index.agent.openai import OpenAIAgentWorker

# OpenAIAgentをツールから初期化
openai_step_engine = OpenAIAgentWorker.from_tools(tools, llm=llm, verbose=True)
agent = AgentRunner(openai_step_engine)

イメージ的には

  • AgentWorker: 個々のサブタスクを実行・処理する。
  • AgentRunner: タスク全体の管理を行い、必要に応じてworkerを使ってタスクを実行する。

ということかな。なので、***AgentというのはWorkerとRunnerのセットを包括したモジュールなのだと思う。カスタムなエージェントを作る場合にもこれはやりやすいということらしい。

低レベルAPIの詳細については以下。これは次の章でやる。
https://docs.llamaindex.ai/en/stable/module_guides/deploying/agents/agent_runner/

エージェントのカスタマイズ

エージェントをカスタムで作成する場合、ステートフルな関数をFnAgentWorkerでラップするのが一番簡単なやり方。

関数に入力・関数から出力されるステート変数には、ツールでも任意の変数でも、好きなものを入れることができる。また、タスクや出力オブジェクトも含まれる。

from llama_index.core.agent import FnAgentWorker
from typing import Tuple, Dict, Any

## これは、入力された数値を毎回2倍する些細な関数の例
## これをエージェントに渡す
def multiply_agent_fn(state: dict) -> Tuple[Dict[str, Any], bool]:
    """Mock agent input function."""
    if "max_count" not in state:
        raise ValueError("max_count must be specified.")

    # __output__ はエージェントの最終出力を示す特別なキー
    # __task__ は、エージェントが関数に渡すタスクオブジェクトを表す特別なキー
    # `task.input` は渡された入力文字列
    if "__output__" not in state:
        state["__output__"] = int(state["__task__"].input)
        state["count"] = 0
    else:
        state["__output__"] = state["__output__"] * 2
        state["count"] += 1

    is_done = state["count"] >= state["max_count"]

    # デバッグ出力
    print("DEBUG(state):", state)
    print("DEBUG(is_done):", is_done)

    # この関数の出力は、state変数とis_doneのタプルでなければならない。
    return state, is_done


agent = FnAgentWorker(
    fn=multiply_agent_fn, initial_state={"max_count": 5}
).as_agent()

agent.query("5")
WARNI [llama_index.core.callbacks.utils] Could not find attribute callback_manager on <class 'llama_index.core.agent.custom.simple_function.FnAgentWorker'>.
WARNI [llama_index.core.callbacks.utils] Could not find attribute callback_manager on <class 'llama_index.core.agent.custom.simple_function.FnAgentWorker'>.
WARNI [llama_index.core.callbacks.utils] Could not find attribute callback_manager on <class 'llama_index.core.agent.custom.simple_function.FnAgentWorker'>.
WARNI [llama_index.core.callbacks.utils] Could not find attribute callback_manager on <class 'llama_index.core.agent.custom.simple_function.FnAgentWorker'>.
WARNI [llama_index.core.callbacks.utils] Could not find attribute callback_manager on <class 'llama_index.core.agent.custom.simple_function.FnAgentWorker'>.
WARNI [llama_index.core.callbacks.utils] Could not find attribute callback_manager on <class 'llama_index.core.agent.custom.simple_function.FnAgentWorker'>.
DEBUG(state): {'max_count': 5, '__task__': Task(task_id='048e10e5-5a32-4df2-910d-56fb390aa127', input='5', memory=ChatMemoryBuffer(chat_store=SimpleChatStore(store={'chat_history': []}), chat_store_key='chat_history', token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all')), callback_manager=<llama_index.core.callbacks.base.CallbackManager object at 0x7e3900a2ee90>, extra_state={}), '__output__': 5, 'count': 0}
DEBUG(is_done): False
DEBUG(state): {'max_count': 5, '__task__': Task(task_id='048e10e5-5a32-4df2-910d-56fb390aa127', input='5', memory=ChatMemoryBuffer(chat_store=SimpleChatStore(store={'chat_history': []}), chat_store_key='chat_history', token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all')), callback_manager=<llama_index.core.callbacks.base.CallbackManager object at 0x7e3900a2ee90>, extra_state={'max_count': 5, '__task__': Task(task_id='048e10e5-5a32-4df2-910d-56fb390aa127', input='5', memory=ChatMemoryBuffer(chat_store=SimpleChatStore(store={'chat_history': []}), chat_store_key='chat_history', token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all')), callback_manager=<llama_index.core.callbacks.base.CallbackManager object at 0x7e3900a2ee90>, extra_state={...}), '__output__': 5, 'count': 0}), '__output__': 10, 'count': 1}
DEBUG(is_done): False
DEBUG(state): {'max_count': 5, '__task__': Task(task_id='048e10e5-5a32-4df2-910d-56fb390aa127', input='5', memory=ChatMemoryBuffer(chat_store=SimpleChatStore(store={'chat_history': []}), chat_store_key='chat_history', token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all')), callback_manager=<llama_index.core.callbacks.base.CallbackManager object at 0x7e3900a2ee90>, extra_state={'max_count': 5, '__task__': Task(task_id='048e10e5-5a32-4df2-910d-56fb390aa127', input='5', memory=ChatMemoryBuffer(chat_store=SimpleChatStore(store={'chat_history': []}), chat_store_key='chat_history', token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all')), callback_manager=<llama_index.core.callbacks.base.CallbackManager object at 0x7e3900a2ee90>, extra_state={...}), '__output__': 10, 'count': 1}), '__output__': 20, 'count': 2}
DEBUG(is_done): False
DEBUG(state): {'max_count': 5, '__task__': Task(task_id='048e10e5-5a32-4df2-910d-56fb390aa127', input='5', memory=ChatMemoryBuffer(chat_store=SimpleChatStore(store={'chat_history': []}), chat_store_key='chat_history', token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all')), callback_manager=<llama_index.core.callbacks.base.CallbackManager object at 0x7e3900a2ee90>, extra_state={'max_count': 5, '__task__': Task(task_id='048e10e5-5a32-4df2-910d-56fb390aa127', input='5', memory=ChatMemoryBuffer(chat_store=SimpleChatStore(store={'chat_history': []}), chat_store_key='chat_history', token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all')), callback_manager=<llama_index.core.callbacks.base.CallbackManager object at 0x7e3900a2ee90>, extra_state={...}), '__output__': 20, 'count': 2}), '__output__': 40, 'count': 3}
DEBUG(is_done): False
DEBUG(state): {'max_count': 5, '__task__': Task(task_id='048e10e5-5a32-4df2-910d-56fb390aa127', input='5', memory=ChatMemoryBuffer(chat_store=SimpleChatStore(store={'chat_history': []}), chat_store_key='chat_history', token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all')), callback_manager=<llama_index.core.callbacks.base.CallbackManager object at 0x7e3900a2ee90>, extra_state={'max_count': 5, '__task__': Task(task_id='048e10e5-5a32-4df2-910d-56fb390aa127', input='5', memory=ChatMemoryBuffer(chat_store=SimpleChatStore(store={'chat_history': []}), chat_store_key='chat_history', token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all')), callback_manager=<llama_index.core.callbacks.base.CallbackManager object at 0x7e3900a2ee90>, extra_state={...}), '__output__': 40, 'count': 3}), '__output__': 80, 'count': 4}
DEBUG(is_done): False
DEBUG(state): {'max_count': 5, '__task__': Task(task_id='048e10e5-5a32-4df2-910d-56fb390aa127', input='5', memory=ChatMemoryBuffer(chat_store=SimpleChatStore(store={'chat_history': []}), chat_store_key='chat_history', token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all')), callback_manager=<llama_index.core.callbacks.base.CallbackManager object at 0x7e3900a2ee90>, extra_state={'max_count': 5, '__task__': Task(task_id='048e10e5-5a32-4df2-910d-56fb390aa127', input='5', memory=ChatMemoryBuffer(chat_store=SimpleChatStore(store={'chat_history': []}), chat_store_key='chat_history', token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all')), callback_manager=<llama_index.core.callbacks.base.CallbackManager object at 0x7e3900a2ee90>, extra_state={...}), '__output__': 80, 'count': 4}), '__output__': 160, 'count': 5}
DEBUG(is_done): True
Response(response='160', source_nodes=[], metadata=None)

なるほど、エージェント内に状態を保持するstate変数が用意され、データが与えられると、これを処理する関数でstateを更新するという感じ。

カスタムなエージェントのサンプルは以下。
https://docs.llamaindex.ai/en/stable/examples/agent/custom_agent/

kun432kun432

OpenAIAgent向けのベータ版)より進んだコンセプト

エージェントは、例えば、クエリ時にインデックスからツールを取得したり、既存のツールセットに対してクエリプランニングを実行したり、といった、より高度な設定でも使用できる。

これらは、OpenAI Function APIに依存しているため、主にOpenAIAgentクラスで実装かつベータという扱いになっていて、他のエージェントモジュール、例えばReActAgentでサポートできるかどうかは今後の調査次第というところらしい。

Function Retreval Agent

Function Retreval Agentは、ツールが多数ある場合に、インデックスを使ってツールの検索を行うというもの。ツールをObjectIndexを使ってツールをインデックス化、クエリ時にエージェントにObjectRetrieverを渡す。これにより、エージェントがツールの選択を行う前に、クエリに関連するツールを検索によって取捨できる。

まずObjectIndexを作成する。

from llama_index.core import VectorStoreIndex
from llama_index.core.objects import ObjectIndex

obj_index = ObjectIndex.from_objects(
    all_tools,
    index_cls=VectorStoreIndex,
)

そしてエージェントを定義する場合に、ObjectRetrieverを指定する。

from llama_index.agent.openai import OpenAIAgent

agent = OpenAIAgent.from_tools(
    tool_retriever=obj_index.as_retriever(similarity_top_k=2), verbose=True
)

ObjectIndexの詳細は以下のnotebookにある。
https://docs.llamaindex.ai/en/stable/examples/objects/object_index/

Context Retrieval Agent

Context Retrieval Agentは、ツールを呼び出す前に、クエリで検索を行い、追加のコンテキストを取得する。つまり、エージェントのツール選択を補完するようなコンテキスト情報を与えて、ツールがより適切に選択されるようにする。

from llama_index.core import Document
from llama_index.agent.openai_legacy import ContextRetrieverOpenAIAgent


# サンプルのインデックス。略語のリストを持たせている。
# toy index - stores a list of Abbreviations
texts = [
    "略語: X = 収益",
    "略語: YZ = リスク要素",
    "略語: Z = 費用",
]
docs = [Document(text=t) for t in texts]
context_index = VectorStoreIndex.from_documents(docs)

# context agentを追加
context_agent = ContextRetrieverOpenAIAgent.from_tools_and_retriever(
    query_engine_tools,
    context_index.as_retriever(similarity_top_k=1),
    verbose=True,
)
response = context_agent.chat("2022年3月のYZは?")

Query Planning

OpenAI Function Agentは、高度なクエリプランニングを行うことができる。このエージェントにQueryPlanToolを提供すると、サブツールのセットに対してクエリプランを表す完全なPydantcスキーマを推論させる。

# クエリプランニングツールを定義
from llama_index.core.tools import QueryPlanTool
from llama_index.core import get_response_synthesizer

response_synthesizer = get_response_synthesizer(
    service_context=service_context
)
query_plan_tool = QueryPlanTool.from_defaults(
    query_engine_tools=[query_tool_sept, query_tool_june, query_tool_march],
    response_synthesizer=response_synthesizer,
)

# エージェントを初期化
agent = OpenAIAgent.from_tools(
    [query_plan_tool],
    max_function_calls=10,
    llm=OpenAI(temperature=0, model="gpt-4-0613"),
    verbose=True,
)

# 3月、6月、9月のツールを呼び出すクエリプランを出力
response = agent.query(
    "Uberの、3月、6月、9月の収益の伸びを分析して。"
)
kun432kun432

あいかわらずUsage Patternというドキュメントの割に、レベル感がバラバラでどこから読み始めればいいのかとてもわかりにくい・・・パターンというところはわかるんだけども、もうちょっと初学者がどういうステップを踏めば理解が進みやすいかというところを意識したドキュメント構成になって欲しい。

kun432kun432

低レベルエージェントAPI

というところで、上でも出てきた低レベルエージェントAPIの話。LlamaIndexにおける「エージェント」がどういう仕組みになっているのか?をプログラム的・モジュール構成的に理解する意味では重要だと思う。

https://docs.llamaindex.ai/en/stable/module_guides/deploying/agents/agent_runner/

  • 高レベルエージェントAPI
    • ユーザクエリをエンドツーエンドで実行できる
    • 抽象化度合いが高め
    • OpenAIAgentとかReActAgentとか
    • エージェントにツールとLLMを与えて、クエリを投げるだけ
  • 低レベルエージェントAPI
    • ユーザクエリをエージェントがどのように処理するかを詳細に設定できる
    • 抽象化度合いが低め
    • より信頼性の高い・制御可能なエージェントを作成するのが目的。
    • AgentWorkerAgentRunnerを使って、エージェントを細かく定義してから、クエリを投げる

この仕組みは以下をベースに設計されているらしい。

AgentProtocol
https://agentprotocol.ai/

OpenAI Assitants API
https://platform.openai.com/docs/assistants/overview

ReWOO: Decoupling Reasoning from Observations for Efficient Augmented Language Models
https://arxiv.org/abs/2305.18323

最後の論文については以下がわかりやすかった。
https://tech.algomatic.jp/entry/2024/03/19/183857

アーキテクチャ

LlamaIndexにおけるエージェントはAgentRunnerAgentWorkersで構成される。

  • AgentRunners
    • 状態(会話メモリを含む)を保存
    • タスクを作成と管理
    • 各タスクを通してステップを実行
    • ユーザが対話するためのユーザ向けの高レベルインターフェースを提供するオーケストレータ
  • AgentWorkers
    • タスクのステップごとの実行をコントロールする
    • 入力されたステップに対して、AgentWorkersは次のステップを生成する。
    • パラメータで初期化することができ、Task/TaskStepオブジェクトから渡された状態に対応する。
    • 外側のAgentRunnerは、AgentWorkerを呼び出し、結果を収集/集約する責任を負う。

以下のような補助クラスもある。

  • Task
    • 高レベルのタスク
    • ユーザーからのクエリーを受け取り、メモリーなど他の情報を渡す。
  • TaskStep
    • 1つのステップを表す。
    • AgentWorkerに対して入力として与えられ、TaskStepOutputを返す。
    • Taskの完了は、複数のTaskStepが関連する場合がありうる。
  • TaskStepOutput
    • 与えられたステップ実行からの出力。
    • タスクが完了したかどうかを出力する。

ドキュメントの図を少し修正してみた。シングルステップで実行した場合。

んー、StructuredPlannerAgentのところでは、なんとなくAgentRunnerがプランニングをしていると思ったのだけど、プランニングもWorkerに行わせてるように見える。この図だとAgentRunnerは単にWorkerの実行管理とインタフェースって感じに思える。実際に動かしてみないとちょっとわからないかも。

利点

低レベルエージェントAPIを使うメリットは以下。

  • タスクの作成と実行を切り離して、タスクを実行するタイミングをコントロールする。
  • 各ステップの実行のデバッグ性を高める。
  • 完了したステップと次のステップを表示して、視認性を上げる。
  • [近日公開予定] 操舵性: 人間のフィードバックを注入することで、中間ステップを直接制御/変更する。
  • タスクの放棄: タスクが実行中に脱線した場合、コアエージェントメモリに影響を与えることなく、タスクを放棄する。
  • [近日公開予定] ステップのやり直し
  • より簡単なカスタマイズ: AgentWorkerを実装することで、新しいエージェントアルゴリズム(ReAct、OpenAIだけでなく、plan+solve、LLMCompilerなど)を簡単にサブクラス化/実装できる。

使用パターン

AgentRunnerAgentWorkerを使った書き方。ちょっとドキュメントのサンプルコードのままだと一部動かなかったので、少し処理の流れがわかりにくいかんがあったので書き換えた。

まず、ツールとAgentRunnerAgentWorkerを定義。

from llama_index.core.agent import AgentRunner
from llama_index.agent.openai import OpenAIAgentWorker
from llama_index.core.tools import FunctionTool
from llama_index.llms.openai import OpenAI

# サンプルのデータベース
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:
    """「商品名」(product_name)から「商品ID」(product_id)を取得する。"""
    return {"product_name": product_name, "product_id": product_list[product_name]}

def get_product_info_from_product_id(product_id: str)->dict:
    """「商品ID」(product_id)から「商品情報(価格、在庫)」(product_info)を取得する。"""
    return {"product_id": product_id, "product_info": product_catalog[product_id]}

# 関数をツールとして使えるように定義
get_product_id_tool = FunctionTool.from_defaults(fn=get_product_id_from_product_name)
get_product_info_tool = FunctionTool.from_defaults(fn=get_product_info_from_product_id)

tools = [get_product_id_tool, get_product_info_tool]

# llmを初期化
llm = OpenAI(model="gpt-3.5-turbo-0125")

# OpenAIAgentをツールから初期化
openai_step_engine = OpenAIAgentWorker.from_tools(tools, llm=llm, verbose=True)
agent = AgentRunner(openai_step_engine)

データベースにアクセスするための関数をツールとして用意しているのだけど、少しひねったツールにしていて、例えば、商品名から在庫を調べようと思うと、マルチステップで推論する必要があるというもの

一番最後のところでAgentRunnerAgentWorkerが定義されている。

で、エージェントにタスクを登録する。

# タスクを作成
task = agent.create_task("タブレットの在庫を調べて。")

タスクを登録するとTaskオブジェクトが生成される。

task
Task(task_id='f8ad17ab-6e70-4718-8db7-8037260171d5',
input='タブレットの在庫を調べて。',
memory=ChatMemoryBuffer(chat_store=SimpleChatStore(store={
}
),
chat_store_key='chat_history',
token_limit=3000,
tokenizer_fn=functools.partial(<boundmethodEncoding.encodeof<Encoding'cl100k_base'>>,
allowed_special='all')),
callback_manager=<llama_index.core.callbacks.base.CallbackManagerobjectat0x7d97bbdb2bc0>,
extra_state={
    'sources': [
    ],
    'n_function_calls': 0,
    'new_memory': ChatMemoryBuffer(chat_store=SimpleChatStore(store={
    }
    ),
    chat_store_key='chat_history',
    token_limit=3000,
    tokenizer_fn=functools.partial(<boundmethodEncoding.encodeof<Encoding'cl100k_base'>>,
    allowed_special='all'))
}
)

Taskオブジェクトはメモリのような役割に見える。

そしてAgentに渡されるのはTaskStepオブジェクトになる。次に渡されるオブジェクトを見てみる。

agent.get_upcoming_steps(task.task_id)
[
    TaskStep(task_id='f8ad17ab-6e70-4718-8db7-8037260171d5',
    step_id='d83a1cfc-a48d-46d9-bc1d-ccd228d6eca1',
    input='タブレットの在庫を調べて。',
    step_state={
    },
    next_steps={
    },
    prev_steps={
    },
    is_ready=True)
]

最初なので特にタスクと違いがあるようには見えないけども、TaskStepとしてのIDが付与されている。

で、このTaskStepをAgentWorkerにわたすことでステップが実行される。

# ステップを実行
step_output = agent.run_step(task.task_id)
Added user message to memory: タブレットの在庫を調べて。
=== Calling Function ===
Calling function: get_product_id_from_product_name with args: {"product_name":"タブレット"}
Got output: {'product_name': 'タブレット', 'product_id': 'E1003'}
========================

ステップの出力も見てみる。

step_output
TaskStepOutput(output=AgentChatResponse(response='None',
sources=[
    ToolOutput(content="{'product_name': 'タブレット', 'product_id': 'E1003'}",
    tool_name='get_product_id_from_product_name',
    raw_input={
        'args': (),
        'kwargs': {
            'product_name': 'タブレット'
        }
    },
    raw_output={
        'product_name': 'タブレット',
        'product_id': 'E1003'
    },
    is_error=False)
],
source_nodes=[
],
is_dummy_stream=False,
metadata=None),
task_step=TaskStep(task_id='f8ad17ab-6e70-4718-8db7-8037260171d5',
step_id='d83a1cfc-a48d-46d9-bc1d-ccd228d6eca1',
input='タブレットの在庫を調べて。',
step_state={
},
next_steps={
},
prev_steps={
},
is_ready=True),
next_steps=[
    TaskStep(task_id='f8ad17ab-6e70-4718-8db7-8037260171d5',
    step_id='7365f053-dced-462b-923f-26184eca2b83',
    input=None,
    step_state={
    },
    next_steps={
    },
    prev_steps={
    },
    is_ready=True)
],
is_last=False)

まず、最初のステップとして商品名から商品IDを検索するために1つ目のツールが実行され、商品IDが取得されたことがわかる。
そして次のステップが指定されているのと、is_lastがFalseになっていて、まだステップが続くことがわかる。

ではステップをさらに進めてみる。

step_output = agent.run_step(task.task_id)
=== Calling Function ===
Calling function: get_product_info_from_product_id with args: {"product_id":"E1003"}
Got output: {'product_id': 'E1003', 'product_info': {'price': 300, 'stock_level': 25}}
========================
step_output
TaskStepOutput(output=AgentChatResponse(response='None',
sources=[
    ToolOutput(content="{'product_name': 'タブレット', 'product_id': 'E1003'}",
    tool_name='get_product_id_from_product_name',
    raw_input={
        'args': (),
        'kwargs': {
            'product_name': 'タブレット'
        }
    },
    raw_output={
        'product_name': 'タブレット',
        'product_id': 'E1003'
    },
    is_error=False),
    ToolOutput(content="{'product_id': 'E1003', 'product_info': {'price': 300, 'stock_level': 25}}",
    tool_name='get_product_info_from_product_id',
    raw_input={
        'args': (),
        'kwargs': {
            'product_id': 'E1003'
        }
    },
    raw_output={
        'product_id': 'E1003',
        'product_info': {
            'price': 300,
            'stock_level': 25
        }
    },
    is_error=False)
],
source_nodes=[
],
is_dummy_stream=False,
metadata=None),
task_step=TaskStep(task_id='f8ad17ab-6e70-4718-8db7-8037260171d5',
step_id='7365f053-dced-462b-923f-26184eca2b83',
input=None,
step_state={
},
next_steps={
},
prev_steps={
},
is_ready=True),
next_steps=[
    TaskStep(task_id='f8ad17ab-6e70-4718-8db7-8037260171d5',
    step_id='e93827b6-1777-423e-802b-f6ca2f6ee4af',
    input=None,
    step_state={
    },
    next_steps={
    },
    prev_steps={
    },
    is_ready=True)
],
is_last=False)

ここで2つ目のツールが実行されて、商品IDから、在庫を含む商品情報を取得できたことがわかる。ただし、is_lastはFalseで次のステップが指定されているのでまだタスクは完了していない。

さらにステップを進める。

step_output = agent.run_step(task.task_id)
step_output
TaskStepOutput(output=AgentChatResponse(response='タブレットの在庫は25個あります。価格は300です。',
sources=[
    ToolOutput(content="{'product_name': 'タブレット', 'product_id': 'E1003'}",
    tool_name='get_product_id_from_product_name',
    raw_input={
        'args': (),
        'kwargs': {
            'product_name': 'タブレット'
        }
    },
    raw_output={
        'product_name': 'タブレット',
        'product_id': 'E1003'
    },
    is_error=False),
    ToolOutput(content="{'product_id': 'E1003', 'product_info': {'price': 300, 'stock_level': 25}}",
    tool_name='get_product_info_from_product_id',
    raw_input={
        'args': (),
        'kwargs': {
            'product_id': 'E1003'
        }
    },
    raw_output={
        'product_id': 'E1003',
        'product_info': {
            'price': 300,
            'stock_level': 25
        }
    },
    is_error=False)
],
source_nodes=[
],
is_dummy_stream=False,
metadata=None),
task_step=TaskStep(task_id='f8ad17ab-6e70-4718-8db7-8037260171d5',
step_id='e93827b6-1777-423e-802b-f6ca2f6ee4af',
input=None,
step_state={
},
next_steps={
},
prev_steps={
},
is_ready=True),
next_steps=[
],
is_last=True)

ツールの実行結果を踏まえて、レスポンスが生成されている。そしてis_lastがTrueになっていてステップがもう存在しない、タスクが完了していることがわかる。

タスクが完了していれば、ファイナライズすることで最終的なレスポンスが取得できる。

# step_output が 完了ならば、レスポンスをファイナライズ
if step_output.is_last:
    response = agent.finalize_response(task.task_id)
    print(response)
タブレットの在庫は25個あります。価格は300です。

最終的な回答が生成された。

上の方に記載した図のとおり、AgentRunnerとAgentWorker間でタスクをステップごとにやり取りしているのがわかる。

あと、完了したステップの一覧は以下で確認ができる。

agent.get_completed_steps(task.task_id)
kun432kun432

Module Guides

複数のエージェントモジュールがある。一覧は以下。

https://docs.llamaindex.ai/en/stable/module_guides/deploying/agents/modules/

  • OpenAI向けエージェント
    • OpenAI Agent
    • OpenAI Agent with Query Engine Tools
    • Retrieval Augmented Agent
    • OpenAI Agent Cookbook
    • Query Planning
    • Context Retrieval Agent
    • Recursive Retriever Agents
    • Multi Document Agents
    • Agent Builder
    • Parallel Function Calling
    • Agent with Planning
  • OpenAI Assistant エージェント
    • OpenAI Assistant
    • OpenAI Assistant Retrieval Benchmark
    • Assistant Query Cookbook
  • 他のFunction Calling エージェント
    • Mistral Agent
  • ReActエージェント
    • ReAct Agent
    • ReAct Agent
    • ReAct Agent with Query Engine Tools
  • LlamaHubで利用できるエージェント
    • LLMCompiler Agent (Cookbook)
    • Chain-of-Abstraction Agent (Cookbook)
    • Language Agent Tree Search Agent (Cookbook)
    • Instrospective Agent (Cookbook)
  • カスタムエージェント
    • Custom Agent
    • Query Pipeline Agent
  • 低レベルエージェントAPI
    • Agent Runner
    • Agent Runner RAG
    • Agent with Planning
    • Controllable Agent Runner

全部は見てられないのでいくつかピックアップして見ていく。

数が多いのでこちらで。

https://zenn.dev/kun432/scraps/0f2da2381099fa

kun432kun432

低レベルAPIを理解した上で、高レベルのモジュールを使えば良いと思う。

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