80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
"""pgvector 索引管理 — LlamaIndex PGVectorStore 封装。
|
||
|
||
Schema 隔离：使用显式表名 `rag_platform_kb_chunks`，不可使用默认 `data_<name>`。
|
||
确认 `create_if_not_exists` 不会触碰 `episodic_memory` 或任何 `memory/` 所属表。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from typing import TYPE_CHECKING
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
if TYPE_CHECKING:
|
||
from llama_index.vector_stores.postgres import PGVectorStore
|
||
|
||
# 显式表名 — 与 memory/episodic.py 的 episodic_memories 表完全隔离
|
||
KB_CHUNKS_TABLE = "rag_platform_kb_chunks"
|
||
|
||
# 默认 embedding 维度（OpenAI text-embedding-3-small）
|
||
DEFAULT_EMBED_DIM = 1536
|
||
|
||
|
||
def _async_to_sync_url(database_url: str) -> str:
|
||
"""将 asyncpg URL 转换为 psycopg2 URL（LlamaIndex PGVectorStore 使用同步驱动）。"""
|
||
return database_url.replace("postgresql+asyncpg://", "postgresql://")
|
||
|
||
|
||
def create_vector_store(
|
||
database_url: str,
|
||
embed_dim: int = DEFAULT_EMBED_DIM,
|
||
hybrid_search: bool = True,
|
||
) -> "PGVectorStore":
|
||
"""创建 PGVectorStore，使用显式表名实现 schema 隔离。
|
||
|
||
Args:
|
||
database_url: SQLAlchemy 数据库 URL（async 或 sync 均可）
|
||
embed_dim: embedding 向量维度
|
||
hybrid_search: 是否启用混合搜索（向量 + 全文）
|
||
|
||
Returns:
|
||
LlamaIndex PGVectorStore 实例
|
||
|
||
Raises:
|
||
ImportError: 如果 llama-index-vector-stores-postgres 未安装
|
||
"""
|
||
from llama_index.vector_stores.postgres import PGVectorStore
|
||
|
||
sync_url = _async_to_sync_url(database_url)
|
||
|
||
logger.info(
|
||
"Creating PGVectorStore: table=%s, embed_dim=%d, hybrid=%s",
|
||
KB_CHUNKS_TABLE,
|
||
embed_dim,
|
||
hybrid_search,
|
||
)
|
||
|
||
return PGVectorStore.from_params(
|
||
database=sync_url,
|
||
table_name=KB_CHUNKS_TABLE,
|
||
embed_dim=embed_dim,
|
||
hybrid_search=hybrid_search,
|
||
text_search_config="english", # U4 将用 jieba 替换为中文分词
|
||
# ponytail: 不设 schema_name，默认 public — 避免创建独立 schema 的运维复杂度
|
||
# 表名前缀 rag_platform_ 已足够隔离
|
||
)
|
||
|
||
|
||
async def ensure_vector_store_schema(database_url: str, embed_dim: int = DEFAULT_EMBED_DIM) -> None:
|
||
"""确保 vector store 表存在（幂等）。
|
||
|
||
在应用启动时调用，创建表结构（如果不存在）。
|
||
不会触碰 episodic_memory 或任何 memory/ 所属表。
|
||
"""
|
||
vs = create_vector_store(database_url, embed_dim=embed_dim)
|
||
# PGVectorStore.__init__ 内部会调用 create_if_not_exists
|
||
# 显式调用 _initialize 来确保表创建
|
||
vs._initialize()
|
||
logger.info("Vector store schema ensured: table=%s", KB_CHUNKS_TABLE)
|