Инструменты для неструктурированного поиска с использованием агентов ИИ

Важный

Эта функция доступна в общедоступной предварительной версии.

В этой статье показано, как создать средства агента ИИ для неструктурированного извлечения данных с помощью Платформы агента ИИ Мозаики. Неструктурированные средства извлечения позволяют агентам запрашивать неструктурированные источники данных, такие как корпус документа, с помощью индексов векторного поиска.

Дополнительные сведения о средствах агента см. в статье средства агента ИИ.

Локальная разработка средств извлечения векторного поиска с помощью AI Bridge

Самый простой способ начать разработку инструмента поиска по векторам Databricks — на локальном компьютере. Используйте пакеты Databricks AI Bridge, такие как databricks-langchain и databricks-openai, чтобы быстро добавить возможности извлечения в агент и экспериментировать с параметрами запроса. Этот подход обеспечивает быструю итерацию во время начальной разработки.

После того как локальное средство будет готово, вы можете напрямую создать его в составе кода агента или перенести его в функцию каталога Unity, которая обеспечивает лучшую возможность обнаружения и управления, но имеет определенные ограничения. См. инструмент для поиска на основе векторов с функциями каталога Unity.

LangChain/LangGraph

Следующий код прототипирует средство извлечения и привязывает его к LLM локально, чтобы вы могли общаться с агентом для проверки поведения вызова инструментов.

Установите последнюю версию databricks-langchain, включающую Мост ИИ Databricks.

%pip install --upgrade databricks-langchain

Заметка

При инициализации VectorSearchRetrieverToolаргументы text_column и embedding обязательны для Delta Sync индексов с самостоятельными встраиваниями и индексов прямого векторного доступа. См. варианты для предоставления вложений.

from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks

# Initialize the retriever tool
vs_tool = VectorSearchRetrieverTool(index_name="catalog.schema.my_index_name")

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

Чтобы настроить вызов инструмента, передайте дополнительные аргументы в VectorSearchRetrieverTool:

from databricks_langchain import VectorSearchRetrieverTool

vs_tool = VectorSearchRetrieverTool(
  index_name, # Index name in the format 'catalog.schema.index'
  num_results, # Max number of documents to return
  columns, # List of columns to include in the search
  filters, # Filters to apply to the query
  query_type, # Query type ("ANN" or "HYBRID").
  tool_name, # Used by the LLM to understand the purpose of the tool
  tool_description, # Used by the LLM to understand the purpose of the tool
  text_column, # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
  embedding # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

OpenAI

Следующий код прототипирует инструмент для поиска по векторам и интегрирует его с моделями GPT от OpenAI.

Дополнительные сведения о рекомендациях OpenAI для инструментов см. в документации по вызову функций OpenAI.

Установите последнюю версию databricks-openai, включающую Мост ИИ Databricks.

%pip install --upgrade databricks-openai

Заметка

При инициализации VectorSearchRetrieverToolаргументы text_column и embedding необходимы для индексов Delta Sync с самостоятельно управляемыми встраиваниями и индексов с прямым доступом к векторам. См. варианты использованиявекторных представлений.

from databricks_openai import VectorSearchRetrieverTool
from openai import OpenAI
import json

# Initialize OpenAI client
client = OpenAI(api_key=<your_API_key>)

# Call model with VectorSearchRetrieverTool defined
dbvs_tool = VectorSearchRetrieverTool(index_name="catalog.schema.my_index_name")
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]
)

Чтобы настроить вызов инструмента, передайте дополнительные аргументы в VectorSearchRetrieverTool:

from databricks_openai import VectorSearchRetrieverTool

vs_tool = VectorSearchRetrieverTool(
    index_name, # Index name in the format 'catalog.schema.index'
    num_results, # Max number of documents to return
    columns, # List of columns to include in the search
    filters, # Filters to apply to the query
    query_type, # Query type ("ANN" or "HYBRID").
    tool_name, # Used by the LLM to understand the purpose of the tool
    tool_description, # Used by the LLM to understand the purpose of the tool
    text_column, # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
    embedding_model_name # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

инструмент поиска по вектору с функциями каталога Unity

В следующем примере создается средство извлечения с помощью функции каталога Unity для запроса данных из индекса Векторного поиска мозаики ИИ.

Функция каталога Unity databricks_docs_vector_search запрашивает гипотетический индекс векторного поиска, содержащий документацию Databricks. Эта функция упаковывает функцию Databricks SQL vector_search() и выравнивает выходные данные с схемой извлечения MLflow. с помощью псевдонимов page_content и metadata.

Заметка

Чтобы соответствовать схеме извлекателя MLflow, все дополнительные столбцы метаданных необходимо добавить в столбец metadata с помощью функции карты SQL, а не в качестве ключей выходных данных верхнего уровня.

Чтобы создать функцию, выполните следующий код в записной книжке или редакторе SQL:

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
  )

