基于知识图谱构建带LLama 3.1的智能药物副作用查询代理

这是用ChatGPT创建的

虽然大多数人更关注于从非结构化文本（如公司文件或文档）中进行提取增强生成（RAG），但我对从结构化信息（特别是知识图谱）中提取系统非常看好。特别是微软的实现，关于GraphRAG引起了很大的关注。然而，在他们的实现中，输入数据是以文档形式的非结构化文本，这些文本通过大型语言模型（LLM）转换成知识图谱。

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

在我看来，目前最好的解决方案是动态生成查询。这种方法不是完全依赖大型语言模型生成完整查询，而是通过一个逻辑层，根据预设的输入参数确定性生成数据库查询。这个解决方案可以通过支持函数调用的大型语言模型实现。使用函数调用功能的优势在于，可以定义大型语言模型如何为函数准备结构化输入。这种方法不仅确保了查询生成过程的可控性和一致性，还允许用户输入的灵活性。

动态查询生成图 — 图片由作者制作

这张图展示了理解用户提问并从中提取特定信息的过程。整个过程分三个主要步骤。

1. 一名用户询问35岁以下人群使用Lyrica药物常见的副作用。
2. 大型语言模型（LLM）决定调用哪个函数及其所需参数。在此示例中，它选择了名为side_effects的函数，参数包括Lyrica药物和最大年龄35。
3. 根据识别出的函数和参数，动态生成一个确定性的Cypher查询语句来检索相关信息。

函数调用功能对于高级LLM用例至关重要，例如允许LLM根据用户意图使用多个检索器或构建多代理流程。我写过一些相关文章，使用商业LLM。然而，我们将使用新发布的开源LLM Llama-3.1，它具有原生的函数调用支持。

代码可在GitHub上找到。

我们将使用Neo4j（一个原生图数据库）来存储不良事件信息。你可以通过这个链接免费设置一个云Sandbox项目，该项目已经预填充了FAERS数据。

实例化数据库实例包含如下模式的图。

作者提供的不良事件图表（图由作者提供）

该架构以病例节点为中心，将药物安全报告的各个部分联系在一起，包括所涉及的药物、产生的反应、结果以及开具的疗法。每种药物被描述为是主要的、次要的、伴随使用的或相互作用的。病例还涉及制造商信息、患者年龄组以及报告来源。此架构允许以结构化的方式追踪和分析药物、其反应和结果之间的相互关系。

让我们从建立一个数据库连接开始，通过实例化一个Neo4j图对象来实现。

    os.environ["NEO4J_URI"] = "bolt://18.206.157.187:7687"  # NEO4J_URI环境变量存储Neo4j数据库的连接字符串  
    os.environ["NEO4J_USERNAME"] = "neo4j"  # NEO4J_USERNAME环境变量存储Neo4j的用户名  
    os.environ["NEO4J_PASSWORD"] = "elevation-reservist-thousands"  # NEO4J_PASSWORD环境变量存储Neo4j的密码  

    graph = Neo4jGraph(refresh_schema=False)  # graph变量被初始化为一个Neo4j图，且不刷新模式

有很多选择可以托管像Llama-3.1这样的开源大型语言模型。我们将使用NVIDIA API目录，它提供了NVIDIA NIM推理微服务并且支持Llama 3.1模型的功能调用功能。注册后，你将获得1,000个代币，足够你跟着教程一起使用。你需要创建一个API密钥，并将其复制到笔记本里。

    os.environ["NVIDIA_API_KEY"] = "nvapi-"  # 设置环境变量NVIDIA_API_KEY，用于NVIDIA的相关API调用
    llm = ChatNVIDIA(model="meta/llama-3.1-70b-instruct")

我们将使用llama-3.1–70b，因为8b版本的可选参数在函数定义里有一些小毛病。

NVIDIA NIM 微服务的其中一个好处是，如果有安全或其他顾虑，你可以非常轻松地本地运行它们， 并且可以轻松替换，只需在 LLM 模型配置中添加一个 URL 参数即可：

# 连接到运行在本地的NIM，地址为localhost:8000，并指定使用特定的模型.
llm = ChatNVIDIA(
  base_url="http://localhost:8000/v1",   
  model="meta/llama-3.1-70b-instruct"  
)

