diff --git a/pipeline_service/init.py b/pipeline_service/init.py index 6bff45c..21f717a 100644 --- a/pipeline_service/init.py +++ b/pipeline_service/init.py @@ -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 diff --git a/pipeline_service/llm_proxy.py b/pipeline_service/llm_proxy.py new file mode 100644 index 0000000..877691b --- /dev/null +++ b/pipeline_service/llm_proxy.py @@ -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:///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