feat(llm-proxy): 内部 LLM 代理 + 短期 token,真 api_key 永不下发运行环境
运行环境(agent run_shell 子进程/bwrap 沙箱/远程目标机/被开发的 AI 应用)需要调 LLM 时, 不下发真实模型 key,改为签发短期 token(绑定 org_id+project_id+task_id,带过期/调用上限/可吊销), 运行环境用 token 当 api_key 调本平台 OpenAI 兼容代理端点;真 key 只在服务进程内存解析。 - create_llm_token / verify_llm_token / revoke_llm_token / revoke_task_tokens - proxy_chat_completion: 按 token 绑定 org_id 解析真 key(机构隔离) + 转发 + 用量回填 - list_llm_tokens: 审计用,不返回 token 明文
This commit is contained in:
parent
b409f18929
commit
106a600c45
@ -536,6 +536,10 @@ def load_pipeline_service():
|
||||
from .workspace import load_workspace
|
||||
load_workspace()
|
||||
|
||||
# 内部 LLM 代理:短期 token 鉴权,真 api_key 永不下发到运行环境
|
||||
from .llm_proxy import load_llm_proxy
|
||||
load_llm_proxy()
|
||||
|
||||
# Register intent classifier + LLM bridge (shared across all pipelines)
|
||||
from .intent_classifier import intent_classify
|
||||
from .llm_bridge import llm_call
|
||||
|
||||
296
pipeline_service/llm_proxy.py
Normal file
296
pipeline_service/llm_proxy.py
Normal file
@ -0,0 +1,296 @@
|
||||
"""内部 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
|
||||
|
||||
|
||||
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 表解析,不下发给调用方。
|
||||
"""
|
||||
ok, info = await verify_llm_token(token)
|
||||
if not ok:
|
||||
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
|
||||
Loading…
x
Reference in New Issue
Block a user