Strumenti per il recupero di informazioni da agenti di intelligenza artificiale non strutturati

Articolo
02/21/2025

Questo articolo illustra come creare strumenti dell'agente di intelligenza artificiale per il recupero di dati non strutturati usando Mosaic AI Agent Framework. I retriever non strutturati consentono agli agenti di eseguire query su origini dati non strutturate, ad esempio un corpus di documenti, usando indici di ricerca vettoriali.

Per altre informazioni sugli strumenti dell'agente, vedere strumenti dell'agente di intelligenza artificiale.

Sviluppare localmente gli strumenti di recupero di ricerca vettoriale con AI Bridge

Il modo più semplice per iniziare a sviluppare uno strumento di recupero di Ricerca vettoriale di Databricks è in locale. Usare i pacchetti di Databricks AI Bridge , come databricks-langchain e databricks-openai, per aggiungere funzionalità di recupero a un agente rapidamente ed eseguire esperimenti con i parametri di query. Questo approccio consente un'iterazione rapida durante lo sviluppo iniziale.

Quando lo strumento locale è pronto, è possibile crearlo direttamente in produzione come parte del codice dell'agente o eseguirne la migrazione a una funzione del catalogo Unity, che offre una migliore individuabilità e governance, ma presenta alcune limitazioni. Consulta lo strumento di recupero ricerca vettoriale con funzioni del Unity Catalog.

Nota

Per usare un indice vettoriale esterno ospitato all'esterno di Databricks, vedere Vector Search retriever usando un indice vettoriale ospitato all'esterno di Databricks.

LangChain/LangGraph

Il seguente codice prototipa uno strumento di recupero e lo collega localmente a un LLM, permettendoti di chattare con l'agente per testarne il comportamento di richiamo degli strumenti.

Installare la versione più recente di databricks-langchain che include Databricks AI Bridge.

%pip install --upgrade databricks-langchain

Nell'esempio seguente viene eseguita una query su un indice di ricerca di vettori ipotetico che recupera il contenuto dalla documentazione del prodotto Databricks.

Fornisci un tool_description chiaro e descrittivo. L'agente LLM usa il tool_description per comprendere lo strumento e determinare quando richiamare lo strumento.

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

Nota

Quando si inizializza il VectorSearchRetrieverTool, gli argomenti text_column e embedding sono necessari per gli indici di sincronizzazione Delta con incorporamenti autogestiti e indici di accesso a vettori diretti. Consultare le opzioni per fornire incorporazioni.

Per altri dettagli, vedere la documentazione dell'API per VectorSearchRetrieverTool

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

OpenAI

Il codice seguente crea prototipi di uno strumento di recupero di ricerca vettoriale e lo integra con i modelli GPT di OpenAI.

Per altre informazioni sulle raccomandazioni openAI per gli strumenti, vedere documentazione relativa alla chiamata di funzioni OpenAI.

Installare la versione più recente di databricks-openai che include Databricks AI Bridge.

%pip install --upgrade databricks-openai

Nell'esempio seguente viene eseguita una query su un indice di ricerca di vettori ipotetico che recupera il contenuto dalla documentazione del prodotto Databricks.

Fornire un tool_description chiaro e descrittivo. L'agente LLM usa il tool_description per comprendere lo strumento e determinare quando richiamare lo strumento.

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

Nota

Per altri dettagli, vedere la documentazione dell'API per VectorSearchRetrieverTool

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

strumento di ricerca vettoriale con funzioni del Catalogo Unity

L'esempio seguente crea uno strumento di recupero usando una funzione del Catalogo Unity per interrogare i dati da un indice di ricerca a vettori di Mosaic AI.

La funzione Catalog di Unity databricks_docs_vector_search esegue una query su un ipotetico indice di ricerca vettoriale contenente la documentazione di Databricks. Questa funzione esegue il wrapping della funzione SQL di Databricks vector_search() e ne allinea l'output con lo schema di recupero MLflow . utilizzando gli alias page_content e metadata.

Nota

Per essere conforme allo schema del retriever MLflow, è necessario aggiungere eventuali colonne di metadati aggiuntive alla colonna metadata usando la funzione mappa SQL , anziché come chiavi di output di primo livello.

Eseguire il codice seguente in un notebook o in un editor SQL per creare la funzione:

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
  )

Per usare questo strumento di recupero nel tuo agente di intelligenza artificiale, racchiudilo con UCFunctionToolkit. In questo modo viene abilitata la traccia automatica tramite MLflow.

MLflow Tracing acquisisce informazioni dettagliate sull'esecuzione per le applicazioni di intelligenza artificiale di generazione. Registra input, output e metadati per ogni passaggio, consentendo di eseguire il debug dei problemi e analizzare le prestazioni.

Quando si usa UCFunctionToolkit, i retriever generano automaticamente tipi di span RETRIEVER nei log MLflow se il loro output è conforme allo schema del retriever di MLflow. Vedi lo schema di tracciamento MLflow .

Per altre informazioni sui UCFunctionToolkit, vedere documentazione di Unity Catalog.

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Questo strumento di recupero presenta le seguenti avvertenze:

