🕸️

LangGraph+Wikipediaを使用したLLMエージェント

2024/05/03に公開

LangChain

はじめに

質問に対して不足している情報をwikipediaから少しずつ取得して、十分に知識が貯まったら回答するエージェントをLangGraphで実装しましたので、ちょっとした情報と感想を共有します。

LangGraph入門者の記事です、ご了承ください<(_ _)>

ここで紹介するコードはgithubに保存しています。

作ったもの

例えば「ジョジョの奇妙な物語の作者の出身小学校の設立年は？」という質問に対して、「まず作者の名前を調べ、次にその作者の出身小学校の名前を調べて、最後に出身小学校の設立年を調べる。」のように、不足する情報を逐一調べるエージェントとなっています。

ノードエッジの構成

langgraphのノードとエッジは以下の構成です。

init_agent
- agentの状態を初期化するノード
generate_task
- ナレッジベースと質問をもとに、回答するために不足している知識を特定し、検索タスクを生成するノード。
search_and_answer
- Wikipediaに検索して、必要な情報をLLMで取得するノード。
final_answer
- ナレッジベースと質問をもとに、回答を生成するノード。
- 回答できないときは回答不可と返答する。

コード解説

エージェントの状態(AgentState)

class AgentState(TypedDict):
    user_question : str #ユーザーの質問
    messages: Annotated[Sequence[BaseMessage], operator.add] #会話履歴
    knowledge_base : Annotated[str, operator.add] #知識ベース(検索した情報を追記していく)
    next_task_search_keyword : str #次に検索するキーワード
    next_task_search_content : str #次に検索して調べる内容
    answer_counter : Annotated[int, operator.add] #最終回答を試みた回数

knowledge_baseは、検索した情報を格納する文字列です。検索ノード(call_search_and_answer)で取得した新しい情報を最後尾に追加していくだけのものです。

検索タスク生成ノード(generate_task)

class SearchTask(BaseModel):
    search_reason: str = Field(description="検索する理由や目的。")
    search_keyword: str = Field(description="Wikipediaでの検索する単語。必ずひとつの単語である必要がある。")
    search_content : str = Field(description="検索したWikipediaのページで取得したい情報。")

gen_task_template ="""
以下の質問に回答するため、不足している情報をwikipediaから検索します。
検索する情報は最小の一つだけです。検索するための単語を記載してください。
検索する項目やその理由を記載してください。
検索する項目は一つだけです。
複数の質問がなされていても、一つの質問に対してのみ回答を生成してください。

{format_instructions}

-----------------
質問例:映画オッペンハイマーの監督が作成した映画の本数を調べてください。
search_reason="情報が不足しているため。とりあえずまず"映画オッペンハイマー"のwikipediaを検索して、監督の名前を調べる必要があります。"
search_keyword="オッペンハイマー"
search_content="映画オッペンハイマーの監督の名前"
-----------------
質問例:中性子と原子核が発見された年に起こった出来事を調べてください。
search_reason="情報が不足しているため、とりあえずまず"中性子"のwikipediaを検索して、とりあえず発見された年を調べる必要があります。"
search_keyword="中性子"
search_content="中性子の発見された年"
-----------------
質問例:中性子と原子核が発見された年に起こった出来事を調べてください。ただし、以下の情報はすでに検索済みです。中性子の発見された年は1932年です。
search_reason="原子核に関する情報が不足しているため、"原子核"のwikipediaを検索して、発見された年を調べる必要があります。"
search_keyword="原子核"
search_content="原子核の発見された年"
-----------------

質問:{user_question}。{knowledge_base}
"""
def call_generate_task(state):
    _parser = PydanticOutputParser(pydantic_object=SearchTask)
    _prompt = PromptTemplate(
        template=gen_task_template,
        input_variables=["user_question", "knowledge_base"],
        partial_variables={"format_instructions": _parser.get_format_instructions()},
    )
    chain = _prompt| model | _parser
    # chain = chain.with_retry(stop_after_attempt=5)
    res = chain.invoke({'user_question': state['user_question'], 'knowledge_base': state['knowledge_base']})

    return {
        "messages": [AIMessage(content=res.search_reason)],
        'next_task_search_keyword':res.search_keyword,
        'next_task_search_content':res.search_content
    }

