Créer un agent d’assistant au savoir sur Databricks avec LangGraph

Un agent personnalisé de base sur Databricks peut évoluer d’une simple appel à un LLM à un assistant de connaissance capable de rechercher dans des documents et de mémoriser les utilisateurs, chaque étape étant progressive. Cet agent final est développé à l’aide du framework langgraph. Il dispose d’un composant de recherche vectorielle pour retrouver du texte ainsi que d’un module de mémoire, lui permettant de se remémorer ce qu’a dit un utilisateur, tant au sein d’une même conversation que entre différentes conversations. Une fois en fonctionnement, vient l’évaluation : mesurer les performances pendant le développement, puis à nouveau en production.

La démarche est expliquée à travers une série de notebooks regroupés en dossiers correspondant aux bases fondamentales, au pipeline RAG et à la mémoire.

Préparation de l’espace de travail

Un agent complet comprend des index vectoriels, des outils et des journaux, qui nécessitent d’abord un espace organisé. Sur Databricks, la couche de gouvernance est Unity Catalog (UC), qui gère tous les actifs au sein d’un espace de noms à trois niveaux. L’exécution dans l’ordre des notebooks situés dans le dossier 00_foundations crée tous les actifs dont ont besoin les étapes suivantes, y compris une combinaison de données structurées et non structurées auxquelles l’agent se connectera.

La mémoire sera stockée dans Lakebase, une base de données auto-échelle compatible avec Postgres. Afin que les tables puissent être créées sans problème par la suite, le notebook de configuration crée également un projet Lakebase ainsi qu’une branche pour les conserver.

Un premier agent minimal

Viennent ensuite les notebooks situés dans 01_rag_pipeline, à commencer par 0_your_first_agent_on_databricks.

Chaque agent est basé sur un modèle de base. Databricks Model Serving énumère ceux qui sont disponibles, et cette liste peut varier en fonction de la configuration de votre espace de travail. Les points d’entrée sont compatibles avec le format API d’OpenAI, mais pour une intégration étroite avec LangGraph, la bibliothèque officielle databricks-langchain est le choix préférable.

Quelques mots sur LangGraph. Un agent en environnement de production n’est rarement constitué d’un seul prompt. Il contient généralement une logique de branchement qui dépend de la question posée et de la réponse de l’utilisateur, ainsi que des boucles permettant de consulter des systèmes externes et de reformuler les requêtes en fonction des réponses reçues. LangGraph organise tout cela en modélisant l’agent sous forme de graphe : les nœuds représentent les fonctions et les appels aux LLM, tandis que les arêtes déterminent la manière dont le contexte circule entre eux. Le graphe de départ est très simple.

On invoque l’agent et un suivi apparaît, listant la question d’entrée, la réponse du modèle ainsi que chaque étape logique intermédiaire. Cela provient de MLflow, qui s’intègre nativement à langgraph. Une seule appel à mlflow.langchain.autolog() en haut du notebook active ces journaux et diagnostics. MLflow fera l’objet d’une description plus détaillée ultérieurement.

Ce premier version se contente de transmettre la question de l’utilisateur au modèle. Un véritable agent ajoute un prompt système pour définir le ton, façonner la réponse et fournir des informations de contexte, comme l’objectif de l’assistant ou les questions auxquelles il doit refuser de répondre. Ainsi, la deuxième itération insère un prompt système dans le graphe.

Ajout de la récupération d’informations

L’un des types d’agents les plus courants est l’agent de récupération de connaissances, également appelé agent RAG. Il aide les gens à rechercher parmi une pile de documents et à y répondre aux questions. Pour conférer cette capacité à un agent, il a besoin d’un index de recherche vectorielle, essentiellement une base de données pour du texte non structuré.

Le notebook 01_building_a_doc_store_on_vector_search traite de la création d’un tel index. Avant qu’un index ne puisse être associé à un agent, il doit déjà exister. Cette partie se concentre sur les étapes fondamentales et omet l’ajustement et l’optimisation.

