Ansluta agenter till ostrukturerade data

AI-agenter behöver ofta fråga ostrukturerade data som dokumentsamlingar, kunskapsbaser eller text corpora för att besvara frågor och tillhandahålla sammanhangsmedvetna svar.

Databricks tillhandahåller flera metoder för att ansluta agenter till ostrukturerade data i Vector Search-index och externa vektorlager. Använd förkonfigurerade MCP-servrar för omedelbar åtkomst till Databricks Vector Search-index, utveckla hämtningsverktyg lokalt med AI Bridge-paket eller skapa anpassade hämtningsfunktioner för specialiserade arbetsflöden.

Fråga ett Databricks Vector Search-index med hjälp av MCP

Om din agent behöver fråga ett Databricks Vector Search-index använder du den Databricks-hanterade MCP-servern. Innan du börjar skapar du ett vektorsökningsindex med databrickshanterade inbäddningar. Se Skapa slutpunkter och index för vektorsökning.

Den hanterade MCP-URL:en för Vector Search är: https://<workspace-hostname>/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}.

I följande exempel visas hur du ansluter din agent till ett Vector Search-index. Ersätt <catalog>, <schema>och <index-name> med namnen på ditt Vector Search-index.

OpenAI Agents SDK (Appar)

from agents import Agent, Runner
from databricks.sdk import WorkspaceClient
from databricks_openai.agents import McpServer

workspace_client = WorkspaceClient()

async with McpServer.from_vector_search(
    catalog="<catalog>",
    schema="<schema>",
    index_name="<index-name>",
    workspace_client=workspace_client,
    name="vector-search",
) as vs_server:
    agent = Agent(
        name="Research assistant",
        instructions="You are a research assistant. Use the vector search tool to find relevant documents and answer questions.",
        model="databricks-claude-sonnet-4-5",
        mcp_servers=[vs_server],
    )
    result = await Runner.run(agent, "What is the return policy?")
    print(result.final_output)

Ge appen åtkomst till Vector Search-indexet i databricks.yml:

resources:
  apps:
    my_agent_app:
      resources:
        - name: 'my_vector_index'
          uc_securable:
            securable_full_name: '<catalog>.<schema>.<index-name>'
            securable_type: 'TABLE'
            permission: 'SELECT'

LangGraph (appar)

from databricks.sdk import WorkspaceClient
from databricks_langchain import ChatDatabricks, DatabricksMCPServer, DatabricksMultiServerMCPClient
from langgraph.prebuilt import create_react_agent

workspace_client = WorkspaceClient()
host = workspace_client.config.host

mcp_client = DatabricksMultiServerMCPClient([
    DatabricksMCPServer(
        name="vector-search",
        url=f"{host}/api/2.0/mcp/vector-search/<catalog>/<schema>/<index-name>",
        workspace_client=workspace_client,
    ),
])

async with mcp_client:
    tools = await mcp_client.get_tools()
    agent = create_react_agent(
        ChatDatabricks(endpoint="databricks-claude-sonnet-4-5"),
        tools=tools,
    )
    result = await agent.ainvoke(
        {"messages": [{"role": "user", "content": "What is the return policy?"}]}
    )
    print(result["messages"][-1].content)

Ge appen åtkomst till Vector Search-indexet i databricks.yml:

resources:
  apps:
    my_agent_app:
      resources:
        - name: 'my_vector_index'
          uc_securable:
            securable_full_name: '<catalog>.<schema>.<index-name>'
            securable_type: 'TABLE'
            permission: 'SELECT'

Modellhantering

from databricks.sdk import WorkspaceClient
from databricks_mcp import DatabricksMCPClient
import mlflow

workspace_client = WorkspaceClient()
host = workspace_client.config.host

# Connect to the Vector Search MCP server
mcp_client = DatabricksMCPClient(
    server_url=f"{host}/api/2.0/mcp/vector-search/<catalog>/<schema>/<index-name>",
    workspace_client=workspace_client,
)

# List available tools from the Vector Search index
tools = mcp_client.list_tools()

# Log the agent with the required resources for deployment
mlflow.pyfunc.log_model(
    "agent",
    python_model=my_agent,
    resources=mcp_client.get_databricks_resources(),
)