Чтобы использовать это средство извлечения в агенте ИИ, оберните его с помощью UCFunctionToolkit. Это обеспечивает автоматическую трассировку через MLflow.

Трассировка MLflow записывает подробные сведения о выполнении для приложений генеративного ИИ. Он регистрирует входные данные, выходные данные и метаданные для каждого шага, помогая выполнять отладку проблем и анализировать производительность.

При использовании UCFunctionToolkitизвлекатели автоматически создают типы диапазонов RETRIEVER в журналах MLflow, если их выходные данные соответствуют схеме извлекателя MLflow. См. схему трассировки MLflow .

Дополнительные сведения о UCFunctionToolkit см. в документации по каталогу Unity .

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Это средство извлечения содержит следующие предостережения:

• Клиенты SQL могут ограничить максимальное количество возвращаемых строк или байтов. Чтобы предотвратить усечение данных, следует усечь значения столбцов, возвращаемые UDF. Например, можно использовать substring(chunked_text, 0, 8192) для уменьшения размера больших столбцов содержимого и предотвращения усечения строк во время выполнения.
• Так как это средство является оболочкой для функции vector_search(), она имеет те же ограничения, что и функция vector_search(). См. Ограничения.

Если этот пример не подходит для вашего варианта использования, создайте инструмент поиска по векторам с использованием пользовательского кода агента.

Извлекатель векторного поиска с кодом агента (PyFunc)

В следующем примере создается извлекатель для векторного поиска, используемый в агенте на основе PyFunc в коде агента.

В этом примере используется databricks-vectorsearch для создания базового извлекателя, выполняющего поиск подобия векторного поиска с фильтрами. В нем используются декораторы MLflow для включения трассировки агента.

Заметка

Чтобы соответствовать схеме извлекателя MLflow, функция извлекателя должна возвращать тип Document и использовать поле metadata в классе Document для добавления дополнительных атрибутов в возвращаемый документ, например like doc_uri и similarity_score.

Используйте следующий код в модуле агента или записной книжке агента.

import mlflow
import json

from mlflow.entities import Document
from typing import List, Dict, Any
from dataclasses import asdict
from databricks.vector_search.client import VectorSearchClient

class VectorSearchRetriever:
    """
    Class using Databricks Vector Search to retrieve relevant documents.
    """
    def __init__(self):
        self.vector_search_client = VectorSearchClient(disable_notice=True)
        # TODO: Replace this with the list of column names to return in the result when querying Vector Search
        self.columns = ["chunk_id", "text_column", "doc_uri"]
        self.vector_search_index = self.vector_search_client.get_index(
            index_name="catalog.schema.chunked_docs_index"
        )
        mlflow.models.set_retriever_schema(
            name="vector_search",
            primary_key="chunk_id",
            text_column="text_column",
            doc_uri="doc_uri"
        )

    @mlflow.trace(span_type="RETRIEVER", name="vector_search")
    def __call__(
        self,
        query: str,
        filters: Dict[Any, Any] = None,
        score_threshold = None
    ) -> List[Document]:
        """
        Performs vector search to retrieve relevant chunks.
        Args:
            query: Search query.
            filters: Optional filters to apply to the search. Filters must follow the Databricks Vector Search filter spec
            score_threshold: Score threshold to use for the query.

        Returns:
            List of retrieved Documents.
        """

        results = self.vector_search_index.similarity_search(
            query_text=query,
            columns=self.columns,
            filters=filters,
            num_results=5,
            query_type="ann"
        )

        documents = self.convert_vector_search_to_documents(
            results, 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 = []
        column_names = [column["name"] for column in vs_results.get("manifest", {}).get("columns", [])]
        result_row_count = vs_results.get("result", {}).get("row_count", 0)

        if result_row_count > 0:
            for item in vs_results["result"]["data_array"]:
                metadata = {}
                score = item[-1]

                if score >= score_threshold:
                    metadata["similarity_score"] = score
                    for i, field in enumerate(item[:-1]):
                        metadata[column_names[i]] = field

                    page_content = metadata.pop("text_column", None)

                    if page_content:
                        doc = Document(
                            page_content=page_content,
                            metadata=metadata
                        )
                        docs.append(doc)

        return docs

Чтобы запустить извлекатель, выполните следующий код на Python. При необходимости можно включить фильтры векторного поиска в запросе для фильтрации результатов.

retriever = VectorSearchRetriever()
query = "What is Databricks?"
filters={"text_column LIKE": "Databricks"},
results = retriever(query, filters=filters, score_threshold=0.1)

Установка схемы извлекателя

Чтобы обеспечить правильное отслеживание и отрисовку в конечных приложениях, вызовите mlflow.models.set_retriever_schema при определении агента. Используйте set_retriever_schema для сопоставления имен столбцов в возвращаемой таблице ожидаемым полям MLflow, таким как primary_key, text_columnи 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"],
)

