Mejora de RAG con gráficos de conocimiento: integración de Llama 3.1, NVIDIA NIM y LangChain para IA dinámica

Mientras que la mayoría de la gente se centra en la generación aumentada por recuperación (RAG) en lugar de texto no estructurado, como documentos o documentación de la empresa, yo soy bastante optimista con respecto a los sistemas de recuperación por encima de la información estructurada, en particular los gráficos de conocimiento . Ha habido mucho entusiasmo en torno a GraphRAG, en concreto la implementación de Microsoft. Sin embargo, en su implementación, los datos de entrada son texto no estructurado en forma de documentos, que se transforma en un gráfico de conocimiento mediante un modelo de lenguaje grande (LLM).

En esta entrada del blog, mostraremos cómo implementar un recuperador sobre un gráfico de conocimiento que contiene información estructurada del Sistema de Notificación de Eventos Adversos de la FDA (FAERS) , que ofrece información sobre eventos adversos de medicamentos. Si alguna vez ha trabajado con gráficos de conocimiento y recuperación, su primera idea podría ser usar un LLM para generar consultas de base de datos para recuperar información relevante de un gráfico de conocimiento para responder una pregunta determinada. Sin embargo, la generación de consultas de base de datos mediante LLM aún está evolucionando y es posible que aún no ofrezca la solución más consistente o sólida. Entonces, ¿cuáles son las alternativas viables en este momento?

En mi opinión, la mejor solución actual es la generación de consultas dinámicas. En lugar de depender completamente de un LLM para generar la consulta completa, este método emplea una capa lógica que genera de manera determinista una consulta de base de datos a partir de parámetros de entrada predefinidos. Esta solución se puede implementar utilizando un LLM con soporte para llamadas de funciones. La ventaja de utilizar una característica de llamadas de funciones radica en la capacidad de definir para un LLM cómo debe preparar una entrada estructurada para una función. Este enfoque garantiza que el proceso de generación de consultas esté controlado y sea consistente, al tiempo que permite la flexibilidad de entrada del usuario.

La imagen ilustra un proceso de comprensión de la pregunta de un usuario para recuperar información específica. El flujo consta de tres pasos principales:

1. Un usuario hace una pregunta sobre los efectos secundarios comunes del medicamento Lyrica para personas menores de 35 años.

   
   

2. El LLM decide qué función llamar y los parámetros necesarios. En este ejemplo, eligió una función denominada side_effects con parámetros que incluyen el medicamento Lyrica y una edad máxima de 35 años.

   
   

3. La función y los parámetros identificados se utilizan para generar de forma determinista y dinámica una declaración de consulta de base de datos (Cypher) para recuperar información relevante.

La compatibilidad con llamadas de funciones es fundamental para los casos de uso avanzados de LLM, como permitir que los LLM utilicen varios recuperadores según la intención del usuario o crear flujos de múltiples agentes. He escrito algunos artículos en los que se utilizan LLM comerciales con compatibilidad nativa con llamadas de funciones. Sin embargo, utilizaremos Llama-3.1, un LLM de código abierto superior con compatibilidad nativa con llamadas de funciones, lanzado recientemente.

El código está disponible en GitHub .

Configuración del gráfico de conocimiento

Usaremos Neo4j, que es una base de datos gráfica nativa, para almacenar la información de eventos adversos. Puedes configurar un proyecto Sandbox gratuito en la nube que viene con FAERS precargado siguiendo este enlace .

La instancia de base de datos instanciada tiene un gráfico con el siguiente esquema.

El esquema se centra en el nodo Caso, que vincula varios aspectos de un informe de seguridad de medicamentos, incluidos los medicamentos involucrados, las reacciones experimentadas, los resultados y las terapias prescritas. Cada medicamento se caracteriza por si es primario, secundario, concomitante o interactuante. Los casos también se asocian con información sobre el fabricante, el grupo de edad del paciente y la fuente del informe. Este esquema permite rastrear y analizar las relaciones entre los medicamentos, sus reacciones y los resultados de una manera estructurada.

Comenzaremos creando una conexión a la base de datos instanciando un 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)

Configuración del entorno LLM

