freedak
641e33b834
- 语义检索(FAISS + embedding)+ 精确查找(法规名+条号) - RAG 问答(SSE 流式,支持 thinking 折叠显示) - 法规浏览(原文阅读) - 历史记录(检索+对话持久化到 SQLite) - 设置页(系统提示词/模板/LLM 参数可配置) - 检索质量评估脚本 Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
80 lines
2.2 KiB
Python
80 lines
2.2 KiB
Python
"""embedding 服务客户端 — 调用 114:8003 Qwen3-Embedding 服务"""
|
|
import asyncio
|
|
import logging
|
|
from typing import List
|
|
|
|
import httpx
|
|
|
|
from app.config import EMBEDDING_URL, EMBEDDING_MODEL, EMBEDDING_BATCH_SIZE
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_client: httpx.AsyncClient | None = None
|
|
|
|
|
|
def _get_client() -> httpx.AsyncClient:
|
|
"""复用 httpx 客户端"""
|
|
global _client
|
|
if _client is None or _client.is_closed:
|
|
_client = httpx.AsyncClient(timeout=60)
|
|
return _client
|
|
|
|
|
|
async def embed_batch(texts: List[str], batch_size: int = EMBEDDING_BATCH_SIZE) -> List[List[float]]:
|
|
"""批量向量化
|
|
|
|
Args:
|
|
texts: 文本列表
|
|
batch_size: 每批大小,默认 32
|
|
|
|
Returns:
|
|
向量列表,顺序与输入一致
|
|
|
|
Raises:
|
|
RuntimeError: embedding 服务不可用或返回错误
|
|
"""
|
|
if not texts:
|
|
return []
|
|
|
|
client = _get_client()
|
|
all_vecs: List[List[float]] = []
|
|
|
|
for i in range(0, len(texts), batch_size):
|
|
batch = texts[i : i + batch_size]
|
|
payload = {
|
|
"model": EMBEDDING_MODEL,
|
|
"input": batch,
|
|
}
|
|
|
|
# 重试 3 次,指数退避
|
|
last_err = None
|
|
for attempt in range(3):
|
|
try:
|
|
resp = await client.post(
|
|
f"{EMBEDDING_URL}/embeddings",
|
|
json=payload,
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
# vLLM embeddings 响应格式: {data: [{embedding: [...]}], ...}
|
|
vecs = [item["embedding"] for item in data["data"]]
|
|
all_vecs.extend(vecs)
|
|
last_err = None
|
|
break
|
|
except Exception as e:
|
|
last_err = e
|
|
wait = 2 ** attempt
|
|
logger.warning(f"embedding 批次 {i//batch_size} 第 {attempt+1} 次失败: {e}, {wait}s 后重试")
|
|
await asyncio.sleep(wait)
|
|
|
|
if last_err is not None:
|
|
raise RuntimeError(f"embedding 服务调用失败(重试 3 次): {last_err}")
|
|
|
|
return all_vecs
|
|
|
|
|
|
async def embed_text(text: str) -> List[float]:
|
|
"""单条文本向量化"""
|
|
vecs = await embed_batch([text])
|
|
return vecs[0]
|