Concevoir et tracer des outils de récupération pour les données non structurées

Utilisez le Framework Agent IA Mosaïque pour créer des outils qui permettent aux agents d’IA d’interroger des données non structurées telles qu’une collection de documents. Cette page montre comment :

• Développer des récupérateurs localement
• Créer un récupérateur à l’aide des fonctions de catalogue Unity
• Interroger des index vectoriels externes
• Ajouter le traçage MLflow afin d’assurer l’observabilité

Pour en savoir plus sur les outils d'agent IA, consultez outils d'agent IA.

développer localement des outils de récupération de recherche vectorielle avec AI Bridge

Le moyen le plus rapide de commencer à créer un outil de récupération de recherche vectorielle Databricks consiste à le développer et à le tester localement à l’aide de packages Databricks AI Bridge comme databricks-langchain et databricks-openai.

LangChain/LangGraph

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

%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 de pouvoir tester son comportement d’appel d’outils.

Fournissez une description tool_description pour aider l’agent à comprendre l’outil et déterminer quand l’appeler.

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 soit des index d'accès direct, soit des index Delta Sync avec des embeddings auto-gérés, vous devez configurer le VectorSearchRetrieverTool et spécifier un modèle d'embedding personnalisé ainsi qu'une colonne de texte. Consultez les options pour fournir des intégrations.

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

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="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 d’informations, consultez la documentation de l’API pour VectorSearchRetrieverTool.

OpenAI

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

%pip install --upgrade databricks-openai

Le code suivant prototype un récupérateur qui interroge un index de recherche vectorielle hypothétique et l’intègre aux modèles GPT d’OpenAI.

Fournissez une description tool_description pour aider l’agent à comprendre l’outil et déterminer quand l’appeler.

Pour plus d’informations sur les suggestions OpenAI pour les outils, consultez la documentation sur les appels de fonction OpenAI.

from databricks_openai import VectorSearchRetrieverTool
from openai import OpenAI
import json

# Initialize OpenAI client
client = OpenAI(api_key=<your_API_key>)

# Initialize the retriever tool
dbvs_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"
)

messages = [
  {"role": "system", "content": "You are a helpful assistant."},
  {
    "role": "user",
    "content": "Using the Databricks documentation, answer what is Spark?"
  }
]
first_response = client.chat.completions.create(
  model="gpt-4o",
  messages=messages,
  tools=[dbvs_tool.tool]
)

# Execute function code and parse the model's response and handle function calls.
tool_call = first_response.choices[0].message.tool_calls[0]
args = json.loads(tool_call.function.arguments)
result = dbvs_tool.execute(query=args["query"])  # For self-managed embeddings, optionally pass in openai_client=client

# Supply model with results – so it can incorporate them into its final response.
messages.append(first_response.choices[0].message)
messages.append({
  "role": "tool",
  "tool_call_id": tool_call.id,
  "content": json.dumps(result)
})
second_response = client.chat.completions.create(
  model="gpt-4o",
  messages=messages,
  tools=[dbvs_tool.tool]
)

Pour les scénarios qui utilisent soit des index d'accès direct, soit des index Delta Sync avec des embeddings auto-gérés, vous devez configurer le VectorSearchRetrieverTool et spécifier un modèle d'embedding personnalisé ainsi qu'une colonne de texte. Consultez les options pour fournir des intégrations.

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

from databricks_openai import VectorSearchRetrieverTool

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="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_model_name="databricks-bge-large-en" # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

Pour plus d’informations, consultez la documentation de l’API pour VectorSearchRetrieverTool.

Une fois que votre outil local est prêt, vous pouvez le produire directement dans le cadre de votre code d’agent ou le migrer vers une fonction de catalogue Unity, qui offre une meilleure détectabilité et gouvernance, mais présente certaines limitations.

La section suivante explique comment migrer le récupérateur vers une fonction du catalogue Unity.

Outil de recherche vectorielle avec les fonctions du catalogue Unity

Vous pouvez créer une fonction du catalogue Unity qui intègre une requête d’index de recherche vectorielle AI Mosaic. Cette approche a les caractéristiques suivantes :

• Prend en charge les cas d'usage en production avec la gouvernance et la découvrabilité
• Utilise la fonction vector_search() SQL sous le capot
• Prend en charge le traçage MLflow automatique
  • Vous devez aligner la sortie de la fonction sur le schéma du récupérateur MLflow à l'aide des alias page_content et metadata.
  • Toutes les colonnes de métadonnées supplémentaires doivent être ajoutées à la colonne à l’aide metadata 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 bloc-notes ou un éditeur SQL pour créer la fonction :

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 Vector 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 IA, encapsulez-le avec UCFunctionToolkit. Cela permet un traçage automatique via MLflow par génération automatique des types de span RETRIEVER dans les logs MLflow.

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 de catalogue Unity ont les avertissements suivants :

