ナレッジグラフによる RAG の強化: 動的 AI のための Llama 3.1、NVIDIA NIM、LangChain の統合

ほとんどの人が、会社の文書やドキュメントなどの非構造化テキストに対する検索拡張生成 (RAG) に注目していますが、私は構造化情報、特にナレッジ グラフに対する検索システムにかなり期待しています。GraphRAG、特に Microsoft の実装については大きな期待が寄せられています。ただし、その実装では、入力データはドキュメント形式の非構造化テキストであり、大規模言語モデル (LLM) を使用してナレッジ グラフに変換されます。

このブログ記事では、医薬品の有害事象に関する情報を提供するFDA 有害事象報告システム (FAERS)からの構造化情報を含むナレッジ グラフに対して、リトリーバーを実装する方法を説明します。ナレッジ グラフと検索を扱ったことがある人なら、まず LLM を使用してデータベース クエリを生成し、ナレッジ グラフから関連情報を取得して特定の質問に答えることを考えます。ただし、LLM を使用したデータベース クエリ生成はまだ進化の途上であり、最も一貫性のある堅牢なソリューションを提供していない可能性があります。では、現時点で実行可能な代替手段は何でしょうか。

私の意見では、現時点で最も優れたソリューションは動的クエリ生成です。この方法では、完全なクエリを生成するために LLM に完全に依存するのではなく、定義済みの入力パラメータからデータベース クエリを決定論的に生成するロジック レイヤーを使用します。このソリューションは、関数呼び出しをサポートする LLM を使用して実装できます。関数呼び出し機能を使用する利点は、LLM に関数への構造化入力を準備する方法を定義できることです。このアプローチにより、クエリ生成プロセスが制御され、一貫性を保ちながら、ユーザー入力の柔軟性を確保できます。

この画像は、ユーザーの質問を理解して特定の情報を取得するプロセスを示しています。フローには主に 3 つのステップがあります。

1. ユーザーが、35 歳未満の人によく見られる薬「リリカ」の副作用について質問します。

   
   

2. LLM は、呼び出す関数と必要なパラメータを決定します。この例では、薬剤 Lyrica と最大年齢 35 歳を含むパラメータを持つ side_effects という名前の関数が選択されました。

   
   

3. 識別された関数とパラメータは、関連情報を取得するためのデータベース クエリ (Cypher) ステートメントを決定論的かつ動的に生成するために使用されます。

関数呼び出しサポートは、LLM がユーザーの意図に基づいて複数のリトリーバーを使用できるようにしたり、マルチエージェント フローを構築したりするなど、高度な LLM ユース ケースにとって不可欠です。私は、ネイティブ関数呼び出しサポートを備えた商用 LLM を使用した記事をいくつか書いています。ただし、ここでは、ネイティブ関数呼び出しサポートを備えた優れたオープン ソース LLM である最近リリースされた Llama-3.1 を使用します。

コードはGitHubで入手できます。

ナレッジグラフの設定

有害事象情報の保存には、ネイティブ グラフ データベースである Neo4j を使用します。このリンクをたどると、FAERS が事前に入力された無料のクラウド サンドボックス プロジェクトを設定できます。

インスタンス化されたデータベース インスタンスには、次のスキーマを持つグラフがあります。

このスキーマは、関連する薬剤、経験した反応、結果、処方された治療法など、薬剤安全性レポートのさまざまな側面をリンクするケース ノードを中心にしています。各薬剤は、主、副、同時、または相互作用のいずれかによって特徴付けられます。ケースは、製造元、患者の年齢層、レポートのソースに関する情報にも関連付けられています。このスキーマにより、薬剤、その反応、結果の関係を構造化された方法で追跡および分析できます。

まず、Neo4jGraph オブジェクトをインスタンス化してデータベースへの接続を作成します。

 os.environ["NEO4J_URI"] = "bolt://18.206.157.187:7687" os.environ["NEO4J_USERNAME"] = "neo4j" os.environ["NEO4J_PASSWORD"] = "elevation-reservist-thousands" graph = Neo4jGraph(refresh_schema=False)