検索タスク生成ノード(call_generate_task)は、質問と知識ベースから、検索する理由、検索ワード、検索したページから取得したい内容を生成するノードです。例えば、質問文と知識ベースが

質問例：中性子と原子核が発見された年に起こった出来事を調べてください。
知識ベース：中性子の発見された年は1932年です。

のとき、以下の項目を生成することを目的としています。

検索する理由：原子核に関する情報が不足しているため、"原子核"のwikipediaを検索して、発見された年を調べる必要があります。
検索ワード：原子核
取得したい内容：原子核の発見された年

gpt-4-truboならば、この程度のfew-shotを書いておくだけでも、そこそこ意図通りに動きました。gpt-3.5-turboだと、例えば上記質問に対して「中性子　原子核　発見された年」といった検索ワードを生成するなど意図通りに動いてくれない傾向がありました。

業務で使うには、もう少し工夫する必要がありそうです。

検索ノード(search_and_answer)


def search_for_wikipedia(query: str) -> str:
    """
    Search for a wikipedia article and return the content of the first article found.
    """
    docs = WikipediaLoader(query=query, load_max_docs=5, lang='ja').load()
    if len(docs) > 0:
        return ''.join([d.page_content for d in docs])
    return ""


search_and_answer_template ="""
以下の参考文書を使用して、知りたい内容に回答してください。
知りたい内容には簡潔に数10文字のテキストで答えてください。
たとえば、「XXXはYYYです。」と短文で回答すること。

不要な情報は削除すること。

参考文書:{searched_content}

知りたい内容:{next_task_search_content}
"""

def call_search_and_answer(state):
    searched_content = search_for_wikipedia(state['next_task_search_keyword'])

    _parser = PydanticOutputParser(pydantic_object=SearchTask)
    _prompt = PromptTemplate(
        template=search_and_answer_template,
        input_variables=["next_task_search_content", "searched_content"],
    )
    chain = _prompt| model | StrOutputParser()

    res = chain.invoke({'next_task_search_content': state['next_task_search_content'], 'searched_content': searched_content})
    new_knowledge = f"検索した結果「{state['next_task_search_content']}」に関する情報は次の通り。\n{res}\n\n\n"
    return {
        "messages": [AIMessage(content=new_knowledge)],
        'knowledge_base':new_knowledge,
    }

ここは特に工夫してません。ここで生成した新しい知識('new_knowledge')をAgentStateのknowledge_baseに追記していきます。

最終回答ノード(final_answer)

final_answer ="""
以下の知識をもとに質問に回答してください。
もし回答に必要な情報が不足している場合は、回答できな理由を述べた後、「回答不可」と回答してください。

知識:{knowledge_base}

質問:{user_question}
"""
def call_final_answer(state):
    prompt = ChatPromptTemplate.from_template(template=final_answer)
    chain = prompt| model | StrOutputParser()
    res = chain.invoke({
        'user_question': state['user_question'],
        'knowledge_base': state['knowledge_base']
    })
    return {"messages": [AIMessage(content=res)], "answer_counter":1}

def router_fa(state):
    res = call_final_answer(state)
    if state['answer_counter'] >= 5:
        # print('回答制限')
        return 'end'
    if '回答不可' in res['messages'][-1].content:
        return 'continue'
    else:
        return 'end'

これは、最終回答ノード(call_final_answer)でAgentStateのknowledge_baseから回答を生成できるか判断します。
knowledge_baseに必要な情報が足りない場合は「回答不可」と出力され、router_faで再度情報検索する分岐がなされます。gpt-4-turboの場合、情報不足のため「回答不可」という文字列を入れてくれるので、この実装でもそこそこ動きました。

感想

Wikipedia情報を2~3回検索して回答するような今回のケースでは、タスクをうまく分割してくれるかどうかがミソでした。
もう少し長期の会話の場合や、サイズの大きい文書から回答する場合はメモリ(知識ベース)も工夫が必要そうな気がします。
gpt-3.5-turboやclaude3-haikuのような軽量モデルでも動かせないか試してみましたが、OutputParserでのエラーや適切な検索ワードを抽出できないなど、あまりうまく動きませんでした。