Nástroje pro nestrukturované vyhledávání AI agenta

Článek
02/21/2025

Důležitý

Tato funkce je ve verzi Public Preview.

V tomto článku se dozvíte, jak vytvořit nástroje agenta AI pro načítání nestrukturovaných dat pomocí architektury Agent Frameworku systému Mosaic AI. Nestrukturované načítáče umožňují agentům dotazovat nestrukturované zdroje dat, jako je korpus dokumentu, pomocí indexů vektorového vyhledávání.

Další informace o nástrojích agenta najdete v tématu nástroje agenta AI.

Místně vyvíjejte nástroje pro vektorové vyhledávání pomocí AI Bridge.

Nejjednodušší způsob, jak začít vyvíjet nástroj pro načítání vektorového vyhledávání Databricks, je místní. Balíčky AI Bridge Databricks, jako databricks-langchain a databricks-openai, umožňují rychle přidat vyhledávací schopnosti agentovi a experimentovat s parametry dotazu. Tento přístup umožňuje rychlou iteraci během počátečního vývoje.

Jakmile je místní nástroj připravený, můžete ho přímo productionizovat jako součást kódu agenta nebo ho migrovat do funkce Katalogu Unity, která poskytuje lepší zjistitelnost a zásady správného řízení, ale má určitá omezení. Viz nástroj pro načítání vektorového vyhledávání s funkcemi katalogu Unity.

Poznámka

Pokud chcete použít externí vektorový index hostovaný mimo Databricks, viz Vektorové vyhledávání pomocí vektorového indexu hostovaného mimo Databricks.

LangChain/LangGraph

Následující kód vytvoří prototyp nástroje retrieveru a místně ho sváže s LLM, abyste mohli chatovat s agentem a otestovat jeho chování při volání nástrojů.

Nainstalujte nejnovější verzi databricks-langchain, která zahrnuje most Databricks AI.

%pip install --upgrade databricks-langchain

Následující příklad dotazuje hypotetický index vyhledávání vektorů, který načte obsah z dokumentace k produktu Databricks.

Zadejte jasné a popisné tool_description. Agent LLM používá tool_description k pochopení nástroje a určení, kdy se má nástroj vyvolat.

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

Poznámka

Při inicializaci VectorSearchRetrieverTooljsou argumenty text_column a embedding vyžadovány pro indexy Delta Sync s vloženými funkcemi a indexy přímého vektoru. Podívejte se na možnosti pro poskytování embeddingů.

Další podrobnosti najdete v dokumentaci rozhraní API pro 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

Následující kód prototypuje nástroj pro načítání vektorového vyhledávání a integruje ho s modely GPT OpenAI.

Další informace o doporučeních OpenAI pro nástroje najdete v dokumentaci k volání funkcí OpenAI.

Nainstalujte nejnovější verzi databricks-openai, která zahrnuje most Databricks AI.

%pip install --upgrade databricks-openai

Následující příklad dotazuje hypotetický index vyhledávání vektorů, který načte obsah z dokumentace k produktu Databricks.

Zadejte jasné a popisné tool_description. Agent LLM používá tool_description k pochopení nástroje a určení, kdy se má nástroj vyvolat.

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

Poznámka

Podrobnosti naleznete v dokumentaci rozhraní API pro 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.
)

nástroj Vector Search retriever s funkcemi katalogu Unity

Následující příklad vytvoří nástroj retriever pomocí funkce Unity Catalog k dotazování dat z indexu Mosaic AI Vector Search.

Funkce Unity Catalog databricks_docs_vector_search dotazuje hypotetický index vektorového vyhledávání obsahující dokumentaci k Databricks. Tato funkce obalí funkci SQL Databricks vector_search() a zarovná její výstup se schématem MLflow retriever. pomocí aliasů page_content a metadata.

Poznámka

Aby bylo možné vyhovovat schématu načítání MLflow, je nutné do sloupce metadata přidat jakékoliv další sloupce metadat pomocí SQL mapovací funkce, a nikoliv jako výstupní klíče vyšší úrovně.

Spuštěním následujícího kódu v poznámkovém bloku nebo editoru SQL vytvořte funkci:

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
  )

Pokud chcete tento nástroj pro vyhledávání použít v agentu AI, obalte jej UCFunctionToolkit. To umožňuje automatické trasování přes MLflow.

Sledování MLflow zaznamenává podrobné informace o spuštění pro aplikace generativní umělé inteligence. Protokoluje vstupy, výstupy a metadata pro každý krok, což vám pomůže ladit problémy a analyzovat výkon.

Při použití UCFunctionToolkitgenerují retrievery automaticky RETRIEVER typy rozsahu v protokolech MLflow, pokud jejich výstup odpovídá schématu načítání MLflow. Podívejte se na schéma trasování MLflow.

Další informace o UCFunctionToolkit naleznete v dokumentaci k katalogu Unity.

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Tento vyhledávací nástroj má následující výhrady:

Klienti SQL můžou omezit maximální počet vrácených řádků nebo bajtů. Chcete-li zabránit zkrácení dat, měli byste zkrátit hodnoty ve sloupcích vrácené uživatelsky definovanou funkcí. Pomocí substring(chunked_text, 0, 8192) můžete například zmenšit velikost velkých sloupců obsahu a vyhnout se zkrácení řádků během provádění.
Vzhledem k tomu, že je tento nástroj obálkou pro funkci vector_search(), podléhá stejným omezením jako funkce vector_search(). Viz Omezení.

