152 lines
4.8 KiB
Python
152 lines
4.8 KiB
Python
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import requests
|
|
|
|
|
|
class LLMConfigManager:
|
|
def __init__(self, config_path: str):
|
|
self.config_path = Path(config_path)
|
|
|
|
def load(self) -> dict:
|
|
if not self.config_path.exists():
|
|
return {}
|
|
with self.config_path.open("r", encoding="utf-8-sig") as file:
|
|
return json.load(file)
|
|
|
|
|
|
def _join_url(base_url: str, path: str) -> str:
|
|
return f"{base_url.rstrip('/')}/{path.lstrip('/')}"
|
|
|
|
|
|
def _timeout(config: dict) -> int:
|
|
section = config.get("request", {}) if isinstance(config, dict) else {}
|
|
try:
|
|
return int(section.get("timeout_seconds", 45))
|
|
except Exception:
|
|
return 45
|
|
|
|
|
|
def _extract_openai_text(payload: Any) -> str:
|
|
if isinstance(payload, dict):
|
|
choices = payload.get("choices")
|
|
if isinstance(choices, list) and choices:
|
|
msg = choices[0].get("message") if isinstance(choices[0], dict) else {}
|
|
if isinstance(msg, dict) and msg.get("content"):
|
|
return str(msg["content"])
|
|
if payload.get("answer"):
|
|
return str(payload["answer"])
|
|
if payload.get("data"):
|
|
return json.dumps(payload["data"], ensure_ascii=False)
|
|
return json.dumps(payload, ensure_ascii=False) if isinstance(payload, (dict, list)) else str(payload)
|
|
|
|
|
|
def ask_fastgpt(config: dict, prompt: str, context: str = "", custom_uid: str = "") -> dict:
|
|
section = config.get("fastgpt", {})
|
|
base_url = section.get("base_url", "").strip()
|
|
api_key = section.get("api_key", "").strip()
|
|
chat_id = str(section.get("chat_id", "111"))
|
|
model = section.get("model", "")
|
|
timeout = _timeout(config)
|
|
|
|
if not base_url or not api_key:
|
|
return {"ok": False, "message": "FastGPT 未配置 base_url 或 api_key"}
|
|
|
|
url = _join_url(base_url, "/v1/chat/completions")
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
content = prompt if not context else f"请基于以下本地知识库内容回答。\n\n知识库:\n{context}\n\n问题: {prompt}"
|
|
messages = []
|
|
system_prompt = section.get("system_prompt")
|
|
if system_prompt:
|
|
messages.append({"role": "system", "content": system_prompt})
|
|
messages.append({"role": "user", "content": content})
|
|
|
|
use_custom_uid = custom_uid or section.get("custom_uid", "")
|
|
payload = {
|
|
"chatId": chat_id,
|
|
"stream": False,
|
|
"detail": False,
|
|
"messages": messages,
|
|
}
|
|
if model:
|
|
payload["model"] = model
|
|
if use_custom_uid:
|
|
payload["customUid"] = use_custom_uid
|
|
|
|
try:
|
|
response = requests.post(url, headers=headers, json=payload, timeout=timeout, allow_redirects=True)
|
|
if response.status_code >= 400:
|
|
return {
|
|
"ok": False,
|
|
"message": f"FastGPT 调用失败: HTTP {response.status_code}",
|
|
"detail": response.text,
|
|
}
|
|
|
|
data = response.json()
|
|
return {
|
|
"ok": True,
|
|
"provider": "fastgpt",
|
|
"raw": data,
|
|
"answer": _extract_openai_text(data),
|
|
}
|
|
except Exception as exc:
|
|
return {
|
|
"ok": False,
|
|
"message": "FastGPT 请求异常",
|
|
"detail": str(exc),
|
|
}
|
|
|
|
|
|
def ask_dify(config: dict, prompt: str, context: str = "", user: str = "anonymous") -> dict:
|
|
section = config.get("dify", {})
|
|
base_url = section.get("base_url", "").strip()
|
|
api_key = section.get("api_key", "").strip()
|
|
timeout = _timeout(config)
|
|
|
|
if not base_url or not api_key:
|
|
return {"ok": False, "message": "Dify 未配置 base_url 或 api_key"}
|
|
|
|
url = _join_url(base_url, "/v1/chat-messages")
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
query = prompt if not context else f"知识库上下文:\n{context}\n\n用户问题:\n{prompt}"
|
|
payload = {
|
|
"inputs": {},
|
|
"query": query,
|
|
"response_mode": "blocking",
|
|
"conversation_id": "",
|
|
"user": user,
|
|
}
|
|
|
|
try:
|
|
response = requests.post(url, headers=headers, json=payload, timeout=timeout, allow_redirects=True)
|
|
if response.status_code >= 400:
|
|
return {
|
|
"ok": False,
|
|
"message": f"Dify 调用失败: HTTP {response.status_code}",
|
|
"detail": response.text,
|
|
}
|
|
|
|
data = response.json()
|
|
answer = data.get("answer") or _extract_openai_text(data)
|
|
return {
|
|
"ok": True,
|
|
"provider": "dify",
|
|
"raw": data,
|
|
"answer": answer,
|
|
}
|
|
except Exception as exc:
|
|
return {
|
|
"ok": False,
|
|
"message": "Dify 请求异常",
|
|
"detail": str(exc),
|
|
}
|