Existen muchas opciones para alojar LLM de código abierto como Llama-3.1. Usaremos el catálogo de API de NVIDIA , que proporciona microservicios de inferencia de NVIDIA NIM y admite llamadas de funciones para modelos de Llama 3.1. Cuando creas una cuenta, obtienes 1000 tokens, que son más que suficientes para seguir adelante. Deberás crear una clave de API y copiarla en el cuaderno:

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

Usaremos llama-3.1–70b porque la versión 8b tiene algunos problemas con los parámetros opcionales en las definiciones de funciones.

Lo bueno de los microservicios NVIDIA NIM es que puedes alojarlos fácilmente de forma local si tienes preocupaciones de seguridad o de otro tipo, por lo que son fácilmente intercambiables y solo necesitas agregar un parámetro de URL a la configuración de 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" )

Definición de herramienta

Configuraremos una única herramienta con cuatro parámetros opcionales. Construiremos una declaración Cypher correspondiente en función de esos parámetros para recuperar la información relevante del gráfico de conocimiento. Nuestra herramienta podrá identificar los efectos secundarios más frecuentes en función del medicamento de entrada, la edad y el fabricante del 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

La función get_side_effects está diseñada para recuperar los efectos secundarios comunes de los medicamentos de un gráfico de conocimiento utilizando criterios de búsqueda específicos. Acepta parámetros opcionales para el nombre del medicamento, el rango de edad del paciente y el fabricante del medicamento para personalizar la búsqueda. Cada parámetro tiene una descripción que se pasa a un LLM junto con la descripción de la función, lo que permite que el LLM comprenda cómo usarlos. Luego, la función construye una consulta Cypher dinámica basada en las entradas proporcionadas, ejecuta esta consulta contra el gráfico de conocimiento y devuelve los datos de efectos secundarios resultantes.

Probemos la función:

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

Nuestra herramienta primero asignó el medicamento Lyrica mencionado en la pregunta a los valores “['LYRICA', 'LYRICA CR']” en el gráfico de conocimiento, luego ejecutó una declaración Cypher correspondiente para encontrar los efectos secundarios más frecuentes.

Agente LLM basado en gráficos

Lo único que queda por hacer es configurar un agente LLM que pueda usar la herramienta definida para responder preguntas sobre los efectos secundarios del medicamento.

La imagen muestra a un usuario interactuando con un agente de Llama 3.1 para consultar sobre los efectos secundarios de un medicamento. El agente accede a una herramienta de efectos secundarios que recupera información de un gráfico de conocimiento para proporcionar al usuario los datos pertinentes.

Comenzaremos definiendo la plantilla del mensaje:

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

La plantilla de solicitud incluye el mensaje del sistema, el historial de chat opcional y la entrada del usuario. El agent_scratchpad está reservado para el LLM, ya que a veces necesita varios pasos para responder la pregunta, como ejecutar y recuperar información de las herramientas.

La biblioteca LangChain facilita la adición de herramientas al LLM mediante el 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 )

El agente procesa la entrada a través de transformaciones y controladores que dan formato al historial de chat, aplican el LLM con las herramientas vinculadas y analizan la salida. Por último, el agente está configurado con un ejecutor que administra el flujo de ejecución, especifica los tipos de entrada y salida e incluye configuraciones de verbosidad para un registro detallado durante la ejecución.

Probemos el agente:

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

Resultados:

El LLM identificó que necesita utilizar la función get_side_effects con los argumentos adecuados. Luego, la función genera dinámicamente una declaración Cypher, obtiene la información relevante y la devuelve al LLM para generar la respuesta final.

Resumen

Las capacidades de llamada de funciones son una poderosa incorporación a los modelos de código abierto como Llama 3.1, que permiten interacciones más estructuradas y controladas con herramientas y fuentes de datos externas. Más allá de simplemente consultar documentos no estructurados, los agentes basados en gráficos ofrecen interesantes posibilidades para interactuar con gráficos de conocimiento y datos estructurados. La facilidad de alojar estos modelos mediante plataformas como los microservicios NVIDIA NIM los hace cada vez más accesibles.

Como siempre, el código está disponible en GitHub .

Para obtener más información sobre este tema, únase a nosotros en NODES 2024 el 7 de noviembre, nuestra conferencia virtual gratuita para desarrolladores sobre aplicaciones inteligentes, gráficos de conocimiento e IA. ¡Regístrese ahora!