90 lines
2.5 KiB
Python
90 lines
2.5 KiB
Python
# core/llm_client.py - LLM 客户端封装
|
||
import json
|
||
from typing import Optional
|
||
|
||
import config
|
||
|
||
|
||
class LLMClient:
|
||
"""
|
||
OpenAI 兼容 LLM 客户端封装。
|
||
支持任何兼容 OpenAI API 格式的服务（OpenAI / Azure / 本地模型等）。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
api_key: str = config.LLM_API_KEY,
|
||
base_url: str = config.LLM_BASE_URL,
|
||
model: str = config.LLM_MODEL,
|
||
temperature: float = config.LLM_TEMPERATURE,
|
||
):
|
||
"""
|
||
初始化 LLM 客户端
|
||
|
||
Args:
|
||
api_key: API 密钥
|
||
base_url: API 基础 URL
|
||
model: 模型名称
|
||
temperature: 生成温度（0~1，越低越确定）
|
||
|
||
Raises:
|
||
ImportError: 未安装 openai 库
|
||
ValueError: api_key 为空
|
||
"""
|
||
try:
|
||
from openai import OpenAI
|
||
except ImportError:
|
||
raise ImportError("请安装 openai: pip install openai")
|
||
|
||
if not api_key:
|
||
raise ValueError("LLM_API_KEY 未配置，请在 .env 文件中设置")
|
||
|
||
self.model = model
|
||
self.temperature = temperature
|
||
self._client = OpenAI(api_key=api_key, base_url=base_url)
|
||
|
||
def chat(self, system_prompt: str, user_prompt: str) -> str:
|
||
"""
|
||
发送对话请求，返回模型回复文本
|
||
|
||
Args:
|
||
system_prompt: 系统提示词
|
||
user_prompt: 用户输入
|
||
|
||
Returns:
|
||
模型回复的文本内容
|
||
|
||
Raises:
|
||
Exception: API 调用失败
|
||
"""
|
||
response = self._client.chat.completions.create(
|
||
model=self.model,
|
||
temperature=self.temperature,
|
||
messages=[
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": user_prompt},
|
||
],
|
||
)
|
||
return response.choices[0].message.content.strip()
|
||
|
||
def chat_json(self, system_prompt: str, user_prompt: str) -> dict:
|
||
"""
|
||
发送对话请求，解析并返回 JSON 结果
|
||
|
||
Args:
|
||
system_prompt: 系统提示词
|
||
user_prompt: 用户输入
|
||
|
||
Returns:
|
||
解析后的 dict 对象
|
||
|
||
Raises:
|
||
json.JSONDecodeError: 模型返回非合法 JSON
|
||
"""
|
||
raw = self.chat(system_prompt, user_prompt)
|
||
# 去除可能的 markdown 代码块包裹
|
||
raw = raw.strip()
|
||
if raw.startswith("```"):
|
||
lines = raw.split("\n")
|
||
raw = "\n".join(lines[1:-1])
|
||
return json.loads(raw) |