La création d’un index nécessite trois actions :

• Le parsing consiste à ouvrir un document et à en lire le contenu.
• Le chunking consiste à utiliser un outil de segmentation pour diviser le document en sections.
• L’embedding consiste à convertir chaque section en un vecteur sémantique stocké dans l’index.

Chacun présente ses propres subtilités, et des erreurs à un stade précoce détériorent silencieusement tout le processus ultérieur. Le tutoriel utilise systématiquement les paramètres par défaut. En environnement de production, différents documents peuvent nécessiter des analyseurs distincts pour extraire correctement le texte, et divers cas d’usage peuvent exiger des tailles de blocs différentes pour le texte, les tableaux et les figures, ainsi que des modèles d’incorporation différents.

Databricks fournit des outils pour faciliter cette tâche. Volumes propose un stockage d’objets natif du cloud, entièrement géré, pour les fichiers bruts. La fonction ai_parse_documents extrait le texte, les figures et les tableaux avec peu de configuration et s’adapte bien à différentes charges de travail. Les UDF de Spark permettent de diviser le texte extrait en blocs à différentes granularités, prêts à être incorporés par n’importe lequel des modèles d’incorporation fournis.

Lorsque les embeddings existent, la recherche de similarité nécessite un index de recherche vectorielle. Un nouveau stockage de données peut sembler impliquer du travail supplémentaire, mais Delta Sync s’en charge en embedding et en indexant automatiquement les blocs au fur et à mesure qu’ils sont ajoutés, modifiés ou supprimés. Il n’y a pas d’infrastructure nouvelle à gérer ni de tâche de synchronisation à entretenir entre la table des blocs et l’index.

Lors des deux premières étapes du notebook, les fichiers Markdown générés lors de la configuration sont ingérés et divisés en blocs superposés de 800 caractères, puis chargés dans Databricks Vector Search. La bibliothèque databricks_langchain comprend un outil prêt à l’emploi, VectorSearchRetrieverTool, qui relie l’index à l’agent langgraph.

Avec l’index créé et testé, la troisième étape redéfinit l’agent. Il ressemble beaucoup aux versions précédentes, avec une seule addition : un ToolNode qui contient une liste d’outils. Dans ce cas, cette liste ne comprend que le VectorSearchRetrieverTool.

Dans le flux résultant, une question est saisie et passe par le modèle de prompt. Le modèle prend connaissance du prompt ainsi que des outils et de leurs descriptions, puis décide s’il faut consulter l’index. Il peut le faire à plusieurs reprises et ne répond au utilisateur qu’une fois qu’il a trouvé ce qu’il cherchait.

C’est un agent RAG fonctionnel. Le graphe reste compact, avec le node d’outil comme seule addition, ce qui montre à quel point un cycle de récupération nécessite peu de structure. Cependant, un assistant à la connaissance a besoin d’un élément supplémentaire : de la mémoire.

Fournir une mémoire à l’agent

Jusqu’à présent, l’agent ne possède absolument aucune mémoire. Une question suivante qui fait référence à la réponse précédente échoue simplement, car le modèle n’a aucune idée de ce qui a été dit il y a un instant.

Afin que l’agent puisse tenir compte des échanges antérieurs, langgraph a besoin d’un checkpointer qui stocke l’historique de la conversation. Les checkpointers peuvent aller de simples fichiers texte à des bases de données complètes. En termes de scalabilité, de latence et de sécurité, une base de données est le choix préférable, et sur Databricks cela signifie Lakebase.

Il existe deux types de mémoire. La mémoire à court terme ne couvre que la session active en cours. La mémoire à long terme conserve des informations entre les sessions. Presque toutes les applications ont besoin de la première, tandis que le besoin de la seconde dépend du cas d’usage.

Mémoire à court terme

