Aller au contenu principal
Dernière mise à jour le

Connectez les agents aux données non structurées

Les agents doivent souvent query des données non structurées, comme des collections de documents, des bases de connaissances ou des corpus de texte, pour répondre aux questions et fournir des réponses contextuelles.

Pour les nouveaux agents Python, utilisez AI Bridge (recommandé) pour query directement les index de recherche AI Search. Pour les intégrations MCP existantes, consultez l’ endpoint de workspace hérité. Vous pouvez également vous connecter à des magasins de vecteurs externes ou lire des fichiers à partir de volumes Unity Catalog.

Développez un retriever avec AI Bridge (recommandé)​

Pour les nouveaux agents Python qui interrogent des index Databricks AI Search, Databricks recommande AI Bridge. Utilisez VectorSearchRetrieverTool à partir de databricks-langchain ou de databricks-openai pour exposer un index en tant qu'outil d'agent.

Avant d’exécuter ces exemples, connectez-vous à votre Workspace et choisissez un index de recherche IA que vous pouvez query. Remplacez le nom de l’index d’exemple par son nom en trois parties. Les premiers exemples utilisent des intégrations gérées par Databricks et un endpoint de modèle qui prend en charge l’appel d’outils.

Exécutez les commandes d’installation dans un terminal. Dans un Notebook, utilisez %pip à la place de pip et redémarrez Python après l’installation.

Installez la dernière version de databricks-langchain et LangGraph :

Bash
pip install --upgrade databricks-langchain langgraph

Créez un outil de récupération et donnez-lui une description qui indique à l’agent quand effectuer une recherche. Cet exemple teste la récupération directement, puis exécute un agent capable d'appeler l'outil et d'utiliser ses résultats.

Python
from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks
from langgraph.prebuilt import create_react_agent

vs_tool = VectorSearchRetrieverTool(
    index_name="catalog.schema.my_databricks_docs_index",
    tool_name="databricks_docs_retriever",
    tool_description="Search Databricks documentation to answer product questions.",
)
print(vs_tool.invoke("What is Databricks Agent Framework?"))

agent = create_react_agent(
    ChatDatabricks(endpoint="databricks-claude-sonnet-4-5"),
    tools=[vs_tool],
)
result = agent.invoke({"messages": [{
    "role": "user",
    "content": "Search the documentation and explain Databricks Agent Framework.",
}]})
for message in result["messages"]:
    print(message)

Veuillez vérifier la conversation enregistrée afin d’y rechercher un appel d’outil, les documents récupérés, ainsi qu’une réponse fondée sur ces derniers.

Pour les scénarios qui utilisent des index à accès direct ou des index Delta Sync utilisant des intégrations autogérées, vous devez configurer le VectorSearchRetrieverTool et spécifier un modèle d'intégration personnalisé et une colonne de texte. Consultez les options pour fournir des intégrations.

L'exemple suivant vous montre comment configurer un VectorSearchRetrieverTool avec des clés columns et embedding.

Python
from databricks_langchain import VectorSearchRetrieverTool
from databricks_langchain import DatabricksEmbeddings

embedding_model = DatabricksEmbeddings(
    endpoint="databricks-bge-large-en",
)

vs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.index_name", # Index name in the format 'catalog.schema.index'
  num_results=5, # Max number of documents to return
  columns=["primary_key", "text_column"], # List of columns to include in the search
  filters={"text_column LIKE": "Databricks"}, # Filters to apply to the query
  query_type="ANN", # Query type ("ANN" or "HYBRID").
  tool_name="databricks_docs_retriever", # Used by the LLM to understand the purpose of the tool
  tool_description="Purpose of the tool", # Used by the LLM to understand the purpose of the tool
  text_column="text_column", # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
  embedding=embedding_model # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

Pour plus de détails, consultez les documents de l'API pour VectorSearchRetrieverTool.

After testing the retriever, add it to your custom agent.

query un index de recherche vectorielle en dehors de Databricks​

Si votre index vectoriel est hébergé en dehors de Databricks, vous pouvez créer une connexion Unity Catalog pour vous connecter au service externe et l'utiliser dans le code de votre agent. Consultez Se connecter à des services HTTP externes.

