361 lines
11 KiB
Python
361 lines
11 KiB
Python
"""
|
||
PGVector 向量检索模块
|
||
====================
|
||
使用 PostgreSQL + pgvector 实现语义向量检索，
|
||
与现有 MySQL BM25 检索配合，实现混合检索。
|
||
|
||
使用方式:
|
||
from app.pgvector_search import vector_search
|
||
results = await vector_search(query, kb_id, top_k=5)
|
||
"""
|
||
|
||
import os
|
||
import logging
|
||
import asyncio
|
||
from typing import List, Dict, Optional
|
||
from datetime import datetime
|
||
|
||
logger = logging.getLogger("aiflow.pgvector")
|
||
|
||
# PostgreSQL 连接配置
|
||
PG_HOST = os.environ.get("POSTGRES_HOST", "pgvector")
|
||
PG_PORT = int(os.environ.get("POSTGRES_PORT", "5432"))
|
||
PG_USER = os.environ.get("POSTGRES_USER", "aiflow")
|
||
PG_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "pwd2026")
|
||
PG_DB = os.environ.get("POSTGRES_DB", "aiflow_vectors")
|
||
|
||
# 嵌入模型配置
|
||
EMBEDDING_API_URL = os.environ.get("SILICONFLOW_BASE_URL", "https://api.siliconflow.cn")
|
||
EMBEDDING_API_KEY = os.environ.get("SILICONFLOW_API_KEY", "")
|
||
EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-large-zh-v1.5")
|
||
|
||
# 连接池
|
||
_pg_pool = None
|
||
_pool_lock = asyncio.Lock()
|
||
|
||
|
||
async def get_pg_pool():
|
||
"""获取 PostgreSQL 异步连接池"""
|
||
global _pg_pool
|
||
if _pg_pool is not None:
|
||
return _pg_pool
|
||
|
||
async with _pool_lock:
|
||
if _pg_pool is not None:
|
||
return _pg_pool
|
||
|
||
try:
|
||
import asyncpg
|
||
except ImportError:
|
||
logger.warning("asyncpg not installed, vector search disabled")
|
||
return None
|
||
|
||
try:
|
||
_pg_pool = await asyncpg.create_pool(
|
||
host=PG_HOST,
|
||
port=PG_PORT,
|
||
user=PG_USER,
|
||
password=PG_PASSWORD,
|
||
database=PG_DB,
|
||
min_size=2,
|
||
max_size=10,
|
||
command_timeout=30,
|
||
)
|
||
logger.info(f"PostgreSQL connection pool created: {PG_HOST}:{PG_PORT}/{PG_DB}")
|
||
|
||
# 初始化表结构
|
||
async with _pg_pool.acquire() as conn:
|
||
await conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS document_embeddings (
|
||
id SERIAL PRIMARY KEY,
|
||
doc_id VARCHAR(64) NOT NULL,
|
||
kb_id VARCHAR(64) NOT NULL,
|
||
chunk_id VARCHAR(64),
|
||
content TEXT NOT NULL,
|
||
embedding vector(1024),
|
||
metadata JSONB DEFAULT '{}',
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||
);
|
||
""")
|
||
await conn.execute("""
|
||
CREATE INDEX IF NOT EXISTS idx_embeddings_kb
|
||
ON document_embeddings(kb_id);
|
||
""")
|
||
await conn.execute("""
|
||
CREATE INDEX IF NOT EXISTS idx_embeddings_doc
|
||
ON document_embeddings(doc_id);
|
||
""")
|
||
await conn.execute("""
|
||
CREATE INDEX IF NOT EXISTS idx_embeddings_vector
|
||
ON document_embeddings USING ivfflat (embedding vector_cosine_ops)
|
||
WITH (lists = 100);
|
||
""")
|
||
logger.info("PGVector tables initialized")
|
||
|
||
return _pg_pool
|
||
except Exception as e:
|
||
logger.error(f"Failed to create PostgreSQL pool: {e}")
|
||
_pg_pool = None
|
||
return None
|
||
|
||
|
||
async def get_embedding(text: str) -> Optional[List[float]]:
|
||
"""调用嵌入API获取文本向量"""
|
||
if not EMBEDDING_API_KEY:
|
||
logger.warning("No embedding API key configured")
|
||
return None
|
||
|
||
try:
|
||
import httpx
|
||
async with httpx.AsyncClient(timeout=30) as client:
|
||
resp = await client.post(
|
||
f"{EMBEDDING_API_URL}/v1/embeddings",
|
||
headers={
|
||
"Authorization": f"Bearer {EMBEDDING_API_KEY}",
|
||
"Content-Type": "application/json",
|
||
},
|
||
json={
|
||
"model": EMBEDDING_MODEL,
|
||
"input": text[:8000], # 截断防止超长
|
||
},
|
||
)
|
||
if resp.status_code == 200:
|
||
data = resp.json()
|
||
embedding = data["data"][0]["embedding"]
|
||
return embedding
|
||
else:
|
||
logger.error(f"Embedding API error: {resp.status_code} {resp.text[:200]}")
|
||
return None
|
||
except Exception as e:
|
||
logger.error(f"Embedding generation failed: {e}")
|
||
return None
|
||
|
||
|
||
async def vector_search(
|
||
query: str,
|
||
kb_id: Optional[str] = None,
|
||
top_k: int = 5,
|
||
threshold: float = 0.5,
|
||
) -> List[Dict]:
|
||
"""
|
||
向量语义检索
|
||
|
||
Args:
|
||
query: 查询文本
|
||
kb_id: 知识库ID (可选，不指定则搜索全部)
|
||
top_k: 返回结果数
|
||
threshold: 相似度阈值 (0-1)
|
||
|
||
Returns:
|
||
[{"doc_id": ..., "content": ..., "score": ..., "metadata": ...}, ...]
|
||
"""
|
||
pool = await get_pg_pool()
|
||
if pool is None:
|
||
return []
|
||
|
||
# 获取查询向量
|
||
query_embedding = await get_embedding(query)
|
||
if query_embedding is None:
|
||
return []
|
||
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
# 将向量转换为 PostgreSQL 格式
|
||
embedding_str = f"[{','.join(str(x) for x in query_embedding)}]"
|
||
|
||
if kb_id:
|
||
rows = await conn.fetch(
|
||
"""
|
||
SELECT doc_id, chunk_id, content, metadata,
|
||
1 - (embedding <=> $1::vector) as score
|
||
FROM document_embeddings
|
||
WHERE kb_id = $2
|
||
ORDER BY embedding <=> $1::vector
|
||
LIMIT $3;
|
||
""",
|
||
embedding_str, kb_id, top_k,
|
||
)
|
||
else:
|
||
rows = await conn.fetch(
|
||
"""
|
||
SELECT doc_id, chunk_id, content, metadata,
|
||
1 - (embedding <=> $1::vector) as score
|
||
FROM document_embeddings
|
||
ORDER BY embedding <=> $1::vector
|
||
LIMIT $2;
|
||
""",
|
||
embedding_str, top_k,
|
||
)
|
||
|
||
results = []
|
||
for row in rows:
|
||
score = float(row["score"])
|
||
if score >= threshold:
|
||
results.append({
|
||
"doc_id": row["doc_id"],
|
||
"chunk_id": row["chunk_id"],
|
||
"content": row["content"],
|
||
"score": score,
|
||
"metadata": row["metadata"] if isinstance(row["metadata"], dict) else {},
|
||
})
|
||
|
||
return results
|
||
|
||
except Exception as e:
|
||
logger.error(f"Vector search failed: {e}")
|
||
return []
|
||
|
||
|
||
async def index_document(
|
||
doc_id: str,
|
||
kb_id: str,
|
||
content: str,
|
||
chunk_id: Optional[str] = None,
|
||
metadata: Optional[Dict] = None,
|
||
) -> bool:
|
||
"""
|
||
将文档内容索引到 PGVector
|
||
|
||
Args:
|
||
doc_id: 文档ID
|
||
kb_id: 知识库ID
|
||
content: 文档内容
|
||
chunk_id: 分块ID (可选)
|
||
metadata: 元数据
|
||
|
||
Returns:
|
||
是否成功
|
||
"""
|
||
pool = await get_pg_pool()
|
||
if pool is None:
|
||
return False
|
||
|
||
# 分块处理（简单按长度分块）
|
||
chunks = []
|
||
chunk_size = 500
|
||
overlap = 50
|
||
if len(content) <= chunk_size:
|
||
chunks = [(content, chunk_id or f"{doc_id}_0")]
|
||
else:
|
||
for i in range(0, len(content), chunk_size - overlap):
|
||
chunk = content[i:i + chunk_size]
|
||
if len(chunk.strip()) > 10:
|
||
chunks.append((chunk, f"{doc_id}_{i // (chunk_size - overlap)}"))
|
||
|
||
success_count = 0
|
||
for chunk_text, cid in chunks:
|
||
embedding = await get_embedding(chunk_text)
|
||
if embedding is None:
|
||
continue
|
||
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
embedding_str = f"[{','.join(str(x) for x in embedding)}]"
|
||
import json
|
||
await conn.execute(
|
||
"""
|
||
INSERT INTO document_embeddings (doc_id, kb_id, chunk_id, content, embedding, metadata)
|
||
VALUES ($1, $2, $3, $4, $5::vector, $6)
|
||
ON CONFLICT DO NOTHING;
|
||
""",
|
||
doc_id, kb_id, cid, chunk_text, embedding_str,
|
||
json.dumps(metadata or {}),
|
||
)
|
||
success_count += 1
|
||
except Exception as e:
|
||
logger.error(f"Failed to index chunk {cid}: {e}")
|
||
|
||
logger.info(f"Indexed document {doc_id}: {success_count}/{len(chunks)} chunks")
|
||
return success_count > 0
|
||
|
||
|
||
async def hybrid_search(
|
||
query: str,
|
||
kb_id: Optional[str] = None,
|
||
top_k: int = 5,
|
||
bm25_results: Optional[List[Dict]] = None,
|
||
) -> List[Dict]:
|
||
"""
|
||
混合检索：BM25 + 向量检索融合
|
||
|
||
Args:
|
||
query: 查询文本
|
||
kb_id: 知识库ID
|
||
top_k: 返回结果数
|
||
bm25_results: BM25 检索结果（来自MySQL）
|
||
|
||
Returns:
|
||
融合后的检索结果
|
||
"""
|
||
# 获取向量检索结果
|
||
vec_results = await vector_search(query, kb_id, top_k)
|
||
|
||
if not vec_results and not bm25_results:
|
||
return []
|
||
|
||
if not vec_results:
|
||
return bm25_results or []
|
||
|
||
if not bm25_results:
|
||
return vec_results
|
||
|
||
# 融合：Reciprocal Rank Fusion (RRF)
|
||
rrf_k = 60
|
||
score_map = {}
|
||
|
||
for rank, item in enumerate(bm25_results):
|
||
doc_id = item.get("doc_id", item.get("id", ""))
|
||
score_map[doc_id] = {
|
||
"content": item.get("content", item.get("text", "")),
|
||
"score": 1.0 / (rrf_k + rank + 1),
|
||
"bm25_rank": rank + 1,
|
||
"metadata": item.get("metadata", {}),
|
||
}
|
||
|
||
for rank, item in enumerate(vec_results):
|
||
doc_id = item["doc_id"]
|
||
if doc_id in score_map:
|
||
score_map[doc_id]["score"] += 1.0 / (rrf_k + rank + 1)
|
||
score_map[doc_id]["vec_score"] = item["score"]
|
||
score_map[doc_id]["vec_rank"] = rank + 1
|
||
else:
|
||
score_map[doc_id] = {
|
||
"content": item["content"],
|
||
"score": 1.0 / (rrf_k + rank + 1),
|
||
"vec_score": item["score"],
|
||
"vec_rank": rank + 1,
|
||
"metadata": item.get("metadata", {}),
|
||
}
|
||
|
||
# 按融合分数排序
|
||
merged = sorted(score_map.values(), key=lambda x: x["score"], reverse=True)
|
||
return merged[:top_k]
|
||
|
||
|
||
async def delete_document_embeddings(doc_id: str) -> bool:
|
||
"""删除文档的所有向量索引"""
|
||
pool = await get_pg_pool()
|
||
if pool is None:
|
||
return False
|
||
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
await conn.execute(
|
||
"DELETE FROM document_embeddings WHERE doc_id = $1;",
|
||
doc_id,
|
||
)
|
||
logger.info(f"Deleted embeddings for document {doc_id}")
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"Failed to delete embeddings: {e}")
|
||
return False
|
||
|
||
|
||
async def close_pool():
|
||
"""关闭连接池"""
|
||
global _pg_pool
|
||
if _pg_pool is not None:
|
||
await _pg_pool.close()
|
||
_pg_pool = None
|
||
logger.info("PostgreSQL connection pool closed")
|