Aller au contenu principal
Dernière mise à jour le

Connectez les agents aux données non structurées

Les agents d'IA doivent souvent interroger des données non structurées, telles que des collections de documents, des bases de connaissances ou des corpus de texte, pour répondre aux questions et fournir des réponses contextuelles.

Databricks propose plusieurs approches pour connecter les agents aux données non structurées dans les index de recherche IA et les magasins de vecteurs externes. Utilisez des serveurs MCP préconfigurés pour un accès immédiat aux index de recherche Databricks AI, développez des outils de récupération localement avec des packages AI Bridge, ou créez des fonctions de récupération personnalisées pour des workflows spécialisés.

Databricks AI Search était anciennement connu sous le nom de Databricks Vector Search. Le préfixe d'URL hérité /api/2.0/mcp/vector-search/ continue de fonctionner pour la rétrocompatibilité.

Interroger un index de recherche IA Databricks à l'aide de MCP​

Utilisez le serveur MCP Databricks AI Search géré par Databricks pour donner à votre agent l'accès à un index Databricks AI Search. Tout d'abord, créez un index en utilisant des intégrations gérées par Databricks. Consultez Créer des Endpoint et des index AI Search.

L'URL MCP gérée pour la recherche IA est https://<workspace-hostname>/api/2.0/mcp/ai-search/{catalog}/{schema}/{index_name}. Connectez-vous à celui-ci et listez les outils qu'il expose :

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 serveurs MCP dans les agents. Accorder à l'agent SELECT sur la ressource sécurisable de Unity Catalog de l'index.

Autres approches​

Interroger un index de recherche vectorielle en dehors de Databricks

Requêter un index de recherche vectorielle hébergé 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 utiliser la connexion dans votre code d'agent. Voir Connecter des agents à des outils tiers avec les services MCP.

L'exemple suivant crée un récupérateur qui appelle un index vectoriel hébergé en dehors de Databricks pour un agent de type PyFunc.

  1. Créez une connexion Unity Catalog au service externe, dans ce cas, Azure.

    SQL
    CREATE CONNECTION ${connection_name}
    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 de l'agent à l'aide de la connexion Unity Catalog. Cet exemple utilise des décorateurs MLflow pour activer le traçage des agents.

remarque

Pour se conformer au schéma du récupérateur MLflow, la fonction de 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. Consultez la documentation MLflow.

Python
import mlflow
import json

from mlflow.entities import Document
from typing import List, Dict, Any
from dataclasses import asdict

class VectorSearchRetriever:
  """
  Class using Databricks AI Search to retrieve relevant documents.
  """

  def __init__(self):
    self.azure_search_index = "hotels_vector_index"

  @mlflow.trace(span_type="RETRIEVER", name="vector_search")
  def __call__(self, query_vector: List[Any], score_threshold=None) -> List[Document]:
    """
    Performs vector search to retrieve relevant chunks.
    Args:
      query: Search query.
      score_threshold: Score threshold to use for the query.

    Returns:
      List of retrieved Documents.
    """
    import requests
    from databricks.sdk import WorkspaceClient

    w = WorkspaceClient()
    json = {
      "count": true,
      "select": "HotelId, HotelName, Description, Category",
      "vectorQueries": [
        {
          "vector": query_vector,
          "k": 7,
          "fields": "DescriptionVector",
          "kind": "vector",
          "exhaustive": true,
        }
      ],
    }

    response = requests.post(
      f"{w.config.host}/api/2.0/unity-catalog/connections/{connection_name}/proxy/indexes/{self.azure_search_index}/docs/search?api-version=2023-07-01-Preview",
      headers={
     **w.config.authenticate(),
     &quot;Content-Type&quot;: &quot;application/json&quot;,
   },
      json=json,
    ).text

    documents = self.convert_vector_search_to_documents(response, score_threshold)
    return [asdict(doc) for doc in documents]

  @mlflow.trace(span_type="PARSER")
  def convert_vector_search_to_documents(
    self, vs_results, score_threshold
  ) -> List[Document]:
    docs = []

    for item in vs_results.get("value", []):
      score = item.get("@search.score", 0)

      if score >= score_threshold:
        metadata = {
          "score": score,
          "HotelName": item.get("HotelName"),
          "Category": item.get("Category"),
        }

        doc = Document(
          page_content=item.get("Description", ""),
          metadata=metadata,
          id=item.get("HotelId"),
        )
        docs.append(doc)

    return docs

  1. Pour exécuter le retriever, exécutez le code Python suivant.

    Python
    retriever = VectorSearchRetriever()
    query = [0.01944167, 0.0040178085 . . .  TRIMMED FOR BREVITY 010858015, -0.017496133]
    results = retriever(query, score_threshold=0.1)

Développer un système de récupération local

Développez un récupérateur localement à l'aide d'AI Bridge​

Pour créer un outil de récupération Databricks AI Search localement, utilisez les packages Databricks AI Bridge comme databricks-langchain et databricks-openai. Ces packages comprennent des fonctions utilitaires telles que from_vector_search et from_uc_function pour créer des extracteurs à partir des Ressources Databricks existantes.

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

Bash
%pip install --upgrade databricks-langchain

Le code suivant prototype un outil de récupération qui interroge un index de recherche vectorielle hypothétique et le lie à un LLM localement afin que vous puissiez tester son comportement d'appel d'outil.

Fournissez une tool_description descriptive pour aider l'agent à comprendre l'outil et à déterminer quand l'invoquer.

Python
from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks

# Initialize the retriever tool.
vs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.my_databricks_docs_index",
  tool_name="databricks_docs_retriever",
  tool_description="Retrieves information about Databricks products from official Databricks documentation."
)

