36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
import logging
|
|
from typing import List, Optional
|
|
|
|
import httpx
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class EmbeddingGenerator:
|
|
"""简单的 OpenAI 兼容嵌入生成器"""
|
|
|
|
def __init__(self, api_base: str, api_key: str, model: str, timeout: float = 30.0) -> None:
|
|
self.api_base = api_base.rstrip("/")
|
|
self.api_key = api_key
|
|
self.model = model
|
|
self.timeout = timeout
|
|
|
|
async def embed_texts(self, texts: List[str]) -> List[Optional[List[float]]]:
|
|
if not texts:
|
|
return []
|
|
payload = {"input": texts, "model": self.model}
|
|
headers = {
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
|
resp = await client.post(f"{self.api_base}/embeddings", json=payload, headers=headers)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
embeddings = []
|
|
for item in data.get("data", []):
|
|
embeddings.append(item.get("embedding"))
|
|
if len(embeddings) != len(texts):
|
|
logger.warning("嵌入结果数量与输入不一致：%s != %s", len(embeddings), len(texts))
|
|
return embeddings
|