Crear un agente asistente de conocimiento en Databricks con LangGraph

Un agente personalizado básico en Databricks puede evolucionar de una simple llamada a un LLM a un asistente de conocimiento que busca en documentos y recuerda a los usuarios, siendo cada paso sencillo. El agente final se construye con el framework langgraph. Cuenta con un componente de búsqueda vectorial para localizar texto y un módulo de memoria, lo que le permite recordar lo que dijo un usuario final tanto dentro de una conversación como en conversaciones separadas. Una vez que está en funcionamiento, sigue la evaluación: se mide el rendimiento durante el desarrollo y nuevamente en producción.

La guía sigue un conjunto de cuadernos de notas, agrupados en carpetas correspondientes a los fundamentos, al pipeline RAG y a la memoria.

Preparación del espacio de trabajo

Un agente completo incluye índices vectoriales, herramientas y registros, y estos necesitan primero un lugar organizado donde almacenarse. En Databricks, la capa de gobernanza es Unity Catalog (UC), que gestiona todos los activos bajo un espacio de nombres de tres niveles. Al ejecutar los cuadernos de notas en el folder 00_foundations en orden, se crean todos los activos de los que dependen los pasos posteriores, incluyendo una combinación de datos estructurados y no estructurados a los que se conectará el agente.

La memoria se almacenará en Lakebase, una base de datos compatible con Postgres y con escalado automático. Para que las tablas puedan crearse sin problemas más adelante, el cuaderno de notas de configuración también crea un proyecto en Lakebase y una rama para guardarlas.

Un primer agente mínimo

A continuación se encuentran los cuadernos de notas en 01_rag_pipeline, comenzando con 0_your_first_agent_on_databricks.

Cada agente se construye sobre un modelo base. Databricks Model Serving enumera los disponibles, y la lista puede variar según cómo esté configurado su espacio de trabajo. Los puntos finales son compatibles con el formato API de OpenAI, pero para una integración más estrecha con LangGraph, la biblioteca oficial databricks-langchain es la mejor opción.

Un breve comentario sobre LangGraph. Un agente en entorno de producción rara vez se basa en un único prompt; generalmente contiene lógica ramificada que depende de la pregunta y de cómo responde el usuario, además de bucles para consultar sistemas externos y reformular las consultas según la respuesta recibida. LangGraph organiza todo esto modelando al agente como un grafo: los nodos albergan funciones y llamadas a modelos de lenguaje grande, mientras que las aristas determinan cómo se transfiere el contexto entre ellos. El grafo inicial es muy simple.

Se invoca al agente y aparece un registro que enumera la pregunta de entrada, la respuesta del modelo y cada paso lógico intermedio. Eso proviene de MLflow, que se integra de forma nativa con langgraph. Una sola llamada a mlflow.langchain.autolog() al principio del cuaderno de notas activa estos registros y diagnósticos. MLflow será tratado con más detalle más adelante.

Esta primera versión simplemente pasa la pregunta del usuario al modelo. Un agente real añade un prompt de sistema para establecer el tono, dar forma a la respuesta y proporcionar contexto adicional, como para qué sirve el asistente y qué preguntas debe negarse a responder. Por lo tanto, en la segunda iteración se inserta un prompt de sistema en el grafo.

Agregar recuperación de información

Uno de los tipos de agentes más comunes es el agente de recuperación de conocimiento, también conocido como agente RAG. Ayuda a las personas a buscar entre una gran cantidad de documentos y a responder preguntas sobre ellos. Para dotar a un agente de esta capacidad, necesita un índice de búsqueda vectorial, esencialmente una base de datos para texto no estructurado.

El cuaderno de notas 01_building_a_doc_store_on_vector_search trata sobre cómo crear uno. Antes de que se pueda asociar un índice al agente, este debe existir ya. Esta sección se centra en los pasos básicos y omite el ajuste y la optimización.

Crear un índice implica tres acciones:

• Analizado significa abrir un documento y leer su contenido.
• Fragmentación significa utilizar un divisor para dividir el documento en secciones.
• Incrustación significa convertir cada fragmento en un vector semántico almacenado en el índice.

Cada uno tiene sus propias sutilezas, y los errores en una etapa inicial degradan silenciosamente todo lo que sigue. El tutorial utiliza valores predeterminados en todo momento. En entornos de producción, diferentes documentos pueden requerir parsers distintos para extraer el texto correctamente, y diferentes casos de uso pueden necesitar tamaños de bloque diferentes para el texto, las tablas y las figuras, además de modelos de incrustación variados.

Databricks proporciona herramientas para facilitar esto. Volumes ofrece un almacén de objetos nativo de la nube y completamente gestionado para archivos en bruto. La función ai_parse_documents extrae texto, figuras y tablas con poca configuración y se escala bien. Las UDF de Spark pueden dividir el texto extraído en bloques con diferentes niveles de detalle, listos para ser incrustados por cualquiera de los modelos de incrustación proporcionados.

Una vez que existen los embeddings, la búsqueda de similitud requiere un índice de búsqueda vectorial. Un nuevo almacén de datos parece implicar trabajo adicional, pero Delta Sync se encarga de ello al incrustar e indexar automáticamente los fragmentos a medida que se añaden, modifican o eliminan. No hay infraestructura nueva que administrar ni tareas de sincronización que mantener entre la tabla de fragmentos y el índice.