Information om hur du distribuerar agenten finns i Distribuera en agent för generativa AI-program (modellservering). Mer information om loggningsagenter med MCP-resurser finns i Använda Databricks-hanterade MCP-servrar.

Andra metoder

%pip install --upgrade databricks-langchain

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

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

%pip install --upgrade databricks-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]
)

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

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
  )

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Lägga till spårning i ett hämtningsverktyg

Lägg till MLflow-spårning för att övervaka och felsöka din retriever. Med spårning kan du visa indata, utdata och metadata för varje körningssteg.

I föregående exempel läggs @mlflow.trace-dekoratören till både __call__ metoderna och parsning. Dekoratören skapar ett span som startar när funktionen anropas och slutar när den returneras. MLflow registrerar automatiskt funktionens indata och utdata och eventuella undantag som genereras.

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]:
    ...

Kontrollera att dekoratören uppfyller följande krav för att verifiera att nedströms applikationer, såsom Agent Evaluation och AI Playground, korrekt återger retriever-spårningen.

• Använd schemat för MLflow retriever span och kontrollera att funktionen returnerar ett List[Document]-objekt.
• Spårningsnamnet och retriever_schema namnet måste matcha för att kunna konfigurera spårningen korrekt. Se följande avsnitt för att lära dig hur du ställer in retriever-schemat.

Ange retriever-schema för att verifiera MLflow-kompatibilitet

Om spårningen som returneras från hämtaren eller span_type="RETRIEVER" inte överensstämmer med MLflows standardåterhämtningsschema måste du mappa det returnerade schemat manuellt till MLflows förväntade fält. Detta verifierar att MLflow korrekt kan spåra din retriever och återge spårningar i underordnade program.

Så här anger du upphämtningsschema manuellt:

1. Anropa mlflow.models.set_retriever_schema när du definierar din agent. Använd set_retriever_schema för att mappa kolumnnamnen i den returnerade tabellen till MLflows förväntade fält, till exempel primary_key, text_columnoch 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. Ange ytterligare kolumner i hämtarens schema genom att ange en lista med kolumnnamn med fältet other_columns .

3. Om du har flera retrievers kan du definiera flera scheman med hjälp av unika namn för varje retrieverschema.

Datahämtningsschemauppsättningen när agenten skapas påverkar underordnade program och arbetsflöden, till exempel granskningsappen och utvärderingsuppsättningarna. Mer specifikt fungerar kolumnen doc_uri som primär identifierare för dokument som returneras av hämtaren.

• granskningsappen visar doc_uri för att hjälpa granskare att utvärdera svar och spåra dokument ursprung. Se över Granska appens användargränssnitt.
• Utvärderingsuppsättningar används doc_uri för att jämföra hämtningsresultat med fördefinierade utvärderingsdatauppsättningar för att fastställa hämtarens träffsäkerhet och precision. Se Utvärderingsuppsättningar (MLflow 2).

Läsa filer från en Unity Catalog-volym

Om din agent behöver läsa ostrukturerade filer (textdokument, rapporter, konfigurationsfiler osv.) som lagras i en Unity Catalog-volym kan du skapa verktyg som använder Databricks SDK Files API för att visa och läsa filer direkt.

I följande exempel skapas två verktyg som agenten kan använda:

• list_volume_files: Visar filer och kataloger i volymen.
• read_volume_file: Läser innehållet i en textfil från volymen.

LangChain/LangGraph

Installera den senaste versionen av databricks-langchain som innehåller Databricks AI Bridge.

%pip install --upgrade databricks-langchain

from databricks.sdk import WorkspaceClient
from langchain_core.tools import tool

VOLUME = "<catalog>.<schema>.<volume>"  # TODO: Replace with your volume
w = WorkspaceClient()


@tool
def list_volume_files(directory: str = "") -> str:
    """Lists files and directories in the Unity Catalog volume.
    Provide a relative directory path, or leave empty to list the volume root."""
    base = f"/Volumes/{VOLUME.replace('.', '/')}"
    path = f"{base}/{directory.lstrip('/')}" if directory else base
    entries = []
    for f in w.files.list_directory_contents(path):
        kind = "dir" if f.is_directory else "file"
        size = f" ({f.file_size} bytes)" if not f.is_directory else ""
        entries.append(f"  [{kind}] {f.name}{size}")
    return "\n".join(entries) if entries else "No files found."