Можно также указать дополнительные столбцы в схеме извлекателя, указав список имен столбцов с полем other_columns.

Если у вас несколько извлекателей, можно определить несколько схем с помощью уникальных имен для каждой схемы извлекателя.

Набор схем ретривера, установленный во время создания агента, влияет на последующие приложения и рабочие процессы, такие как приложение для обзора и наборы для оценки. В частности, столбец doc_uri служит основным идентификатором документов, возвращаемых извлекателем.

• приложение для обзора отображает doc_uri, чтобы помочь рецензентам оценивать ответы и отслеживать источники документов. См. пользовательский интерфейс Review App.
• наборы оценки используются doc_uri для сравнения результатов извлекателя с предопределенными наборами данных оценки для определения полноты и точности извлекателя. См. наборы оценок.

Поиск извлекателя

Отслеживание MLflow повышает наблюдаемость, фиксируя подробную информацию о выполнении агента. Он предоставляет способ записи входных, выходных данных и метаданных, связанных с каждым промежуточным шагом запроса, что позволяет быстро определить источник ошибок и непредвиденных действий.

В этом примере используется декоратор @mlflow.trace для создания трейса для извлекателя и парсера. Другие варианты настройки методов трассировки можно найти в разделе MLflow Трассировка для агентов.

Декоратор создает интервал, который начинается при вызове функции и заканчивается, когда функция возвращается. MLflow автоматически записывает входные и выходные данные функции и все исключения, вызванные.

Заметка

Пользователи библиотек LangChain, LlamaIndex и OpenAI могут использовать автоматическое логирование MLflow вместо определения трассировок вручную с помощью декоратора. См. раздел Использование автологирования для добавления трассировок в агентов.

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

Чтобы последующие приложения, такие как оценка агентов и платформа ИИ, правильно отображали трассировку извлекателя, убедитесь, что декоратор соответствует следующим требованиям.

• Используйте span_type="RETRIEVER" и убедитесь, что функция возвращает объект List[Document]. См. Retriever охватывает.
• Имя трассировки и имя retriever_schema должны совпадать для правильной настройки трассировки.

Фильтрация результатов векторного поиска

Область поиска можно ограничить подмножеством данных с помощью фильтра векторного поиска.

Параметр filters в VectorSearchRetriever определяет условия фильтрации с помощью спецификации фильтра Databricks Vector Search.

filters = {"text_column LIKE": "Databricks"}

В методе __call__ словарь фильтров передается непосредственно в функцию similarity_search:

results = self.vector_search_index.similarity_search(
    query_text=query,
    columns=self.columns,
    filters=filters,
    num_results=5,
    query_type="ann"
)

После первоначальной фильтрации параметр score_threshold предоставляет дополнительную фильтрацию, задав минимальную оценку сходства.

if score >= score_threshold:
    metadata["similarity_score"] = score

Окончательный результат включает документы, соответствующие условиям filters и score_threshold.

Дальнейшие действия

После того как вы создадите инструмент для функции каталога Unity, добавьте его в ИИ-агент. См. статью Добавление средств каталога Unity в агенты.