LlamaIndex PostgresML 托管索引集成：PostgresMLIndex 与 PostgresMLRetriever 深度解析

本文基于 LlamaIndex 官方 API 参考页 indices/postgresml 展开，系统讲解 LlamaIndex 的 PostgresML 托管索引集成（llama-index-indices-managed-postgresml 包）：PostgresMLIndex 的创建与写入机制、PostgresMLRetriever 的检索流程、以及配套的 PostgresMLQueryEngine 端到端 RAG 查询能力。读完后，你可以直接在 LlamaIndex 应用中接入 PostgresML，理解其默认的 Pipeline 配置（文本切分、向量化、Rerank）如何被序列化下发到服务端，并掌握检索、增量写入与流式回答等实战用法。

一、什么是 PostgresML 托管索引

该集成对应源码位于 llama-index-integrations/indices/llama-index-indices-managed-postgresml，核心入口为 llama_index/indices/managed/postgresml/init.py：

from llama_index.indices.managed.postgresml.base import PostgresMLIndex
from llama_index.indices.managed.postgresml.retriever import PostgresMLRetriever

__all__ = ["PostgresMLIndex", "PostgresMLRetriever"]

与 LlamaIndex 中把 chunk、embedding、存储都留在本地的向量索引不同，PostgresMLIndex 属于托管索引（Managed Index）——它继承自 llama_index/core/indices/managed/base.py 中的 BaseManagedIndex。按 base.py 中类文档字符串的说法，PostgresML 把传统索引的多项职能放到了服务端：

• 将文档切分为 chunks（nodes）；
• 为每个 chunk 生成 embedding；
• 对查询执行 top-k 相似度检索；
• 可选地在服务端完成文本生成或 chat completion。

BaseManagedIndex 在基类层面声明了四个抽象方法（见 managed/base.py#L50-L67）：_insert、delete_ref_doc、update_ref_doc 和 as_retriever，PostgresMLIndex 全部实现了它们。测试文件 tests/test_indices_managed_postgresml.py 也验证了 PostgresMLIndex 的 MRO 中包含 BaseManagedIndex、PostgresMLRetriever 的 MRO 中包含 BaseRetriever。

从版本约束看（见 pyproject.toml），该包当前版本为 0.6.0，要求 Python >=3.10,<4.0，依赖 pgml>=1.1.0,<2（PostgresML 官方 Python SDK，基于 pyo3 绑定）与 llama-index-core>=0.13.0,<0.15。

安装方式（来自 README.md）：

pip install llama-index-indices-managed-postgresml

前提是你已有一个 PostgresML 数据库的连接字符串（PostgresML 托管服务或自建实例）。

二、PostgresMLIndex：构造、默认 Pipeline 与写入

2.1 构造函数参数

PostgresMLIndex 定义在 base.py#L52-L105，完整参数如下：

参数	类型	默认值	说明
collection_name	str	必填	PostgresML 集合（Collection）名称，同时被用作索引结构的 index_id
pipeline_name	Optional[str]	"v1"	服务端 Pipeline 名称，未指定时默认为 v1
pipeline_schema	Optional[Dict[str, Any]]	见下	Pipeline 的 JSON 配置，定义切分与语义搜索模型
pgml_database_url	Optional[str]	None	数据库连接串，也可通过环境变量 PGML_DATABASE_URL 提供
show_progress	bool	True	是否显示 tqdm 进度条
upsert_parallel_batches	int	1	写入时的并行批次数，会作为 parallel_batches 参数传给 upsert_documents
nodes	Optional[Sequence[BaseNode]]	None	构造时即写入的节点序列

若未传入 pipeline_schema，源码会套用如下默认配置（base.py#L82-L94）：

pipeline_schema = {
    "content": {
        "splitter": {
            "model": "recursive_character",
            "parameters": {"chunk_size": 1500},
        },
        "semantic_search": {
            "model": "intfloat/e5-small-v2",
            "parameters": {"prompt": "passage: "},
        },
    }
}

也就是说，默认使用 recursive_character 分词器（chunk 大小为 1500 字符），并用 intfloat/e5-small-v2 模型做语义搜索，入库端 prompt 前缀为 "passage: "。这些配置通过 pgml.Pipeline(pipeline_name, pipeline_schema) 构造，随后调用 collection.add_pipeline(pipeline) 注册到服务端。

一个值得注意的实现细节：add_pipeline、upsert_documents、delete_documents 这类 SDK 调用都被包在一个 async def 里，再经由 llama_index.core.async_utils.run_async_tasks 桥接同步/异步（base.py#L97-L102）。源码注释明确说明这是 pyo3 异步实现的一个限制——pgml SDK 的方法本身是 async 的，而 PostgresMLIndex 的对外接口是同步的，因此需要这层 run_async_tasks 包装。

2.2 写入路径：from_documents、add_documents 与 _insert

from_documents 类方法（base.py#L169-L193）是最常用的入口，参数与构造函数一致，但 collection_name 为必填（缺失时抛出 Exception("collection_name is a required argument")）。它把 Document 序列转换为 TextNode 序列后交给构造函数写入：

