64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
"""Embedding 服务：调用硅基流动 bge-m3"""
|
||
|
||
import asyncio
|
||
import httpx
|
||
from config import (
|
||
SILICONFLOW_API_KEY,
|
||
SILICONFLOW_BASE_URL,
|
||
EMBEDDING_MODEL,
|
||
EMBEDDING_DIM,
|
||
)
|
||
|
||
BATCH_SIZE = 50
|
||
MAX_RETRIES = 2
|
||
|
||
|
||
async def embed_single(text: str) -> list[float]:
|
||
"""单条文本 embedding"""
|
||
async with httpx.AsyncClient(timeout=30) as client:
|
||
resp = await client.post(
|
||
f"{SILICONFLOW_BASE_URL}/embeddings",
|
||
headers={"Authorization": f"Bearer {SILICONFLOW_API_KEY}"},
|
||
json={
|
||
"model": EMBEDDING_MODEL,
|
||
"input": [text],
|
||
},
|
||
)
|
||
if resp.status_code != 200:
|
||
raise RuntimeError(f"Embedding API error: {resp.status_code} {resp.text}")
|
||
data = resp.json()
|
||
return data["data"][0]["embedding"]
|
||
|
||
|
||
async def embed_batch(texts: list[str]) -> list[list[float]]:
|
||
"""批量 embedding，自动分批 + 重试"""
|
||
all_embeddings = []
|
||
|
||
for i in range(0, len(texts), BATCH_SIZE):
|
||
batch = texts[i:i + BATCH_SIZE]
|
||
for attempt in range(MAX_RETRIES + 1):
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60) as client:
|
||
resp = await client.post(
|
||
f"{SILICONFLOW_BASE_URL}/embeddings",
|
||
headers={"Authorization": f"Bearer {SILICONFLOW_API_KEY}"},
|
||
json={
|
||
"model": EMBEDDING_MODEL,
|
||
"input": batch,
|
||
},
|
||
)
|
||
if resp.status_code == 200:
|
||
data = resp.json()
|
||
all_embeddings.extend([d["embedding"] for d in data["data"]])
|
||
break
|
||
else:
|
||
err = f"Status {resp.status_code}"
|
||
if attempt == MAX_RETRIES:
|
||
raise RuntimeError(f"Embedding batch failed after {MAX_RETRIES} retries: {err}")
|
||
except Exception as e:
|
||
if attempt == MAX_RETRIES:
|
||
raise RuntimeError(f"Embedding batch failed: {e}")
|
||
await asyncio.sleep(2 ** attempt)
|
||
|
||
return all_embeddings
|