L’exemple suivant query un index Azure AI Search via une connexion HTTP. Vous avez besoin d’une connexion disposant des autorisations nécessaires pour query cet index et de USE CONNECTION sur la connexion. Installez requests, databricks-sdk et mlflow avant d’exécuter l’exemple Python.

  1. Créer une connexion Unity Catalog vers Azure AI Search à l'aide de la configuration de connexion HTTP. Cet exemple utilise un jeton porteur ; choisissez la méthode d'authentification requise par votre déploiement Azure AI Search.

    SQL
    CREATE CONNECTION azure_search_connection
    TYPE HTTP
    OPTIONS (
      host 'https://example.search.windows.net',
      base_path '/',
      bearer_token secret ('<secret-scope>','<secret-key>')
    );

  2. Définissez l’outil de récupération dans le code d’agent à l’aide de la connexion Unity Catalog. Cet exemple utilise les décorateurs MLflow pour activer le traçage d’agent.

remarque

Pour se conformer au schéma du récupérateur MLflow, la fonction du récupérateur doit renvoyer un objet List[Document] et utiliser le champ metadata dans la classe Document pour ajouter des attributs supplémentaires au document renvoyé, tels que doc_uri et similarity_score. Voir la documentation MLflow.

Python
import mlflow
import requests
from databricks.sdk import WorkspaceClient
from mlflow.entities import Document

class VectorSearchRetriever:
    def __init__(self, connection_name, index_name):
        self.workspace = WorkspaceClient()
        self.connection_name = connection_name
        self.index_name = index_name

    @mlflow.trace(span_type="RETRIEVER", name="vector_search")
    def __call__(self, query_vector: list[float], score_threshold: float = 0.0) -> list[Document]:
        response = requests.post(
            f"{self.workspace.config.host}/api/2.0/unity-catalog/connections/"
            f"{self.connection_name}/proxy/indexes/{self.index_name}/docs/search",
            params={"api-version": "2024-07-01"},
            headers=self.workspace.config.authenticate(),
            json={
                "select": "HotelId,HotelName,Description,Category",
                "vectorQueries": [{
                    "vector": query_vector,
                    "k": 7,
                    "fields": "DescriptionVector",
                    "kind": "vector",
                    "exhaustive": True,
                }],
            },
            timeout=30,
        )
        response.raise_for_status()
        return self.parse_results(response.json(), score_threshold)

    @mlflow.trace(span_type="PARSER")
    def parse_results(self, results, score_threshold) -> list[Document]:
        return [
            Document(
                page_content=item.get("Description", ""),
                id=item.get("HotelId"),
                metadata={
                    "similarity_score": item.get("@search.score", 0),
                    "HotelName": item.get("HotelName"),
                    "Category": item.get("Category"),
                },
            )
            for item in results.get("value", [])
            if item.get("@search.score", 0) >= score_threshold
        ]

  1. Générez un embedding de query avec le même modèle et les mêmes dimensions que ceux utilisés par votre index. Enregistrez le tableau numérique sous la forme query-vector.json, puis exécutez le récupérateur :

    Python
    import json
    from pathlib import Path

    query_vector = json.loads(Path("query-vector.json").read_text())
    retriever = VectorSearchRetriever("azure_search_connection", "hotels_vector_index")
    results = retriever(query_vector, score_threshold=0.1)
    for document in results:
        print(document.page_content)

Ajouter le traçage à un outil de recherche​

Ajouter le traçage MLflow pour surveiller et déboguer votre récupérateur. Le traçage vous permet de visualiser les entrées, les sorties et les métadonnées pour chaque étape d'exécution.

L'exemple précédent ajoute le décorateur @mlflow.trace aux méthodes __call__ et d'analyse. Le décorateur crée un span qui start lorsque la fonction est invoquée et se termine lorsqu'elle renvoie un résultat. MLflow enregistre automatiquement les entrées et les sorties de la fonction, ainsi que toutes les exceptions levées.

remarque

LangChain, LlamaIndex, and OpenAI library users can use MLflow auto logging in addition to manually defining traces with the decorator. Consultez la vue d’ensemble du traçage.

Renvoyez une liste d'objets Document avec page_content et metadata, comme le fait l'exemple ci-dessus. Cela suit le schéma d'étendue de récupérateur MLflow, afin que les outils d'évaluation et de révision puissent reconnaître les documents récupérés.

Lisez les fichiers d'un volume Unity Catalog​

Si votre agent doit lire des fichiers non structurés (documents texte, rapports, fichiers de configuration, etc.) stockés dans un volume Unity Catalog, vous pouvez créer des outils qui utilisent l’ API Files du SDK Databricks pour lister et lire les fichiers directement.

Choisissez un volume existant et accordez à l’identité de votre agent USE CATALOG, USE SCHEMA et READ VOLUME. Remplacez <catalog>.<schema>.<volume> dans le code. Pour la tâche d’exemple, utilisez un fichier texte UTF-8 nommé about_databricks.txt ou modifiez l’invite pour pointer vers un fichier de votre volume.