@tool
def read_volume_file(file_path: str) -> str:
    """Reads a text file from the Unity Catalog volume.
    Provide the path relative to the volume root, for example 'reports/q1_summary.txt'."""
    base = f"/Volumes/{VOLUME.replace('.', '/')}"
    full_path = f"{base}/{file_path.lstrip('/')}"
    resp = w.files.download(full_path)
    return resp.contents.read().decode("utf-8")

Binda verktygen till en LLM och kör en verktygsanropsloop:

from databricks_langchain import ChatDatabricks
from langchain_core.messages import HumanMessage, ToolMessage

llm = ChatDatabricks(endpoint="databricks-claude-sonnet-4-5")
llm_with_tools = llm.bind_tools([list_volume_files, read_volume_file])

messages = [HumanMessage(content="What files are in the volume? Can you read about_databricks.txt and summarize it in 2 sentences?")]
tool_map = {"list_volume_files": list_volume_files, "read_volume_file": read_volume_file}

for _ in range(5):  # max iterations
    response = llm_with_tools.invoke(messages)
    messages.append(response)
    if not response.tool_calls:
        break
    for tc in response.tool_calls:
        result = tool_map[tc["name"]].invoke(tc["args"])
        messages.append(ToolMessage(content=result, tool_call_id=tc["id"]))

print(response.content)

OpenAI (på engelska)

Installera den senaste versionen av databricks-openai som innehåller Databricks AI Bridge.

%pip install --upgrade databricks-openai

from databricks.sdk import WorkspaceClient
from databricks_openai import DatabricksOpenAI
import json

VOLUME = "<catalog>.<schema>.<volume>"  # TODO: Replace with your volume
w = WorkspaceClient()
client = DatabricksOpenAI()

# Define the tool specifications
tools = [
    {
        "type": "function",
        "function": {
            "name": "list_volume_files",
            "description": "Lists files and directories in the Unity Catalog volume. Provide a relative directory path, or leave empty to list the volume root.",
            "parameters": {
                "type": "object",
                "properties": {
                    "directory": {
                        "type": "string",
                        "description": "Relative directory path within the volume. Leave empty for root.",
                    }
                },
                "required": [],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "read_volume_file",
            "description": "Reads a text file from the Unity Catalog volume. Provide the path relative to the volume root, for example 'reports/q4_summary.txt'.",
            "parameters": {
                "type": "object",
                "properties": {
                    "file_path": {
                        "type": "string",
                        "description": "Path to the file relative to the volume root.",
                    }
                },
                "required": ["file_path"],
            },
        },
    },
]


def execute_tool(name: str, args: dict) -> str:
    base = f"/Volumes/{VOLUME.replace('.', '/')}"
    if name == "list_volume_files":
        directory = args.get("directory", "")
        path = f"{base}/{directory.lstrip('/')}" if directory else base
        entries = []
        for f in w.files.list_directory_contents(path):
            kind = "dir" if f.is_directory else "file"
            size = f" ({f.file_size} bytes)" if not f.is_directory else ""
            entries.append(f"[{kind}] {f.name}{size}")
        return "\n".join(entries) if entries else "No files found."
    elif name == "read_volume_file":
        full_path = f"{base}/{args['file_path'].lstrip('/')}"
        resp = w.files.download(full_path)
        return resp.contents.read().decode("utf-8")
    return f"Unknown tool: {name}"


# Call the model with tools
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "List the files in the volume, then read about_databricks.txt and summarize it."},
]

response = client.chat.completions.create(
    model="databricks-claude-sonnet-4-5", messages=messages, tools=tools
)

# Execute tool calls and send results back
while response.choices[0].finish_reason == "tool_calls":
    messages.append(response.choices[0].message)
    for tool_call in response.choices[0].message.tool_calls:
        args = json.loads(tool_call.function.arguments)
        result = execute_tool(tool_call.function.name, args)
        messages.append(
            {"role": "tool", "tool_call_id": tool_call.id, "content": result}
        )
    response = client.chat.completions.create(
        model="databricks-claude-sonnet-4-5", messages=messages, tools=tools
    )

print(response.choices[0].message.content)