Nástroje agenta umělé inteligence pro vyhledávání nestrukturovaných dat
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, například korpus dokumentu pomocí různých metod načítání, včetně vektorového vyhledávání.
Další informace o nástrojích agenta najdete v tématu nástroje agenta AI.
Lokálně vyvíjet nástroje pro vyhledávání vektorů s využitím AI Bridge
Nejjednodušší způsob, jak začít vyvíjet nástroj pro vyhledávání vektorů v prostředí Databricks, je na lokálním počítači. 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 vektorové vyhledávání používaný 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 VectorSearchRetrieverTool
jsou argumenty text_column
a embedding
vyžadovány pro indexy Delta Sync se samostatně spravovaným vkládáním a indexy pro přímý přístup k vektorům. 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
Při inicializaci VectorSearchRetrieverTool
jsou argumenty text_column
a embedding
vyžadovány pro indexy Delta Sync se samostatně spravovaným vkládáním a indexy pro přímý přístup k vektorům. Podívejte se na možnosti pro poskytování embeddingů.
Další podrobnosti najdete 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í aplikací generativní umělé inteligence. Zaznamenává vstupy, výstupy a metadata pro každý krok, což pomáhá při ladění problémů a analýze výkonu.
Při použití UCFunctionToolkit
generují retrievery automaticky RETRIEVER
typy rozsahu v protokolech MLflow, pokud jejich výstup odpovídá schématu retrievery 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 oříznutí dat, měli byste zkrátit hodnoty ve sloupcích vrácené uživatelsky definovanou funkcí (UDF). 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 funkcevector_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ří vyhledávač vektorového vyhledávání, který volá vektorový index hostovaný mimo Databricks pro agenta s podporou 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_column
a 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 používají
doc_uri
k porovnání výsledků retrieveru s předdefinovanými hodnoticími datovými sadami za účelem určení zpětné dolovatelnosti a přesnosti retrieveru. Viz sady vyhodnocení.
Sledování retrívra
Trasování MLflow přidává možnost sledování tím, že zachycuje podrobné informace o provádě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í trasování pro retriever a parser. Další možnosti nastavení trasovacích metod najdete v tématu Pozorovatelnost agenta pomocí trasování MLflow.
Dekorátor vytvoří časový interval, 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. Viz Automatické trasování.
...
@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í objektList[Document]
. Podívejte se na rozsah Retrieveru. - Název sledování a název retriever_schema musí být shodné, aby bylo sledování správně nakonfigurováno.
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ů.