Les exemples suivants créent deux outils que votre agent peut utiliser :

  • list_volume_files : Répertorie les fichiers et les répertoires dans le volume.
  • read_volume_file : Lit le contenu d'un fichier texte depuis le volume.

Installez la dernière version de databricks-langchain qui inclut Databricks AI Bridge.

Bash
pip install --upgrade databricks-langchain langgraph
Python
from pathlib import PurePosixPath
from databricks.sdk import WorkspaceClient
from langchain_core.tools import tool

VOLUME = "<catalog>.<schema>.<volume>"
w = WorkspaceClient()


def volume_path(relative_path: str = "") -> str:
    path = PurePosixPath(relative_path)
    if path.is_absolute() or ".." in path.parts:
        raise ValueError("Use a path relative to the configured volume.")
    return str(PurePosixPath("/Volumes", *VOLUME.split(".")) / path)


@tool
def list_volume_files(directory: str = "") -> str:
    """Lists files and directories in the Unity Catalog volume.
    Provide a relative directory path, or leave empty to list the volume root."""
    path = volume_path(directory)
    entries = []
    for f in w.files.list_directory_contents(path):
        kind = "dir" if f.is_directory else "file"
        size = f" ({f.file_size} bytes)" if not f.is_directory else ""
        entries.append(f"  [{kind}] {f.name}{size}")
    return "\n".join(entries) if entries else "No files found."


@tool
def read_volume_file(file_path: str) -> str:
    """Reads a text file from the Unity Catalog volume.
    Provide the path relative to the volume root, for example 'reports/q1_summary.txt'."""
    full_path = volume_path(file_path)
    resp = w.files.download(full_path)
    return resp.contents.read().decode("utf-8")

Liez les outils à un LLM et exécutez une boucle d'appel d'outil :

Python
from databricks_langchain import ChatDatabricks
from langchain_core.messages import HumanMessage, ToolMessage

llm = ChatDatabricks(endpoint="databricks-claude-sonnet-4-5")
llm_with_tools = llm.bind_tools([list_volume_files, read_volume_file])

messages = [HumanMessage(content="What files are in the volume? Can you read about_databricks.txt and summarize it in 2 sentences?")]
tool_map = {"list_volume_files": list_volume_files, "read_volume_file": read_volume_file}

for _ in range(5):  # max iterations
    response = llm_with_tools.invoke(messages)
    messages.append(response)
    if not response.tool_calls:
        break
    for tc in response.tool_calls:
        result = tool_map[tc["name"]].invoke(tc["args"])
        messages.append(ToolMessage(content=result, tool_call_id=tc["id"]))

else:
    raise RuntimeError("The agent reached the tool-calling limit.")

print(response.content)

Query un index de recherche AI à l’aide du protocole MCP (ancienne version)​

Le serveur MCP AI Search utilise un endpoint MCP de workspace hérité. Il nécessite un index avec des embeddings gérés par Databricks. Consultez Créer des Endpoint et des index de recherche AI.

The legacy workspace MCP URL for AI Search is https://<workspace-hostname>/api/2.0/mcp/ai-search/{catalog}/{schema}/{index_name}. Connect to it and list the tools it exposes:

Python
from databricks.sdk import WorkspaceClient
from databricks_mcp import DatabricksMCPClient

workspace_client = WorkspaceClient()
host = workspace_client.config.host

mcp_client = DatabricksMCPClient(
    server_url=f"{host}/api/2.0/mcp/ai-search/<catalog>/<schema>/<index-name>",
    workspace_client=workspace_client,
)
tools = mcp_client.list_tools()

Pour créer et déployer un agent qui utilise ce serveur, consultez Utiliser des outils MCP dans un agent Python. Accordez à l'agent SELECT sur l'élément sécurisable Unity Catalog de l'index.

AI Search s’appelait auparavant Vector Search. Le préfixe d’URL /api/2.0/mcp/vector-search/ précédent fonctionne toujours.

Query Databricks AI Search à l’aide de fonctions UC (déprécié)​

remarque

Pour les nouveaux agents Python, utilisez AI Bridge. L'exemple suivant sert de référence pour les intégrations existantes.

