Partager via

Outils d'agent IA pour la recherche non structurée

  • Article

Important

Cette fonctionnalité est disponible en préversion publique.

Cet article explique comment créer des outils d’assistant IA pour l’extraction de données non structurées à l’aide de l’infrastructure de l’assistant AI de Mosaic. Les récupérateurs non structurés permettent aux agents d’interroger des sources de données non structurées, telles qu’un corpus de documents, à l’aide d’index de recherche vectorielle.

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 simple de commencer à développer un outil de récupération de recherche vectorielle Databricks est localement. Utilisez packages Databricks AI Bridge comme databricks-langchain et databricks-openai pour ajouter rapidement des fonctionnalités de récupération à un agent et expérimenter des paramètres de requête. Cette approche permet une itération rapide pendant le développement initial.

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. Consultez l’outil de recherche vectorielle avec les fonctions de Unity Catalog.

LangChain/LangGraph

Le code suivant prototype un outil de récupération et le lie à un LLM localement afin de pouvoir discuter avec l’agent pour tester son comportement d’appel d’outil.

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

%pip install --upgrade databricks-langchain

Remarque

Lors de l’initialisation du VectorSearchRetrieverTool, les arguments text_column et embedding sont requis pour les index delta sync avec des incorporations auto-gérées et des index d’accès vectoriel direct. Consultez les options pour fournir des intégrations.

from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks

# Initialize the retriever tool
vs_tool = VectorSearchRetrieverTool(index_name="catalog.schema.my_index_name")

# 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-meta-llama-3-1-70b-instruct")
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 personnaliser l’appel d’outils, passez des arguments supplémentaires à VectorSearchRetrieverTool:

from databricks_langchain import VectorSearchRetrieverTool

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

OpenAI

Le code suivant prototype un outil de récupération de recherche vectorielle et l’intègre aux modèles GPT d’OpenAI.

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

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

%pip install --upgrade databricks-openai

Remarque

Lors de l’initialisation du VectorSearchRetrieverTool, les arguments text_column et embedding sont requis pour les index delta sync avec des incorporations auto-gérées et des index d’accès vectoriel direct. Consultez les options pour fournir des intégrations.

from databricks_openai import VectorSearchRetrieverTool
from openai import OpenAI
import json

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

# Call model with VectorSearchRetrieverTool defined
dbvs_tool = VectorSearchRetrieverTool(index_name="catalog.schema.my_index_name")
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 personnaliser l’appel d’outils, passez des arguments supplémentaires à VectorSearchRetrieverTool:

from databricks_openai import VectorSearchRetrieverTool

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

Outil de recherche vectorielle avec les fonctions du catalogue Unity

L’exemple suivant crée un outil de récupération à l’aide d’une fonction du catalogue Unity pour interroger des données à partir d’un index de recherche vectorielle Mosaic AI .

La fonction Unity Catalog databricks_docs_vector_search interroge un index de recherche vectorielle hypothétique contenant la documentation Databricks. Cette fonction encapsule la fonction DATAbricks SQL vector_search() et aligne sa sortie avec le schéma du récupérateur MLflow . à l’aide des alias page_content et metadata.

Remarque

Pour se conformer au schéma du récupérateur MLflow, 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 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 le suivi automatique via MLflow.

Le traçage MLflow capture des informations d’exécution détaillées pour les applications d'IA générative. Il journalise les entrées, les sorties et les métadonnées pour chaque étape, ce qui vous aide à déboguer des problèmes et à analyser les performances.

Lorsque vous utilisez UCFunctionToolkit, les extracteurs génèrent automatiquement des types de segments RETRIEVER dans les enregistrements MLflow si leur sortie est conforme au schéma de l’extracteur MLflow. Consultez schéma de traçage MLflow.

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

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Cet outil de récupération présente 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.

Si cet exemple n’est pas adapté à votre cas d’usage, créez un outil de récupération de recherche vectorielle à l’aide du code d’agent personnalisé à la place.

Récupérateur de recherche vectorielle avec le code de l’agent (PyFunc)

L'exemple suivant crée un récupérateur de recherche vectorielle pour un agent au format PyFunc dans le code agent.

Cet exemple utilise databricks-vectorsearch pour créer un récupérateur de base qui effectue une recherche de similarité Recherche vectorielle avec des filtres. Il utilise des décorateurs MLflow pour permettre le suivi de l’agent.

Remarque

Pour se conformer au schéma du récupérateur MLflow, la fonction retriever doit retourner un type Document et utiliser le champ metadata de la classe Document pour ajouter des attributs supplémentaires au document retourné, comme like doc_uri et similarity_score.

Utilisez le code suivant dans le module agent ou le notebook de l’agent.

import mlflow
import json

from mlflow.entities import Document
from typing import List, Dict, Any
from dataclasses import asdict
from databricks.vector_search.client import VectorSearchClient