# Run a query against the vector search index locally for testing
vs_tool.invoke("Databricks Agent Framework?")

# Bind the retriever tool to your Langchain LLM of choice
llm = ChatDatabricks(endpoint="databricks-claude-sonnet-4-5")
llm_with_tools = llm.bind_tools([vs_tool])

# Chat with your LLM to test the tool calling functionality
llm_with_tools.invoke("Based on the Databricks documentation, what is Databricks Agent Framework?")

Pour les scénarios qui utilisent des index à accès direct ou des index Delta Sync à l'aide d'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 les 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={&quot;text_column LIKE&quot;: &quot;Databricks&quot;}, # Filters to apply to the query
  query_type="ANN", # Query type ("ANN" or "HYBRID").
  tool_name="name of the tool", # 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 la documentation API pour VectorSearchRetrieverTool.

Une fois votre outil local prêt, vous pouvez le mettre directement en production dans le cadre du code de votre agent, ou le migrer vers une fonction Unity Catalog, ce qui offre une meilleure visibilité et gouvernance, mais présente certaines limitations.

Interroger Databricks AI Search à l’aide des fonctions UC (déprécié).

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

remarque

Databricks recommends MCP servers for most agent tools, but defining tools with Unity Catalog functions remains available for prototyping.

Vous pouvez créer une fonction Unity Catalog qui enveloppe 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 automatique de MLflow
    • Vous devez aligner la sortie de la fonction avec le schéma de récupérateur MLflow en utilisant les alias page_content et metadata.
    • Les colonnes de métadonnées supplémentaires doivent être ajoutées à la colonne metadata à l'aide de la fonction map SQL, plutôt qu'en tant que 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
  )

Pour utiliser cet outil de récupération dans votre agent d'IA, encapsulez-le avec UCFunctionToolkit. Ceci active le traçage automatique via MLflow en générant automatiquement RETRIEVER types d'étendue dans les logs MLflow.

Python
from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

toolkit = UCFunctionToolkit(
    function_names=[
        "main.default.databricks_docs_vector_search"
    ]
)
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 pourriez utiliser substring(chunked_text, 0, 8192) pour réduire la taille des colonnes de contenu volumineuses et éviter le tronquage 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, reportez-vous à la documentation Unity Catalog.

Ajoutez le traçage à un outil de récupération​

Ajoutez 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 quand elle retourne. MLflow enregistre automatiquement l'entrée et la sortie de la fonction ainsi que toute exception levée.

remarque

Les utilisateurs des bibliothèques LangChain, LlamaIndex et OpenAI peuvent utiliser le log automatique MLflow en plus de définir manuellement les traces avec le décorateur. Voir Ajouter des traces aux applications : traçage automatique et manuel.

Python
import mlflow
from mlflow.entities import Document

# This code snippet has been truncated for brevity. See the full retriever example above.
class VectorSearchRetriever:
  ...

  # Create a RETRIEVER span. The span name must match the retriever schema name.
  @mlflow.trace(span_type="RETRIEVER", name="vector_search")
  def __call__(...) -> List[Document]:
    ...

  # Create a PARSER span.
  @mlflow.trace(span_type="PARSER")
  def parse_results(...) -> List[Document]:
    ...

Pour vérifier que les applications en aval, telles que l'Agent Evaluation et l'AI Playground, rendent correctement la trace du récupérateur, assurez-vous que le décorateur respecte les exigences suivantes :

  • Utilisez le schéma de portée du récupérateur MLflow et vérifiez que la fonction renvoie un objet List[Document].
  • Le nom de la trace et le nom retriever_schema doivent correspondre pour configurer correctement la trace. Veuillez consulter la section suivante pour savoir comment configurer le schéma du récupérateur.

Définir le schéma du récupérateur pour vérifier la compatibilité MLflow​

Si la trace renvoyée par l'extracteur ou span_type="RETRIEVER" n'est pas conforme au schéma d'extracteur standard de MLflow, vous devez mapper manuellement le schéma renvoyé aux champs attendus de MLflow. Cela vérifie que MLflow peut tracer correctement votre extracteur et afficher les traces dans les applications en aval.

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
    # 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 du récupérateur défini lors de la création de l'agent affecte les applications et flux de travail en aval, tels que l'application de révision et les jeux d'évaluation. Plus précisément, la colonne doc_uri sert d'identifiant principal pour les documents renvoyés par le récupérateur.

  • L'**application d'avis** affiche les doc_uri pour aider les évaluateurs à examiner les réponses et à retracer l'origine des documents. Consultez l'interface utilisateur de l'application d'avis.
  • Les ensembles d'évaluation utilisent doc_uri pour comparer les résultats de l'outil d'extraction avec des datasets d'évaluation prédéfinis afin de déterminer le rappel et la précision de l'outil d'extraction. Consultez les ensembles d'évaluation (MLflow 2).

Lire les fichiers à partir 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 la Files API du SDK Databricks pour répertorier et lire directement les fichiers.

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

  • list_volume_files ** ** : Liste les fichiers et les répertoires dans le volume.
  • read_volume_file : Lit le contenu d'un fichier texte à partir du volume.

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

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

VOLUME = "<catalog>.<schema>.<volume>"  # TODO: Replace with your volume
w = WorkspaceClient()


@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."""
    base = f"/Volumes/{VOLUME.replace('.', '/')}"
    path = f"{base}/{directory.lstrip('/')}" if directory else base
    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'."""
    base = f"/Volumes/{VOLUME.replace('.', '/')}"
    full_path = f"{base}/{file_path.lstrip('/')}"
    resp = w.files.download(full_path)
    return resp.contents.read().decode("utf-8")

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

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

print(response.content)
Sur cette page