LLM 環境の設定

Llama-3.1 のようなオープンソース LLM をホストするには、多くのオプションがあります。ここでは、 NVIDIA NIM 推論マイクロサービスを提供し、Llama 3.1 モデルの関数呼び出しをサポートするNVIDIA API カタログを使用します。アカウントを作成すると、1,000 トークンが付与されます。このトークンは、このチュートリアルを進めるには十分すぎるほどです。API キーを作成し、ノートブックにコピーする必要があります。

 os.environ["NVIDIA_API_KEY"] = "nvapi-" llm = ChatNVIDIA(model="meta/llama-3.1-70b-instruct")

8b バージョンでは関数定義のオプション パラメータに関して問題があるため、 llama-3.1–70bを使用します。

NVIDIA NIM マイクロサービスの優れた点は、セキュリティやその他の懸念がある場合でもローカルで簡単にホストできるため、簡単に交換でき、LLM 構成に URL パラメータを追加するだけで済むことです。

 # connect to an local NIM running at localhost:8000, # specifying a specific model llm = ChatNVIDIA( base_url="http://localhost:8000/v1", model="meta/llama-3.1-70b-instruct" )

ツールの定義

4 つのオプション パラメータを使用して単一のツールを構成します。これらのパラメータに基づいて対応する Cypher ステートメントを作成し、ナレッジ グラフから関連情報を取得します。このツールは、入力された薬剤、年齢、および薬剤製造元に基づいて、最も頻繁に発生する副作用を特定できます。

 @tool def get_side_effects( drug: Optional[str] = Field( description="disease mentioned in the question. Return None if no mentioned." ), min_age: Optional[int] = Field( description="Minimum age of the patient. Return None if no mentioned." ), max_age: Optional[int] = Field( description="Maximum age of the patient. Return None if no mentioned." ), manufacturer: Optional[str] = Field( description="manufacturer of the drug. Return None if no mentioned." ), ): """Useful for when you need to find common side effects.""" params = {} filters = [] side_effects_base_query = """ MATCH (c:Case)-[:HAS_REACTION]->(r:Reaction), (c)-[:IS_PRIMARY_SUSPECT]->(d:Drug) """ if drug and isinstance(drug, str): candidate_drugs = [el["candidate"] for el in get_candidates(drug, "drug")] if not candidate_drugs: return "The mentioned drug was not found" filters.append("d.name IN $drugs") params["drugs"] = candidate_drugs if min_age and isinstance(min_age, int): filters.append("c.age > $min_age ") params["min_age"] = min_age if max_age and isinstance(max_age, int): filters.append("c.age < $max_age ") params["max_age"] = max_age if manufacturer and isinstance(manufacturer, str): candidate_manufacturers = [ el["candidate"] for el in get_candidates(manufacturer, "manufacturer") ] if not candidate_manufacturers: return "The mentioned manufacturer was not found" filters.append( "EXISTS {(c)<-[:REGISTERED]-(:Manufacturer {manufacturerName: $manufacturer})}" ) params["manufacturer"] = candidate_manufacturers[0] if filters: side_effects_base_query += " WHERE " side_effects_base_query += " AND ".join(filters) side_effects_base_query += """ RETURN d.name AS drug, r.description AS side_effect, count(*) AS count ORDER BY count DESC LIMIT 10 """ print(f"Using parameters: {params}") data = graph.query(side_effects_base_query, params=params) return data

get_side_effects 関数は、指定された検索条件を使用してナレッジ グラフから薬の一般的な副作用を取得するように設計されています。検索をカスタマイズするために、薬の名前、患者の年齢範囲、薬の製造元などのオプション パラメータを受け入れます。各パラメータには関数の説明とともに LLM に渡される説明があり、LLM がそれらの使用方法を理解できるようにします。次に、関数は提供された入力に基づいて動的な Cypher クエリを作成し、このクエリをナレッジ グラフに対して実行して、結果として得られた副作用データを返します。

