simonkoson
69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
"""
|
||
Embedding 调用服务 — 封装 MiniMax embo-01
|
||
|
||
请求格式（确认自探路脚本）：
|
||
POST /v1/embeddings
|
||
Body: {"model": "embo-01", "texts": [...], "type": "db"|"query"}
|
||
响应格式：
|
||
{"vectors": [[...1536 floats...]], "total_tokens": N, "base_resp": {"status_code": 0, "status_msg": "success"}}
|
||
"""
|
||
|
||
import httpx
|
||
from typing import List
|
||
|
||
from app.core.config import settings
|
||
|
||
|
||
class EmbeddingService:
|
||
"""MiniMax embo-01 embedding 调用封装"""
|
||
|
||
def __init__(self):
|
||
self.api_key = settings.MINIMAX_EMBED_API_KEY
|
||
self.group_id = settings.MINIMAX_GROUP_ID
|
||
self.endpoint = "https://api.minimax.chat/v1/embeddings"
|
||
|
||
def embed(self, texts: List[str], embed_type: str = "db") -> List[List[float]]:
|
||
"""
|
||
调用 embo-01 将文本列表转为向量
|
||
|
||
Args:
|
||
texts: 文本列表（支持批量）
|
||
embed_type: "db" = 存入库，"query" = 查询
|
||
Returns:
|
||
List[List[float]]，每个元素是一组 1536 维向量
|
||
"""
|
||
if not self.api_key or self.api_key == "your_api_key_here":
|
||
raise RuntimeError("MINIMAX_EMBED_API_KEY not configured in .env")
|
||
if not self.group_id or self.group_id == "your_group_id_here":
|
||
raise RuntimeError("MINIMAX_GROUP_ID not configured in .env")
|
||
|
||
headers = {
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"GroupId": self.group_id,
|
||
"Content-Type": "application/json",
|
||
}
|
||
payload = {
|
||
"model": "embo-01",
|
||
"texts": texts,
|
||
"type": embed_type,
|
||
}
|
||
|
||
resp = httpx.post(self.endpoint, headers=headers, json=payload, timeout=60.0)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
|
||
# 检查业务错误
|
||
base_resp = data.get("base_resp", {})
|
||
if base_resp.get("status_code", 0) != 0:
|
||
raise RuntimeError(f"Embedding API error: {base_resp.get('status_msg', 'unknown')}")
|
||
|
||
vectors = data.get("vectors", [])
|
||
if not vectors:
|
||
raise RuntimeError("No vectors returned from embedding API")
|
||
|
||
return vectors
|
||
|
||
def embed_single(self, text: str, embed_type: str = "db") -> List[float]:
|
||
"""单文本 embedding，返回 1536 维向量列表（Python list）"""
|
||
vectors = self.embed([text], embed_type=embed_type)
|
||
return vectors[0] |