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]
|