技术动态 | 构建基于知识图谱的医药智能体Agent

转载公众号 | 知识图谱科技

---

Build a knowledge graph-based agent with Llama 3.1, NVIDIA NIM, and LangChain by Tomaz Bratanic

利用Llama-3.1的原生函数调用能力，从知识图谱中检索结构化数据，以增强RAG应用程序

使用ChatGPT创建

虽然大多数人关注RAG与非结构化文本的关系，比如公司文档或技术文档，但我对基于结构化信息的检索系统持乐观态度，特别是知识图谱。关于GraphRAG，特别是微软的实现，重磅 - 微软官宣正式在GitHub开源GraphRAG，令人非常兴奋。然而，在他们的实现中，输入数据是以文档形式存在的非结构化文本，这些文本通过大型语言模型（LLM）转换为知识图谱。

在这篇博客文章中，我们将展示如何在包含来自FDA不良事件报告系统（FAERS-https://www.fda.gov/drugs/questions-and-answers-fdas-adverse-event-reporting-system-faers/fda-adverse-event-reporting-system-faers-public-dashboard）的结构化信息的知识图谱上实现检索器，该系统提供关于药物不良事件AE的信息。如果你曾经尝试过知识图谱和检索，你的第一反应可能是使用LLM生成数据库查询，从知识图谱中检索相关信息以回答给定的问题。然而，使用LLM生成数据库查询仍在发展中，可能尚未提供最一致或最强大的解决方案。那么，目前可行的替代方案是什么呢？

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

动态查询生成流程 作者提供的图像

该图示展示了理解用户提问以获取特定信息的过程。该流程包括三个主要步骤：

1. 用户询问关于35岁以下人群使用药物Lyrica的常见副作用的问题。

2. 大型语言模型决定调用哪个函数以及所需的参数。在这个例子中，它选择了一个名为“side_effects”的函数，参数包括药物“Lyrica”和最大年龄35岁。

3. 识别出的函数和参数被用来确定性和动态地生成数据库查询（Cypher）语句，以检索相关信息。

函数调用支持对高级大型语言模型用例至关重要，例如允许大型语言模型根据用户意图使用多个检索器或构建多代理流程。我写了一些关于使用商业大型语言模型及其原生函数调用支持的文章。然而，在这篇博客文章中，我们将使用Llama-3.1，这是一种出色的开源大型语言模型，具备原生函数调用支持，并且刚刚发布-开源引领新时代 -Meta重磅发布Llama3.1 405B，可媲美GPT-4o级 及 社评。

代码可在GitHub上获得 

https://github.com/tomasonjo/blogs/blob/master/llm/nvidia_neo4j_langchain.ipynb

建立知识图谱

我们将使用Neo4j，这是一个本地图形数据库，用于存储不良事件信息。您可以通过以下链接（https://sandbox.neo4j.com/?usecase=healthcare-analytics）设置一个免费的云Sandbox项目，该项目配备了预先填充的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环境

有许多选项可以托管开源的LLM，比如Llama-3.1。在这篇博客中，我们将使用NVIDIA API目录，它提供NVIDIA NIM推理微服务，并支持Llama 3.1模型的函数调用。当您创建一个帐户时，您将获得1000个令牌，这足够完成这篇博客。您需要创建一个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"
)

工具定义

在这个例子中，我们将配置一个具有四个可选参数的工具。基于这些参数，我们将构建一个相应的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函数旨在通过指定的搜索条件从知识图谱中检索药物的常见副作用。它接受可选参数，包括药物名称、患者年龄范围和药物制造商，以自定义搜索。每个参数都有一个描述，这些描述与函数描述一起传递给大型语言模型，使其理解如何使用这些参数。然后该函数根据提供的输入构建动态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语句以查找最常见的副作用。

基于图谱的大模型智能体

唯一需要做的就是配置一个可以使用定义工具回答有关药物副作用问题的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?"
    }
)

结果

Agent execution. Image by author.

LLM识别到需要使用get_side_effects函数并提供适当的参数。该函数动态生成Cypher语句，获取相关信息，并将其返回给LLM以生成最终答案。

Summary摘要

函数调用能力是对开源模型如Llama 3.1的强大补充，使与外部数据源和工具的互动更加结构化和可控。除了查询非结构化文档之外，基于图形的代理为与知识图谱和结构化数据的互动提供了令人兴奋的可能性。使用像NVIDIA NIM(https://build.nvidia.com/explore/discover)微服务这样的平台轻松托管这些模型，使它们变得越来越易于获取

如常，代码在GitHub上可用

https://github.com/tomasonjo/blogs/blob/master/llm/nvidia_neo4j_langchain.ipynb

---

OpenKG

OpenKG（中文开放知识图谱）旨在推动以中文为核心的知识图谱数据的开放、互联及众包，并促进知识图谱算法、工具及平台的开源开放。

点击阅读原文，进入 OpenKG 网站。