104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
"""统一的 OpenAI 兼容 LLM client：文本 + 多模态。
|
||
|
||
配置来自用户的 UserSettings（base_url / api_key / 模型名）。api_key 在库里
|
||
是 Fernet 加密的，这里解密后使用。
|
||
"""
|
||
import base64
|
||
from dataclasses import dataclass
|
||
|
||
from openai import AsyncOpenAI
|
||
|
||
from app.models.user import UserSettings
|
||
from app.services.crypto import decrypt
|
||
|
||
|
||
class LLMConfigError(RuntimeError):
|
||
"""用户尚未正确配置 AI 接口。"""
|
||
|
||
|
||
@dataclass
|
||
class LLMResult:
|
||
content: str
|
||
model: str | None
|
||
prompt_tokens: int | None
|
||
completion_tokens: int | None
|
||
|
||
|
||
class LLMClient:
|
||
def __init__(self, settings: UserSettings):
|
||
if not settings or not settings.llm_base_url or not settings.llm_api_key_encrypted:
|
||
raise LLMConfigError("请先在设置里配置 AI 接口（base_url 与 api_key）")
|
||
api_key = decrypt(settings.llm_api_key_encrypted)
|
||
if not api_key:
|
||
raise LLMConfigError("api_key 解密失败，请在设置里重新填写")
|
||
self._client = AsyncOpenAI(base_url=settings.llm_base_url, api_key=api_key)
|
||
self._text_model = settings.llm_text_model
|
||
self._vision_model = settings.llm_vision_model or settings.llm_text_model
|
||
|
||
async def chat_text(
|
||
self,
|
||
system: str,
|
||
user: str,
|
||
*,
|
||
json_mode: bool = False,
|
||
model: str | None = None,
|
||
) -> LLMResult:
|
||
m = model or self._text_model
|
||
if not m:
|
||
raise LLMConfigError("未配置文本模型名")
|
||
kwargs = {}
|
||
if json_mode:
|
||
kwargs["response_format"] = {"type": "json_object"}
|
||
resp = await self._client.chat.completions.create(
|
||
model=m,
|
||
messages=[
|
||
{"role": "system", "content": system},
|
||
{"role": "user", "content": user},
|
||
],
|
||
**kwargs,
|
||
)
|
||
return self._to_result(resp, m)
|
||
|
||
async def chat_vision(
|
||
self,
|
||
system: str,
|
||
user_text: str,
|
||
image_bytes: bytes,
|
||
mime: str = "image/jpeg",
|
||
*,
|
||
json_mode: bool = False,
|
||
model: str | None = None,
|
||
) -> LLMResult:
|
||
m = model or self._vision_model
|
||
if not m:
|
||
raise LLMConfigError("未配置多模态模型名")
|
||
data_url = f"data:{mime};base64,{base64.b64encode(image_bytes).decode()}"
|
||
kwargs = {}
|
||
if json_mode:
|
||
kwargs["response_format"] = {"type": "json_object"}
|
||
resp = await self._client.chat.completions.create(
|
||
model=m,
|
||
messages=[
|
||
{"role": "system", "content": system},
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": user_text},
|
||
{"type": "image_url", "image_url": {"url": data_url}},
|
||
],
|
||
},
|
||
],
|
||
**kwargs,
|
||
)
|
||
return self._to_result(resp, m)
|
||
|
||
@staticmethod
|
||
def _to_result(resp, model: str) -> LLMResult:
|
||
usage = getattr(resp, "usage", None)
|
||
return LLMResult(
|
||
content=resp.choices[0].message.content or "",
|
||
model=model,
|
||
prompt_tokens=getattr(usage, "prompt_tokens", None),
|
||
completion_tokens=getattr(usage, "completion_tokens", None),
|
||
)
|