import os

os.environ["PGML_DATABASE_URL"] = "..."  # 也可在构造函数中传 pgml_database_url

from llama_index.core import Document
from llama_index.indices.managed.postgresml import PostgresMLIndex

# 创建索引（同时写入文档）
index = PostgresMLIndex.from_documents(
    "llama-index-test-1", [Document.example()]
)

# 连接已有索引（不写入）
index = PostgresMLIndex("llama-index-test-1")

底层 _insert 方法（base.py#L107-L129）把每个 BaseNode 转成 PostgresML 文档结构：

documents = [
    {
        "id": node.node_id,        # 沿用 LlamaIndex 的 node_id，保证 ID 体系一致
        "content": node.get_content(),
        "metadata": node.metadata,
    }
    for node in nodes
]
args = {"parallel_batches": self.upsert_parallel_batches, **insert_kwargs}

随后以 parallel_batches 为并行度调用 collection.upsert_documents(documents, args)。这里使用 upsert 语义意味着相同 id 的文档会被覆盖，这也是 update_ref_doc 能以 merge=True 方式更新文档的原因：

def update_ref_doc(self, document: Document) -> None:
    node = TextNode(**document.dict())
    self._insert([node], merge=True)

删除则走 base.py#L139-L145 的 delete_ref_doc，按 {"id": {"$eq": ref_doc_id}} 条件删除，id 即写入时使用的 node_id。add_documents 则是把 Sequence[Document] 转成 TextNode 后追加写入。

三、PostgresMLRetriever：检索参数与查询结构

PostgresMLRetriever 定义在 retriever.py#L18-L103，继承自 BaseRetriever。构造函数参数：

参数	类型	默认值	说明
index	PostgresMLIndex	必填	所属索引，通常由 index.as_retriever(**kwargs) 自动传入
callback_manager	Optional[CallbackManager]	None	LlamaIndex 回调管理器
pgml_query	Optional[Dict[str, Any]]	None	直接透传给 PostgresML 的原始查询 DSL；提供后检索完全由该 DSL 决定
limit	Optional[int]	5	返回结果条数上限
rerank	Optional[Dict[str, Any]]	None	重排配置，检索时会自动合并 {"query": query_str}

从 base.py#L151-L157 看，index.as_retriever() 就是把 self 与 **kwargs 转发给 PostgresMLRetriever 构造器的工厂方法，因此你可以直接传参定制检索行为：

# 创建 retriever，并指定检索条数与重排模型
retriever = index.as_retriever(limit=8, rerank={"model": "cross-encoder/ms-marco-MiniLM-L-6-v2"})

results = retriever.retrieve("What managed index is the best?")
print(results)

3.1 检索的两种路径

_aretrieve 的核心逻辑（retriever.py#L50-L103）区分两种情况：

1. 提供了 pgml_query：原样调用 collection.vector_search(self._pgml_query, self._index.pipeline)，检索行为完全由你提供的 PostgresML 查询 DSL 决定，LlamaIndex 不再拼装查询体。
2. 未提供 pgml_query：要求必须有 query_bundle（否则抛出 Exception("Must provide either query or query_bundle...")），并按如下结构组装 vector_search 请求：

{
    "query": {
        "fields": {
            "content": {
                "query": query_bundle.query_str,
                "parameters": {"prompt": "query: "},  # 查询端 prompt 前缀，与入库端 "passage: " 对应
            }
        }
    },
    "rerank": self._rerank,  # 若提供，会先合并 {"query": query_bundle.query_str}
    "limit": self._limit,
}

结果随后被还原为 LlamaIndex 的 NodeWithScore 列表，每个节点的 id_、text、metadata 分别取自服务端的 document.id、chunk 与 document.metadata。分数取值有一个细节：

• 未启用 rerank 时，score 取 r["score"]（向量相似度分）；
• 启用 rerank 时，score 取 r["rerank_score"]（重排分）。

同步方法 _retrieve 同样通过 run_async_tasks 桥接异步实现，调用方无需感知 asyncio。

四、PostgresMLQueryEngine：服务端一站式 RAG

API 参考页只列出 PostgresMLIndex 与 PostgresMLRetriever 两个成员，但包内还有第三个关键组件 PostgresMLQueryEngine（query.py），由 index.as_query_engine() 创建（base.py#L159-L167）。它利用 PostgresML 的服务端 RAG 能力，把“向量检索 + 上下文聚合 + LLM 生成”一次性下推到数据库完成。

构造参数（query.py#L80-L102）：