関数をテストしてみましょう:

 get_side_effects("lyrica") # Using parameters: {'drugs': ['LYRICA', 'LYRICA CR']} # [{'drug': 'LYRICA', 'side_effect': 'Pain', 'count': 32}, # {'drug': 'LYRICA', 'side_effect': 'Fall', 'count': 21}, # {'drug': 'LYRICA', 'side_effect': 'Intentional product use issue', 'count': 20}, # {'drug': 'LYRICA', 'side_effect': 'Insomnia', 'count': 19}, # ...

私たちのツールは、まず質問で言及されたリリカという薬をナレッジグラフ内の「['LYRICA', 'LYRICA CR']」の値にマッピングし、対応する Cypher ステートメントを実行して最も頻繁な副作用を見つけ出しました。

グラフベースの LLM エージェント

残っているのは、定義されたツールを使用して薬の副作用に関する質問に答えることができる LLM エージェントを構成することだけです。

この画像は、ユーザーが Llama 3.1 エージェントと対話して薬の副作用について問い合わせている様子を示しています。エージェントは副作用ツールにアクセスし、ナレッジ グラフから情報を取得して、関連データをユーザーに提供します。

まずプロンプト テンプレートを定義します。

 prompt = ChatPromptTemplate.from_messages( [ ( "system", "You are a helpful assistant that finds information about common side effects. " "If tools require follow up questions, " "make sure to ask the user for clarification. Make sure to include any " "available options that need to be clarified in the follow up questions " "Do only the things the user specifically requested. ", ), MessagesPlaceholder(variable_name="chat_history"), ("user", "{input}"), MessagesPlaceholder(variable_name="agent_scratchpad"), ] )

プロンプト テンプレートには、システム メッセージ、オプションのチャット履歴、およびユーザー入力が含まれます。agent_scratchpad は、ツールから情報を実行したり取得したりするなど、質問に答えるために複数の手順が必要になることがあるため、LLM 用に予約されています。

LangChain ライブラリでは、bind_tools メソッドを使用して LLM にツールを簡単に追加できます。

 tools = [get_side_effects] llm_with_tools = llm.bind_tools(tools=tools) agent = ( { "input": lambda x: x["input"], "chat_history": lambda x: _format_chat_history(x["chat_history"]) if x.get("chat_history") else [], "agent_scratchpad": lambda x: format_to_openai_function_messages( x["intermediate_steps"] ), } | prompt | llm_with_tools | OpenAIFunctionsAgentOutputParser() ) agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True).with_types( input_type=AgentInput, output_type=Output )

エージェントは、チャット履歴をフォーマットし、バインドされたツールで LLM を適用し、出力を解析する変換とハンドラーを介して入力を処理します。最後に、エージェントは、実行フローを管理し、入力と出力のタイプを指定し、実行中の詳細なログ記録のための冗長設定を含むエグゼキュータを使用してセットアップされます。

エージェントをテストしてみましょう:

 agent_executor.invoke( { "input": "What are the most common side effects when using lyrica for people below 35 years old?" } )

結果：

LLM は、適切な引数を持つ get_side_effects 関数を使用する必要があることを特定しました。次に、関数は Cypher ステートメントを動的に生成し、関連情報を取得して、それを LLM に返して最終的な回答を生成します。

まとめ

関数呼び出し機能は、Llama 3.1 などのオープンソース モデルに強力に追加され、外部データ ソースやツールとのより構造化され制御されたインタラクションを可能にします。グラフ ベースのエージェントは、非構造化ドキュメントのクエリだけでなく、ナレッジ グラフや構造化データとインタラクションするための魅力的な可能性を提供します。NVIDIA NIM マイクロサービスなどのプラットフォームを使用してこれらのモデルを簡単にホストできるため、アクセスしやすくなっています。

いつものように、コードはGitHubで入手できます。

このトピックについて詳しく知るには、11 月 7 日に開催されるインテリジェント アプリ、ナレッジ グラフ、AI に関する無料の仮想開発者会議 NODES 2024 にご参加ください。今すぐ登録してください。