Le dossier 02_memory s’ouvre avec 01_short_term_memory, qui ajoute un point de contrôle. Les étapes 1 à 3 créent un StateGraph identique au précédent, et l’étape 4 montre à quel point une conversation devient difficile en l’absence de mémoire.

La création d’un point de contrôle lié à Lakebase restaure la continuité, et databricks_langchain propose une classe native CheckpointSaver pour l’instancier et le rattacher au graphique.

Un agent est un service partagé par de nombreux utilisateurs, donc les conversations enregistrées nécessitent un propriétaire. Le thread_id remplit cette fonction. Dans une application en production, l’interface utilisateur doit générer un thread_id unique par session active, afin que chaque utilisateur ne voie que son propre historique de chat.

Au terme de l’étape 6, le point de contrôle est en place et la conversation se déroule naturellement.

Mémoire à long terme

La mémoire à court terme disparaît à la fin de la session. Pour se souvenir d’un utilisateur et de ses préférences, l’agent a besoin d’une mémoire à long terme, et la décision de l’intégrer dépend du produit. Un assistant personnel à long terme devrait se souvenir des conversations précédentes, ainsi que de faits tels que le nom de l’utilisateur et de ses préférences, comme des réponses courtes et directes plutôt que des explications détaillées. Un service d’assistance sur un portail client peut préférer ne pas conserver ces détails, car cela signifierait assumer la charge de sécuriser, chiffrer et gérer des données personnelles sur une longue période.

La mémoire à court terme ne nécessitait qu’un CheckpointSaver. La mémoire à long terme exige une certaine conception. Enregistrer chaque conversation pour l’éternité n’est généralement pas souhaité. Ce qui vaut la peine d’être conservé, ce sont les faits, tels qu’un nom ou une ville natale, ainsi que les préférences.

Le cahier de notes 02_long_term_memory décrit les mécanismes en détail. Les étapes 1 à 4 reconstruisent l’agent RAG avec une mémoire à court terme, tandis que l’étape 5 montre que les préférences disparaissent lorsque la session et le thread_id changent.

L’étape 6 introduit un objet DatabricksStore qui persiste les données dans Lakebase. Cette partie est simple ; la vraie question de conception réside dans le choix des données à persister. L’étape 7 adopte une approche basique : quelques faits codés en dur sont écrits dans le stockage.

Pour les utiliser, l’étape 8 recrée l’agent. Auparavant, le nœud de l’agent n’était qu’une fonction simple qui appelait invoke sur le modèle de base. Désormais, il s’agit d’une fonction plus complexe : elle prend un ID d’utilisateur (codé en dur pour la démonstration), récupère les faits correspondants de DatabricksStore et les concatène dans le contexte envoyé au modèle. Ce schéma nécessite clairement une conception minutieuse pour être fiable.

Une approche plus flexible permet au modèle lui-même de décider quand enregistrer, supprimer et récupérer des mémoires. L’étape 12 met en œuvre une version de base : trois nouveaux outils interagissent avec le module mémoire, et l’instruction système est modifiée pour informer le modèle à leur sujet et encourager leur utilisation.

Le résultat est un assistant complet, dont chaque fonctionnalité peut être testée indépendamment, car elle a été intégrée via un notebook distinct ainsi qu’une modification séparée du graphe. Il recherche des documents à l’aide d’un index vectoriel, suit la conversation en cours grâce à une mémoire à court terme et conserve des détails sur les utilisateurs entre différentes sessions grâce à une mémoire à long terme.

Quelles sont les prochaines étapes ?

Le parcours a évolué, allant d’un bot qui transmettait directement les questions à un LLM à un système doté d’un modèle de prompt, de mémoires à court et long terme ainsi que de recherche vectorielle pour le contexte. Les travaux ultérieurs visent à le mettre en production : évaluer la qualité de l’agent, y ajouter davantage de fonctionnalités et le déployer en tant que service pour les utilisateurs finaux.