Ligue agentes a dados não estruturados

Consultar um índice de pesquisa vetorial alojado fora dos Databricks

Se o seu índice vetorial estiver alojado fora do Azure Databricks, pode criar uma ligação ao Unity Catalog para se ligar ao serviço externo e usar a ligação no seu código de agente. Consulte Ligue agentes a ferramentas de terceiros com os Serviços MCP.

O exemplo seguinte cria um retriever que chama um índice vetorial alojado fora do Databricks para um agente com a funcionalidade PyFunc.

1. Crie uma conexão de catálogo Unity com o serviço externo, neste caso, o Azure.

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

2. Defina a ferramenta retriever no código do agente usando a conexão do Unity Catalog. Este exemplo usa decoradores MLflow para habilitar o rastreamento de agentes.

   Note

   Para estar em conformidade com o esquema MLflow retriever, a função retriever deve retornar um List[Document] objeto e usar o metadata campo na classe Document para adicionar atributos adicionais ao documento retornado, como doc_uri e similarity_score. Consulte Documento 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 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(),
           "Content-Type": "application/json",
         },
         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. Para executar o retriever, execute o seguinte código Python.

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

Desenvolver um retriever localmente usando o AI Bridge

Para criar localmente uma ferramenta de recuperação do Databricks AI Search, use os pacotes Databricks AI Bridge, como databricks-langchain e databricks-openai. Estes pacotes incluem funções auxiliares como from_vector_search e from_uc_function para criar retrievers a partir de recursos Databricks existentes.

LangChain/LangGraph

Instale a versão mais recente de databricks-langchain que inclui o Databricks AI Bridge.

%pip install --upgrade databricks-langchain

O código a seguir prototipa uma ferramenta retriever que consulta um índice de pesquisa vetorial hipotético e o vincula a um LLM localmente para que você possa testar seu comportamento de chamada de ferramenta.

Forneça um descritivo tool_description para ajudar o agente a entender a ferramenta e determinar quando invocá-la.

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

Para cenários que usam índices de acesso direto ou índices Delta Sync usando representações autogeridas, deve configurar o VectorSearchRetrieverTool e deve especificar um modelo de representação personalizado e uma coluna de texto. Consulte as opções para fornecer embeddings.

O exemplo a seguir mostra como configurar um VectorSearchRetrieverTool com as chaves columns e 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.
)

Para obter detalhes adicionais, consulte os documentos da API para VectorSearchRetrieverTool.

OpenAI

Instale a versão mais recente de databricks-openai que inclui o Databricks AI Bridge.

%pip install --upgrade databricks-openai

O código a seguir prototipa um retriever que consulta um índice de pesquisa vetorial hipotético e o integra com os modelos GPT da OpenAI.

Forneça um descritivo tool_description para ajudar o agente a entender a ferramenta e determinar quando invocá-la.

Para obter mais informações sobre recomendações do OpenAI para ferramentas, consulte documentação do OpenAI Function Calling.

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

Para cenários que usam índices de acesso direto ou índices Delta Sync usando representações autogeridas, deve configurar o VectorSearchRetrieverTool e deve especificar um modelo de representação personalizado e uma coluna de texto. Consulte as opções para fornecer embeddings.

O exemplo a seguir mostra como configurar um VectorSearchRetrieverTool com as chaves columns e 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.
)

Para obter detalhes adicionais, consulte os documentos da API para VectorSearchRetrieverTool.

Depois de a sua ferramenta local estar pronta, pode colocá-la em produção diretamente como parte do código do seu agente, ou migrá-la para uma função no Unity Catalog, que oferece melhor capacidade de descoberta e governação, mas tem certas limitações.

Consultar o Databricks AI Search com funções UC (preterido)

Pode criar uma função do Unity Catalog que encapsule uma consulta a um índice do AI Search. Esta abordagem:

• Suporta casos de uso de produção com governança e capacidade de descoberta
• Usa a função vetor_search() SQL nos bastidores
• Suporta rastreamento automático de MLflow
  • Você deve alinhar a saída da função ao esquema MLflow retriever usando os aliases page_content e metadata.
  • Quaisquer colunas de metadados adicionais devem ser adicionadas à coluna metadata usando a função SQL map, em vez de como chaves de saída de nível superior.

Execute o seguinte código em um bloco de anotações ou editor SQL para criar a função:

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
  )

Para usar essa ferramenta retriever em seu agente de IA, envolva-a com UCFunctionToolkit. Isso permite o rastreamento automático através do MLflow, criando automaticamente tipos de span nos logs do MLflow.

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

As ferramentas de recuperação do Unity Catalog têm as seguintes advertências:

• Os clientes SQL podem limitar o número máximo de linhas ou bytes devolvidos. Para evitar truncamento de dados, trunque os valores das colunas devolvidos pelo UDF. Por exemplo, você pode usar substring(chunked_text, 0, 8192) para reduzir o tamanho de colunas de conteúdo grandes e evitar o truncamento de linha durante a execução.
• Como esta ferramenta é um wrapper para a vector_search() função, está sujeita às mesmas limitações da vector_search() função. Consulte Limitações.

Para mais informações sobre UCFunctionToolkit, consulte a documentação do Unity Catalog.