""" 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", "qihu-aiflow-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") # ---- 本地嵌入缓存（SQLite 持久化）：破解"每次嵌入消耗外部 API 额度"的隐性成本障碍 ---- import sqlite3, hashlib, json as _json _EMBED_CACHE_DB = os.environ.get("EMBED_CACHE_DB", "/app/backend/python/.embed_cache.db") def _embed_cache_get(text, model): try: key = hashlib.sha256(f"{model}::{text}".encode("utf-8")).hexdigest() conn = sqlite3.connect(_EMBED_CACHE_DB, timeout=5) row = conn.execute("SELECT vec FROM embed_cache WHERE k=?", (key,)).fetchone() conn.close() if row: return _json.loads(row[0]) except Exception: return None return None def _embed_cache_put(text, model, vec): try: key = hashlib.sha256(f"{model}::{text}".encode("utf-8")).hexdigest() conn = sqlite3.connect(_EMBED_CACHE_DB, timeout=5) conn.execute("CREATE TABLE IF NOT EXISTS embed_cache (k TEXT PRIMARY KEY, vec TEXT, ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") conn.execute("INSERT OR REPLACE INTO embed_cache (k, vec) VALUES (?, ?)", (key, _json.dumps(vec))) conn.commit() conn.close() except Exception: pass # 连接池 _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 cached = _embed_cache_get(text, EMBEDDING_MODEL) if cached: return cached try: # 收敛到共享标准模型层（app.core.llm_client.aembed），与聊天共用代理/超时/failover 语义 from app.core.llm_client import aembed embedding = await aembed( EMBEDDING_API_URL, EMBEDDING_API_KEY, text[:8000], EMBEDDING_MODEL, timeout=30 ) if embedding: _embed_cache_put(text, EMBEDDING_MODEL, embedding) return embedding logger.error("Embedding API returned empty") 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")