• Les clients SQL peuvent limiter le nombre maximal de lignes ou d’octets retournés. Pour empêcher la troncation des données, vous devez tronquer les valeurs de colonne retournées par la fonction 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 de ligne 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 Limitations.

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

Récupérateur qui interroge un index vectoriel hébergé en dehors de Databricks

Si votre index vectoriel est hébergé en dehors d’Azure Databricks, vous pouvez créer une connexion de catalogue Unity pour vous connecter au service externe et utiliser la connexion dans votre code d’agent. Consultez Connecter les outils d'agent IA aux services externes.

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 de catalogue Unity au service externe, dans ce cas, Azure.

   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 du catalogue Unity. 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 retriever doit retourner un List[Document] objet et utiliser le metadata champ de la classe Document pour ajouter des attributs supplémentaires au document retourné, tels que doc_uri et similarity_score. Consultez le document MLflow.

   import mlflow
   import json
   
   from mlflow.entities import Document
   from typing import List, Dict, Any
   from dataclasses import asdict
   
   class VectorSearchRetriever:
     """
     Class using Databricks Vector 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.
       """
       from databricks.sdk import WorkspaceClient
       from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod
   
       json = {
         "count": true,
         "select": "HotelId, HotelName, Description, Category",
         "vectorQueries": [
           {
             "vector": query_vector,
             "k": 7,
             "fields": "DescriptionVector",
             "kind": "vector",
             "exhaustive": true,
           }
         ],
       }
   
       response = (
         WorkspaceClient()
         .serving_endpoints.http_request(
           conn=connection_name,
           method=ExternalFunctionRequestHttpMethod.POST,
           path=f"indexes/{self.azure_search_index}/docs/search?api-version=2023-07-01-Preview",
           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
   

3. Pour exécuter l’extracteur, exécutez le code Python suivant. Vous pouvez éventuellement inclure les filtres Recherche vectorielle dans la requête pour filtrer les résultats.

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

Ajouter le traçage à un récupérateur

Ajoutez le traçage MLflow pour surveiller et déboguer votre récupérateur. Le suivi vous permet d’afficher 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 de parsing. Le décorateur crée une étendue qui démarre lorsque la fonction est appelée et se termine à son retour. MLflow enregistre automatiquement l’entrée et la sortie de la fonction et toutes les exceptions déclenchées.

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 garantir que les applications en aval telles que Agent Evaluation et AI Playground affichent correctement la trace du récupérateur, assurez-vous que le décorateur répond aux exigences suivantes :

• Utilisez (https://mlflow.org/docs/latest/tracing/tracing-schema.html#retriever-spans) et vérifiez que la fonction renvoie un objet List[Document].
• Le nom de la trace et le retriever_schema nom doivent correspondre pour configurer la trace correctement. Consultez la section suivante pour savoir comment définir le schéma du récupérateur.

Définir le schéma du récupérateur pour garantir la compatibilité de MLflow

Si la trace retournée à partir du 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 retourné aux champs attendus de MLflow. Cela garantit que MLflow peut suivre correctement votre récupérateur et afficher les traces dans les applications en aval.

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

1. Appelez mlflow.models.set_retriever_schema lorsque vous définissez votre agent. Permet set_retriever_schema de mapper les noms de colonnes dans la table retournée aux champs attendus de MLflow tels que primary_key, text_columnet doc_uri.

   # 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 other_columns champ.

3. Si vous avez plusieurs extracteurs, vous pouvez définir plusieurs schémas à l’aide de noms uniques pour chacun d’entre eux.

Le schéma du récupérateur défini lors de la création de l'agent affecte les applications et les flux de travail en aval, comme l'application de révision et les ensembles d'évaluation. Plus précisément, la colonne doc_uri sert d’identificateur principal pour les documents retournés par le récupérateur.

• L’application de révision affiche les doc_uri pour faciliter l'évaluation des réponses et la traçabilité des origines des documents. Consultez Interface utilisateur de l’application de révision.
• Les jeux d’évaluation utilisent doc_uri pour comparer les résultats de l'extracteur avec les jeux de données d’évaluation prédéfinis afin de déterminer le rappel et la précision de l'extracteur. Consultez les jeux d’évaluation (MLflow 2).

Étapes suivantes

Après avoir créé un outil de catalogue Unity, ajoutez-le à votre agent. Consultez Créer un outil d’agent.