非結構化擷取 AI 代理程式工具

發行項
12/27/2024

重要

本文說明如何使用馬賽克 AI 代理程式架構建立非結構化數據擷取的 AI 代理程式工具。非結構化擷取程式可讓代理程式使用向量搜尋索引來查詢非結構化數據源，例如文件主體。

若要深入瞭解代理程式工具，請參閱建立 AI 代理程式工具。

使用 Unity Catalog 功能的向量搜尋檢索工具

下列範例將為一個檢索工具建立 Unity Catalog 函數，以便從 Mosaic AI 向量搜尋索引中查詢數據。

Unity Catalog 函式 databricks_docs_vector_search 查詢包含 Databricks 檔的假設向量搜尋索引。它會封裝 Databricks SQL 函式 vector_search()，並使用別名 page_content 和 metadata，使其輸出對應到 MLflow 擷取器 schema。

注意

若要符合 MLflow 擷取器 schema，必須將任何其他元數據 columns 使用 SQL 映射函式新增至 metadatacolumn，而不是作為最上層輸出索引鍵。

在筆記本或 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
  )

此擷取程式工具有下列注意事項：

MLflow 會將此 Unity Catalog 函式追蹤為 TOOL 範圍類型，而不是 RETRIEVER 範圍類型。因此，下游的代理檢閱應用程式和 AI 遊樂場等 Agent Framework 應用程式不會顯示與擷取工具相關的詳細資料，例如區塊的連結。如需範圍類型的詳細資訊，請參閱 MLflow 追蹤 Schema。
SQL 用戶端可能會限制傳回的資料列數目或位元組數量的上限 limit。為了防止數據截斷，您應該截斷 UDF 傳回的 columnvalues。例如，您可以使用 substring(chunked_text, 0, 8192) 來減少大型內容的大小 columns，並避免執行期間數據列截斷。
由於此工具是 vector_search() 函式的包裝函式，因此受限於與 vector_search() 函式相同的限制。請參閱限制。

如果此範例不適合您的使用案例，請改為使用自定義代理程式代碼建立向量搜尋擷取器工具。

具有代理程式代碼（PyFunc）的向量搜尋檢索器

下列範例會在代理代碼中為 PyFunc 風格的代理程序建立向量搜尋擷取器。

此範例使用 databricks-vectorsearch 來建立基本擷取程式，以使用篩選執行向量搜尋相似度搜尋。它會使用 MLflow 裝飾項目來啟用代理程式追蹤。

注意

若要符合 MLflow 擷取器 schema，擷取器函式應該會傳回 Document 類型，並使用 Document 類別中的 metadata 欄位，將其他屬性新增至傳回的檔，例如 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)

Set 擷取器 schema

若要確保檢索器在下游應用程式中正確追蹤並呈現，請在定義代理程式時呼叫 mlflow.models.set_retriever_schema。使用 set_retriever_schema，將傳回 table 中的 column 名稱對應至 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"],
)

您也可以在擷取器的 schema 中指定其他 columns，方法是提供含有 [other_columns] 欄位的 column 名稱 list。

如果您有多個擷取器，您可以使用每個擷取器的唯一名稱來定義多個架構，schema。

在代理程式建立期間，擷取元件 schemaset 會影響下游應用和工作流程，例如審查應用程式和評估集。具體而言，doc_uricolumn 可作為檢索工具所傳回檔案的主要 identifier。

檢閱應用程式 會顯示 doc_uri，以協助檢閱者評估回應和追蹤檔來源。請參閱檢查應用程式使用者介面。
評估集 使用 doc_uri 來比較擷取器結果與預先定義的評估數據集，以判斷擷取器的召回率和精確度。請參閱評估集。

追蹤資料檢索器

MLflow 追蹤藉由擷取代理程式執行的詳細資訊來增加可觀察性。它提供一種方式來記錄與要求每個中繼步驟相關聯的輸入、輸出和元數據，讓您能夠快速找出 Bug 和非預期行為的來源。

這個範例會使用 @mlflow.trace 裝飾器來為擷取器和剖析器建立追蹤。如需設定追蹤方法的其他選項，請參閱 MLflow 追蹤代理程式。

裝飾器會建立範圍，該範圍會在函式被呼叫時啟動，並在函式回傳時結束。 MLflow 會自動記錄函式的輸入和輸出，以及引發的任何例外狀況。

注意

LangChain、LlamaIndex 和 OpenAI 函式庫的使用者可以使用 MLflow 自動記錄功能，而不需要使用裝飾器手動定義追蹤。請參閱使用自動記錄將追蹤新增至代理程式。

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

若要確保代理程式評估和 AI 遊樂場等後續應用程式正確呈現擷取器追蹤，請確定裝飾器符合下列需求：

使用 span_type="RETRIEVER" 並確定函式會傳回 List[Document] 物件。請參閱擷取器範圍。
追蹤名稱和 retriever_schema 名稱必須相符，才能正確設定追蹤。

篩選向量搜尋結果

您可以使用向量搜尋篩選器，將搜尋範圍 limit 至數據子集。

filters 中的 VectorSearchRetriever 參數會使用 Databricks 向量搜尋篩選來定義篩選條件。

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 Catalog 函式代理程式工具之後，請將此工具新增至 AI 代理程式。請參閱將 Unity Catalog 工具新增至代理程式。

共用方式為