参数	类型	默认值	说明
retriever	PostgresMLRetriever	必填	检索器，由 as_query_engine() 内部构造
streaming	Optional[bool]	False	是否流式输出 token
pgml_query	Optional[Dict[str, Any]]	None	原始 PostgresML 查询 DSL，提供时跳过默认组装
vector_search_limit	Optional[int]	4	RAG 上下文中检索的 chunk 数量
vector_search_rerank	Optional[Dict[str, Any]]	None	重排配置，运行时自动合并当前 query
vector_search_document	Optional[Dict[str, Any]]	{"keys": ["id", "metadata"]}	控制服务端返回哪些文档字段
model	Optional[str]	"meta-llama/Meta-Llama-3-8B-Instruct"	服务端生成使用的 LLM 模型
model_parameters	Optional[Dict[str, Any]]	{"max_tokens": 2048}	模型参数（如 max_tokens）

非流式路径 _do_query 的组装逻辑（query.py#L136-L216）值得细看：

1. 用内置 text_qa_template 提示词（SYSTEM 角色为 "You are a helpful chatbot"，USER 消息中 {context_str} 占位为 {CONTEXT}）格式化出 messages；
2. 组装一个包含 CONTEXT 与 chat 两大块的查询：

query = {
    "CONTEXT": {
        "vector_search": {
            "query": {"fields": {"content": {"query": query_bundle.query_str,
                                              "parameters": {"prompt": "query: "}}}},
            "document": self._vector_search_document,
            "limit": self._vector_search_limit,
            "rerank": self._vector_search_rerank,
        },
        "aggregate": {"join": "\n"},  # 多个 chunk 用换行拼接成上下文
    },
    "chat": model_parameters,  # {"model": ..., "max_tokens": ..., "messages": [...]}
}

源码注释指出，{CONTEXT} 占位符会被 pgml SDK 生成的 SQL 替换为真实的检索上下文；

3. 非流式时调用 collection.rag(query, pipeline)，最终返回 Response(response=results["rag"][0], source_nodes=source_nodes)，其中 source_nodes 由 results["sources"]["CONTEXT"] 中的每个 chunk 还原为 NodeWithScore（rerank 场景同样取 rerank_score 作为 score）；
4. 流式时（streaming=True）调用 collection.rag_stream，并用自定义的 AsyncJsonGenerator（同时实现 Generator 与 AsyncGenerator）把异步 token 流适配为 LlamaIndex 的 AsyncStreamingResponse/StreamingResponse。源码中有一条明确注释：pgml SDK 在流式模式下目前不返回 sources，因此流式响应只携带文本 token。

README 中的用法示例：

query_engine = index.as_query_engine()
response = query_engine.query("What managed index is the best?")
print(response)

五、关键调用链小结

把三个组件串起来，一次典型的 PostgresML 托管索引工作流如下（均有源码对应）：

• 建库写入：PostgresMLIndex.from_documents(collection_name, documents) → TextNode 化 → _insert 组装 {id, content, metadata} → collection.upsert_documents(documents, {"parallel_batches": n})；Pipeline 在构造时经 add_pipeline 注册，默认 recursive_character（chunk 1500）+ intfloat/e5-small-v2。
• 检索：index.as_retriever(limit=..., rerank=...) → PostgresMLRetriever._aretrieve → collection.vector_search(query, pipeline) → 还原为 NodeWithScore 列表。
• RAG 问答：index.as_query_engine(...) → PostgresMLQueryEngine._do_query 组装 CONTEXT（vector_search + aggregate）与 chat（模型 + messages）→ collection.rag / collection.rag_stream → 返回带 source_nodes 的 Response 或流式响应。

六、使用注意与限制

结合源码可以归纳出几点实际使用中的边界条件：

1. 依赖前提：需要有效的 PostgresML 连接串（pgml_database_url 参数或 PGML_DATABASE_URL 环境变量），以及 pgml SDK（>=1.1.0,<2）；包版本约束 llama-index-core>=0.13.0,<0.15，在更新的大版本 core 上需自行验证兼容性。
2. pyo3 异步桥接：所有服务端调用都经过 run_async_tasks 同步包装，同步 API 下无需自建事件循环，但也意味着同步方法内部会临时运行 asyncio 任务。
3. rerank 分数字段切换：启用 rerank 后分数来源从 score 变为 rerank_score，若下游逻辑依赖相似度绝对值，需留意语义差异。
4. 流式无 sources：streaming=True 时流式响应不含来源节点（源码注释明确 pgml SDK 当前限制）。
5. 更新语义：update_ref_doc 依赖 node id 不变进行 merge upsert；删除依赖写入时保留的 node_id，因此外部自行构造 node 时保证 node_id 稳定可追踪很重要。
6. 索引结构类型：PostgresMLIndexStruct.get_type() 返回字符串 "POSTGRESML"（base.py#L33-L37），注释掉的 IndexStructType.POSTGRESML 表明该类型尚未进入 core 的枚举体系。

参考路径汇总

内容	路径
API 参考入口	docs/api_reference/api_reference/indices/postgresml.md
索引实现	llama_index/indices/managed/postgresml/base.py
检索器实现	llama_index/indices/managed/postgresml/retriever.py
查询引擎实现	llama_index/indices/managed/postgresml/query.py
使用文档	README.md
包配置	pyproject.toml
继承测试	tests/test_indices_managed_postgresml.py
托管索引基类	llama_index/core/indices/managed/base.py