ai_prep_search 函数

适用于： Databricks SQL Databricks Runtime

该 ai_prep_search() 函数将结构化输出 ai_parse_document 转换为针对 RAG 矢量搜索和信息检索系统优化的格式。 对于每个输入文档，函数将内容拆分为语义区块，使用文档级上下文（如文档标题、节标题、页面引用）扩充每个区块，并生成嵌入就绪表示形式。

要求

• Databricks Runtime 18.2 或更高版本。
• 如果使用无服务器计算，则还需要满足以下条件：
  • 必须将无服务器环境的版本设置为 3 或更高，因为这会启用诸如 VARIANT 之类的功能。
  • 必须使用 Python 或 SQL。 有关其他无服务器功能和限制，请参阅 无服务器计算限制。
• ai_prep_search 函数可通过 Databricks 笔记本、SQL 编辑器、Databricks 工作流、作业或 Lakeflow Spark 声明性管道使用。

Syntax

ai_prep_search(
    parsed VARIANT,
    [options MAP<STRING, STRING>]
) RETURNS VARIANT

Arguments

• parsed：表示 VARIANT 的结构化输出的 ai_parse_document表达式。
• options：可选 MAP<STRING, STRING>。 支持的密钥：
  • 'version'：要使用的输出架构的版本。

退货

VARIANT包含为矢量搜索索引设置格式的文档区块。 输出中的每个行表示一个输入文档。

输出架构为：

{
  "document": {
    "contents": [
      {
        "chunk_id": STRING,       // Unique identifier composed of the document ID and chunk position
        "chunk_position": INT,    // 0-based position of the chunk within the document
        "chunk_to_retrieve": STRING,  // Raw text content of the chunk
        "chunk_to_embed": STRING, // Context-enriched text prepared for embedding; see chunk_to_embed format
        "pages": [
          {
            "page_id": INT,       // Page index that this chunk appears on
            "image_uri": STRING   // Path to the page image for multi-modal retrieval
          }
        ]
      }
    ],
    "pages": [
      {
        "id": INT,           // 0-based page index
        "image_uri": STRING  // Path to the rendered page image, populated when
                             // imageOutputPath is set in ai_parse_document
      }
    ],
    "source_uri": STRING   // Source document URI
  },
  "error_status": {...}
}

chunk_to_embed格式

该 chunk_to_embed 字段将文档级上下文与区块内容相结合，以提高语义搜索期间的检索质量。 格式为：

The following passage represents a chunk of content from a document.
- 'Content' contains raw document text
- All other fields describe document context and hierarchical information
- For visual elements like images/charts, a summary is generated as part of 'Content'

Document Title: {doc_title}
Page Header: {page_header}
Page Footer: {page_footer}
Section Header: {section_header}
Caption: {caption}
Footnote: {footnote}
Page Number: {page_number}

Content:
{chunk_to_retrieve}

示例

带ai_parse_document的链

以下示例链接 ai_prep_search 用于 ai_parse_document 从存储在 Unity 目录卷中的原始文档生成搜索就绪区块：

WITH parsed_documents AS (
  SELECT ai_parse_document(content) AS parsed
  FROM READ_FILES('/Volumes/mydata/documents/', format => 'binaryFile')
)
SELECT ai_prep_search(parsed) AS result
FROM parsed_documents;

生成矢量搜索源表

以下示例将输出平展为单个区块行，并将其写入 Delta 表。 然后，该表可用作 Databricks 矢量搜索 索引的源，用作 chunk_to_embed 嵌入列和 chunk_id 主键。

WITH parsed_documents AS (
  SELECT ai_parse_document(content) AS parsed
  FROM READ_FILES('/Volumes/mydata/documents/', format => 'binaryFile')
),
prepped_documents AS (
  SELECT ai_prep_search(parsed) AS result
  FROM parsed_documents
)
SELECT
  chunk.value:chunk_id::STRING AS chunk_id,
  chunk.value:chunk_position::INT AS chunk_position,
  chunk.value:chunk_to_retrieve::STRING AS chunk_to_retrieve,
  chunk.value:chunk_to_embed::STRING AS chunk_to_embed,
  prepped_documents.result:document.source_uri::STRING AS source_uri
FROM
  prepped_documents,
  LATERAL variant_explode(prepped_documents.result:document.contents) AS chunk;

生成的行具有以下架构：

启用多模式检索

使用ai_parse_document选项调用时imageOutputPath，呈现的页面图像将保存到 Unity 目录卷，并image_uri填充每个区块pages数组中的字段。 这些图像引用可以在查询时传递给支持视觉的模型，以回答需要视觉上下文的问题，例如块图、图表或文本中未完全表示的表。

WITH parsed_documents AS (
  SELECT ai_parse_document(
    content,
    map(
      'imageOutputPath', '/Volumes/catalog/schema/volume/page_images/',
      'descriptionElementTypes', '*'
    )
  ) AS parsed
  FROM READ_FILES('/Volumes/mydata/documents/', format => 'binaryFile')
),
prepped_documents AS (
  SELECT ai_prep_search(parsed) AS result
  FROM parsed_documents
)
SELECT
  chunk.value:chunk_id::STRING AS chunk_id,
  chunk.value:chunk_to_embed::STRING AS chunk_to_embed,
  chunk.value:pages AS pages
FROM
  prepped_documents,
  LATERAL variant_explode(prepped_documents.result:document.contents) AS chunk;

局限性

• 该 ai_prep_search 函数需要有效的 ai_parse_document 输出作为输入。 传递其他 VARIANT 数据或不支持的架构版本可能会产生意外的结果或错误。
• 最大输入大小与最大输出大小 ai_parse_document一致。

相关函数

• ai_classify 函数
• ai_extract 函数
• ai_parse_document 函数