Aprimorando o RAG com gráficos de conhecimento: integrando Llama 3.1, NVIDIA NIM e LangChain para IA dinâmica

Enquanto a maioria das pessoas foca na geração aumentada de recuperação (RAG) em vez de texto não estruturado, como documentos ou documentação da empresa, eu sou bem otimista em relação aos sistemas de recuperação em vez de informações estruturadas, particularmente gráficos de conhecimento . Tem havido muita empolgação sobre o GraphRAG, especificamente a implementação da Microsoft. No entanto, em sua implementação, os dados de entrada são textos não estruturados na forma de documentos, que são transformados em um gráfico de conhecimento usando um modelo de linguagem grande (LLM).

Nesta postagem do blog, mostraremos como implementar um recuperador sobre um gráfico de conhecimento contendo informações estruturadas do FDA Adverse Event Reporting System (FAERS) , que oferece informações sobre eventos adversos a medicamentos. Se você já mexeu com gráficos de conhecimento e recuperação, seu primeiro pensamento pode ser usar um LLM para gerar consultas de banco de dados para recuperar informações relevantes de um gráfico de conhecimento para responder a uma determinada pergunta. No entanto, a geração de consultas de banco de dados usando LLMs ainda está evoluindo e pode ainda não oferecer a solução mais consistente ou robusta. Então, quais são as alternativas viáveis no momento?

Na minha opinião, a melhor solução atual é a geração dinâmica de consultas. Em vez de depender inteiramente de um LLM para gerar a consulta completa, esse método emprega uma camada lógica que gera deterministicamente uma consulta de banco de dados a partir de parâmetros de entrada predefinidos. Essa solução pode ser implementada usando um LLM com suporte a chamada de função. A vantagem de usar um recurso de chamada de função está na capacidade de definir para um LLM como ele deve preparar uma entrada estruturada para uma função. Essa abordagem garante que o processo de geração de consultas seja controlado e consistente, ao mesmo tempo em que permite flexibilidade de entrada do usuário.

A imagem ilustra um processo de compreensão da pergunta de um usuário para recuperar informações específicas. O fluxo envolve três etapas principais:

1. Um usuário faz uma pergunta sobre os efeitos colaterais comuns do medicamento Lyrica para pessoas com menos de 35 anos.

   
   

2. O LLM decide qual função chamar e os parâmetros necessários. Neste exemplo, ele escolheu uma função chamada side_effects com parâmetros incluindo o medicamento Lyrica e uma idade máxima de 35.

   
   

3. A função e os parâmetros identificados são usados para gerar de forma determinística e dinâmica uma instrução de consulta ao banco de dados (Cypher) para recuperar informações relevantes.

O suporte a chamadas de função é vital para casos de uso avançados de LLM, como permitir que LLMs usem vários recuperadores com base na intenção do usuário ou criem fluxos multiagentes. Escrevi alguns artigos usando LLMs comerciais com suporte nativo a chamadas de função. No entanto, usaremos o Llama-3.1 lançado recentemente, um LLM de código aberto superior com suporte nativo a chamadas de função.

O código está disponível no GitHub .

Configurando o Knowledge Graph

Usaremos o Neo4j, que é um banco de dados de gráficos nativo, para armazenar as informações de eventos adversos. Você pode configurar um projeto Sandbox gratuito na nuvem que vem com FAERS pré-preenchido seguindo este link .

A instância do banco de dados instanciada tem um gráfico com o seguinte esquema.

O esquema é centralizado no nó Case, que vincula vários aspectos de um relatório de segurança de medicamentos, incluindo os medicamentos envolvidos, reações experimentadas, resultados e terapias prescritas. Cada medicamento é caracterizado por ser primário, secundário, concomitante ou interativo. Os casos também são associados a informações sobre o fabricante, a faixa etária do paciente e a fonte do relatório. Este esquema permite rastrear e analisar as relações entre medicamentos, suas reações e resultados de forma estruturada.

Começaremos criando uma conexão com o banco de dados instanciando um objeto 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)

Configurando o ambiente LLM

