P1+P2 complete: pgvector, Casdoor SSO, CLI, KB router, K8s, monitoring, UI fixes (Promotion rewrite, Storage/Automation btn fix, membership configs mojibake fix)

2026-07-08 12:10:32 +08:00
1 parent 1f47951895
commit 177e345539
@@ -0,0 +1,360 @@
+"""
+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")