9bb1c5ce65
窗口 300s 内鉴权失败 >= 20 次即拒绝，不再查库(避免枚举放大 DB 压力)。 限速键取 token 前 12 字符，内存不留全量密钥；条目超 5000 清理过期窗口。
347 lines
13 KiB
Python
347 lines
13 KiB
Python
"""内部 LLM 代理：短期 token 鉴权，真 api_key 永不出服务进程。
|
||
|
||
背景（设计铁律）：
|
||
运行环境（角色 agent 的 run_shell 子进程、bwrap 沙箱、远程 SSH 目标机、被开发的应用）
|
||
里的东西需要调 LLM 时，**绝不下发真 api_key**——运行环境恰恰是 agent 能自由执行
|
||
shell 的地方，key 落进去等于可被 `env | grep KEY` 读出、写进交付件、commit 进 git、
|
||
打进日志；落磁盘（env/*.json 会 git 提交）更会泄露给所有能读该仓库的人，并破坏
|
||
llm 表的机构隔离。
|
||
|
||
做法：
|
||
1. 服务端签发短期 token（绑定 org_id + project_id + task_id，有过期时间/调用上限，可吊销）。
|
||
2. 运行环境拿 token 当 api_key，base_url 指向本平台的 OpenAI 兼容代理端点。
|
||
3. 代理端点校验 token → 按 token 绑定的 org_id 走 llm_bridge 解析真 key（机构隔离）
|
||
→ 转发上游 → 回填用量。真 key 只在服务进程内存里出现。
|
||
|
||
适配 browser-use / LangChain / OpenAI SDK 等一切 OpenAI 兼容客户端：
|
||
ChatOpenAI(base_url="http://<host>/pipeline_core/api/llm_proxy_v1/chat/completions" 的父路径,
|
||
api_key="<短期 token>")
|
||
"""
|
||
|
||
import json
|
||
import logging
|
||
import secrets
|
||
from datetime import datetime, timedelta
|
||
|
||
from sqlor.dbpools import DBPools
|
||
from appPublic.uniqueID import getID
|
||
|
||
DBNAME = "pipeline"
|
||
logger = logging.getLogger("pipeline.llm_proxy")
|
||
|
||
TOKEN_PREFIX = "plk-" # pipeline llm key
|
||
S_ACTIVE = "active"
|
||
S_REVOKED = "revoked"
|
||
|
||
DEFAULT_TTL_HOURS = 8
|
||
DEFAULT_MAX_CALLS = 500
|
||
|
||
# 失败限速：防 token 枚举。按来源标识（token 前缀/调用方）计数，
|
||
# 窗口内失败次数超阈值则直接拒绝，不再查库（避免枚举放大 DB 压力）。
|
||
_FAIL_WINDOW_SEC = 300
|
||
_FAIL_MAX = 20
|
||
_fail_counters: dict = {}
|
||
|
||
|
||
def _fail_key(token):
|
||
"""限速键：取 token 前 12 字符（足够区分调用方，又不在内存留全量密钥）。"""
|
||
t = (token or '').strip()
|
||
if t.lower().startswith('bearer '):
|
||
t = t[7:].strip()
|
||
return t[:12] or 'anon'
|
||
|
||
|
||
def _rate_limited(token):
|
||
"""是否已被限速。顺带清理过期窗口。"""
|
||
import time
|
||
now = time.time()
|
||
k = _fail_key(token)
|
||
rec = _fail_counters.get(k)
|
||
if not rec:
|
||
return False
|
||
start, cnt = rec
|
||
if now - start > _FAIL_WINDOW_SEC:
|
||
_fail_counters.pop(k, None)
|
||
return False
|
||
return cnt >= _FAIL_MAX
|
||
|
||
|
||
def _record_fail(token):
|
||
"""记一次鉴权失败。"""
|
||
import time
|
||
now = time.time()
|
||
k = _fail_key(token)
|
||
rec = _fail_counters.get(k)
|
||
if not rec or now - rec[0] > _FAIL_WINDOW_SEC:
|
||
_fail_counters[k] = (now, 1)
|
||
else:
|
||
_fail_counters[k] = (rec[0], rec[1] + 1)
|
||
# 防内存膨胀：条目过多时清理已过期窗口
|
||
if len(_fail_counters) > 5000:
|
||
for kk in [x for x, v in _fail_counters.items() if now - v[0] > _FAIL_WINDOW_SEC]:
|
||
_fail_counters.pop(kk, None)
|
||
|
||
|
||
def _get_db():
|
||
db = DBPools()
|
||
if not db.databases:
|
||
from appPublic.jsonConfig import getConfig
|
||
config = getConfig()
|
||
if config.databases:
|
||
db.databases = config.databases
|
||
return db, DBNAME
|
||
|
||
|
||
def _now():
|
||
return datetime.now()
|
||
|
||
|
||
async def create_llm_token(org_id, project_id="", task_id="", model_name="",
|
||
purpose="", ttl_hours=DEFAULT_TTL_HOURS,
|
||
max_calls=DEFAULT_MAX_CALLS, created_by=""):
|
||
"""签发短期 LLM 代理 token。返回 (True, token) 或 (False, 错误信息)。
|
||
|
||
org_id 决定机构隔离边界：代理调用时按此 org_id 查 llm 表解析真 key，
|
||
运行环境无法越权使用其他机构的模型。
|
||
"""
|
||
if not org_id:
|
||
return False, "缺少 org_id（机构隔离必需）"
|
||
try:
|
||
ttl = int(ttl_hours)
|
||
except (TypeError, ValueError):
|
||
ttl = DEFAULT_TTL_HOURS
|
||
if ttl <= 0 or ttl > 72:
|
||
ttl = DEFAULT_TTL_HOURS
|
||
try:
|
||
mc = int(max_calls)
|
||
except (TypeError, ValueError):
|
||
mc = DEFAULT_MAX_CALLS
|
||
if mc <= 0:
|
||
mc = DEFAULT_MAX_CALLS
|
||
|
||
token = TOKEN_PREFIX + secrets.token_urlsafe(32)
|
||
expires_at = _now() + timedelta(hours=ttl)
|
||
db, dbname = _get_db()
|
||
async with db.sqlorContext(dbname) as sor:
|
||
await sor.C('pipeline_llm_tokens', {
|
||
'id': getID(),
|
||
'token': token,
|
||
'org_id': org_id,
|
||
'project_id': project_id or '',
|
||
'task_id': task_id or '',
|
||
'model_name': model_name or '',
|
||
'purpose': purpose or '',
|
||
'status': S_ACTIVE,
|
||
'expires_at': expires_at.strftime('%Y-%m-%d %H:%M:%S'),
|
||
'max_calls': mc,
|
||
'call_count': 0,
|
||
'prompt_tokens': 0,
|
||
'completion_tokens': 0,
|
||
'created_by': created_by or '',
|
||
})
|
||
await sor.sqlExe("COMMIT", {})
|
||
logger.info("create_llm_token: org=%s project=%s task=%s purpose=%s ttl=%dh max_calls=%d",
|
||
org_id, project_id, task_id, purpose, ttl, mc)
|
||
return True, token
|
||
|
||
|
||
async def verify_llm_token(token):
|
||
"""校验 token。返回 (True, 绑定信息 dict) 或 (False, 错误信息)。
|
||
|
||
校验项：存在、status=active、未过期、未超调用上限。
|
||
"""
|
||
if not token:
|
||
return False, "缺少 token"
|
||
token = token.strip()
|
||
if token.lower().startswith("bearer "):
|
||
token = token[7:].strip()
|
||
db, dbname = _get_db()
|
||
async with db.sqlorContext(dbname) as sor:
|
||
recs = await sor.sqlExe(
|
||
"SELECT id, org_id, project_id, task_id, model_name, status, expires_at, "
|
||
"max_calls, call_count FROM pipeline_llm_tokens WHERE token=${t}$",
|
||
{"t": token})
|
||
await sor.sqlExe("COMMIT", {})
|
||
if not recs:
|
||
return False, "token 无效"
|
||
r = recs[0]
|
||
status = getattr(r, 'status', '') or ''
|
||
if status != S_ACTIVE:
|
||
return False, f"token 已{('吊销' if status == S_REVOKED else status)}"
|
||
exp = getattr(r, 'expires_at', None)
|
||
if exp:
|
||
if isinstance(exp, str):
|
||
try:
|
||
exp = datetime.strptime(exp[:19], '%Y-%m-%d %H:%M:%S')
|
||
except ValueError:
|
||
exp = None
|
||
if exp and _now() > exp:
|
||
return False, "token 已过期"
|
||
mc = int(getattr(r, 'max_calls', 0) or 0)
|
||
cc = int(getattr(r, 'call_count', 0) or 0)
|
||
if mc and cc >= mc:
|
||
return False, f"token 调用次数已达上限（{mc}）"
|
||
return True, {
|
||
"id": getattr(r, 'id', ''),
|
||
"org_id": getattr(r, 'org_id', '') or '',
|
||
"project_id": getattr(r, 'project_id', '') or '',
|
||
"task_id": getattr(r, 'task_id', '') or '',
|
||
"model_name": getattr(r, 'model_name', '') or '',
|
||
}
|
||
|
||
|
||
async def revoke_llm_token(token=None, token_id=None):
|
||
"""吊销 token（按 token 明文或 id）。返回 (True, 消息) 或 (False, 错误)。"""
|
||
if not token and not token_id:
|
||
return False, "须指定 token 或 token_id"
|
||
db, dbname = _get_db()
|
||
async with db.sqlorContext(dbname) as sor:
|
||
if token_id:
|
||
await sor.sqlExe(
|
||
"UPDATE pipeline_llm_tokens SET status=${s}$ WHERE id=${i}$",
|
||
{"s": S_REVOKED, "i": token_id})
|
||
else:
|
||
t = token.strip()
|
||
if t.lower().startswith("bearer "):
|
||
t = t[7:].strip()
|
||
await sor.sqlExe(
|
||
"UPDATE pipeline_llm_tokens SET status=${s}$ WHERE token=${t}$",
|
||
{"s": S_REVOKED, "t": t})
|
||
await sor.sqlExe("COMMIT", {})
|
||
return True, "已吊销"
|
||
|
||
|
||
async def revoke_task_tokens(task_id):
|
||
"""任务结束时吊销该任务签发的全部 token（防止 token 生命周期超出任务）。"""
|
||
if not task_id:
|
||
return False, "缺少 task_id"
|
||
db, dbname = _get_db()
|
||
async with db.sqlorContext(dbname) as sor:
|
||
await sor.sqlExe(
|
||
"UPDATE pipeline_llm_tokens SET status=${s}$ "
|
||
"WHERE task_id=${t}$ AND status=${a}$",
|
||
{"s": S_REVOKED, "t": task_id, "a": S_ACTIVE})
|
||
await sor.sqlExe("COMMIT", {})
|
||
return True, "已吊销该任务全部 token"
|
||
|
||
|
||
async def _record_usage(token_id, usage):
|
||
"""回填用量（调用次数 + token 数）。usage 为上游返回的 usage dict。"""
|
||
if not token_id:
|
||
return
|
||
pt = 0
|
||
ct = 0
|
||
if isinstance(usage, dict):
|
||
try:
|
||
pt = int(usage.get('prompt_tokens') or 0)
|
||
ct = int(usage.get('completion_tokens') or 0)
|
||
except (TypeError, ValueError):
|
||
pt, ct = 0, 0
|
||
db, dbname = _get_db()
|
||
try:
|
||
async with db.sqlorContext(dbname) as sor:
|
||
await sor.sqlExe(
|
||
"UPDATE pipeline_llm_tokens SET call_count=call_count+1, "
|
||
"prompt_tokens=prompt_tokens+${pt}$, completion_tokens=completion_tokens+${ct}$, "
|
||
"last_used_at=NOW() WHERE id=${i}$",
|
||
{"pt": pt, "ct": ct, "i": token_id})
|
||
await sor.sqlExe("COMMIT", {})
|
||
except Exception as e:
|
||
logger.warning("_record_usage failed: %s", e)
|
||
|
||
|
||
async def proxy_chat_completion(token, payload):
|
||
"""OpenAI 兼容代理转发。
|
||
|
||
token 运行环境持有的短期 token（当 api_key 用）
|
||
payload 客户端原始请求体（含 model/messages/tools/temperature 等）
|
||
|
||
返回 (True, 上游响应 dict) 或 (False, 错误信息)。
|
||
真 api_key 由 llm_bridge 按 token 绑定的 org_id 从 llm 表解析，不下发给调用方。
|
||
"""
|
||
# 失败限速前置：窗口内鉴权失败过多直接拒绝（防 token 枚举）
|
||
if _rate_limited(token):
|
||
logger.warning("proxy_chat_completion: 限速拒绝 key=%s", _fail_key(token))
|
||
return False, "鉴权失败次数过多，请稍后再试"
|
||
ok, info = await verify_llm_token(token)
|
||
if not ok:
|
||
_record_fail(token)
|
||
return False, info
|
||
if not isinstance(payload, dict) or not payload.get('messages'):
|
||
return False, "请求体缺 messages"
|
||
|
||
org_id = info['org_id']
|
||
# 模型选择：token 绑定了 model_name 则强制用它（防运行环境越权指定贵模型）；
|
||
# 否则用请求里的 model；都没有则由 llm_bridge 取该机构第一个 active 模型。
|
||
model_name = info.get('model_name') or payload.get('model') or None
|
||
|
||
from .llm_bridge import _get_model_config, _post_chat_completion
|
||
cfg = await _get_model_config(model_name, org_id=org_id)
|
||
if not (cfg.get('api_key') and cfg.get('api_base')):
|
||
return False, f"机构 {org_id} 未配置可用模型（llm 表 status=active）"
|
||
|
||
api_base = cfg['api_base']
|
||
api_key = cfg['api_key'] # 只在本进程内存中使用，不返回给调用方
|
||
model_id = cfg.get('model_id') or model_name or 'default'
|
||
|
||
# 透传客户端参数（messages/tools/temperature/max_tokens 等），但 model 换成真实 model_id，
|
||
# 且不透传 stream（代理暂不支持流式）。
|
||
upstream = {k: v for k, v in payload.items() if k not in ('model', 'stream')}
|
||
upstream['model'] = model_id
|
||
|
||
headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
|
||
url = api_base.rstrip('/') + '/chat/completions'
|
||
try:
|
||
data = await _post_chat_completion(url, headers, upstream)
|
||
except Exception as e:
|
||
logger.warning("proxy_chat_completion upstream failed: %s", e)
|
||
return False, f"上游调用失败: {str(e)[:300]}"
|
||
|
||
await _record_usage(info['id'], data.get('usage'))
|
||
logger.info("proxy_chat_completion: org=%s project=%s model=%s ok",
|
||
org_id, info.get('project_id'), model_id)
|
||
return True, data
|
||
|
||
|
||
async def list_llm_tokens(org_id=None, project_id=None, status=None, limit=50):
|
||
"""列出已签发 token（不返回 token 明文，只返回元信息 + 用量，供审计/管理）。"""
|
||
db, dbname = _get_db()
|
||
async with db.sqlorContext(dbname) as sor:
|
||
conds = ["1=1"]
|
||
params = {}
|
||
if org_id:
|
||
conds.append("org_id=${o}$")
|
||
params["o"] = org_id
|
||
if project_id:
|
||
conds.append("project_id=${p}$")
|
||
params["p"] = project_id
|
||
if status:
|
||
conds.append("status=${s}$")
|
||
params["s"] = status
|
||
where = " AND ".join(conds)
|
||
recs = await sor.sqlExe(
|
||
f"SELECT id, org_id, project_id, task_id, model_name, purpose, status, "
|
||
f"expires_at, max_calls, call_count, prompt_tokens, completion_tokens, "
|
||
f"created_by, created_at, last_used_at FROM pipeline_llm_tokens "
|
||
f"WHERE {where} ORDER BY created_at DESC LIMIT {int(limit)}", params)
|
||
await sor.sqlExe("COMMIT", {})
|
||
out = []
|
||
for r in (recs or []):
|
||
try:
|
||
out.append(dict(r))
|
||
except (TypeError, ValueError):
|
||
pass
|
||
return out
|
||
|
||
|
||
def load_llm_proxy():
|
||
"""注册到 ServerEnv，供 dspy 端点调用。"""
|
||
from ahserver.serverenv import ServerEnv
|
||
env = ServerEnv()
|
||
env.create_llm_token = create_llm_token
|
||
env.verify_llm_token = verify_llm_token
|
||
env.revoke_llm_token = revoke_llm_token
|
||
env.revoke_task_tokens = revoke_task_tokens
|
||
env.proxy_chat_completion = proxy_chat_completion
|
||
env.list_llm_tokens = list_llm_tokens
|