Votre agent RAG oublie tout après un seul message – Voici comment vous l’avez résolu avec Databricks…
Guide pas à pas pour comprendre pourquoi votre agent RAG oublie tout après un seul message – Voici comment vous l’avez résolu avec Databricks… : contrats, vérifications et emplacements de code prêts à l’emploi pour les équipes.
Ce guide reconstitue le parcours allant des matières premières à un système fonctionnel pour le cas suivant : Votre agent RAG oublie tout après un seul message – voici comment je l’ai résolu avec Databricks Lakebase. L’accent est mis sur des étapes opérationnelles, des vérifications explicites, ainsi que du code que vous pouvez intégrer directement dans un dépôt sans devoir deviner l’intention derrière les actions. Afin d’obtenir une vue d’ensemble, définissez les entrées, le responsable de chaque étape et les critères d’achèvement avant de modifier le code. Les opérateurs doivent pouvoir relancer l’étape à partir d’un point de contrôle connu, sans avoir à deviner l’état caché du système. Considérez cette étape comme un contrat entre les entrées et les sorties validées. Donnez des noms aux artefacts, définissez des vérifications de succès et refusez toute complétion partielle silencieuse.
L’architecture
Lorsque vous travaillez sur The Architecture, notez d’abord le contrat : les entrées requises, le signal de succès et ce qui se passe en cas d’échec partiel. Cette liste de contrôle permet de rester honnête lors des modifications ultérieures du code. Enregistrez les temps d’exécution ainsi que le coût des tokens ou des requêtes à côté des résultats fonctionnels. Une visibilité précoce des coûts évite les factures inattendues lorsque le processus passe de l’environnement de démonstration à des environnements partagés. Mesurez le taux de rappel sur un ensemble fixe de questions avant d’ajuster les prompts. Un changement fréquent des prompts ne résout que rarement un système de récupération insuffisant.
PDFs (text + images + diagrams)
↓ ai_parse_document() (Version 2.0)
Parsed elements (text, tables, figure descriptions)
↓ RecursiveCharacterTextSplitter
Chunks table (Delta, with Change Data Feed)
↓ Delta Sync + GTE-Large
Vector Search Index
↓ VectorSearchRetrieverTool
LangChain Agent + PostgresSaver
↓ ↓
LLM (Foundation Model) Lakebase (conversation memory)
↓
Model Serving Endpoint (MLflow ResponsesAgent)
Étape 1 : Parser les documents avec ai_parse_document()
Lorsque vous travaillez sur l’Étape 1 : Analyser les documents avec ai_parse_document(), notez d’abord le contrat : les entrées requises, le signal de succès, ainsi que ce qui se passe en cas d’échec partiel. Cette liste de contrôle permet de rester honnête lors des modifications ultérieures du code.
Conservez la configuration en dehors du code de l’application. Les fichiers d’environnement, les bases de données secrètes et les indicateurs fonctionnels doivent être regroupés en un seul endroit que les administrateurs peuvent auditer sans devoir lire l’ensemble du système.
Mesurez le taux de rappel sur un ensemble fixe de questions avant d’ajuster les prompts. Un changement fréquent des prompts ne résout que rarement un système de récupération insuffisant.
from pyspark.sql.functions import expr
# Volume path where PDFs are stored
docs_path = "/Volumes/<YOUR_CATALOG>/<YOUR_SCHEMA>/source_docs/"
# Read all files as binary
docs_df = spark.read.format("binaryFile").load(docs_path)
# Parse each document using ai_parse_document v2.0
parsed_df = docs_df.withColumn(
"parsed_content",
expr(f"""ai_parse_document(content, map(
"version", "2.0",
"imageOutputPath", "{docs_path}/parsed_images/"
))""")
)
# Drop binary content (too large to display)
parsed_df = parsed_df.drop("content")
# Save to Delta table
output_table = "<YOUR_CATALOG>.<YOUR_SCHEMA>.docs_parsed"
parsed_df.write.format("delta").mode("overwrite").saveAsTable(output_table)
print(f"✅ Parsed results saved to: {output_table}")
Étape 2 : Nettoyer, transformer et découper
Lors de l’exécution de l’Étape 2 : Nettoyer, transformer et découper, notez d’abord le contrat : les entrées requises, le signal de succès et ce qui se passe en cas d’échec partiel. Cette liste de contrôle permet de garantir l’honnêteté des modifications ultérieures du code. Documentez ensemble le parcours normal et le parcours de récupération. Les tentatives répétées, les contrôles humains et la gestion des messages non traités font partie intégrante du produit, et non d’améliorations ultérieures. Mesurez le taux de rappel sur un ensemble fixe de questions avant d’ajuster les prompts. Le remplacement fréquent des prompts résout rarement un système de récupération insuffisant. Lors de l’exécution de l’Étape 2 : Nettoyer, transformer et découper, notez d’abord le contrat : les entrées requises, le signal de succès et ce qui se passe en cas d’échec partiel. Cette liste de contrôle permet de garantir l’honnêteté des modifications ultérieures du code. Considérez cette étape comme un contrat entre les entrées et les sorties validées. Donnez des noms aux artefacts, définez des contrôles de succès et refusez les complétions partielles silencieuses.
Extraction rapide de texte brut
L’extraction rapide de texte brut fonctionne le mieux lorsqu’elle est considérée comme une surface mesurable. Capturez un transcript parfait, un cas d’échec et la note de réversion avant d’élargir le périmètre. Enregistrez les temps d’exécution ainsi que le coût des tokens ou des requêtes à côté des résultats fonctionnels. Une visibilité précoce des coûts évite les factures inattendues lorsque le processus passe d’un environnement de démonstration à des environnements partagés. Séparez la politique de segmentation de la politique de récupération : modifier l’une ne doit pas obliger à réécrire l’autre lorsque les métriques de qualité changent.
from pyspark.sql import functions as F
# Convert VARIANT to JSON string, then extract text content
safe_json_col = F.coalesce(
F.to_json(F.col("parsed_content")),
F.col("parsed_content").cast("string")
)
plain_text_df = parsed_df.withColumn(
"plain_text",
extract_contents_udf()(safe_json_col) # Custom UDF to join text elements
)
Segmenter avec LangChain
Le chunking avec LangChain fonctionne le mieux lorsqu’il est considéré comme une surface mesurable. Capturez un transcript parfait, un cas d’échec et la note de rollback avant d’élargir le périmètre. Conservez les configurations en dehors du code de l’application. Les fichiers d’environnement, les bases de données secrètes et les flags fonctionnels doivent être regroupés en un seul endroit que les administrateurs peuvent auditer sans devoir lire l’ensemble du système. Séparez la politique de chunking de la politique de récupération. Modifier l’une ne doit pas obliger à réécrire l’autre lorsque les métriques de qualité évoluent.
from langchain_text_splitters import RecursiveCharacterTextSplitter
from pyspark.sql.types import StructType, StructField, StringType
import pandas as pd
CHUNK_SIZE = 2000
CHUNK_OVERLAP = 200
splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP,
separators=["\n== page ==\n", "== page ==", "\n\n", "\n", " ", ""]
)
schema = StructType([
StructField("path", StringType(), True),
StructField("chunk", StringType(), True),
])
def split_rows(iterator):
for pdf in iterator:
out = []
for _, row in pdf.iterrows():
path, text = row["path"], row["plain_text"]
if isinstance(text, str) and text.strip():
for c in splitter.split_text(text):
if c and c.strip():
out.append((path, c))
yield pd.DataFrame(out, columns=["path", "chunk"])
df_chunks = (
plain_text_df.select("path", "plain_text")
.mapInPandas(split_rows, schema=schema)
)
# Add unique IDs and save
df_chunks = df_chunks.withColumn("id", F.monotonically_increasing_id())
chunked_table = "<YOUR_CATALOG>.<YOUR_SCHEMA>.docs_chunked"
df_chunks.write.format("delta") \
.mode("overwrite") \
.option("mergeSchema", "true") \
.saveAsTable(chunked_table)
Étape 3 : Construire une recherche vectorielle
Étape 3 : La mise en place de la recherche vectorielle fonctionne le mieux lorsqu’elle est considérée comme une surface mesurable. Capturez un exemple réussi, un cas d’échec ainsi que la note de réversion avant d’élargir le périmètre. Documentez conjointement le parcours optimal et celui de récupération. Les tentatives répétées, les contrôles humains et la gestion des messages non livrés font partie intégrante du produit, et non d’une amélioration ultérieure. Séparez la politique de segmentation des données de la politique de récupération ; modifier l’une ne doit pas obliger à réécrire l’autre lorsque les métriques de qualité évoluent.
Activer l’alimentation en données de modification
Activer le flux de données de modification fonctionne le mieux lorsqu’il est considéré comme une entité mesurable. Capturez un exemplaire idéal, un cas d’échec et la note de réversion avant d’élargir le périmètre. Préférez des unités petites et testables aux scripts complexes. Lorsqu’une étape échoue, l’échec doit pointer vers une seule responsabilité plutôt que vers un processus embrouillé. Séparez la politique de segmentation de la politique de récupération. Modifier l’une ne doit pas obliger à réécrire l’autre lorsque les métriques de qualité changent.
ALTER TABLE <YOUR_CATALOG>.<YOUR_SCHEMA>.docs_chunked
SET TBLPROPERTIES (delta.enableChangeDataFeed = true);
Créer l’index de synchronisation Delta
La création de l’index Delta Sync fonctionne le mieux lorsqu’elle est considérée comme une surface mesurable. Capturez un transcript idéal, un cas d’échec et la note de réversion avant d’élargir le périmètre. Considérez cette étape comme un contrat entre les entrées et les sorties validées. Nommez les artefacts, définites des critères de succès et refusez toute mise en œuvre partielle silencieuse. Séparez la politique de segmentation des données de la politique de récupération. Modifier l’une ne doit pas obliger à réécrire l’autre lorsque les métriques de qualité évoluent.
from databricks.vector_search.client import VectorSearchClient
vsc = VectorSearchClient(disable_notice=True)
index_name = "<YOUR_CATALOG>.<YOUR_SCHEMA>.docs_chunked_index"
vsc.create_delta_sync_index_and_wait(
endpoint_name="<YOUR_VS_ENDPOINT>",
index_name=index_name,
source_table_name="<YOUR_CATALOG>.<YOUR_SCHEMA>.docs_chunked",
primary_key="id",
embedding_source_column="chunk",
embedding_model_endpoint_name="databricks-gte-large-en",
pipeline_type="TRIGGERED",
)
Test de récupération
Le test de récupération fonctionne le mieux lorsqu’il est considéré comme une surface mesurable. Capturez un exemple réussi, un cas d’échec ainsi que la note de réversion avant d’élargir le périmètre. Enregistrez les temps d’exécution ainsi que le coût des tokens ou des requêtes à côté des résultats fonctionnels. Une visibilité précoce des coûts permet d’éviter des factures inattendues lorsque le système passe de l’environnement de démonstration à des environnements partagés. Séparez la politique de segmentation de la politique de récupération : modifier l’une ne doit pas obliger à réécrire l’autre lorsque les métriques de qualité évoluent.
index = vsc.get_index(index_name=index_name)
results = index.similarity_search(
query_text="How does the system prevent overheating?",
columns=["path", "chunk"],
num_results=5,
)
display(results)
Étape 4 : Configurer Lakebase pour la mémoire de conversation
Étape 4 : Configurer Lakebase pour la mémoire de conversation fonctionne le mieux lorsqu’elle est considérée comme une surface mesurable. Capturez un transcript parfait, un cas d’échec et la note de réversion avant d’élargir le périmètre. Conservez la configuration en dehors du code de l’application. Les fichiers d’environnement, les stocks de secrets et les flags fonctionnels doivent être regroupés en un seul endroit que les opérateurs peuvent auditer sans devoir lire l’ensemble du système. Séparez la politique de segmentation des données de la politique de récupération. Modifier l’une ne doit pas obliger à réécrire l’autre lorsque les métriques de qualité évoluent.
Mettre en place un projet d’auto-échelle pour Lakebase
La mise en place d’un projet de mise à l’échelle automatique Lakebase fonctionne le mieux lorsqu’il est considéré comme une entité mesurable. Capturez un exemple idéal, un cas d’échec et des notes de réversion avant d’élargir le périmètre. Documentez à la fois le parcours normal et celui de récupération. Les tentatives répétées, les contrôles humains et le traitement des messages non livrés font partie intégrante du produit, et non d’une étape de finition ultérieure. Séparez la politique de segmentation de la politique de récupération. Modifier l’une ne doit pas obliger à réécrire l’autre lorsque les métriques de qualité évoluent. La mise en place d’un projet de mise à l’échelle automatique Lakebase fonctionne le mieux lorsqu’il est considéré comme une entité mesurable. Capturez un exemple idéal, un cas d’échec et des notes de réversion avant d’élargir le périmètre. Considérez cette étape comme un contrat entre les entrées et les sorties validées. Donnez des noms aux artefacts, définez des critères de succès et refusez toute complétion partielle silencieuse.
Récupérer les détails de connexion de manière programmée
Pour obtenir les détails de la connexion de manière programmée, définitz les entrées, le responsable de l’étape et les critères d’arrêt avant de modifier le code. Les opérateurs doivent pouvoir relancer l’étape à partir d’un point de contrôle connu sans deviner l’état caché. Enregistrez les temps d’exécution ainsi que le coût des tokens ou des requêtes à côté des résultats fonctionnels. Une visibilité précoce du coût permet d’éviter des factures inattendues lorsque le parcours passe de l’environnement de démonstration à des environnements partagés. Citez les passages qui ont réellement servi de base à la réponse. Sans citations, les opérateurs ne peuvent pas distinguer une hallucination d’un manque dans l’indexation.
from databricks.sdk import WorkspaceClient
w = WorkspaceClient()
project_id = "<YOUR_PROJECT_NAME>"
# Get branch and endpoint
branches = list(w.postgres.list_branches(parent=f"projects/{project_id}"))
branch_name = branches[0].name
endpoints = list(w.postgres.list_endpoints(parent=branch_name))
ep = endpoints[0]
HOST = ep.status.hosts.host
ENDPOINT = ep.name
USERNAME = "<YOUR_DATABRICKS_EMAIL>"
print(f"Host: {HOST}")
print(f"Endpoint: {ENDPOINT}")
Tester la connexion
Pour tester la connexion, définissez les entrées, le responsable de l’étape et les critères de fin avant de modifier le code. Les opérateurs doivent pouvoir relancer l’étape à partir d’un point de contrôle connu sans deviner l’état caché. Conservez la configuration en dehors du code de l’application. Les fichiers d’environnement, les bases de données secrètes et les indicateurs fonctionnels doivent être regroupés en un seul endroit que les opérateurs peuvent auditer sans avoir à lire l’ensemble du système. Citez les passages qui justifient réellement la réponse. Sans citations, les opérateurs ne peuvent pas distinguer une hallucination d’un manque d’indexation.
import psycopg2
cred = w.postgres.generate_database_credential(endpoint=ENDPOINT)
conn = psycopg2.connect(
host=HOST,
dbname="databricks_postgres",
user=USERNAME,
password=cred.token,
port=5432,
sslmode="require"
)
with conn.cursor() as cur:
cur.execute("SELECT version()")
print(cur.fetchone()[0])
conn.close()
print("✅ Connected to Lakebase!")
Créer des tables de points de contrôle
Pour créer des tables de points de contrôle, définissez les entrées, le responsable de l’étape et les critères d’arrêt avant de modifier le code. Les opérateurs doivent pouvoir relancer l’étape à partir d’un point de contrôle connu sans deviner l’état caché. Documentez conjointement le parcours normal et le parcours de récupération. Les tentatives répétées, les contrôles humains et la gestion des messages non traités font partie intégrante du produit, et non d’une mise en forme ultérieure. Citez les passages qui justifient réellement la réponse. Sans citations, les opérateurs ne peuvent pas distinguer une hallucination d’un manque d’indexation.
from urllib.parse import quote
from langgraph.checkpoint.postgres import PostgresSaver
cred = w.postgres.generate_database_credential(endpoint=ENDPOINT)
DB_URI = (
f"postgresql://{quote(USERNAME, safe='')}:{quote(cred.token, safe='')}"
f"@{HOST}:5432/databricks_postgres"
f"?sslmode=require"
)
with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
checkpointer.setup()
print("✅ Checkpoint tables created!")
Pour créer des tables de points de contrôle, définissez les entrées, le responsable de l’étape et les critères d’arrêt avant de modifier le code. Les opérateurs doivent pouvoir relancer l’étape à partir d’un point de contrôle connu sans deviner l’état caché. Considérez cette étape comme un contrat entre les entrées et les sorties validées. Donnez des noms aux artefacts, définissez des vérifications de succès et refusez les terminations partielles silencieuses.
Étape 5 : Créer l’agent conscient du contexte
Lors de l’exécution de l’Étape 5 : Créer l’agent conscient du contexte, notez d’abord les exigences : entrées requises, signal de succès et conséquences en cas d’échec partiel. Cette liste de contrôle permet de rester honnête lors des modifications ultérieures du code. Enregistrez les temps d’exécution ainsi que le coût en tokens ou requêtes à côté des résultats fonctionnels. Une visibilité précoce des coûts évite les factures inattendues lorsque le système passe de l’environnement de démonstration à des environnements partagés. Mesurez le taux de rappel sur un ensemble fixe de questions avant d’ajuster les prompts. Changer fréquemment les prompts ne résout que rarement un système de récupération insuffisant.
Agent interactif (Notebook)
Lorsque vous travaillez sur l’Interactive Agent (Notebook), notez d’abord les exigences : entrées requises, signal de succès, et ce qui se passe en cas d’échec partiel. Cette liste de contrôle permet de rester honnête lors des modifications ultérieures du code. Conservez la configuration en dehors du code de l’application. Les fichiers d’environnement, les bases de données secrètes et les indicateurs fonctionnels doivent être regroupés en un seul endroit que les administrateurs peuvent auditer sans devoir lire l’ensemble du système. Évaluez le taux de rappel sur un ensemble fixe de questions avant d’ajuster les prompts. Changer fréquemment les prompts ne résout que rarement un système de récupération insuffisant.
from urllib.parse import quote
from langchain.agents import create_agent
from databricks_langchain import ChatDatabricks, VectorSearchRetrieverTool
from langgraph.checkpoint.postgres import PostgresSaver
from databricks.sdk import WorkspaceClient
import psycopg
from psycopg.rows import dict_row
def get_lakebase_checkpointer(host: str, endpoint: str, username: str):
"""Create a PostgresSaver backed by Lakebase Autoscaling."""
w = WorkspaceClient()
cred = w.postgres.generate_database_credential(endpoint=endpoint)
db_uri = (
f"postgresql://{quote(username, safe='')}:{quote(cred.token, safe='')}"
f"@{host}:5432/databricks_postgres"
f"?sslmode=require"
)
# IMPORTANT: Use psycopg.connect directly, not from_conn_string
# from_conn_string returns a context manager, not a persistent instance
conn = psycopg.connect(db_uri, autocommit=True, row_factory=dict_row)
checkpointer = PostgresSaver(conn=conn)
checkpointer.setup()
return checkpointer
def build_agent(llm_endpoint: str, index_name: str, num_results: int = 3):
model = ChatDatabricks(endpoint=llm_endpoint, max_tokens=500)
vs_tool = VectorSearchRetrieverTool(
name="knowledge_search",
index_name=index_name,
description="Search knowledge base for relevant information",
num_results=num_results,
)
# Lakebase-backed checkpointer instead of InMemorySaver
checkpointer = get_lakebase_checkpointer(HOST, ENDPOINT, USERNAME)
system_prompt = """You are a Knowledge Assistant. Respond in a clear,
professional tone. Use only verified information from the provided documents.
If the answer cannot be found, clearly state that."""
return create_agent(
model=model,
tools=[vs_tool],
system_prompt=system_prompt,
checkpointer=checkpointer,
)
Tester une conversation à plusieurs tours
Lorsque vous travaillez sur le test de conversation à plusieurs tours, notez d’abord les conditions prévues : les entrées requises, le signal de succès et ce qui se passe en cas d’échec partiel. Cette liste de contrôle permet de garantir l’honnêteté des modifications ultérieures du code. Documentez en même temps le parcours normal et celui de récupération. Les tentatives répétées, les contrôles humains et la gestion des messages non traités font partie intégrante du produit, et non d’améliorations apportées ultérieurement. Mesurez le taux de rappel sur un ensemble fixe de questions avant d’ajuster les prompts. Un changement fréquent des prompts résout rarement un système de récupération insuffisant.
agent = build_agent("<YOUR_LLM_ENDPOINT>", "<YOUR_INDEX_NAME>", 3)
# STABLE thread_id - this is what enables context awareness
config = {"configurable": {"thread_id": "demo-session-001"}}
# Turn 1
r1 = agent.invoke(
{"messages": [{"role": "user", "content": "What is the Orion system?"}]},
config=config
)
print("Turn 1:", r1['messages'][-1].content)
# Turn 2 - agent should know "it" = Orion
r2 = agent.invoke(
{"messages": [{"role": "user", "content": "How does it handle overheating?"}]},
config=config
)
print("Turn 2:", r2['messages'][-1].content)
Lorsque vous travaillez sur le test de conversation à plusieurs tours, notez d’abord les conditions prévues : les entrées requises, le signal de succès et ce qui se passe en cas d’échec partiel. Cette liste de contrôle permet de garantir l’honnêteté des modifications ultérieures du code. Considérez cette étape comme un contrat entre les entrées et les sorties validées. Donnez des noms aux éléments concernés, définez des critères de succès et refusez les complétions partielles silencieuses.
Étape 6 : Code de l’agent de production (agent.py)
L’Étape 6 : Code de l’agent de production (agent.py) fonctionne le mieux lorsqu’elle est considérée comme une surface mesurable. Capturez un enregistrement exemplaire, un cas d’échec et la note de rollback avant d’élargir le périmètre. Notez les temps d’exécution ainsi que le coût des tokens ou des requêtes à côté des résultats fonctionnels. Une visibilité précoce des coûts évite les factures inattendues lorsque le processus passe de l’environnement de démonstration à des environnements partagés. Séparez la politique de segmentation des données de la politique de récupération ; modifier l’une ne doit pas obliger à réécrire l’autre lorsque les métriques de qualité changent.
# agent.py
import os
from uuid import uuid4
from typing import Any, Dict, List
from urllib.parse import quote
import yaml
import mlflow
import psycopg
from psycopg.rows import dict_row
from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import ResponsesAgentRequest, ResponsesAgentResponse
from langchain.agents import create_agent
from databricks_langchain import ChatDatabricks, VectorSearchRetrieverTool
from langgraph.checkpoint.postgres import PostgresSaver
from databricks.sdk import WorkspaceClient
def _load_config(path: str = "agent-config.yaml") -> Dict[str, Any]:
if not os.path.exists(path):
raise FileNotFoundError(f"Config file not found at '{path}'")
with open(path, "r", encoding="utf-8") as f:
cfg = yaml.safe_load(f) or {}
llm_endpoint = cfg.get("llm_endpoint_name")
vs = cfg.get("vector_search", {}) or {}
index_name = vs.get("index_name")
num_results = int(vs.get("num_results", 3))
lakebase = cfg.get("lakebase", {}) or {}
return {
"llm_endpoint_name": llm_endpoint,
"vs_index_name": index_name,
"vs_num_results": num_results,
"lakebase_host": lakebase.get("host"),
"lakebase_endpoint": lakebase.get("endpoint"),
"lakebase_user": lakebase.get("user"),
}
def get_lakebase_checkpointer(host, endpoint, user):
w = WorkspaceClient()
cred = w.postgres.generate_database_credential(endpoint=endpoint)
db_uri = (
f"postgresql://{quote(user, safe='')}:{quote(cred.token, safe='')}"
f"@{host}:5432/databricks_postgres?sslmode=require"
)
conn = psycopg.connect(db_uri, autocommit=True, row_factory=dict_row)
checkpointer = PostgresSaver(conn=conn)
checkpointer.setup()
return checkpointer
def build_agent(llm_endpoint, index_name, num_results,
lakebase_host, lakebase_endpoint, lakebase_user):
model = ChatDatabricks(endpoint=llm_endpoint, max_tokens=500)
vs_tool = VectorSearchRetrieverTool(
name="knowledge_search",
index_name=index_name,
description="Search knowledge base for relevant information",
num_results=num_results,
)
checkpointer = get_lakebase_checkpointer(
lakebase_host, lakebase_endpoint, lakebase_user
)
system_prompt = (
"You are a Knowledge Assistant. Respond in a clear, professional tone. "
"Use only verified information from the provided documents. "
"If the answer cannot be found, clearly state that."
)
return create_agent(
model=model, tools=[vs_tool],
system_prompt=system_prompt, checkpointer=checkpointer,
)
def _last_user_text(messages):
user_msgs = [m for m in messages if m.get("role") == "user"]
return str(user_msgs[-1].get("content", "")) if user_msgs else ""
class LangChainResponsesAgent(ResponsesAgent):
def __init__(self):
cfg = _load_config()
self._agent = build_agent(
cfg["llm_endpoint_name"], cfg["vs_index_name"],
cfg["vs_num_results"], cfg["lakebase_host"],
cfg["lakebase_endpoint"], cfg["lakebase_user"],
)
def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
msgs = [m.model_dump() for m in request.input]
custom_inputs = dict(request.custom_inputs or {})
thread_id = custom_inputs.get("thread_id", f"session-{uuid4()}")
result = self._agent.invoke(
{"messages": msgs},
config={"configurable": {"thread_id": thread_id}},
)
try:
text = result["messages"][-1].content
except Exception:
text = str(result)
return ResponsesAgentResponse(
output=[self.create_text_output_item(text, str(uuid4()))],
custom_outputs={"thread_id": thread_id},
)
AGENT = LangChainResponsesAgent()
mlflow.models.set_model(AGENT)
Configuration (agent-config.yaml)
La configuration (agent-config.yaml) fonctionne le mieux lorsqu’elle est considérée comme une entité mesurable. Capturez un exemple idéal, un cas d’échec et la note de réversion avant d’élargir le périmètre.
Conservez la configuration en dehors du code de l’application. Les fichiers d’environnement, les bases de données secrètes et les indicateurs fonctionnels doivent être regroupés en un seul endroit que les opérateurs peuvent auditer sans devoir lire l’ensemble du système.
Séparez la politique de segmentation de la politique de récupération. Modifier l’une ne doit pas obliger à réécrire l’autre lorsque les métriques de qualité évoluent.
llm_endpoint_name: <YOUR_LLM_ENDPOINT>
vector_search:
index_name: <YOUR_CATALOG>.<YOUR_SCHEMA>.docs_chunked_index
num_results: 3
lakebase:
host: <YOUR_LAKEBASE_HOST>
endpoint: projects/<YOUR_PROJECT>/branches/production/endpoints/primary
user: <YOUR_DATABRICKS_EMAIL>
Étape 7 : Journaliser, enregistrer et déployer
Étape 7 : Journaliser, enregistrer et déployer fonctionne le mieux lorsqu’elle est considérée comme une entité mesurable. Capturez un exemple idéal, un cas d’échec et la note de réversion avant d’élargir le périmètre.
Journaliser dans MLflow
import mlflow
from importlib.metadata import version as get_version
from mlflow.models.resources import DatabricksVectorSearchIndex, DatabricksServingEndpoint
resources = [
DatabricksVectorSearchIndex(index_name="<YOUR_CATALOG>.<YOUR_SCHEMA>.docs_chunked_index"),
DatabricksServingEndpoint(endpoint_name="<YOUR_LLM_ENDPOINT>"),
]
with mlflow.start_run():
mlflow.set_tags({
"model_type": "retrieval_agent",
"framework": "langchain",
"memory": "lakebase_autoscaling",
})
logged_agent_info = mlflow.pyfunc.log_model(
name="knowledge_assistant",
python_model="agent.py",
code_paths=["agent-config.yaml"],
input_example={"input": [{"role": "user", "content": "What is Orion?"}]},
pip_requirements=[
f"databricks-vectorsearch=={get_version('databricks-vectorsearch')}",
f"databricks-langchain=={get_version('databricks-langchain')}",
f"langchain=={get_version('langchain')}",
f"mlflow=={get_version('mlflow')}",
"langgraph-checkpoint-postgres",
"psycopg[binary]",
"databricks-sdk>=0.89.0",
],
resources=resources,
)
model_uri = logged_agent_info.model_uri
Enregistrer dans Unity Catalog
mlflow.set_registry_uri("databricks-uc")
UC_MODEL_NAME = "<YOUR_CATALOG>.<YOUR_SCHEMA>.knowledge_assistant"
uc_info = mlflow.register_model(model_uri=model_uri, name=UC_MODEL_NAME)
print(f"✅ Registered: {UC_MODEL_NAME} v{uc_info.version}")
Déployer
from databricks import agents
deployment = agents.deploy(
model_name=UC_MODEL_NAME,
model_version=uc_info.version,
scale_to_zero_enabled=True,
)
print(f"✅ Endpoint: {deployment.query_endpoint}")
Le résultat
ws = WorkspaceClient()
client = ws.serving_endpoints.get_open_ai_client()
session = "user-session-042"
# Turn 1
r1 = client.responses.create(
model="knowledge_assistant",
input=[{"role": "user", "content": "What is the Orion motion controller?"}],
extra_body={"custom_inputs": {"thread_id": session}}
)
# Turn 2 - "it" resolves correctly to Orion
r2 = client.responses.create(
model="knowledge_assistant",
input=[{"role": "user", "content": "How does it prevent overheating?"}],
extra_body={"custom_inputs": {"thread_id": session}}
)
Vérification rapide de la persistance de la mémoire :
import json
from uuid import uuid4
thread_id = f"memory-test-{uuid4()}"
# --- 1. Use custom input instead of input_example ---
custom_input = {
"input": [{"role": "user", "content": "What are the main components of Orion?"}],
"custom_inputs": {"thread_id": thread_id},
}
print("=" * 50)
print("TEST 1: Custom input (no input_example needed)")
print("=" * 50)
# Use output_path to capture results (without it, mlflow.models.predict returns None)
mlflow.models.predict(
model_uri=model_uri,
input_data=custom_input,
env_manager="uv",
output_path="/tmp/result_1.json",
)
with open("/tmp/result_1.json", "r") as f:
result_1 = json.load(f)
thread_id_1 = result_1["custom_outputs"]["thread_id"]
response_1 = result_1["output"][0]["content"][0]["text"]
print(f"Thread ID: {thread_id_1}")
print(f"Response: {response_1[:300]}...")
# --- 2. Test memory persistence with same thread_id ---
follow_up_input = {
"input": [
{
"role": "user",
"content": "Can you elaborate more on the first component you mentioned?",
}
],
"custom_inputs": {"thread_id": thread_id}, # Reuse same thread
}
print("\n" + "=" * 50)
print("TEST 2: Follow-up on same thread (memory test)")
print("=" * 50)
mlflow.models.predict(
model_uri=model_uri,
input_data=follow_up_input,
env_manager="uv",
output_path="/tmp/result_2.json",
)
with open("/tmp/result_2.json", "r") as f:
result_2 = json.load(f)
thread_id_2 = result_2["custom_outputs"]["thread_id"]
response_2 = result_2["output"][0]["content"][0]["text"]
print(f"Thread ID: {thread_id_2}")
print(f"Response: {response_2[:300]}...")
# --- 3. Verify memory with actual conditions ---
print("\n" + "=" * 50)
print("MEMORY CHECK")
print("=" * 50)
# Check 1: Thread IDs match
if thread_id_1 == thread_id_2:
print(f"✅ Thread ID match: {thread_id_1}")
else:
print(f"❌ Thread ID mismatch! Call 1: {thread_id_1}, Call 2: {thread_id_2}")
# Check 2: Follow-up response references context from the first response
follow_up_lower = response_2.lower()
if len(response_2) > 50 and any(
keyword in follow_up_lower
for keyword in [
"motion",
"vision",
"cognition",
"communication",
"subsystem",
"component",
]
):
print(
"✅ Follow-up response references components from the first answer — memory is intact!"
)
else:
print(
"⚠️ Follow-up response may not reference the first answer. Manual review recommended."
)
print(f" Follow-up preview: {response_2[:200]}")
print(
"\n✅ Lakebase Postgres checkpointing is working correctly!"
if thread_id_1 == thread_id_2
else "\n❌ Memory persistence test FAILED."
)