En los dos primeros pasos del cuaderno de notas, los archivos en formato markdown generados durante la configuración se ingieren y se dividen en fragmentos superpuestos de 800 caracteres, para luego cargarse en Databricks Vector Search. La biblioteca databricks_langchain incluye una herramienta VectorSearchRetrieverTool ya preparada que conecta el índice con el agente langgraph.

Con el índice listo y probado, la tercera fase redefine al agente. Se parece mucho a las versiones anteriores, con una adición: un ToolNode que recibe una lista de herramientas. En este caso, dicha lista contiene únicamente el VectorSearchRetrieverTool.

En el flujo resultante, una pregunta entra y pasa por el plantilla de prompt. El modelo ve el prompt junto con las herramientas y sus descripciones, y decide si debe consultar el índice. Puede hacerlo repetidamente, y responde al usuario solo una vez que ha encontrado lo que buscaba.

Ese es un agente RAG funcional. El grafo sigue siendo compacto, con el nodo de herramienta como única adición, lo que demuestra cuán poca estructura necesita un bucle de recuperación. Sin embargo, un asistente de conocimiento necesita algo más: memoria.

Proporcionar memoria al agente

Hasta ahora, el agente no tiene ninguna memoria en absoluto. Una pregunta posterior que hace referencia a la respuesta anterior simplemente falla, porque el modelo no tiene idea de lo que se dijo hace un momento.

Para que el agente tenga en cuenta las interacciones anteriores, langgraph necesita un checkpoint que almacene el historial de la conversación. Los checkpoints pueden ser desde archivos de texto simples hasta bases de datos completas. Por razones de escalabilidad, latencia y seguridad, una base de datos es la opción más adecuada, y en Databricks eso significa Lakebase.

La memoria existe en dos tipos. La memoria a corto plazo abarca solo la sesión activa actual. La memoria a largo plazo conserva datos a través de distintas sesiones. Casi todas las aplicaciones necesitan la primera, mientras que la necesidad de la segunda depende del caso de uso.

Memoria a corto plazo

La carpeta 02_memory se abre junto con 01_short_term_memory, lo que agrega un punto de verificación. Los pasos del 1 al 3 crean un StateGraph idéntico al anterior, y el paso 4 muestra cuán complicada se vuelve una conversación sin memoria.

Crear un punto de verificación vinculado a Lakebase restablece la continuidad, y databricks_langchain ofrece una clase nativa CheckpointSaver para instanciarla y adjuntarla al grafo.

Un agente es un servicio compartido por muchos usuarios, por lo que las conversaciones guardadas necesitan un propietario. El thread_id cumple esa función. En una aplicación en producción, el front end debe generar un thread_id único por sesión activa, de modo que cada usuario vea solo su propio historial de chat.

Para el paso 6, el punto de verificación ya está en su lugar y la conversación fluye de manera natural.

Memoria a largo plazo

La memoria a corto plazo desaparece cuando finaliza la sesión. Para recordar al usuario y sus preferencias, el agente necesita memoria a largo plazo, y la decisión de incluirla depende del producto. Un asistente personal de larga duración debería recordar conversaciones anteriores, datos como el nombre del usuario y preferencias como respuestas breves y directas frente a explicaciones detalladas. Un servicio de ayuda en un portal de clientes podría ser mejor si no conserva esos detalles, ya que eso implicaría asumir la carga de proteger, encriptar y gestionar datos personales durante mucho tiempo.

La memoria a corto plazo solo requería un CheckpointSaver. La memoria a largo plazo necesita algún tipo de diseño. Por lo general, no se desea registrar cada conversación para el futuro. Lo que vale la pena conservar son datos como un nombre o ciudad natal, además de las preferencias.

El cuaderno de notas 02_long_term_memory explica los mecanismos involucrados. Los pasos del 1 al 4 reconstruyen el agente RAG con memoria a corto plazo, y el paso 5 muestra que las preferencias desaparecen cuando cambia la sesión o el thread_id.

El paso 6 introduce un objeto DatabricksStore que guarda los datos en Lakebase. Esa parte es sencilla; decidir qué datos persistir es la verdadera pregunta de diseño. El paso 7 adopta un enfoque básico: se escriben unos pocos datos fijos directamente en el almacén.

Para utilizarlos, el paso 8 recrea el agente. Anteriormente, el nodo del agente era una función simple que llamaba a invoke en el modelo base. Ahora es una función más compleja: toma un ID de usuario (fijado para la demostración), obtiene los datos de ese usuario desde DatabricksStore y los concatena en el contexto que se envía al modelo. Este patrón requiere claramente un diseño cuidadoso para funcionar adecuadamente.

Un enfoque más flexible permite que el modelo mismo decida cuándo guardar, eliminar y recuperar memorias. El paso 12 implementa una versión básica: tres nuevas herramientas interactúan con el módulo de memoria, y el prompt del sistema cambia para informar al modelo sobre ellas e incentivar su uso.

El resultado es un asistente completo, y cada capacidad puede probarse por separado ya que se implementó a través de un cuaderno de notas distinto y de un cambio independiente en el grafo. Busca documentos mediante un índice vectorial, sigue la conversación actual con memoria a corto plazo y recuerda detalles sobre los usuarios entre sesiones con memoria a largo plazo.

Hacia dónde va ahora

El camino evolucionó desde un bot que enviaba las preguntas directamente a un LLM, hasta uno que cuenta con un modelo de prompt, memoria a corto y largo plazo y búsqueda vectorial para obtener contexto. El trabajo posterior consiste en ponerlo en producción: evaluar la calidad del agente, añadir más funcionalidades y desplegarlo como servicio para los usuarios finales.