219 lines
9.1 KiB
Python
219 lines
9.1 KiB
Python
"""向量检索服务（B 档）。
|
||
|
||
- embedding：调用 OpenAI 兼容 /embeddings API（如硅基流动 `BAAI/bge-m3`，1024 维）
|
||
- rerank：调用 /v1/rerank API（如硅基流动 `BAAI/bge-reranker-v2-m3`）
|
||
- 检索：书本段规模几千~几万，直接全量 numpy 余弦扫描，无需引入向量数据库
|
||
- 融合：关键词结果 + 向量结果做 RRF（Reciprocal Rank Fusion）
|
||
- 降级：未配置 API Key / 调用失败时自动退回纯关键词模式，绝不阻塞主流程
|
||
"""
|
||
import json
|
||
import logging
|
||
import threading
|
||
|
||
import numpy as np
|
||
import httpx
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.core.config import settings
|
||
from app.models.book import Segment
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
EMBEDDING_ENABLED = bool(settings.embedding_api_key)
|
||
RERANK_ENABLED = bool(settings.rerank_api_key)
|
||
|
||
# book_id -> 向量生成状态："done" / "running" / "failed"
|
||
_vector_status: dict[str, str] = {}
|
||
_status_lock = threading.Lock()
|
||
|
||
|
||
# ======================= API 客户端 =======================
|
||
|
||
async def embed_texts(texts: list[str]) -> list[list[float] | None] | None:
|
||
"""生成一批文本的向量，返回与 texts 同序的列表（单条失败置 None，其余正常）。
|
||
|
||
单次请求过大 / 服务端失败时自动对半拆开递归重试，直到单条；
|
||
单条超模型输入上限时截断重试一次，仍失败则跳过该条（置 None）。
|
||
未配置时返回 None；鉴权/配置类错误（API 不可用）直接抛出。
|
||
"""
|
||
if not EMBEDDING_ENABLED or not texts:
|
||
return None
|
||
url = settings.embedding_base_url.rstrip("/") + "/embeddings"
|
||
headers = {"Authorization": f"Bearer {settings.embedding_api_key}"}
|
||
out: list[list[float] | None] = [None] * len(texts)
|
||
# 空文本无法向量化（OpenRouter 对空串直接 400 too_small），预过滤不送请求
|
||
tasks: list[tuple[int, str]] = [(i, t) for i, t in enumerate(texts) if t and t.strip()]
|
||
if not tasks:
|
||
return out
|
||
|
||
async with httpx.AsyncClient(timeout=httpx.Timeout(120.0, connect=15.0)) as client:
|
||
async def _call(items: list[str], idxs: list[int]) -> None:
|
||
if not items:
|
||
return
|
||
try:
|
||
resp = await client.post(
|
||
url, headers=headers,
|
||
json={"model": settings.embedding_model, "input": items},
|
||
)
|
||
resp.raise_for_status()
|
||
rows = resp.json().get("data") or []
|
||
for j, row in enumerate(rows):
|
||
idx = row.get("index", j)
|
||
if 0 <= idx < len(idxs) and isinstance(row.get("embedding"), list):
|
||
out[idxs[idx]] = [float(v) for v in row["embedding"]]
|
||
except Exception as e: # noqa: BLE001
|
||
# 鉴权/配置错误拆分也救不了，直接整体失败
|
||
if isinstance(e, httpx.HTTPStatusError) and e.response.status_code in (401, 403, 404, 405):
|
||
raise
|
||
if len(items) > 1:
|
||
mid = len(items) // 2
|
||
await _call(items[:mid], idxs[:mid])
|
||
await _call(items[mid:], idxs[mid:])
|
||
elif len(items[0]) > 1500:
|
||
# 单条超模型输入上限：截断再试一次（超长段向量化本就意义有限）
|
||
try:
|
||
resp = await client.post(
|
||
url, headers=headers,
|
||
json={"model": settings.embedding_model, "input": [items[0][:1500]]},
|
||
)
|
||
resp.raise_for_status()
|
||
rows = resp.json().get("data") or []
|
||
if rows and isinstance(rows[0].get("embedding"), list):
|
||
out[idxs[0]] = [float(v) for v in rows[0]["embedding"]]
|
||
except Exception as e2: # noqa: BLE001
|
||
logger.warning("embedding 单条失败（跳过）: %s", e)
|
||
else:
|
||
logger.warning("embedding 单条失败（跳过）: %s", e)
|
||
|
||
await _call([t for _, t in tasks], [i for i, _ in tasks])
|
||
|
||
return out
|
||
|
||
|
||
async def rerank_documents(query: str, documents: list[str], top_n: int | None = None) -> list[int] | None:
|
||
"""调用 rerank API。返回按相关性降序的原始文档索引；未配置或失败返回 None。"""
|
||
if not RERANK_ENABLED or not documents:
|
||
return None
|
||
url = settings.rerank_base_url.rstrip("/") + "/rerank"
|
||
headers = {"Authorization": f"Bearer {settings.rerank_api_key}"}
|
||
payload = {"model": settings.rerank_model, "query": query, "documents": documents}
|
||
if top_n:
|
||
payload["top_n"] = top_n
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60) as client:
|
||
resp = await client.post(url, headers=headers, json=payload)
|
||
resp.raise_for_status()
|
||
results = resp.json()["results"]
|
||
ordered = sorted(results, key=lambda r: -float(r["relevance_score"]))
|
||
return [int(r["index"]) for r in ordered]
|
||
except Exception as e: # noqa: BLE001
|
||
logger.warning("rerank API 调用失败: %s", e)
|
||
return None
|
||
|
||
|
||
# ======================= 向量检索 =======================
|
||
|
||
async def load_vectors(db: AsyncSession, book_id: str) -> tuple[np.ndarray | None, list[Segment]]:
|
||
"""读取该书所有已生成向量的段。返回 (向量矩阵, 对应段列表)。"""
|
||
result = await db.execute(
|
||
select(Segment)
|
||
.where(Segment.book_id == book_id, Segment.embedding.is_not(None))
|
||
.order_by(Segment.order_idx)
|
||
)
|
||
segs = list(result.scalars().all())
|
||
vecs = []
|
||
kept = []
|
||
for s in segs:
|
||
try:
|
||
vecs.append(json.loads(s.embedding))
|
||
kept.append(s)
|
||
except (json.JSONDecodeError, TypeError):
|
||
continue
|
||
if not vecs:
|
||
return None, []
|
||
return np.asarray(vecs, dtype=np.float32), kept
|
||
|
||
|
||
def vector_search(query_vec: list[float], mat: np.ndarray, top_k: int) -> list[int]:
|
||
"""余弦相似度 top_k，返回 mat 行索引（降序）。"""
|
||
q = np.asarray(query_vec, dtype=np.float32)
|
||
q_norm = np.linalg.norm(q)
|
||
if q_norm == 0:
|
||
return []
|
||
sims = (mat @ q) / (np.linalg.norm(mat, axis=1) * q_norm + 1e-9)
|
||
top = np.argsort(-sims)[:top_k]
|
||
return [int(i) for i in top]
|
||
|
||
|
||
def rrf_fuse(*rankings: list[str], k: int = 60) -> list[str]:
|
||
"""Reciprocal Rank Fusion：合并多个按相关性降序的 segment_id 列表。"""
|
||
scores: dict[str, float] = {}
|
||
for ranking in rankings:
|
||
for rank, sid in enumerate(ranking):
|
||
scores[sid] = scores.get(sid, 0.0) + 1.0 / (k + rank + 1)
|
||
return [sid for sid, _ in sorted(scores.items(), key=lambda kv: -kv[1])]
|
||
|
||
|
||
# ======================= 向量生成（上传后后台 + 检索前懒补）=======================
|
||
|
||
def _set_status(book_id: str, status: str) -> None:
|
||
with _status_lock:
|
||
_vector_status[book_id] = status
|
||
|
||
|
||
def get_status(book_id: str) -> str | None:
|
||
with _status_lock:
|
||
return _vector_status.get(book_id)
|
||
|
||
|
||
async def ensure_embeddings(db: AsyncSession, book_id: str) -> str:
|
||
"""确保 book 的每个 Segment 都有向量。返回状态 done / running / failed / disabled。
|
||
|
||
- 后台任务与请求可能并发调用：running 状态会直接返回，由另一方补齐
|
||
- 只对缺少向量的段生成，重复调用是增量补差
|
||
"""
|
||
if not EMBEDDING_ENABLED:
|
||
_set_status(book_id, "failed")
|
||
return "disabled"
|
||
|
||
status = get_status(book_id)
|
||
if status == "running":
|
||
return "running"
|
||
|
||
result = await db.execute(
|
||
select(Segment)
|
||
.where(Segment.book_id == book_id, Segment.embedding.is_(None))
|
||
.order_by(Segment.order_idx)
|
||
)
|
||
missing = list(result.scalars().all())
|
||
# 空文本段无法向量化（API 对空串 400），也永远无检索价值，直接排除
|
||
missing = [s for s in missing if (s.text or "").strip()]
|
||
if not missing:
|
||
_set_status(book_id, "done")
|
||
return "done"
|
||
|
||
_set_status(book_id, "running")
|
||
try:
|
||
batch = settings.embedding_batch_size
|
||
for start in range(0, len(missing), batch):
|
||
segs = missing[start:start + batch]
|
||
vecs = await embed_texts([s.text for s in segs])
|
||
if vecs is None:
|
||
raise RuntimeError("embedding API 不可用")
|
||
# 单条失败时 vec 为 None，跳过不落库；其余照常持久化
|
||
for seg, vec in zip(segs, vecs):
|
||
if vec is not None:
|
||
seg.embedding = json.dumps([float(v) for v in vec], ensure_ascii=False)
|
||
await db.commit()
|
||
# 已尝试全部缺失段。个别被模型拒绝的段（内容/超限）由关键词档兜底，
|
||
# 标记 done 避免每次检索都重复重试；API 级故障会在 except 里落为 failed
|
||
_set_status(book_id, "done")
|
||
return "done"
|
||
except Exception as e: # noqa: BLE001
|
||
logger.warning("向量生成失败 book=%s: %s", book_id, e)
|
||
_set_status(book_id, "failed")
|
||
# 已提交的批次不受影响；rollback 仅丢弃本批次未提交变更
|
||
await db.rollback()
|
||
return "failed"
|