Aller au contenu principal
Dernière mise à jour le

Connectez les agents aux données non structurées

Les agents d'IA 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.

Databricks propose plusieurs approches pour connecter les agents aux données non structurées dans les index AI Search et les magasins de vecteurs externes. Utilisez des serveurs MCP préconfigurés pour un accès immédiat aux index Databricks AI Search, 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 auparavant 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 Databricks AI Search à l'aide de MCP​

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

L'URL MCP gérée pour AI Search est https://<workspace-hostname>/api/2.0/mcp/ai-search/{catalog}/{schema}/{index_name}. Connectez-vous-y 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. Octroyez à l'agent SELECT sur l'élément sécurisable Unity Catalog de l'index.

Autres approches​

Interroger un index de recherche vectorielle en dehors de Databricks

Query 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. Consultez Connectez les agents aux 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 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 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 récupérateur, 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éveloppez un extracteur local

Développez un agent de récupération 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 d'assistance telles que from_vector_search et from_uc_function pour créer des récupérateurs à partir de 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 présente un prototype d'outil de récupération qui interroge un index de recherche vectorielle hypothétique et le lie localement à un LLM afin que vous puissiez tester son comportement d'appel d'outil.

Fournissez un tool_description descriptif 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 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={&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 les documents de l'API pour VectorSearchRetrieverTool.

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

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

Query Databricks AI Search à l’aide de 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 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
  )

Pour utiliser cet outil de récupération dans votre agent d'IA, encapsulez-le avec UCFunctionToolkit. Cela permet le traçage automatique via MLflow en générant automatiquement RETRIEVER types de spans 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 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.

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

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

Les utilisateurs des bibliothèques LangChain, LlamaIndex et OpenAI peuvent utiliser l'autologging MLflow en plus de définir manuellement des traces avec le décorateur. Consultez 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 Agent Evaluation et l'AI Playground, affichent correctement la trace de récupération, assurez-vous que le décorateur respecte les exigences suivantes :

  • Utilisez le schéma de span de 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. Consultez la section suivante pour apprendre à définir le schéma du récupérateur.

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

Si la trace renvoyée par le récupérateur ou span_type="RETRIEVER" n'est pas conforme au schéma de récupérateur 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 récupérateur 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 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 jeux d’évaluation utilisent doc_uri pour comparer les résultats du récupérateur avec des datasets d’évaluation prédéfinis afin de déterminer le rappel et la précision du récupérateur. Consultez les jeux d'évaluation (MLflow 2).

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.

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

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

print(response.content)
Sur cette page