药物副作用问答系统

虽然大多数人关注的是检索增强生成 (RAG) 对非结构化文本（例如公司文档或文件）的检索，但我对检索系统对结构化信息（尤其是知识图谱）的检索非常看好。GraphRAG 引起了很多关注，尤其是微软的实现。然而，在他们的实现中，输入数据是文档形式的非结构化文本，使用大型语言模型 (LLM) 将其转换为知识图谱。

在这篇博文中，我们将展示如何在包含来自 FDA 不良事件报告系统 (FAERS) 的结构化信息的知识图谱上实现检索器，该系统提供有关药物不良事件的信息。如果你曾经摆弄过知识图谱和检索，你的第一个想法可能是使用 LLM 生成数据库查询，以从知识图谱中检索相关信息来回答给定的问题。然而，使用 LLM 生成数据库查询仍在发展中，可能还不能提供最一致或最强大的解决方案。那么，目前有哪些可行的替代方案呢？

我认为，目前最好的解决方案是动态查询生成。这种方法不是完全依赖 LLM 来生成完整的查询，而是采用逻辑层，从预定义的输入参数确定性地生成数据库查询。可以使用具有函数调用支持的 LLM 来实现此解决方案。使用函数调用功能的优势在于能够向 LLM 定义它应该如何准备函数的结构化输入。这种方法确保查询生成过程是可控且一致的，同时允许用户输入灵活性。

动态查询生成流程

该图说明了理解用户问题以检索特定信息的过程。该流程涉及三个主要步骤：

用户询问有关 35 岁以下人群服用 Lyrica 药物的常见副作用的问题。
LLM 决定调用哪个函数以及所需的参数。在这个例子中，它选择了一个名为 side_effects 的函数，其参数包括药物 Lyrica 和最大年龄 35 岁。
已识别的函数和参数用于确定性和动态地生成数据库查询 (Cypher) 语句以检索相关信息。

函数调用支持对于高级 LLM 用例至关重要，例如允许 LLM 根据用户意图使用多个检索器或构建多代理流。我写了一些使用具有本机函数调用支持的商业 LLM 的文章。但是，我们将使用最近发布的 Llama-3.1，这是一款具有本机函数调用支持的高级开源 LLM。

代码可在 GitHub 上找到。

1、设置知识图谱

我们将使用 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)

2、设置 LLM 环境

有许多选项可以托管开源 LLM，例如 Llama-3.1。我们将使用 NVIDIA API 目录，它提供 NVIDIA NIM 推理微服务并支持 Llama 3.1 模型的函数调用。创建帐户后，你将获得 1,000 个token，这足以让你继续学习。你需要创建一个 API 密钥并将其复制到笔记本中：

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

我们将使用 llama-3.1–70b，因为 8b 版本在函数定义中的可选参数方面存在一些问题。

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"
)

3、工具定义

我们将配置一个具有四个可选参数的单一工具。我们将根据这些参数构建相应的 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', 'LYRICA CR']"值，然后执行相应的 Cypher 语句来查找最常见的副作用。

4、基于图的 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 以生成最终答案。

5、结束语

函数调用功能是 Llama 3.1 等开源模型的强大补充，可实现与外部数据源和工具的更结构化和受控的交互。除了查询非结构化文档之外，基于图的代理还提供了与知识图和结构化数据交互的令人兴奋的可能性。使用 NVIDIA NIM 微服务等平台托管这些模型的便利性使它们越来越易于访问。

与往常一样，代码可在 GitHub 上找到。

原文链接：Enhancing RAG with Knowledge Graphs: Integrating Llama 3.1, NVIDIA NIM, and LangChain for Dynamic AI

汇智网翻译整理，转载请标明出处