陈世睿
74caa9d112
- crawler/:豆果(主)/下厨房/Bing 爬虫 + dishes.jsonl(3056 菜→9167 图映射)+ verify_repair 按 URL 重下 - retrieval/:三路检索(BM25 + 本地 BGE-M3 向量 + RRF 融合),FastAPI + 前端;写死图片目录(默认 crawler/images,可 IMAGE_DIR 覆盖)、绑 0.0.0.0 局域网访问、启动自动建索引、服务器 serve 图片 - 图片(1.7G)与向量模型(2.3G)不进 git Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
232 lines
9.8 KiB
Python
232 lines
9.8 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""通义千问 DashScope 文本向量客户端(text-embedding-v3,1024 维)。
|
||
|
||
要点:
|
||
- Key 复用 PriceBot 的 QWEN_API_KEY;可用环境变量 DASHSCOPE_API_KEY / QWEN_API_KEY 覆盖。
|
||
- 接口走 OpenAI 兼容端点;单次批量上限 10(硬限),用线程池并行多批。
|
||
- name -> 向量 落盘缓存(cache/emb_<model>.jsonl):重复加载同一批菜名零成本、零计费。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import sys
|
||
import threading
|
||
import time
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
|
||
import numpy as np
|
||
import requests
|
||
|
||
DASHSCOPE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings"
|
||
# 默认复用 PriceBot 的通义千问 Key(本工具在 dish-images 下、不进 git);优先用环境变量覆盖。
|
||
DEFAULT_KEY = "sk-f66e689a9b0c43239e299137f68c453c"
|
||
BATCH = 10 # text-embedding-v3 单次最多 10 条
|
||
# 免费额度对"并发齐发"很敏感(冷却后多线程同时重试会互撞、永远挤不进窗口),默认单线程顺序发,
|
||
# 由下方自适应间隔 _delay 控速最稳。Key 额度充足时可设环境变量 EMBED_WORKERS=4 之类提速。
|
||
MAX_WORKERS = max(1, int(os.environ.get("EMBED_WORKERS", "1")))
|
||
|
||
|
||
def get_key(override: str | None = None) -> str:
|
||
return (override or os.environ.get("DASHSCOPE_API_KEY")
|
||
or os.environ.get("QWEN_API_KEY") or DEFAULT_KEY)
|
||
|
||
|
||
_tls = threading.local()
|
||
|
||
|
||
def _sess() -> requests.Session:
|
||
"""线程本地 session,trust_env=False:忽略系统代理(HTTP(S)_PROXY),直连 dashscope。
|
||
Clash 等代理偶尔把请求甩到海外节点,会被阿里云按来源风控拒(HTTP 400 Access denied);
|
||
dashscope.aliyuncs.com 是大陆公网端点,直连最稳(与 PriceBot 调千问一致)。"""
|
||
s = getattr(_tls, "s", None)
|
||
if s is None:
|
||
s = _tls.s = requests.Session()
|
||
s.trust_env = False
|
||
return s
|
||
|
||
|
||
# ---------- 全局自适应限速:稳态间隔 + 冷却闸,收敛到免费额度可持续的速率 ----------
|
||
# dashscope 免费额度是「滚动每分钟预算」式:瞬时小量能过,持续高频会返回 400,且文案写成
|
||
# Arrearage/overdue(其实是"免费窗口用尽且无余额续费",约 1 分钟自动回血)。故这是【可重试】
|
||
# 信号:命中就放大间隔(降速)+ 拉长全员冷却(等回血),成功就缓慢提速,自动逼近可持续速率。
|
||
_pace_lock = threading.Lock()
|
||
_next_at = 0.0
|
||
_resume_at = 0.0
|
||
_delay = 0.6 # 当前稳态请求间隔(秒),单线程顺序发
|
||
_cool = 5.0 # 当前冷却时长(秒),命中放大、成功回落
|
||
DELAY_MIN, DELAY_MAX = 0.3, 4.0
|
||
COOL_MIN, COOL_MAX = 5.0, 45.0
|
||
|
||
|
||
def _pace() -> None:
|
||
"""按当前稳态间隔 _delay 给每个请求排队发牌(有效速率 ≈ 1/_delay)。"""
|
||
global _next_at
|
||
with _pace_lock:
|
||
now = time.time()
|
||
slot = max(now, _next_at)
|
||
_next_at = slot + _delay
|
||
wait = slot - now
|
||
if wait > 0:
|
||
time.sleep(wait)
|
||
|
||
|
||
def _gate_wait() -> None:
|
||
"""命中限速后的全员冷却,期间所有线程一起等(等额度回血)。"""
|
||
while True:
|
||
with _pace_lock:
|
||
wait = _resume_at - time.time()
|
||
if wait <= 0:
|
||
return
|
||
time.sleep(min(wait, 2.0) + 0.05)
|
||
|
||
|
||
def _rate_hit() -> None:
|
||
"""命中免费额度限速:降速 + 拉长冷却等回血。"""
|
||
global _delay, _cool, _resume_at
|
||
with _pace_lock:
|
||
_delay = min(_delay * 1.5, DELAY_MAX)
|
||
_cool = min(_cool * 1.4, COOL_MAX)
|
||
_resume_at = max(_resume_at, time.time() + _cool)
|
||
|
||
|
||
def _ok() -> None:
|
||
"""成功:缓慢提速、缩短冷却,收敛到刚好不撞墙。"""
|
||
global _delay, _cool
|
||
with _pace_lock:
|
||
_delay = max(_delay * 0.9, DELAY_MIN)
|
||
_cool = max(_cool * 0.9, COOL_MIN)
|
||
|
||
|
||
class ArrearsError(RuntimeError):
|
||
"""重试多轮仍被免费额度拒(疑似额度耗尽/欠费),供上层提示"稍后重试/换 Key"。"""
|
||
|
||
|
||
def _is_rate_signal(status: int, text: str) -> bool:
|
||
"""免费额度限速信号(可重试):含真限速(429/503/throttling)与"伪欠费"(Arrearage/overdue,
|
||
实为免费窗口用尽、约 1 分钟回血)。两者都靠退避+等待解决。"""
|
||
if status in (429, 503):
|
||
return True
|
||
low = text.lower()
|
||
return ("throttling" in low or "rate limit" in low or "requests rate" in low
|
||
or "arrearage" in low or "overdue" in low or "good standing" in low
|
||
or "欠费" in text or "余额不足" in text)
|
||
|
||
|
||
class Embedder:
|
||
def __init__(self, model: str = "text-embedding-v4", cache_dir: str = "cache",
|
||
key: str | None = None):
|
||
self.model = model
|
||
self.key = get_key(key)
|
||
os.makedirs(cache_dir, exist_ok=True)
|
||
self.cache_path = os.path.join(cache_dir, f"emb_{model}.jsonl")
|
||
self.cache: dict[str, list[float]] = {}
|
||
self.last_failed = 0
|
||
self.last_error = ""
|
||
self._lock = threading.Lock()
|
||
self._load_cache()
|
||
|
||
# ---------- 缓存 ----------
|
||
def _load_cache(self) -> None:
|
||
if not os.path.exists(self.cache_path):
|
||
return
|
||
with open(self.cache_path, encoding="utf-8") as f:
|
||
for line in f:
|
||
try:
|
||
o = json.loads(line)
|
||
self.cache[o["t"]] = o["v"]
|
||
except Exception:
|
||
pass
|
||
|
||
def _append_cache(self, items: list[tuple[str, list[float]]]) -> None:
|
||
with self._lock:
|
||
with open(self.cache_path, "a", encoding="utf-8") as f:
|
||
for t, v in items:
|
||
f.write(json.dumps({"t": t, "v": v}, ensure_ascii=False) + "\n")
|
||
self.cache[t] = v
|
||
|
||
# ---------- 接口 ----------
|
||
def _embed_batch(self, texts: list[str]) -> list[list[float]]:
|
||
body = {"model": self.model, "input": texts, "encoding_format": "float"}
|
||
last = None
|
||
rate_hits = 0
|
||
for _ in range(7): # 配合冷却升级,总等待可超 1 分钟,足以熬过免费窗口回血
|
||
_gate_wait()
|
||
_pace()
|
||
try:
|
||
r = _sess().post(
|
||
DASHSCOPE_URL,
|
||
headers={"Authorization": f"Bearer {self.key}",
|
||
"Content-Type": "application/json"},
|
||
json=body, timeout=30)
|
||
if r.status_code == 200:
|
||
_ok()
|
||
data = sorted(r.json()["data"], key=lambda d: d["index"])
|
||
return [d["embedding"] for d in data]
|
||
txt = r.text[:200]
|
||
last = f"HTTP {r.status_code}: {txt}"
|
||
if _is_rate_signal(r.status_code, txt):
|
||
rate_hits += 1
|
||
_rate_hit() # 降速 + 拉长冷却等回血后重试
|
||
else:
|
||
time.sleep(1.0) # 其它错误轻退避后重试
|
||
except requests.RequestException as e:
|
||
last = f"{type(e).__name__}: {e}"
|
||
time.sleep(2)
|
||
if rate_hits >= 4: # 多轮仍被额度拒 → 让上层提示换Key/稍后重试
|
||
raise ArrearsError(f"模型 {self.model} 免费额度反复受限(可能已用尽);"
|
||
f"可稍后重试,或在 embed.py 换 DASHSCOPE_API_KEY")
|
||
raise RuntimeError(f"embed batch failed: {last}")
|
||
|
||
def embed_corpus(self, names: list[str], progress=None) -> None:
|
||
"""把 names 全部编码进缓存。progress(done, total) 回调可选。失败的批跳过(不致命)。"""
|
||
todo = [n for n in dict.fromkeys(names) if n not in self.cache]
|
||
total = len(todo)
|
||
self.last_failed = 0
|
||
self.last_error = ""
|
||
if progress:
|
||
progress(0, total)
|
||
if not total:
|
||
return
|
||
batches = [todo[i:i + BATCH] for i in range(0, total, BATCH)]
|
||
done = 0
|
||
ok_count = 0
|
||
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as ex:
|
||
futs = {ex.submit(self._embed_batch, b): b for b in batches}
|
||
for fut in as_completed(futs):
|
||
b = futs[fut]
|
||
try:
|
||
self._append_cache(list(zip(b, fut.result())))
|
||
ok_count += len(b)
|
||
except Exception as e: # noqa: BLE001
|
||
self.last_failed += len(b)
|
||
self.last_error = str(e)[:200]
|
||
if self.last_failed <= len(b) * 2: # 只打前一两次,避免刷屏
|
||
print(f"[embed] 批失败({self.last_failed}): {e}", file=sys.stderr, flush=True)
|
||
# 一次都没成功 + 已连续失败 ≥2 批 → 判定 Key 不可用,中止省额度;
|
||
# 若前面已有成功(只是滚动额度临时受限),则继续磨,失败的批留待下次续抓。
|
||
if ok_count == 0 and self.last_failed >= 20:
|
||
print("[embed] 持续失败且零成功,中止向量编码(BM25 仍可用)",
|
||
file=sys.stderr, flush=True)
|
||
ex.shutdown(wait=False, cancel_futures=True)
|
||
break
|
||
done += len(b)
|
||
if progress:
|
||
progress(done, total)
|
||
|
||
# ---------- 取用 ----------
|
||
def matrix(self, names: list[str]) -> tuple[list[str], np.ndarray]:
|
||
"""返回 (有向量的 name 列表, L2 归一化矩阵 N×dim)。"""
|
||
keep = [n for n in names if n in self.cache]
|
||
if not keep:
|
||
return [], np.zeros((0, 0), dtype=np.float32)
|
||
m = np.array([self.cache[n] for n in keep], dtype=np.float32)
|
||
norm = np.linalg.norm(m, axis=1, keepdims=True)
|
||
norm[norm == 0] = 1.0
|
||
return keep, m / norm
|
||
|
||
def embed_query(self, q: str) -> np.ndarray:
|
||
v = np.array(self._embed_batch([q])[0], dtype=np.float32)
|
||
n = np.linalg.norm(v)
|
||
return v / (n if n else 1.0)
|