Verbinden von Agents mit unstrukturierten Daten

Abfragen eines Databricks-Vektorsuchindex mithilfe der Vektorsuche MCP

Wenn Ihr Agent einen Databricks Vector Search-Index abfragen muss, verwenden Sie dafür den vom Databricks verwalteten MCP-Server:

1. Erstellen Sie einen Vektorsuchindex mit verwalteten Databricks-Embeddings. Siehe Erstellen von Vektorsuchendpunkten und Indizes.
2. Erstellen Sie einen Agent, und verbinden Sie ihn mit der vorkonfigurierten verwalteten MCP-URL für den Vektorsuchindex: https://<workspace-hostname>/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}.

Informationen zum Erstellen eines Agents, der eine Verbindung mit verwalteten MCP-Servern herstellt, finden Sie unter Verwenden von verwalteten MCP-Servern von Databricks.

Abfragen eines Vektorsuchindex, der außerhalb von Databricks gehostet wird

Wenn Ihr Vektorindex außerhalb von Azure Databricks gehostet wird, können Sie eine Unity-Katalogverbindung erstellen, um eine Verbindung mit dem externen Dienst herzustellen und die Verbindung in Ihrem Agentcode zu verwenden. Siehe Verbinden von KI-Agent-Tools mit externen Diensten.

Im folgenden Beispiel wird ein Retriever erstellt, der einen Vektorindex aufruft, der außerhalb von Databricks für einen PyFunc-Flavor-Agent gehostet wird.

1. Erstellen Sie eine Unity-Katalogverbindung mit dem externen Dienst, in diesem Fall Azure.

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

2. Definieren Sie das Retriever-Tool im Agentcode mithilfe der Unity-Katalogverbindung. In diesem Beispiel werden MLflow-Decorator-Elemente verwendet, um die Agent-Nachverfolgung zu aktivieren.

   Hinweis

   Um dem MLflow-Retriever-Schema zu entsprechen, sollte die Retriever-Funktion ein List[Document]-Objekt zurückgeben und das metadata-Feld in der Document-Klasse verwenden, um dem zurückgegebenen Dokument zusätzliche Attribute hinzuzufügen, wie z. B. doc_uri und similarity_score. Siehe MLflow-Dokument.

   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. Führen Sie zum Ausführen des Retrievers den folgenden Python-Code aus. Sie können optional Vektorsuchfilter in die Anforderung einschließen, um Ergebnisse zu filtern.

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

Entwickeln eines Retrievers lokal mithilfe der AI-Brücke

Um ein Databricks Vector Search Retriever-Tool lokal zu erstellen, verwenden Sie Databricks AI Bridge-Pakete wie databricks-langchain und databricks-openai. Zu diesen Paketen gehören Hilfsfunktionen wie from_vector_search und from_uc_function zum Erstellen von Retrievern aus vorhandenen Databricks-Ressourcen.

LangChain/LangGraph

Installieren Sie die neueste Version davon databricks-langchain , die Databricks AI Bridge enthält.

%pip install --upgrade databricks-langchain

Der folgende Codeprototyp stellt ein Retriever-Tool bereit, das einen hypothetischen Vektorsuchindex abfragt und lokal an eine LLM bindet, damit Sie das Toolaufrufverhalten testen können.

Stellen Sie einen beschreibenden Text tool_description bereit, der dem Agent hilft, das Tool zu verstehen und zu bestimmen, wann es aufgerufen werden soll.

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

Für Szenarien, die entweder Direktzugriffsindizes oder Delta-Synchronisierungsindizes mit selbstverwalteten Einbettungen verwenden, müssen Sie das VectorSearchRetrieverTool benutzerdefinierte Einbettungsmodell und eine Textspalte konfigurieren und angeben. Weitere Informationen finden Sie unter Optionen zum Bereitstellen von Einbettungen.

Das folgende Beispiel zeigt, wie Sie ein VectorSearchRetrieverTool mit columns- und embedding-Schlüsseln konfigurieren.

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

Weitere Details finden Sie in den API-Dokumenten für VectorSearchRetrieverTool.

OpenAI

Installieren Sie die neueste Version davon databricks-openai , die Databricks AI Bridge enthält.

%pip install --upgrade databricks-openai

Der folgende Codeprototyp veranschaulicht einen Retriever, der einen hypothetischen Vektorsuchindex abfragt und in die GPT-Modelle von OpenAI integriert.

Stellen Sie einen beschreibenden Text tool_description bereit, der dem Agent hilft, das Tool zu verstehen und zu bestimmen, wann es aufgerufen werden soll.

Weitere Informationen zu OpenAI-Empfehlungen für Tools finden Sie in OpenAI Function Calling-Dokumentation.

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

Für Szenarien, die entweder Direktzugriffsindizes oder Delta-Synchronisierungsindizes mit selbstverwalteten Einbettungen verwenden, müssen Sie das VectorSearchRetrieverTool benutzerdefinierte Einbettungsmodell und eine Textspalte konfigurieren und angeben. Weitere Informationen finden Sie unter Optionen zum Bereitstellen von Einbettungen.

Das folgende Beispiel zeigt, wie Sie ein VectorSearchRetrieverTool mit columns- und embedding-Schlüsseln konfigurieren.

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

Weitere Details finden Sie in den API-Dokumenten für VectorSearchRetrieverTool.

Nachdem Ihr lokales Tool bereit ist, können Sie es direkt als Teil Ihres Agentcodes produzieren oder zu einer Unity-Katalogfunktion migrieren, die eine bessere Auffindbarkeit und Governance bietet, aber bestimmte Einschränkungen hat.

Databricks-Vektorsuche mit UC-Funktionen (veraltet)

Sie können eine Unity-Katalogfunktion erstellen, die eine Indexabfrage für die Mosaik-AI-Vektorsuche umschließt. Dieser Ansatz:

• Unterstützung für produktive Anwendungsfälle mit Governance und Auffindbarkeit
• Verwendet die SQL-Funktion vector_search() unter der Haube.
• Unterstützt automatisches MLflow-Tracing
  • Sie müssen die Ausgabe der Funktion mithilfe der Aliase und page_content an das metadata ausrichten.
  • Alle zusätzlichen Metadatenspalten müssen der Spalte mithilfe der metadataSQL-Zuordnungsfunktion und nicht als Ausgabeschlüssel der obersten Ebene hinzugefügt werden.

Führen Sie den folgenden Code in einem Notizbuch oder SQL-Editor aus, um die Funktion zu erstellen:

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
  )

Um dieses Retriever-Tool in Ihrem KI-Agent zu verwenden, schließen Sie es mit UCFunctionToolkitum. Dies ermöglicht automatisches Tracing durch MLflow, indem automatisch RETRIEVER Span-Typen in MLflow-Protokollen generiert werden.

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Unity Catalog Retriever-Tools weisen die folgenden Vorbehalte auf:

• SQL-Clients können die maximale Anzahl zurückgegebener Zeilen oder Bytes einschränken. Um das Abschneiden von Daten zu verhindern, kürzen Sie spaltenwerte ab, die von der UDF zurückgegeben werden. Sie können z. B. substring(chunked_text, 0, 8192) verwenden, um die Größe großer Inhaltsspalten zu verringern und zeilenabkürzungen während der Ausführung zu vermeiden.
• Da es sich bei diesem Tool um einen Wrapper für die vector_search() Funktion handelt, unterliegt es den gleichen Einschränkungen wie die vector_search() Funktion. Informationen finden Sie unter Einschränkungen.

Weitere Informationen zu UCFunctionToolkit finden Sie in Unity-Katalogdokumentation.