구조화되지 않은 검색 AI 에이전트 도구

이 문서에서는 Mosaic AI 에이전트 프레임워크를 사용하여 구조화되지 않은 데이터 검색을 위한 AI 에이전트 도구를 만드는 방법을 보여 줍니다. 구조화되지 않은 검색기를 사용하면 에이전트가 벡터 검색 인덱스를 사용하여 문서 모음과 같은 구조화되지 않은 데이터 원본을 쿼리할 수 있습니다.

에이전트 도구에 대한 자세한 내용은 AI 에이전트 도구 만들기참조하세요.

Unity Catalog 함수가 포함된 벡터 검색 도구

다음 예제는 Mosaic AI Vector Search 인덱스에서 데이터를 쿼리할 수 있는 검색 도구를 위한 Unity 함수 Catalog을(를) 생성하는 방법을 보여줍니다.

Unity Catalog 함수 databricks_docs_vector_search Databricks 설명서를 포함하는 가상 벡터 검색 인덱스를 쿼리합니다. Databricks SQL 함수 vector_search() 를 래핑하여 별칭 page_content 및 metadata 을 사용하여 출력을 MLflow 리트리버 schema에 일치시킵니다.

Notebook 또는 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 함수를 RETRIEVER 범위 형식이 아닌 TOOL 범위 형식으로 추적합니다. 따라서 에이전트 검토 앱 및 AI Playground와 같은 다운스트림 에이전트 프레임워크 애플리케이션은 청크에 대한 링크와 같은 리트리버 관련 세부 정보를 표시하지 않습니다. 범위 유형에 대한 자세한 내용은 MLflow 추적 Schema참조하세요.
• SQL 클라이언트는 반환되는 최대 행 수 또는 바이트 수를 limit 수 있습니다. 데이터 잘림을 방지하려면 UDF에서 반환된 columnvalues 잘라야 합니다. 예를 들어 substring(chunked_text, 0, 8192) 사용하여 큰 콘텐츠 columns 크기를 줄이고 실행 중에 행 잘림을 방지할 수 있습니다.
• 이 도구는 vector_search() 함수의 래퍼이므로 vector_search() 함수와 동일한 제한 사항이 적용됩니다. 제한 사항참조하세요.

이 예제가 사용 사례에 적합하지 않은 경우, 사용자 지정 에이전트 코드를 사용하여 벡터 검색 도구를 만들어 보세요.

에이전트 코드가 있는 벡터 검색 검색기(PyFunc)

다음 예제에서는 에이전트 코드에서 PyFunc 스타일의 에이전트를 위한 Vector Search 리트리버를 생성합니다.

이 예제에서는 databricks-vectorsearch 사용하여 필터사용하여 벡터 검색 유사성 검색을 수행하는 기본 검색기를 만듭니다. MLflow 데코레이터를 사용하여 에이전트 추적사용하도록 .

에이전트 모듈 또는 에이전트 Notebook에서 다음 코드를 사용합니다.

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를 표시합니다. 앱 UI을 검토하세요.
• 있습니다. 평가 집합을 참조하세요.

리트리버 추적

MLflow 추적은 에이전트 실행에 대한 자세한 정보를 캡처하여 관찰 가능성을 추가합니다. 요청의 각 중간 단계와 연결된 입력, 출력 및 메타데이터를 기록하여 버그 및 예기치 않은 동작의 원인을 신속하게 파악할 수 있는 방법을 제공합니다.

이 예제에서는 @mlflow.trace 데코레이터를 사용하여 검색기 및 파서에 대한 트레이스를 생성합니다. 추적 메서드를 설정하는 다른 옵션은에이전트에 대한 MLflow 추적을 참조하세요.

데코레이터는 함수가 호출될 때 시작되고 반환될 때 끝나는 범위 만듭니다. MLflow는 함수의 입력 및 출력 및 발생한 예외를 자동으로 기록합니다.

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

에이전트 평가 및 AI Playground와 같은 다운스트림 애플리케이션이 검색기 추적을 올바르게 렌더링하도록 하려면 데코레이터가 다음 요구 사항을 충족하는지 확인합니다.

• span_type="RETRIEVER" 사용하여 함수가 List[Document] 개체를 반환하는지 확인합니다. 리트리버의 범위를에서 참조하세요.
• 추적을 올바르게 구성하려면 추적 이름과 retriever_schema 이름이 일치해야 합니다.

벡터 검색 결과 필터링

Vector Search 필터를 사용하여 검색 범위를 데이터의 하위 집합으로 limit 수 있습니다.

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 Catalog 함수 에이전트 도구를 만든 후 AI 에이전트에 도구를 추가합니다. 에이전트에 Unity 도구 추가 참조하세요.