执行摘要

• 一句话：为 ColBERT 添加 query/document 嵌入 IO 处理器
• 推荐动作：对于使用 ColBERT 模型的团队，强烈建议采用此插件以获取正确嵌入。对于其他开发者，可作为实现自定义 IOProcessor 插件的参考范例，展示了插件系统的使用方式、entry point 注册、类型定义和集成测试的模式。

功能与动机

根据 PR body，现有代码无法为 ColBERT 模型输出正确的 query 嵌入（应为 32x128）和 document 嵌入。该 PR 弥补了这一空白，确保输出与 Jina AI API 和原始 ColBERT 论文预期一致。关联 issue #42234。

实现拆解

1. 插件脚手架：在 tests/plugins/colbert_query_plugin 下创建可安装 Python 包，通过 setup.py 注册 vllm.io_processor_plugins entry point，将 colbert_query_plugin 映射到 ColBERTQueryEmbeddingProcessor 类。
2. 数据模型定义：types.py 定义了 InputType (query|document)、ColBERTEmbeddingCompletionRequestMixin（添加 input_type 字段）、ColBERTEmbeddingResponseData 和 ColBERTEmbeddingResponse，使用 Pydantic 模型。
3. 核心处理器：ColBERTQueryEmbeddingProcessor 继承 IOProcessor，实现 parse_data 方法。parse_data 根据 input_type 分别调用 _build_query_prompt 或 _build_document_prompt 构建 TokensPrompt，然后执行推理，并将输出解析为 ColBERTEmbeddingResponse。
4. 提示构建：_build_query_prompt 插入 [CLS]、[QueryMarker]、内容 token、[SEP]，然后填充 [MASK] 至 QUERY_MAXLEN=32；_build_document_prompt 仅插入 [CLS]、[DocumentMarker]、内容 token、[SEP]。
5. CI 集成：在 .buildkite/test-amd.yaml 和 .buildkite/test_areas/plugins.yaml 添加步骤，安装插件并运行集成测试 test_colbert_query_io_processor_plugins.py。

关键文件：

• tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py（模块 嵌入处理器；类别 source；类型 core-logic；符号 ColBERTQueryEmbeddingProcessor, init, repr, _resolve_marker_ids）: 核心插件实现，包含 ColBERTQueryEmbeddingProcessor 类，处理 query 和 document 提示构建及推理。
• tests/plugins/colbert_query_plugin/colbert_query_processor/types.py（模块 数据模型；类别 source；类型 data-contract；符号 ColBERTEmbeddingCompletionRequestMixin, ColBERTEmbeddingResponseData, ColBERTEmbeddingResponse）: 定义了输入输出数据模型，包括 input_type 字段和响应结构。
• tests/plugins_tests/test_colbert_query_io_processor_plugins.py（模块 集成测试；类别 test；类型 test-coverage；符号 ModelConfig, _get_attr_or_val, _check_token_embeddings, server）: 集成测试，验证 query/document 嵌入的正确性及错误处理。
• tests/plugins/colbert_query_plugin/setup.py（模块 打包脚本；类别 infra；类型 configuration）: 通过 entry_points 注册插件，使 vllm 能够发现 colbert_query_plugin。
• tests/plugins/colbert_query_plugin/colbert_query_processor/__init__.py（模块 插件入口；类别 infra；类型 entrypoint；符号 register_colbert_query_embedding_processor）: 暴露插件注册函数，返回处理器类的完全限定名。
• .buildkite/test-amd.yaml（模块 CI配置；类别 infra；类型 configuration）: CI 配置中添加 colbert_query 插件的安装和测试步骤。
• .buildkite/test_areas/plugins.yaml（模块 CI配置；类别 infra；类型 configuration）: 将 colbert_query 插件纳入插件测试区域。

关键符号：ColBERTQueryEmbeddingProcessor.parse_data, ColBERTQueryEmbeddingProcessor._build_query_prompt, ColBERTQueryEmbeddingProcessor._build_document_prompt, ColBERTQueryEmbeddingProcessor._resolve_marker_ids, ColBERTQueryEmbeddingProcessor._iter_content_token_ids, register_colbert_query_embedding_processor

关键源码片段

tests/plugins/colbert_query_plugin/colbert_query_processor/types.py

定义了输入输出数据模型，包括 input_type 字段和响应结构。

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
​
from typing import Literal, get_args
​
from pydantic import BaseModel, Field
​
from vllm.entrypoints.openai.engine.protocol import UsageInfo
from vllm.entrypoints.pooling.base.protocol import CompletionRequestMixin
​
# 输入类型：query 或 document
InputType = Literal["query", "document"]
INPUT_TYPES: tuple[InputType, ...] = get_args(InputType)
# query 嵌入向量的固定长度（ColBERT 论文中的 Q* 超参数）
QUERY_MAXLEN = 32
​
​
class ColBERTEmbeddingCompletionRequestMixin(CompletionRequestMixin):
    """请求扩展，增加 input_type 字段，描述输入应如何处理。"""
    input_type: InputType = Field(
        description="Whether to encode the input as a ColBERT 'query' "
        f"(query marker + [mask] expansion to {QUERY_MAXLEN} tokens) or as a "
        "'document' (document marker only). Required.",
    )
​
​
class ColBERTEmbeddingResponseData(BaseModel):
    """单个嵌入响应条目，包含 token 级别的嵌入向量列表。"""
    index: int
    object: str = "embedding"
    input_type: InputType
    embedding: list[list[float]] # shape: (num_tokens, embedding_dim)
​
​
class ColBERTEmbeddingResponse(BaseModel):
    """嵌入响应顶层结构，兼容 vLLM 的 pooling 响应格式。"""
    data: list[ColBERTEmbeddingResponseData]
    usage: UsageInfo

评论区精华

Review 中 bennethinz 确认插件适用于 jinaai/jina-colbert-v2，但指出不适用于 colbert-ir/colbertv2.0 等变体。作者承诺在文档中明确说明支持范围。该讨论影响了插件的使用建议。

• 插件兼容性范围 (design): 作者同意并将在文档中明确说明只支持 jinaai/jina-colbert-v2。

风险与影响

• 风险：低风险。插件仅通过插件系统加载，不影响核心路径。主要风险包括：1）仅测试了 jinaai/jina-colbert-v2，其他 ColBERT 变体可能因 tokenizer 缺少特殊标记而失败；2）插件位于 tests/ 目录，可能被当作仅测试代码，但实际可通过 --io-processor-plugin 在生产中使用，需注意稳定性；3）CI 测试仅在 AMD 上运行，未覆盖 NVIDIA 环境。
• 影响：影响较小。仅使用 ColBERT 模型的用户受益，通过 /pooling 端点能正确获得 query/document 嵌入。对现有其他模型无影响。团队需注意该插件非官方模型支持，而是用户贡献的额外插件。
• 风险标记：插件特定模型限制, 仅 AMD CI 测试, 代码位于 tests 目录可能被忽视

关联脉络

• 暂无明显关联 PR