zqf
b1f93588f1
- 排除 .env / *.bak / 内部运维文档（README_INTERNAL.html） - config.py 默认密码已替换为占位符 CHANGE_ME_* - init_db.sql 移除生产数据库用户 GRANT 段 - README.html 数据库用户名已脱敏 - 保留：源码 + 公网 API 文档 + 建表 SQL（无授权语句）
68 lines
2.4 KiB
Python
68 lines
2.4 KiB
Python
"""BGE Embedding Service Client"""
|
|
import requests
|
|
import logging
|
|
from typing import List
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# BGE-small-zh-v1.5: 512 tokens, roughly ~1500 Chinese chars
|
|
# Conservative limit: 2000 chars (most text under this is safe)
|
|
MAX_TEXT_CHARS = 2000
|
|
|
|
|
|
class EmbeddingService:
|
|
"""Client for BGE-small-zh-v1.5 embedding service."""
|
|
|
|
def __init__(self, base_url: str = "http://127.0.0.1:8080", timeout: int = 30):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.timeout = timeout
|
|
|
|
def embed(self, text: str) -> List[float]:
|
|
"""Get embedding for a single text.
|
|
|
|
Text is truncated to MAX_TEXT_CHARS if too long.
|
|
"""
|
|
original_len = len(text)
|
|
if len(text) > MAX_TEXT_CHARS:
|
|
logger.warning(f"Text too long ({original_len} chars), truncating to {MAX_TEXT_CHARS}")
|
|
text = text[:MAX_TEXT_CHARS]
|
|
|
|
try:
|
|
resp = requests.post(
|
|
f"{self.base_url}/embed",
|
|
json={"text": text},
|
|
timeout=self.timeout,
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
return data["embedding"]
|
|
except requests.exceptions.ConnectionError:
|
|
msg = "Embedding service unavailable (connection refused)"
|
|
logger.error(msg)
|
|
raise RuntimeError(msg)
|
|
except requests.exceptions.Timeout:
|
|
msg = f"Embedding request timeout ({self.timeout}s). Text may be too long ({original_len} chars, max {MAX_TEXT_CHARS})"
|
|
logger.error(msg)
|
|
raise RuntimeError(msg)
|
|
except requests.exceptions.HTTPError as e:
|
|
status = e.response.status_code if e.response is not None else "unknown"
|
|
msg = f"Embedding service error (HTTP {status}): {e}"
|
|
logger.error(msg)
|
|
raise RuntimeError(msg)
|
|
except Exception as e:
|
|
msg = f"Embedding failed: {e}"
|
|
logger.error(msg)
|
|
raise RuntimeError(msg)
|
|
|
|
def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
|
"""Get embeddings for multiple texts."""
|
|
return [self.embed(t) for t in texts]
|
|
|
|
def health_check(self) -> bool:
|
|
"""Check if embedding service is alive."""
|
|
try:
|
|
resp = requests.get(f"{self.base_url}/health", timeout=5)
|
|
return resp.status_code == 200
|
|
except Exception:
|
|
return False
|