我们将配置一个单一工具，该工具拥有四个可选参数。我们将根据这些参数构建相应的Cypher语句如下配置工具参数，从知识图谱中检索相关信息如下。我们的工具将能够根据输入的药物、年龄和药物制造商来识别最常见的副作用，如下配置。

@tool
def get_side_effects(
    drug: Optional[str] = Field(
        description="问题中提到的药物。如果没有提到，则返回None。"
    ),
    min_age: Optional[int] = Field(
        description="患者的最小年龄。如果没有提到，则返回None。"
    ),
    max_age: Optional[int] = Field(
        description="患者的最大年龄。如果没有提到，则返回None。"
    ),
    manufacturer: Optional[str] = Field(
        description="药物的制造商。如果没有提到，则返回None。"
    ),
):
    """当需要查找常见副作用时很有用。"""
    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 "未找到提到的药物"
        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 "未找到提到的制造商"
        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"使用参数: {params}")
    data = graph.query(side_effects_base_query, params=params)
    return data  

get_side_effects根据指定的搜索条件从知识图谱中检索药物的常见副作用信息。它接受药物名称、患者年龄范围和药物制造商等可选参数，以便定制搜索。每个参数的描述都会传递给LLM，连同函数描述，使LLM能够理解如何使用这些参数。然后，该函数根据提供的参数构建动态的Cypher查询，执行此查询来从知识图谱中获取数据，并返回相应的副作用信息。

我们来试试这个功能吧

    get_side_effects("lyrica")  
    # 使用参数：{'drugs': ['LYRICA', 'LYRICA CR']}  
    # [{'drug': 'LYRICA', 'side_effect': '疼痛', 'count': 32},  
    #  {'drug': 'LYRICA', 'side_effect': '跌倒', 'count': 21},  
    #  {'drug': 'LYRICA', 'side_effect': '有意滥用产品', 'count': 20},  
    #  {'drug': 'LYRICA', 'side_effect': '失眠', 'count': 19},  
    #  ...

我们的工具首先将问题中提到的Lyrica药物映射为知识图中的“[‘LYRICA’，‘LYRICA CR’]”值，然后执行相应的Cypher查询来查找侧效应。

最后要做的就是配置一个LLM代理程序，让它能用定义的工具回答关于这种药的副作用的问题。

代理数据流图 — 图由作者提供

该图展示了一个用户与Llama 3.1 agent互动，询问药物副作用。代理访问一个用于查询副作用的工具，从知识图谱中获取信息，为用户提供相关的信息。

我们先来定义一下提示模板：

    prompt = ChatPromptTemplate.from_messages(  
        [  
            (  
                "system",  
                "你是一个乐于助人的助手，帮助用户了解常见副作用的信息。"  
                "如果有后续问题，记得向用户提问以澄清这些选项。"  
                "只根据用户的明确要求行事。",  
            ),  
            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)  
    # llm_with_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 是用于执行 agent 的执行器，指定了输入输出类型，并开启了详细日志  
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True).with_types(  
        input_type=AgentInput, output_type=Output  
    )

代理通过转换和处理程序来处理输入，这些程序会格式化聊天历史记录，使用绑定工具的应用LLM，并解析输出。最终，代理会被设置一个执行器来管理执行流程，指定输入和输出类型，并包含详细日志设置以在执行期间记录详细信息。

我们来试一下这个代理吧

    agent_executor.invoke(  
        {  
            "input": "使用利克林时，35岁以下的人常见的副作用有哪些?"  
        }  
    )

结果如下：

代理执行图 — 作者的插图

LLM 发现它需要使用带有适当参数的 get_side_effects 函数。该函数随后动态生成 Cypher 语句，获取相关信息，并将这些信息返回给 LLM 以生成最终答案。

函数调用能力大大增强了开源模型如Llama 3.1的功能，使其能更有效地与外部数据源和工具互动。除了能查询非结构化文档，图代理提供了与知识图和结构化数据互动的多种可能性。通过使用类似NVIDIA NIM微服务的平台，可以很方便地托管这些模型，使得它们越来越容易使用。

一如既往，代码可在 GitHub 上获得。