I client SQL possono limitare il numero massimo di righe o byte restituiti. Per evitare il troncamento dei dati, è necessario troncare i valori delle colonne restituiti dalla UDF (funzione definita dall'utente). Ad esempio, è possibile usare substring(chunked_text, 0, 8192) per ridurre le dimensioni delle colonne di contenuto di grandi dimensioni ed evitare il troncamento delle righe durante l'esecuzione.
Poiché questo strumento è un wrapper per la funzione vector_search(), è soggetto alle stesse limitazioni della funzione vector_search(). Vedere Limitazioni.

Recuperatore di ricerca vettoriale utilizzando un indice vettoriale ospitato al di fuori di Databricks

Se l'indice vettoriale è ospitato all'esterno di Azure Databricks, è possibile creare una connessione del catalogo Unity per connettersi al servizio esterno e usare il codice dell'agente di connessione. Per altre informazioni, vedere Connettere gli strumenti dell'agente di intelligenza artificiale ai servizi esterni.

Nell'esempio seguente viene creato un retriever di ricerca vettoriale che chiama un indice vettoriale ospitato esternamente a Databricks per un agente di tipo PyFunc.

Creare una connessione del catalogo Unity al servizio esterno, in questo caso Azure.

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

Definisci lo strumento di recupero nel codice dell'agente utilizzando la connessione al Catalogo Unity che hai creato. Questo esempio usa i decoratori MLflow per abilitare il tracciamento dell'agente.

Nota

Per essere conforme allo schema del retriever MLflow, la funzione retriever deve restituire un tipo di document e utilizzare il campo metadata nella classe Document per aggiungere altri attributi al documento restituito, ad esempio like doc_uri e similarity_score.

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

Per eseguire il retriever, eseguire il codice Python seguente. Facoltativamente, è possibile includere filtri di ricerca vettoriale nella richiesta per filtrare i risultati.
```
retriever = VectorSearchRetriever()
query = [0.01944167, 0.0040178085 . . .  TRIMMED FOR BREVITY 010858015, -0.017496133]
results = retriever(query, score_threshold=0.1)
```

Impostare il modello retriever

Se la traccia restituita dal retriever o span_type="RETRIEVER" non è conforme allo schema del retriever standard di MLflow , è necessario eseguire manualmente il mapping dello schema restituito ai campi previsti di MLflow. In questo modo, MLflow può tracciare correttamente il retriever ed eseguire correttamente il rendering delle tracce nelle applicazioni downstream.

Per impostare manualmente lo schema del retriever, chiamare mlflow.models.set_retriever_schema quando si definisce l'agente. Usare set_retriever_schema per eseguire il mapping dei nomi delle colonne nella tabella restituita ai campi previsti di MLflow, ad esempio primary_key, text_columne 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"],
)

È anche possibile specificare colonne aggiuntive nello schema del retriever fornendo un elenco di nomi di colonna con il campo other_columns.

Se si dispone di più retriever, è possibile definire più schemi usando nomi univoci per ogni schema di recupero.

Lo schema del retriever impostato durante la creazione dell'agente influisce su applicazioni e flussi di lavoro downstream, ad esempio l'app di revisione e i set di valutazione. In particolare, la colonna doc_uri funge da identificatore primario per i documenti restituiti dal retriever.

L'app di revisione visualizza il doc_uri per aiutare i revisori a valutare le risposte e tracciare le origini dei documenti. Vedi Revisione dell'interfaccia utente dell'app.
Set di valutazione utilizzano doc_uri per confrontare i risultati del retriever con set di dati di valutazione predefiniti al fine di determinare il richiamo e la precisione del retriever. Vedere Set di valutazione.

Traccia il retriever

Il monitoraggio MLflow aggiunge osservabilità catturando informazioni dettagliate sull'esecuzione dell'agente. Offre un modo per registrare gli input, gli output e i metadati associati a ogni passaggio intermedio di una richiesta, consentendo di individuare rapidamente l'origine dei bug e comportamenti imprevisti.

In questo esempio viene utilizzato il decoratore @mlflow.trace per creare una traccia per il retriever e il parser. Per altre opzioni per la configurazione dei metodi di tracciamento, vedere Osservabilità degli agenti con MLflow Tracing.

Il decoratore crea un intervallo che inizia quando la funzione viene richiamata e termina al suo ritorno. MLflow registra automaticamente l'input e l'output della funzione ed eventuali eccezioni generate.

Nota

Gli utenti delle librerie LangChain, LlamaIndex e OpenAI possono usare la registrazione automatica di MLflow anziché definire manualmente le tracce con il decoratore. Vedere tracciamento automatico.

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

Per garantire che le applicazioni downstream come Agent Evaluation e AI Playground eseguano correttamente la rappresentazione della traccia del retriever, assicurarsi che il decoratore soddisfi i seguenti requisiti:

Usare span_type="RETRIEVER" e verificare che la funzione restituisca List[Document] oggetto . Vedi spans di Retriever.
Il nome della traccia e il nome retriever_schema devono corrispondere per configurare correttamente la traccia.

Passaggi successivi

Dopo aver creato uno strumento funzione di Unity Catalog, aggiungi lo strumento a un agente di intelligenza artificiale. Vedere Aggiungere strumenti del catalogo Unity agli agenti.

Condividi tramite

Strumenti per il recupero di informazioni da agenti di intelligenza artificiale non strutturati

Sviluppare localmente gli strumenti di recupero di ricerca vettoriale con AI Bridge

LangChain/LangGraph

OpenAI

strumento di ricerca vettoriale con funzioni del Catalogo Unity

Recuperatore di ricerca vettoriale utilizzando un indice vettoriale ospitato al di fuori di Databricks

Impostare il modello retriever

Traccia il retriever

Passaggi successivi

Commenti e suggerimenti

Risorse aggiuntive