124 lines
3.6 KiB
Python
124 lines
3.6 KiB
Python
"""
|
|
Kimi VLM Provider
|
|
使用 Moonshot AI 的 Kimi 1.5 视觉模型
|
|
"""
|
|
|
|
import base64
|
|
import time
|
|
from typing import List
|
|
|
|
from .base import VLMProvider
|
|
|
|
|
|
class KimiProvider(VLMProvider):
|
|
"""
|
|
Kimi (Moonshot AI) 视觉模型提供商
|
|
使用 OpenAI SDK 兼容接口
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
api_key: str,
|
|
base_url: str = "https://api.moonshot.cn/v1",
|
|
model: str = "moonshot-v1-32k-vision-preview",
|
|
):
|
|
"""
|
|
Args:
|
|
api_key: Moonshot AI API Key
|
|
base_url: API 端点地址
|
|
model: 模型名称
|
|
"""
|
|
if not api_key:
|
|
raise ValueError("Kimi API Key 不能为空")
|
|
|
|
self.api_key = api_key
|
|
self.base_url = base_url
|
|
self.model = model
|
|
|
|
def chat_vision(
|
|
self,
|
|
prompt: str,
|
|
image_paths: List[str],
|
|
temperature: float = 0.1,
|
|
timeout_sec: int = 120,
|
|
) -> str:
|
|
"""调用 Kimi API 进行视觉推理"""
|
|
try:
|
|
from openai import OpenAI
|
|
except ImportError:
|
|
raise ImportError(
|
|
"Kimi Provider 需要 openai 库，请安装: pip install openai"
|
|
)
|
|
|
|
print(
|
|
f"[Kimi] Prepare chat model={self.model} temp={temperature} "
|
|
f"images={len(image_paths)} timeout={timeout_sec}s"
|
|
)
|
|
|
|
# 初始化 OpenAI 客户端
|
|
client = OpenAI(
|
|
api_key=self.api_key,
|
|
base_url=self.base_url,
|
|
timeout=timeout_sec,
|
|
)
|
|
|
|
# 构建消息内容（文本 + 图片）
|
|
# Kimi 支持 image_url 格式，使用 base64 编码
|
|
content_parts = [{"type": "text", "text": prompt}]
|
|
|
|
for img_path in image_paths:
|
|
try:
|
|
with open(img_path, "rb") as f:
|
|
img_data = f.read()
|
|
b64_str = base64.b64encode(img_data).decode("utf-8")
|
|
|
|
# 根据文件扩展名判断 MIME 类型
|
|
if img_path.lower().endswith((".jpg", ".jpeg")):
|
|
mime = "image/jpeg"
|
|
elif img_path.lower().endswith(".png"):
|
|
mime = "image/png"
|
|
else:
|
|
mime = "image/jpeg" # 默认
|
|
|
|
data_uri = f"data:{mime};base64,{b64_str}"
|
|
content_parts.append({
|
|
"type": "image_url",
|
|
"image_url": {"url": data_uri}
|
|
})
|
|
except Exception as e:
|
|
print(f"[Kimi] Warning: Failed to load image {img_path}: {e}")
|
|
continue
|
|
|
|
# 调用 API
|
|
t0 = time.time()
|
|
try:
|
|
completion = client.chat.completions.create(
|
|
model=self.model,
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": content_parts,
|
|
}
|
|
],
|
|
temperature=temperature,
|
|
)
|
|
|
|
dt = time.time() - t0
|
|
text = completion.choices[0].message.content or ""
|
|
|
|
result = text.replace("\n", " ")
|
|
print(f"[Kimi] Response elapsed={dt:.2f}s length={len(text)} preview={result}")
|
|
|
|
if not text:
|
|
raise RuntimeError("Kimi 返回内容为空")
|
|
|
|
return text
|
|
|
|
except Exception as e:
|
|
dt = time.time() - t0
|
|
print(f"[Kimi][Error] during request elapsed={dt:.2f}s: {e}")
|
|
raise
|
|
|
|
def __repr__(self):
|
|
return f"<KimiProvider model={self.model} base_url={self.base_url}>"
|