Vous pouvez créer une fonction Unity Catalog qui encapsule une query d’ index de recherche IA. Cette approche :

  • Prend en charge les cas d'usage de production avec gouvernance et découvrabilité.
  • Utilise la fonction SQL vector_search() sous le capot.
  • Prend en charge le suivi MLflow automatique
    • Vous devez aligner la sortie de la fonction sur le schéma de l'extracteur MLflow en utilisant les alias page_content et metadata.
    • Toutes les colonnes de métadonnées supplémentaires doivent être ajoutées à la colonne metadata à l'aide de la fonction de mappage SQL, plutôt que comme clés de sortie de niveau supérieur.

Exécutez le code suivant dans un Notebook ou un éditeur SQL pour créer la fonction :

SQL
CREATE OR REPLACE FUNCTION main.default.databricks_docs_vector_search (
  -- The agent uses this comment to determine how to generate the query string parameter.
  query STRING
  COMMENT 'The query string for searching Databricks documentation.'
) RETURNS TABLE
-- The agent uses this comment to determine when to call this tool. It describes the types of documents and information contained within the index.
COMMENT 'Executes a search on Databricks documentation to retrieve text documents most relevant to the input query.' RETURN
SELECT
  chunked_text as page_content,
  map('doc_uri', url, 'chunk_id', chunk_id) as metadata
FROM
  vector_search(
    -- Specify your AI Search index name here
    index => 'catalog.schema.databricks_docs_index',
    query => query,
    num_results => 5
  )

Suivez la configuration de l'outil de fonction, puis enveloppez ce récupérateur avec UCFunctionToolkit pour l'utiliser dans votre agent. Cela permet le traçage automatique via MLflow en générant automatiquement des types de portée RETRIEVER dans les logs de MLflow.

Python
from databricks_langchain import UCFunctionToolkit
from unitycatalog.ai.core.databricks import DatabricksFunctionClient

toolkit = UCFunctionToolkit(
    function_names=[
        "main.default.databricks_docs_vector_search"
    ],
    client=DatabricksFunctionClient(),
)
tools = toolkit.tools

Les outils de récupération Unity Catalog présentent les mises en garde suivantes :

  • Les clients SQL peuvent limiter le nombre maximal de lignes ou d’octets renvoyés. Pour éviter la troncation des données, tronquez les valeurs de colonne renvoyées par l’UDF. Par exemple, vous pouvez utiliser substring(chunked_text, 0, 8192) pour réduire la taille des colonnes de contenu volumineuses et éviter la troncation des lignes pendant l’exécution.
  • Étant donné que cet outil est un wrapper pour la fonction vector_search(), il est soumis aux mêmes limitations que la fonction vector_search(). Consultez les Limitations.

Pour plus d'informations sur UCFunctionToolkit, consultez la documentation Unity Catalog.

Schémas de récupérateur personnalisés (hérité)​

mlflow.models.set_retriever_schema est obsolète. Pour les nouveaux récupérateurs, renvoyez des objets Document en utilisant le schéma de récupérateur standard. La configuration suivante concerne les intégrations existantes qui renvoient des champs personnalisés.

Pour définir le schéma du récupérateur manuellement :

  1. Appelez mlflow.models.set_retriever_schema. lorsque vous définissez votre agent. Utilisez set_retriever_schema pour mapper les noms de colonnes de la table renvoyée aux champs attendus de MLflow, tels que primary_key, text_column et doc_uri.

    Python
    import mlflow

    # Define the retriever's schema by providing your column names
    mlflow.models.set_retriever_schema(
      name="vector_search",
      primary_key="chunk_id",
      text_column="text_column",
      doc_uri="doc_uri"
      # other_columns=["column1", "column2"],
    )

  2. Spécifiez des colonnes supplémentaires dans le schéma de votre récupérateur en fournissant une liste de noms de colonnes avec le champ other_columns.

  3. Si vous avez plusieurs extracteurs, vous pouvez définir plusieurs schémas en utilisant des noms uniques pour chaque schéma d'extracteur.

Le schéma de l'outil de récupération défini lors de la création de l'agent affecte les applications et les workflows en aval, tels que l'application de révision et les ensembles d'évaluation. Plus précisément, la colonne doc_uri sert d'identifiant principal pour les documents renvoyés par l'outil de récupération.

  • L' application d'examen affiche le doc_uri pour aider les évaluateurs à évaluer les réponses et à retracer les origines des documents. Consultez l'interface utilisateur de l'application d'examen.
  • Les evaluation sets utilisent doc_uri pour comparer les résultats du système de recherche (retriever) aux datasets d’évaluation prédéfinis afin de déterminer le rappel et la précision du système. Consultez la référence des datasets d’évaluation.
Sur cette page