class VectorSearchRetriever:
    """
    Class using Databricks Vector Search to retrieve relevant documents.
    """
    def __init__(self):
        self.vector_search_client = VectorSearchClient(disable_notice=True)
        # TODO: Replace this with the list of column names to return in the result when querying Vector Search
        self.columns = ["chunk_id", "text_column", "doc_uri"]
        self.vector_search_index = self.vector_search_client.get_index(
            index_name="catalog.schema.chunked_docs_index"
        )
        mlflow.models.set_retriever_schema(
            name="vector_search",
            primary_key="chunk_id",
            text_column="text_column",
            doc_uri="doc_uri"
        )

    @mlflow.trace(span_type="RETRIEVER", name="vector_search")
    def __call__(
        self,
        query: str,
        filters: Dict[Any, Any] = None,
        score_threshold = None
    ) -> List[Document]:
        """
        Performs vector search to retrieve relevant chunks.
        Args:
            query: Search query.
            filters: Optional filters to apply to the search. Filters must follow the Databricks Vector Search filter spec
            score_threshold: Score threshold to use for the query.

        Returns:
            List of retrieved Documents.
        """

        results = self.vector_search_index.similarity_search(
            query_text=query,
            columns=self.columns,
            filters=filters,
            num_results=5,
            query_type="ann"
        )

        documents = self.convert_vector_search_to_documents(
            results, 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 = []
        column_names = [column["name"] for column in vs_results.get("manifest", {}).get("columns", [])]
        result_row_count = vs_results.get("result", {}).get("row_count", 0)

        if result_row_count > 0:
            for item in vs_results["result"]["data_array"]:
                metadata = {}
                score = item[-1]

                if score >= score_threshold:
                    metadata["similarity_score"] = score
                    for i, field in enumerate(item[:-1]):
                        metadata[column_names[i]] = field

                    page_content = metadata.pop("text_column", None)

                    if page_content:
                        doc = Document(
                            page_content=page_content,
                            metadata=metadata
                        )
                        docs.append(doc)

        return docs

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 = "What is Databricks?"
filters={"text_column LIKE": "Databricks"},
results = retriever(query, filters=filters, score_threshold=0.1)

Définir le schéma de l’extracteur

Pour vous assurer que les récupérateurs sont suivis correctement et s’affichent correctement dans les applications en aval, appelez mlflow.models.set_retriever_schema lorsque vous définissez votre agent. Utilisez set_retriever_schema pour 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"],
)

Vous pouvez également spécifier des colonnes supplémentaires dans le schéma de votre extracteur en fournissant une liste de noms de colonnes avec le champ other_columns.

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 ensembles d’évaluation utilisent doc_uri pour comparer les résultats du récupérateur aux ensembles de données d’évaluation prédéfinis afin de déterminer le rappel et la précision du récupérateur. Consultez ensembles d'évaluation.

Suivre l’extracteur

Le traçage MLflow améliore l’observabilité en capturant des informations détaillées sur l’exécution de votre agent. Il permet d’enregistrer les entrées, sorties et métadonnées associées à chaque étape intermédiaire d’une requête, ce qui vous permet d’identifier rapidement la source des bogues et des comportements inattendus.

Cet exemple utilise le décorateur @mlflow.trace pour créer une trace pour l’extracteur et l’analyseur. Pour obtenir d’autres options de configuration des méthodes de suivi, consultez traçage MLflow pour les agents.

Le décorateur crée une étendue de qui commence lorsque la fonction est appelée et se termine lorsqu'elle renvoie. MLflow enregistre automatiquement l’entrée et la sortie de la fonction et toutes les exceptions déclenchées.

Remarque

Les utilisateurs de bibliothèque LangChain, LlamaIndex et OpenAI peuvent utiliser l’enregistrement automatique MLflow au lieu de définir manuellement des traces avec le décorateur. Consultez Utiliser l'autologging pour ajouter des traces à vos agents.

...
@mlflow.trace(span_type="RETRIEVER", name="vector_search")
def __call__(self, query: str) -> List[Document]:
  ...

Pour garantir que les applications en aval, telles que l'Évaluation de l'Agent et l'IA Playground, affichent correctement la trace du récupérateur, assurez-vous que le décorateur répond aux exigences suivantes :

  • Utilisez span_type="RETRIEVER" et vérifiez que la fonction retourne l’objet List[Document]. Consultez les étendues Récupérateur.
  • Le nom de la trace et le nom retriever_schema doivent correspondre pour configurer la trace correctement.

Filtrer les résultats de la recherche vectorielle

Vous pouvez limiter l’étendue de recherche à un sous-ensemble de données à l’aide d’un filtre Recherche vectorielle.

Le paramètre filters dans VectorSearchRetriever définit les conditions de filtre à l’aide de la spécification de filtre Recherche vectorielle Databricks.

filters = {"text_column LIKE": "Databricks"}

À l’intérieur de la méthode __call__, le dictionnaire de filtres est transmis directement à la fonction similarity_search :

results = self.vector_search_index.similarity_search(
    query_text=query,
    columns=self.columns,
    filters=filters,
    num_results=5,
    query_type="ann"
)

Après le filtrage initial, le paramètre score_threshold fournit un filtrage supplémentaire en définissant un score de similarité minimal.

if score >= score_threshold:
    metadata["similarity_score"] = score

Le résultat final comprend des documents qui répondent aux conditions filters et score_threshold.

Étapes suivantes

Après avoir créé un outil d'agent de fonction Unity Catalog, ajoutez l'outil à un agent IA. Consultez Ajouter des outils de catalogue Unity aux agents.