다음을 통해 공유

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

  • 아티클

중요하다

이 기능은 공개 프리뷰.

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

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

AI Bridge를 사용하여 벡터 검색 검색 도구 로컬 개발

Databricks Vector Search 검색 도구 개발을 시작하는 가장 쉬운 방법은 로컬입니다. 및 databricks-langchain 같은 databricks-openai Databricks AI Bridge 패키지를 사용하여 검색 기능을 에이전트에 빠르게 추가하고 쿼리 매개 변수를 실험합니다. 이 방법을 사용하면 초기 개발 중에 빠르게 반복할 수 있습니다.

로컬 도구가 준비되면 에이전트 코드의 일부로 직접 프로덕션화하거나 Unity 카탈로그 함수로 마이그레이션하여 더 나은 검색 가능성과 거버넌스를 제공하지만 특정 제한 사항이 있습니다. Unity 카탈로그 함수가 포함된벡터 검색 도구를 참조하십시오.

메모

Databricks 외부에서 호스트되는 외부 벡터 인덱스를 사용하려면 Databricks 외부에서 호스트되는 벡터 인덱스를 사용하여벡터 검색 검색기를 참조하세요.

LangChain/LangGraph

다음 코드는 리트리버 도구를 프로토타입하고 로컬로 LLM에 바인딩하므로 에이전트와 채팅하여 도구 호출 동작을 테스트할 수 있습니다.

Databricks AI Bridge를 포함하는 최신 버전의 databricks-langchain 설치합니다.

%pip install --upgrade databricks-langchain

다음 예제에서는 Databricks 제품 설명서에서 콘텐츠를 가져오는 가상의 벡터 검색 인덱스를 쿼리합니다.

tool_description를 명확하고 설명적으로 작성하십시오. 에이전트 LLM은 tool_description 사용하여 도구를 이해하고 도구를 호출할 시기를 결정합니다.

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

메모

VectorSearchRetrieverTool을 초기화할 때, 자체 관리형 임베딩 및 직접 벡터 액세스 인덱스를 사용하는 델타 동기화 인덱스에는 text_column 및 embedding 인수가 필요합니다. 임베딩을 제공하기 위한옵션을 참조하세요.

자세한 내용은 VectorSearchRetrieverTool 대한 API 문서 참조하세요.

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

다음 코드는 벡터 검색 검색 도구의 프로토타입을 만들고 OpenAI의 GPT 모델과 통합합니다.

도구에 대한 OpenAI 권장 사항에 대한 자세한 내용은 OpenAI 함수 호출 설명서참조하세요.

Databricks AI Bridge를 포함하는 최신 버전의 databricks-openai 설치합니다.

%pip install --upgrade databricks-openai

다음 예제에서는 Databricks 제품 설명서에서 콘텐츠를 가져오는 가상의 벡터 검색 인덱스를 쿼리합니다.

명확하고 설명적인 tool_description를 제공하세요. 에이전트 LLM은 tool_description 사용하여 도구를 이해하고 도구를 호출할 시기를 결정합니다.

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

메모

VectorSearchRetrieverTool을 초기화할 때, 자체 관리형 임베딩 및 직접 벡터 액세스 인덱스를 사용하는 델타 동기화 인덱스에는 text_column 및 embedding 인수가 필요합니다. 임베딩을 제공하기 위한옵션을 참조하세요.

자세한 내용은 VectorSearchRetrieverTool 대한 API 문서 참조하세요.

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

Unity 카탈로그 함수가 있는 벡터 검색 도구

다음 예제에서는 Unity 카탈로그 함수를 사용하여 Mosaic AI Vector Search 인덱스데이터를 쿼리하는 검색 도구가 만들어집니다.

Unity 카탈로그 함수 databricks_docs_vector_search Databricks 설명서를 포함하는 가상 벡터 검색 인덱스를 쿼리합니다. 이 함수는 Databricks SQL 함수 vector_search() 래핑하고 출력을 MLflow 검색기 스키마맞춥니다. page_content 및 metadata 별칭을 사용하여

메모

MLflow 검색기 스키마를 준수하려면 최상위 출력 키가 아닌 metadata사용하여 추가 메타데이터 열을 열에 추가해야 합니다.

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
  )

AI 에이전트에서 이 검색기 도구를 사용하려면 UCFunctionToolkit래핑합니다. 이렇게 하면 MLflow를 통한 자동 추적이 가능합니다.

MLflow 추적은 Gen AI 애플리케이션에 대한 자세한 실행 정보를 캡처합니다. 각 단계에 대한 입력, 출력 및 메타데이터를 기록하여 문제를 디버그하고 성능을 분석하는 데 도움이 됩니다.

UCFunctionToolkit사용하는 경우 출력이 MLflow 검색기 스키마를 준수하는 경우 검색기는 MLflow 로그에서 RETRIEVER 범위 형식을 자동으로 생성합니다. 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() 함수와 동일한 제한 사항이 적용됩니다. 제한 사항참조하세요.

Databricks 외부에서 호스팅되는 벡터 인덱스를 사용하는 벡터 검색기

벡터 인덱스가 Azure Databricks 외부에서 호스트되는 경우 Unity 카탈로그 연결을 만들어 외부 서비스에 연결하고 연결 에이전트 코드를 사용할 수 있습니다. 자세한 내용은 AI 에이전트 도구를 외부 서비스에 연결하는 방법에 대한 정보를 보려면을 참조하세요.

다음 예제에서는 PyFunc 버전 에이전트에 대해 Databricks 외부에서 호스트되는 벡터 인덱스를 호출하는 벡터 검색 검색기를 만듭니다.

  1. 외부 서비스(이 경우 Azure)에 대한 Unity 카탈로그 연결을 만듭니다.

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

  2. 만든 Unity 카탈로그 연결을 사용하여 에이전트 코드에서 검색기 도구를 정의합니다. 이 예제에서는 MLflow 데코레이터를 사용하여 에이전트 추적을 사용하도록 설정합니다.

    메모

    MLflow 검색기 스키마를 준수하기 위해 검색기 함수는 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

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

  3. 검색기를 실행하려면 다음 Python 코드를 실행합니다. 요청 시 필요에 따라 벡터 검색 필터를 포함하여 결과를 제한할 수 있습니다.

    retriever = VectorSearchRetriever()
query = [0.01944167, 0.0040178085 . . .  TRIMMED FOR BREVITY 010858015, -0.017496133]
results = retriever(query, score_threshold=0.1)

검색자 스키마 설정

검색기 또는 span_type="RETRIEVER"에서 반환된 추적이 MLflow의 표준 검색기 스키마를 준수하지 않는 경우, 반환된 스키마를 MLflow가 기대하는 필드에 수동으로 매핑해야 합니다. 이렇게 하면 MLflow가 리트리버를 제대로 추적하고 다운스트림 애플리케이션에서 추적을 올바르게 렌더링할 수 있습니다.

리트리버 스키마를 수동으로 설정하려면 에이전트를 정의할 때 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를 표시합니다. 앱 UI을 검토하세요.
  • 있습니다. 평가 집합을 참조하세요.

리트리버 추적

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

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

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

메모

LangChain, LlamaIndex 및 OpenAI 라이브러리 사용자는 데코레이터를 사용하여 추적을 수동으로 정의하는 대신 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 이름이 일치해야 합니다.

다음 단계

Unity 카탈로그 함수 에이전트 도구를 만든 후 AI 에이전트에 도구를 추가합니다. Unity 카탈로그 도구를 에이전트에 추가하는 방법은 을 참조하세요.