Há muitas opções para hospedar LLMs de código aberto como o Llama-3.1. Usaremos o catálogo de API da NVIDIA , que fornece microsserviços de inferência NVIDIA NIM e suporta chamadas de função para modelos Llama 3.1. Ao criar uma conta, você ganha 1.000 tokens, o que é mais do que suficiente para acompanhar. Você precisará criar uma chave de API e copiá-la para o notebook:

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

Usaremos o llama-3.1–70b porque a versão 8b tem alguns problemas com parâmetros opcionais em definições de funções.

O melhor dos microsserviços NVIDIA NIM é que você pode hospedá-los localmente facilmente se tiver preocupações com segurança ou outros problemas, então eles são facilmente trocáveis, e você só precisa adicionar um parâmetro de URL à configuração do LLM:

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

Definição de ferramenta

Configuraremos uma única ferramenta com quatro parâmetros opcionais. Construiremos uma declaração Cypher correspondente com base nesses parâmetros para recuperar as informações relevantes do gráfico de conhecimento. Nossa ferramenta será capaz de identificar os efeitos colaterais mais frequentes com base no medicamento de entrada, idade e fabricante do medicamento.

 @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

A função get_side_effects foi projetada para recuperar efeitos colaterais comuns de medicamentos de um gráfico de conhecimento usando critérios de pesquisa especificados. Ela aceita parâmetros opcionais para nome do medicamento, faixa etária do paciente e fabricante do medicamento para personalizar a pesquisa. Cada parâmetro tem uma descrição passada para um LLM junto com a descrição da função, permitindo que o LLM entenda como usá-los. A função então constrói uma consulta Cypher dinâmica com base nas entradas fornecidas, executa essa consulta no gráfico de conhecimento e retorna os dados de efeitos colaterais resultantes.

Vamos testar a função:

 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}, # ...

Nossa ferramenta primeiro mapeou o medicamento Lyrica mencionado na pergunta para os valores “['LYRICA', 'LYRICA CR']” no gráfico de conhecimento e, em seguida, executou uma instrução Cypher correspondente para encontrar os efeitos colaterais mais frequentes.

Agente LLM baseado em gráfico

A única coisa que resta a fazer é configurar um agente LLM que possa usar a ferramenta definida para responder a perguntas sobre os efeitos colaterais do medicamento.

A imagem descreve um usuário interagindo com um agente Llama 3.1 para perguntar sobre efeitos colaterais de medicamentos. O agente acessa uma ferramenta de efeitos colaterais que recupera informações de um gráfico de conhecimento para fornecer ao usuário os dados relevantes.

Começaremos definindo o modelo de prompt:

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

O modelo de prompt inclui a mensagem do sistema, histórico de bate-papo opcional e entrada do usuário. O agent_scratchpad é reservado para o LLM, pois às vezes ele precisa de várias etapas para responder à pergunta, como executar e recuperar informações de ferramentas.

A biblioteca LangChain facilita a adição de ferramentas ao LLM usando o método bind_tools:

 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 )

O agente processa a entrada por meio de transformações e manipuladores que formatam o histórico de bate-papo, aplicam o LLM com as ferramentas vinculadas e analisam a saída. Finalmente, o agente é configurado com um executor que gerencia o fluxo de execução, especifica os tipos de entrada e saída e inclui configurações de verbosidade para registro detalhado durante a execução.

Vamos testar o agente:

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

Resultados:

O LLM identificou que precisa usar a função get_side_effects com argumentos apropriados. A função então gera dinamicamente uma declaração Cypher, busca as informações relevantes e as retorna ao LLM para gerar a resposta final.

Resumo

Os recursos de chamada de função são uma adição poderosa aos modelos de código aberto como o Llama 3.1, permitindo interações mais estruturadas e controladas com fontes de dados e ferramentas externas. Além de apenas consultar documentos não estruturados, os agentes baseados em gráficos oferecem possibilidades interessantes para interagir com gráficos de conhecimento e dados estruturados. A facilidade de hospedar esses modelos usando plataformas como os microsserviços NVIDIA NIM os torna cada vez mais acessíveis.

Como sempre, o código está disponível no GitHub .

Para saber mais sobre este tópico, junte-se a nós no NODES 2024 em 7 de novembro, nossa conferência virtual gratuita para desenvolvedores sobre aplicativos inteligentes, gráficos de conhecimento e IA. Registre-se agora!