Modul pro vyhledávání vektoru využívající vektorový index hostovaný mimo Databricks

Pokud je váš vektorový index hostovaný mimo Azure Databricks, můžete vytvořit připojení Unity Catalog Connection pro připojení k externí službě a použít kód agenta připojení. Další informace najdete v tématu Připojení nástrojů agenta AI k externím službám.

Následující příklad vytvoří retriever vektorového vyhledávání, který volá vektorový index hostovaný mimo Databricks pro agenta ve stylu PyFunc.

Vytvořte připojení katalogu Unity k externí službě, v tomto případě Azure.

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

Definujte nástroj pro vyhledávání v kódu agenta pomocí připojení k Unity Catalogu, které jste vytvořili. Tento příklad používá dekorátory MLflow pro umožnění sledování agenta.

Poznámka

Aby bylo možné postupovat podle schématu načítání MLflow, měla by funkce retrieveru vrátit typ dokumentu a pomocí pole ve třídě Dokumentu přidat do vráceného dokumentu další atributy, jako jsou a

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

Pokud chcete spustit retriever, spusťte následující kód Pythonu. Volitelně můžete do požadavku zahrnout filtry vektorového vyhledávání k filtrování výsledků.
```
retriever = VectorSearchRetriever()
query = [0.01944167, 0.0040178085 . . .  TRIMMED FOR BREVITY 010858015, -0.017496133]
results = retriever(query, score_threshold=0.1)
```

Nastavení schématu retrieveru

Pokud trasování vrácené z retrieveru nebo span_type="RETRIEVER" neodpovídá standardnímu schématu retrieveru MLflow, musíte ručně namapovat vrácené schéma na očekávaná pole MLflow. Tím zajistíte, že MLflow dokáže správně trasovat váš vyhledávač a vykreslit trasování správně v následných aplikacích.

Pokud chcete schéma retrieveru nastavit ručně, zavolejte mlflow.models.set_retriever_schema při definování agenta. Pomocí set_retriever_schema namapujte názvy sloupců ve vrácené tabulce na očekávaná pole MLflow, jako jsou primary_key, text_columna 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"],
)

Můžete také zadat další sloupce ve schématu vašeho načítáče tak, že zadáte seznam názvů sloupců s polem other_columns.

Pokud máte více vyhledávačů, můžete definovat více schémat pomocí jedinečných názvů pro každé schéma vyhledávače.

Sada schématu retrieveru během vytváření agenta má vliv na podřízené aplikace a pracovní postupy, jako je například revizní aplikace a sady vyhodnocení. Konkrétně sloupec doc_uri slouží jako primární identifikátor dokumentů vrácených retrieverem.

Aplikace pro kontrolu zobrazí doc_uri, která revidujícím pomáhá vyhodnotit odpovědi a sledovat původy dokumentů. Viz Kontrola uživatelského rozhraní aplikace.
Vyhodnocovací sady jsou použity doc_uri k porovnání výsledků retrieveru s předdefinovanými vyhodnocovacími datovými sadami za účelem určení jeho úplnosti a přesnosti. Viz sady vyhodnocení.

Sledování retrívra

Trasování MLflow přidává sledovatelnost tím, že zachycuje podrobné informace během spuštění vašeho agenta. Poskytuje způsob, jak zaznamenávat vstupy, výstupy a metadata přidružená ke každému přechodnému kroku požadavku, což vám umožní rychle určit zdroj chyb a neočekávané chování.

Tento příklad používá dekorátor @mlflow.trace k vytvoření sledování pro retriever a analyzátor. Další možnosti nastavení trasovacích metod najdete v tématu trasování MLflow pro agenty.

Dekorátor vytvoří úsek, který začíná, když je funkce vyvolána, a končí, když se vrátí. MLflow automaticky zaznamenává vstup a výstup funkce a všechny výjimky vyvolané.

Poznámka

Uživatelé knihoven LangChain, LlamaIndex a OpenAI mohou využít automatické protokolování v MLflow namísto ručního definování trasování pomocí dekorátoru. Vizte Použití automatického logování k přidání stop do agentů.

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

Aby se zajistilo, že následné aplikace, například Vyhodnocení Agenta a AI Hřiště, vykreslují trasování načítače správně, ujistěte se, že dekorátor splňuje následující požadavky:

Použijte span_type="RETRIEVER" a ujistěte se, že funkce vrací objekt List[Document]. Podívejte se na rozsahy Retrieveru.
Název sledování a název retriever_schema se musí shodovat, aby se sledování správně nakonfigurovalo.

Další kroky

Po vytvoření nástroje agenta pro funkci Unity Catalog přidejte tento nástroj do AI agenta. Viz Přidejte nástroje katalogu Unity do agentů.

Sdílet prostřednictvím

Nástroje pro nestrukturované vyhledávání AI agenta

Místně vyvíjejte nástroje pro vektorové vyhledávání pomocí AI Bridge.

LangChain/LangGraph

OpenAI

nástroj Vector Search retriever s funkcemi katalogu Unity

Modul pro vyhledávání vektoru využívající vektorový index hostovaný mimo Databricks

Nastavení schématu retrieveru

Sledování retrívra

Další kroky

Váš názor

Další materiály