chore(align): sync Git with running container as authoritative source
One-time alignment to close drift between the running container (qihu-aiflow / image qihuaiflow:with-deps) and this Git repository. Container hotfixes previously persisted only via `docker commit` are now captured in Git so Git becomes the single source of truth. Backend: - chat_service.py: accumulated fixes incl. EigenFlux README guidance - chat_context.py / chat_media.py: URL/web-repo context injection - api/url_reader.py: NEW SSRF-safe web/README reader (Issue #725) - skills_admin.py & Skills.tsx: 14-scene capability unification - pgvector_search.py, orchestrator.py, platform_knowledge.py, market.py, admin/models.py: drift sync Frontend: - ChatPage.tsx: collapsible AI think/act block (Issue #724) - MarketPage / EvolutionDashboard / TokenDashboard styling sync Migrations: - 018_normalize_marketapp_scene.sql: fixed capability allowlist (prevents capability wipe on restart) Alignment performed 2026-07-23 05:42:47. No deletions; secrets/venv/node_modules/ caches were not touched (tracked-source-only copy).
This commit is contained in:
1 parent
49f7503978
commit
a08ec87b22
19 files changed
+7474
-7196
No files matched your search
@@ -932,7 +932,13 @@ async def test_model_connection(
|
||||
"n": 1,
|
||||
"size": "1024x1024",
|
||||
}
|
||||
resp = await client.post(image_url, json=payload, headers=headers)
|
||||
# 与运行时媒体生成一致:用 _proxy_for_url 解析代理并 trust_env=False;
|
||||
# 图像端点会被 _proxy_for_url 判定为直连(公网 API 直连比代理更稳定)
|
||||
from ..chat_utils import _proxy_for_url
|
||||
async with httpx.AsyncClient(
|
||||
timeout=timeout, proxy=_proxy_for_url(image_url), trust_env=False
|
||||
) as img_client:
|
||||
resp = await img_client.post(image_url, json=payload, headers=headers)
|
||||
latency = int((time.time() - start) * 1000)
|
||||
|
||||
if resp.status_code == 200:
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -1121,7 +1121,7 @@ def _build_system_prompt(
|
||||
|
||||
all_mcp: list = None,
|
||||
|
||||
|
||||
db: Session = None,
|
||||
|
||||
) -> tuple:
|
||||
|
||||
@@ -1195,7 +1195,7 @@ def _build_system_prompt(
|
||||
|
||||
|
||||
|
||||
pk = get_platform_knowledge()
|
||||
pk = get_platform_knowledge(db=db)
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1566,11 +1566,16 @@ async def _generate_media_task(ai_msg_id: str, ai_model_id: str, prompt: str, us
|
||||
|
||||
|
||||
except Exception as e:
|
||||
|
||||
|
||||
|
||||
logger.warning(f"[_generate_media_task] 后台生成异常: {e}")
|
||||
|
||||
logger.warning(f"[_generate_media_task] 后台生成异常: {e}", exc_info=True)
|
||||
# 兜底回写:任何意外异常都不能让消息永久停在「⏳ 生成中…」
|
||||
try:
|
||||
db.rollback()
|
||||
_err = db.query(Message).filter(Message.id == ai_msg_id).first()
|
||||
if _err and str(_err.content or "").startswith("⏳"):
|
||||
_err.content = "⚠️ 媒体生成失败,请稍后重试"
|
||||
db.commit()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
finally:
|
||||
|
||||
+4304
-4273
File diff suppressed because it is too large.
Load diff
@@ -1,387 +1,387 @@
|
||||
# Auto-extracted from chat.py during Phase 1 refactor (Strangler Fig).
|
||||
# Source of truth: scripts/phase1_extract.py. Do not hand-merge back.
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Response
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy import and_, desc, asc, func
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Optional
|
||||
import uuid
|
||||
from datetime import datetime, timezone, timedelta
|
||||
import re
|
||||
import json
|
||||
import asyncio
|
||||
import math
|
||||
import httpx
|
||||
import os
|
||||
import logging
|
||||
from ..database import get_db, SessionLocal
|
||||
from ..models import (
|
||||
|
||||
|
||||
|
||||
User, Conversation, Message, AIModel, ModelProvider, TokenUsageLog,
|
||||
|
||||
|
||||
|
||||
KnowledgeBase, KnowledgeDocument, KnowledgeDocumentChunk, EvolutionSkill,
|
||||
|
||||
|
||||
|
||||
EvolutionTrajectory, MarketApp, AppPurchase, AppInstallation,
|
||||
|
||||
|
||||
|
||||
MCPServer,
|
||||
|
||||
|
||||
|
||||
gen_uuid
|
||||
|
||||
|
||||
|
||||
)
|
||||
from ..auth import get_current_user
|
||||
from ..platform_knowledge import get_platform_knowledge
|
||||
from ..mcp_client import list_mcp_tools, call_mcp_tool
|
||||
from ..middleware.log_service import LogService
|
||||
from ..capability_registry import resolve_auto_model
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CST_TZ = timezone(timedelta(hours=8))
|
||||
|
||||
def cst_now():
|
||||
|
||||
|
||||
|
||||
"""Get current time in CST timezone"""
|
||||
|
||||
|
||||
|
||||
return datetime.now(CST_TZ)
|
||||
|
||||
def _html_escape(s: str) -> str:
|
||||
|
||||
|
||||
|
||||
return (str(s).replace("&", "&").replace("<", "<").replace(">", ">")
|
||||
|
||||
|
||||
|
||||
.replace('"', """).replace("'", "'"))
|
||||
|
||||
def _outbound_proxy() -> str:
|
||||
|
||||
|
||||
|
||||
"""出网代理地址(兜底)。默认经 172.18.0.1:7890 智能路由代理;
|
||||
|
||||
|
||||
|
||||
运行时若未注入代理环境变量则回退到该地址。"""
|
||||
|
||||
|
||||
|
||||
import os
|
||||
|
||||
|
||||
|
||||
return os.environ.get("HTTPS_PROXY") or os.environ.get("HTTP_PROXY") or "http://172.18.0.1:7890"
|
||||
|
||||
_HK_PROXY = "http://qyyfw:XHR521yhf@103.242.14.208:11495"
|
||||
|
||||
def _proxy_for_url(url: str) -> str | None:
|
||||
|
||||
|
||||
|
||||
"""按目标域名选择是否走代理。
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
2026-07-15 调整:AGNES(apihub.agnes-ai.com / platform-outputs.agnes-ai.space)
|
||||
|
||||
|
||||
|
||||
经用户实测走香港代理 103.242.14.208:11495 延迟更低、更稳定,故改为走 HK 代理
|
||||
|
||||
|
||||
|
||||
(带鉴权),不再直连,也不经智能路由代理(172.18.0.1:7890 对 AGNES 返 502)。
|
||||
|
||||
|
||||
|
||||
其余国内/平台域名直连,真正的境外域名走 _outbound_proxy()。
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
注意:httpx 默认 trust_env=True 会忽略 proxy 参数而改用环境变量 HTTP_PROXY,
|
||||
|
||||
|
||||
|
||||
故调用方必须同时传 trust_env=False 才能真正走此处指定的代理。
|
||||
|
||||
|
||||
|
||||
"""
|
||||
|
||||
|
||||
|
||||
import os
|
||||
|
||||
|
||||
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
||||
|
||||
host = (urlparse(url).hostname or "").lower()
|
||||
|
||||
|
||||
|
||||
# 文生图 / 视频生成走直连:智能路由代理(172.18.0.1:7890)与 HK 代理对这类长耗时、
|
||||
# 大响应体的生成连接会偶发重置(Server disconnected),而各图像/视频供应商均为公网
|
||||
# API,直连更稳定(已实测 gpt-image-2 / agnes 直连均返回 200)。故媒体端点统一直连。
|
||||
_lower = (url or "").lower()
|
||||
if "/images/" in _lower or "/videos/" in _lower or "/v1/images" in _lower or "/v1/videos" in _lower:
|
||||
return None
|
||||
|
||||
|
||||
|
||||
if "agnes-ai.com" in host or "agnes-ai.space" in host:
|
||||
|
||||
|
||||
|
||||
return _HK_PROXY
|
||||
|
||||
|
||||
|
||||
# 平台自有域名(含自建 Dify/n8n 等)必须直连,不经出网代理。
|
||||
|
||||
|
||||
|
||||
if "qihuyun.com" in host:
|
||||
|
||||
|
||||
|
||||
return None
|
||||
|
||||
|
||||
|
||||
# 尊重 NO_PROXY:国内域名(硅基流动/DeepSeek/小米MoE 等)直连更稳定,
|
||||
|
||||
|
||||
|
||||
# 避免经易波动的智能路由代理 172.18.0.1:7890(对国内域名偶发连接重置)。
|
||||
|
||||
|
||||
|
||||
_no_proxy = [h.strip().lower() for h in os.environ.get("NO_PROXY", "").split(",") if h.strip()]
|
||||
|
||||
|
||||
|
||||
for _np in _no_proxy:
|
||||
|
||||
|
||||
|
||||
if host == _np or host.endswith("." + _np) or _np in host:
|
||||
|
||||
|
||||
|
||||
return None
|
||||
|
||||
|
||||
|
||||
return _outbound_proxy()
|
||||
|
||||
def _extract_text_from_json(obj) -> str:
|
||||
|
||||
|
||||
|
||||
"""从 n8n / 通用 webhook 的任意 JSON 响应中尽力抽取文本回答。"""
|
||||
|
||||
|
||||
|
||||
if obj is None:
|
||||
|
||||
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
|
||||
if isinstance(obj, str):
|
||||
|
||||
|
||||
|
||||
return obj
|
||||
|
||||
|
||||
|
||||
if isinstance(obj, (int, float, bool)):
|
||||
|
||||
|
||||
|
||||
return str(obj)
|
||||
|
||||
|
||||
|
||||
if isinstance(obj, list):
|
||||
|
||||
|
||||
|
||||
# 取第一个可抽取到文本的元素
|
||||
|
||||
|
||||
|
||||
for it in obj:
|
||||
|
||||
|
||||
|
||||
t = _extract_text_from_json(it)
|
||||
|
||||
|
||||
|
||||
if t:
|
||||
|
||||
|
||||
|
||||
return t
|
||||
|
||||
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
|
||||
if isinstance(obj, dict):
|
||||
|
||||
|
||||
|
||||
# 常见回答字段优先级
|
||||
|
||||
|
||||
|
||||
for k in ("output", "text", "answer", "response", "reply",
|
||||
|
||||
|
||||
|
||||
"message", "content", "result", "data"):
|
||||
|
||||
|
||||
|
||||
if k in obj:
|
||||
|
||||
|
||||
|
||||
t = _extract_text_from_json(obj[k])
|
||||
|
||||
|
||||
|
||||
if t:
|
||||
|
||||
|
||||
|
||||
return t
|
||||
|
||||
|
||||
|
||||
# 兜底:整体序列化
|
||||
|
||||
|
||||
|
||||
try:
|
||||
|
||||
|
||||
|
||||
return json.dumps(obj, ensure_ascii=False)
|
||||
|
||||
|
||||
|
||||
except Exception:
|
||||
|
||||
|
||||
|
||||
return str(obj)
|
||||
|
||||
|
||||
|
||||
return str(obj)
|
||||
|
||||
class CreateConversationRequest(BaseModel):
|
||||
|
||||
|
||||
|
||||
title: str = Field(default="新对话")
|
||||
|
||||
|
||||
|
||||
model_id: Optional[str] = None
|
||||
|
||||
class UpdateConversationRequest(BaseModel):
|
||||
|
||||
|
||||
|
||||
title: Optional[str] = None
|
||||
|
||||
|
||||
|
||||
status: Optional[str] = None # active, archived
|
||||
|
||||
|
||||
|
||||
pinned: Optional[bool] = None
|
||||
|
||||
|
||||
|
||||
agent_ids: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
skill_ids: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
knowledge_base_ids: Optional[list] = None
|
||||
|
||||
class SendMessageRequest(BaseModel):
|
||||
|
||||
|
||||
|
||||
content: str
|
||||
|
||||
|
||||
|
||||
model_id: Optional[str] = None
|
||||
|
||||
|
||||
|
||||
skill_ids: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
agent_ids: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
knowledge_base_ids: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
mcp_ids: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
attachments: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
use_tools: Optional[bool] = False
|
||||
# Auto-extracted from chat.py during Phase 1 refactor (Strangler Fig).
|
||||
# Source of truth: scripts/phase1_extract.py. Do not hand-merge back.
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Response
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy import and_, desc, asc, func
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Optional
|
||||
import uuid
|
||||
from datetime import datetime, timezone, timedelta
|
||||
import re
|
||||
import json
|
||||
import asyncio
|
||||
import math
|
||||
import httpx
|
||||
import os
|
||||
import logging
|
||||
from ..database import get_db, SessionLocal
|
||||
from ..models import (
|
||||
|
||||
|
||||
|
||||
User, Conversation, Message, AIModel, ModelProvider, TokenUsageLog,
|
||||
|
||||
|
||||
|
||||
KnowledgeBase, KnowledgeDocument, KnowledgeDocumentChunk, EvolutionSkill,
|
||||
|
||||
|
||||
|
||||
EvolutionTrajectory, MarketApp, AppPurchase, AppInstallation,
|
||||
|
||||
|
||||
|
||||
MCPServer,
|
||||
|
||||
|
||||
|
||||
gen_uuid
|
||||
|
||||
|
||||
|
||||
)
|
||||
from ..auth import get_current_user
|
||||
from ..platform_knowledge import get_platform_knowledge
|
||||
from ..mcp_client import list_mcp_tools, call_mcp_tool
|
||||
from ..middleware.log_service import LogService
|
||||
from ..capability_registry import resolve_auto_model
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CST_TZ = timezone(timedelta(hours=8))
|
||||
|
||||
def cst_now():
|
||||
|
||||
|
||||
|
||||
"""Get current time in CST timezone"""
|
||||
|
||||
|
||||
|
||||
return datetime.now(CST_TZ)
|
||||
|
||||
def _html_escape(s: str) -> str:
|
||||
|
||||
|
||||
|
||||
return (str(s).replace("&", "&").replace("<", "<").replace(">", ">")
|
||||
|
||||
|
||||
|
||||
.replace('"', """).replace("'", "'"))
|
||||
|
||||
def _outbound_proxy() -> str:
|
||||
|
||||
|
||||
|
||||
"""出网代理地址(兜底)。默认经 172.18.0.1:7890 智能路由代理;
|
||||
|
||||
|
||||
|
||||
运行时若未注入代理环境变量则回退到该地址。"""
|
||||
|
||||
|
||||
|
||||
import os
|
||||
|
||||
|
||||
|
||||
return os.environ.get("HTTPS_PROXY") or os.environ.get("HTTP_PROXY") or "http://172.18.0.1:7890"
|
||||
|
||||
_HK_PROXY = "http://qyyfw:XHR521yhf@103.242.14.208:11495"
|
||||
|
||||
def _proxy_for_url(url: str) -> str | None:
|
||||
|
||||
|
||||
|
||||
"""按目标域名选择是否走代理。
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
2026-07-15 调整:AGNES(apihub.agnes-ai.com / platform-outputs.agnes-ai.space)
|
||||
|
||||
|
||||
|
||||
经用户实测走香港代理 103.242.14.208:11495 延迟更低、更稳定,故改为走 HK 代理
|
||||
|
||||
|
||||
|
||||
(带鉴权),不再直连,也不经智能路由代理(172.18.0.1:7890 对 AGNES 返 502)。
|
||||
|
||||
|
||||
|
||||
其余国内/平台域名直连,真正的境外域名走 _outbound_proxy()。
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
注意:httpx 默认 trust_env=True 会忽略 proxy 参数而改用环境变量 HTTP_PROXY,
|
||||
|
||||
|
||||
|
||||
故调用方必须同时传 trust_env=False 才能真正走此处指定的代理。
|
||||
|
||||
|
||||
|
||||
"""
|
||||
|
||||
|
||||
|
||||
import os
|
||||
|
||||
|
||||
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
||||
|
||||
host = (urlparse(url).hostname or "").lower()
|
||||
|
||||
|
||||
|
||||
# 文生图 / 视频生成走直连:智能路由代理(172.18.0.1:7890)与 HK 代理对这类长耗时、
|
||||
# 大响应体的生成连接会偶发重置(Server disconnected),而各图像/视频供应商均为公网
|
||||
# API,直连更稳定(已实测 gpt-image-2 / agnes 直连均返回 200)。故媒体端点统一直连。
|
||||
_lower = (url or "").lower()
|
||||
if "/images/" in _lower or "/videos/" in _lower or "/v1/images" in _lower or "/v1/videos" in _lower:
|
||||
return None
|
||||
|
||||
|
||||
|
||||
if "agnes-ai.com" in host or "agnes-ai.space" in host:
|
||||
|
||||
|
||||
|
||||
return _HK_PROXY
|
||||
|
||||
|
||||
|
||||
# 平台自有域名(含自建 Dify/n8n 等)必须直连,不经出网代理。
|
||||
|
||||
|
||||
|
||||
if "qihuyun.com" in host:
|
||||
|
||||
|
||||
|
||||
return None
|
||||
|
||||
|
||||
|
||||
# 尊重 NO_PROXY:国内域名(硅基流动/DeepSeek/小米MoE 等)直连更稳定,
|
||||
|
||||
|
||||
|
||||
# 避免经易波动的智能路由代理 172.18.0.1:7890(对国内域名偶发连接重置)。
|
||||
|
||||
|
||||
|
||||
_no_proxy = [h.strip().lower() for h in os.environ.get("NO_PROXY", "").split(",") if h.strip()]
|
||||
|
||||
|
||||
|
||||
for _np in _no_proxy:
|
||||
|
||||
|
||||
|
||||
if host == _np or host.endswith("." + _np) or _np in host:
|
||||
|
||||
|
||||
|
||||
return None
|
||||
|
||||
|
||||
|
||||
return _outbound_proxy()
|
||||
|
||||
def _extract_text_from_json(obj) -> str:
|
||||
|
||||
|
||||
|
||||
"""从 n8n / 通用 webhook 的任意 JSON 响应中尽力抽取文本回答。"""
|
||||
|
||||
|
||||
|
||||
if obj is None:
|
||||
|
||||
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
|
||||
if isinstance(obj, str):
|
||||
|
||||
|
||||
|
||||
return obj
|
||||
|
||||
|
||||
|
||||
if isinstance(obj, (int, float, bool)):
|
||||
|
||||
|
||||
|
||||
return str(obj)
|
||||
|
||||
|
||||
|
||||
if isinstance(obj, list):
|
||||
|
||||
|
||||
|
||||
# 取第一个可抽取到文本的元素
|
||||
|
||||
|
||||
|
||||
for it in obj:
|
||||
|
||||
|
||||
|
||||
t = _extract_text_from_json(it)
|
||||
|
||||
|
||||
|
||||
if t:
|
||||
|
||||
|
||||
|
||||
return t
|
||||
|
||||
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
|
||||
if isinstance(obj, dict):
|
||||
|
||||
|
||||
|
||||
# 常见回答字段优先级
|
||||
|
||||
|
||||
|
||||
for k in ("output", "text", "answer", "response", "reply",
|
||||
|
||||
|
||||
|
||||
"message", "content", "result", "data"):
|
||||
|
||||
|
||||
|
||||
if k in obj:
|
||||
|
||||
|
||||
|
||||
t = _extract_text_from_json(obj[k])
|
||||
|
||||
|
||||
|
||||
if t:
|
||||
|
||||
|
||||
|
||||
return t
|
||||
|
||||
|
||||
|
||||
# 兜底:整体序列化
|
||||
|
||||
|
||||
|
||||
try:
|
||||
|
||||
|
||||
|
||||
return json.dumps(obj, ensure_ascii=False)
|
||||
|
||||
|
||||
|
||||
except Exception:
|
||||
|
||||
|
||||
|
||||
return str(obj)
|
||||
|
||||
|
||||
|
||||
return str(obj)
|
||||
|
||||
class CreateConversationRequest(BaseModel):
|
||||
|
||||
|
||||
|
||||
title: str = Field(default="新对话")
|
||||
|
||||
|
||||
|
||||
model_id: Optional[str] = None
|
||||
|
||||
class UpdateConversationRequest(BaseModel):
|
||||
|
||||
|
||||
|
||||
title: Optional[str] = None
|
||||
|
||||
|
||||
|
||||
status: Optional[str] = None # active, archived
|
||||
|
||||
|
||||
|
||||
pinned: Optional[bool] = None
|
||||
|
||||
|
||||
|
||||
agent_ids: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
skill_ids: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
knowledge_base_ids: Optional[list] = None
|
||||
|
||||
class SendMessageRequest(BaseModel):
|
||||
|
||||
|
||||
|
||||
content: str
|
||||
|
||||
|
||||
|
||||
model_id: Optional[str] = None
|
||||
|
||||
|
||||
|
||||
skill_ids: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
agent_ids: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
knowledge_base_ids: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
mcp_ids: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
attachments: Optional[list] = None
|
||||
|
||||
|
||||
|
||||
use_tools: Optional[bool] = False
|
||||
@@ -691,7 +691,12 @@ def list_apps(
|
||||
"capability": app.capability or None,
|
||||
"category": app.category or None,
|
||||
"tags": app.tags or None,
|
||||
"installed": (app.id in installed_app_ids or is_user_owned) if user_id else False,
|
||||
# 官方导入并发布的技能默认全局已安装(见 #716:无论前台/后台导入,默认全局安装)
|
||||
"installed": (
|
||||
app.id in installed_app_ids
|
||||
or is_user_owned
|
||||
or (app.source == "official" and app.status == "published")
|
||||
),
|
||||
"version": app.version or "1.0.0",
|
||||
"security_scan": app.security_scan,
|
||||
"verified": bool(app.verified),
|
||||
@@ -756,7 +761,8 @@ def get_app(
|
||||
"tags": app.tags or None,
|
||||
"created_at": str(app.created_at),
|
||||
"purchased": purchased,
|
||||
"installed": installed,
|
||||
# 官方导入并发布的技能默认全局已安装(见 #716)
|
||||
"installed": installed or (app.source == "official" and app.status == "published"),
|
||||
"version": app.version or "1.0.0",
|
||||
"security_scan": app.security_scan,
|
||||
"verified": bool(app.verified),
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
"""服务端 URL 读取器(Issue 2 修复)
|
||||
|
||||
当用户输入包含 URL 且明确要求『读取 / 加入 / 安装 / 分析 / 总结』等意图时,
|
||||
服务端抓取该链接的真实内容并作为参考上下文注入对话,使模型能基于页面真实信息
|
||||
作答(例如按 GitHub 仓库 README 的 Quick Start 引导用户完成安装/加入),
|
||||
而不是凭空编造。
|
||||
|
||||
SSRF 防护:仅允许 http/https;解析域名 IP 后阻断内网/回环/链路本地/云元数据地址;
|
||||
限制跳转、响应体大小与超时。
|
||||
"""
|
||||
import ipaddress
|
||||
import logging
|
||||
import re
|
||||
import socket
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_URL_RE = re.compile(r'https?://[^\s"\)\\\]>]+', re.IGNORECASE)
|
||||
# 触发意图:显式要求读取/加入/安装/分析/总结/打开/访问 等
|
||||
_INTENT_RE = re.compile(
|
||||
r'(read|join|install|fetch|summar|analys|open|visit|access|browse|'
|
||||
r'读取|加入|安装|打开|访问|浏览|总结|分析|抓取|看下这个|查看该|这个网页|这个页面|这个链接|help me join)',
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
# 常见私有/保留网段(阻断内网扫描与云元数据)
|
||||
_BLOCKED_NETS = [
|
||||
ipaddress.ip_network('10.0.0.0/8'),
|
||||
ipaddress.ip_network('172.16.0.0/12'),
|
||||
ipaddress.ip_network('192.168.0.0/16'),
|
||||
ipaddress.ip_network('127.0.0.0/8'),
|
||||
ipaddress.ip_network('169.254.0.0/16'),
|
||||
ipaddress.ip_network('::1/128'),
|
||||
ipaddress.ip_network('fc00::/7'),
|
||||
ipaddress.ip_network('fe80::/10'),
|
||||
ipaddress.ip_network('0.0.0.0/8'),
|
||||
]
|
||||
|
||||
_MAX_CHARS = 16000
|
||||
_TIMEOUT = 30.0
|
||||
|
||||
|
||||
def _is_blocked_host(host: str) -> bool:
|
||||
host = (host or '').strip().lower()
|
||||
if not host:
|
||||
return True
|
||||
if re.fullmatch(r'(\d{1,3}\.){3}\d{1,3}', host):
|
||||
try:
|
||||
ip = ipaddress.ip_address(host)
|
||||
except ValueError:
|
||||
return True
|
||||
return any(ip in net for net in _BLOCKED_NETS)
|
||||
try:
|
||||
infos = socket.getaddrinfo(host, None)
|
||||
for info in infos:
|
||||
ipstr = info[4][0]
|
||||
try:
|
||||
ip = ipaddress.ip_address(ipstr)
|
||||
except ValueError:
|
||||
continue
|
||||
if any(ip in net for net in _BLOCKED_NETS):
|
||||
return True
|
||||
except Exception:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
async def _make_client():
|
||||
try:
|
||||
return httpx.AsyncClient(trust_env=True, follow_redirects=True, timeout=_TIMEOUT)
|
||||
except Exception:
|
||||
return httpx.AsyncClient(trust_env=False, follow_redirects=True, timeout=_TIMEOUT)
|
||||
|
||||
|
||||
async def _github_text(url: str):
|
||||
p = urlparse(url)
|
||||
if p.netloc.lower() not in ('github.com', 'www.github.com'):
|
||||
return None
|
||||
parts = [x for x in p.path.split('/') if x]
|
||||
if len(parts) < 2:
|
||||
return None
|
||||
owner, repo = parts[0], parts[1]
|
||||
if len(parts) >= 4 and parts[2] == 'blob':
|
||||
ref = parts[3]
|
||||
fpath = '/'.join(parts[4:])
|
||||
raw = f'https://raw.githubusercontent.com/{owner}/{repo}/{ref}/{fpath}'
|
||||
return await _fetch_text(raw)
|
||||
api = f'https://api.github.com/repos/{owner}/{repo}/readme'
|
||||
try:
|
||||
async with await _make_client() as c:
|
||||
r = await c.get(api, headers={
|
||||
'Accept': 'application/vnd.github+json',
|
||||
'User-Agent': 'Mozilla/5.0',
|
||||
})
|
||||
if r.status_code == 200:
|
||||
import base64
|
||||
d = r.json()
|
||||
return base64.b64decode(d.get('content', '')).decode('utf-8', 'replace')
|
||||
except Exception as e:
|
||||
logger.warning('github readme fetch failed: %s', e)
|
||||
return None
|
||||
|
||||
|
||||
async def _fetch_text(url: str) -> str:
|
||||
async with await _make_client() as c:
|
||||
r = await c.get(url, headers={'User-Agent': 'Mozilla/5.0'})
|
||||
r.raise_for_status()
|
||||
data = r.content[:2_000_000]
|
||||
ctype = r.headers.get('content-type', '')
|
||||
text = data.decode('utf-8', 'replace')
|
||||
if 'html' in ctype:
|
||||
return _html_to_text(text)
|
||||
return text
|
||||
|
||||
|
||||
def _html_to_text(html: str) -> str:
|
||||
html = re.sub(r'(?is)<(script|style|head|noscript|svg|template)[^>]*>.*?</\1>', ' ', html)
|
||||
html = re.sub(r'(?is)<!--.*?-->', ' ', html)
|
||||
html = re.sub(r'(?is)<[^>]+>', ' ', html)
|
||||
import html as _html
|
||||
text = _html.unescape(html)
|
||||
text = re.sub(r'[ \t]+', ' ', text)
|
||||
text = re.sub(r'\n{3,}', '\n\n', text)
|
||||
return text.strip()
|
||||
|
||||
|
||||
async def fetch_url_context(text: str):
|
||||
"""若有 URL 且命中读取意图,返回 {'url','text'};否则返回 None。"""
|
||||
if not text:
|
||||
return None
|
||||
m = _URL_RE.search(text)
|
||||
if not m:
|
||||
return None
|
||||
url = m.group(0).rstrip('.,);]')
|
||||
if not _INTENT_RE.search(text):
|
||||
return None
|
||||
try:
|
||||
p = urlparse(url)
|
||||
if p.scheme not in ('http', 'https'):
|
||||
return None
|
||||
if _is_blocked_host(p.hostname):
|
||||
logger.warning('blocked host for url fetch: %s', p.hostname)
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
try:
|
||||
if 'github.com' in (p.hostname or '').lower():
|
||||
body = await _github_text(url)
|
||||
if not body:
|
||||
body = await _fetch_text(url)
|
||||
else:
|
||||
body = await _fetch_text(url)
|
||||
except Exception as e:
|
||||
logger.warning('url fetch failed: %s', e)
|
||||
return None
|
||||
|
||||
if not body or not body.strip():
|
||||
return None
|
||||
return {'url': url, 'text': body[:_MAX_CHARS]}
|
||||
@@ -1,395 +1,395 @@
|
||||
"""app/core/orchestrator.py — Phase 3.3 多 Agent 编排图引擎。
|
||||
|
||||
设计目标:
|
||||
- 以轻量 DAG 描述多 Agent 协作:节点(llm/local/external/reduce) + 有向边(数据流)。
|
||||
- 拓扑分层后按层并行执行(asyncio.gather),互相独立的节点真正并发。
|
||||
- 失败隔离:每个节点异常不影响其它分支;per-node on_error=halt|continue。
|
||||
- 节点执行复用既有原语(_build_system_prompt / _execute_skill_tool / _call_external_agent /
|
||||
3.1 工具沙箱),完全向后兼容——agent.orchestration 为 null 时 chat_service 走原逻辑。
|
||||
- 并发安全:每个节点用独立 SessionLocal 会话做 DB 读,避免共享请求会话跨协程竞态。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger("aiflow.orchestrator")
|
||||
|
||||
VALID_NODE_TYPES = ("llm", "local", "external", "reduce")
|
||||
|
||||
|
||||
class OrchestrationError(ValueError):
|
||||
"""图结构非法或执行被 halt。"""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 校验 + 拓扑
|
||||
# ---------------------------------------------------------------------------
|
||||
def validate_graph(graph: Any) -> None:
|
||||
if not isinstance(graph, dict):
|
||||
raise OrchestrationError("orchestration 必须为对象")
|
||||
nodes = graph.get("nodes")
|
||||
edges = graph.get("edges")
|
||||
if not isinstance(nodes, list) or not nodes:
|
||||
raise OrchestrationError("nodes 必须为非空数组")
|
||||
if not isinstance(edges, list):
|
||||
raise OrchestrationError("edges 必须为数组")
|
||||
|
||||
ids = set()
|
||||
for i, n in enumerate(nodes):
|
||||
if not isinstance(n, dict):
|
||||
raise OrchestrationError(f"节点 #{i} 必须为对象")
|
||||
nid = n.get("id")
|
||||
if not nid or not isinstance(nid, str):
|
||||
raise OrchestrationError(f"节点 #{i} 缺少字符串 id")
|
||||
if nid in ids:
|
||||
raise OrchestrationError(f"节点 id 重复: {nid}")
|
||||
ids.add(nid)
|
||||
ntype = n.get("type")
|
||||
if ntype not in VALID_NODE_TYPES:
|
||||
raise OrchestrationError(
|
||||
f"节点 {nid} 类型非法: {ntype} (允许 {', '.join(VALID_NODE_TYPES)})"
|
||||
)
|
||||
if ntype in ("local", "external") and not n.get("agent_id"):
|
||||
raise OrchestrationError(f"节点 {nid} 类型 {ntype} 必须指定 agent_id")
|
||||
|
||||
for i, e in enumerate(edges):
|
||||
if not isinstance(e, dict):
|
||||
raise OrchestrationError(f"边 #{i} 必须为对象")
|
||||
f = e.get("from")
|
||||
t = e.get("to")
|
||||
if f not in ids:
|
||||
raise OrchestrationError(f"边 #{i} 的 from={f} 指向不存在的节点")
|
||||
if t not in ids:
|
||||
raise OrchestrationError(f"边 #{i} 的 to={t} 指向不存在的节点")
|
||||
if f == t:
|
||||
raise OrchestrationError(f"边 #{i} 存在自环: {f}")
|
||||
|
||||
|
||||
def topological_layers(graph: dict) -> list:
|
||||
"""Kahn 算法返回分层(每层内节点可并行);含环则抛错。"""
|
||||
nodes = graph["nodes"]
|
||||
edges = graph.get("edges", [])
|
||||
indeg = {n["id"]: 0 for n in nodes}
|
||||
adj = {n["id"]: [] for n in nodes}
|
||||
for e in edges:
|
||||
f, t = e["from"], e["to"]
|
||||
adj[f].append(t)
|
||||
indeg[t] += 1
|
||||
order = {n["id"]: i for i, n in enumerate(nodes)}
|
||||
layers: list = []
|
||||
remaining = set(indeg)
|
||||
while remaining:
|
||||
layer = sorted([nid for nid in remaining if indeg[nid] == 0], key=lambda x: order[x])
|
||||
if not layer:
|
||||
cycle = sorted(remaining, key=lambda x: order[x])
|
||||
raise OrchestrationError(f"检测到环路,无法拓扑排序: {cycle}")
|
||||
layers.append(layer)
|
||||
for nid in layer:
|
||||
remaining.discard(nid)
|
||||
for nxt in adj[nid]:
|
||||
indeg[nxt] -= 1
|
||||
return layers
|
||||
|
||||
|
||||
def _node_by_id(graph: dict, nid: str) -> dict:
|
||||
for n in graph["nodes"]:
|
||||
if n["id"] == nid:
|
||||
return n
|
||||
return {}
|
||||
|
||||
|
||||
def _predecessors(graph: dict, nid: str) -> list:
|
||||
return [e["from"] for e in graph.get("edges", []) if e["to"] == nid]
|
||||
|
||||
|
||||
def _successors(graph: dict, nid: str) -> list:
|
||||
return [e["to"] for e in graph.get("edges", []) if e["from"] == nid]
|
||||
|
||||
|
||||
def _find_output_node(graph: dict) -> str:
|
||||
if graph.get("output_node"):
|
||||
return graph["output_node"]
|
||||
sinks = [n["id"] for n in graph["nodes"] if not _successors(graph, n["id"])]
|
||||
if len(sinks) == 1:
|
||||
return sinks[0]
|
||||
for s in sinks:
|
||||
if _node_by_id(graph, s).get("type") == "reduce":
|
||||
return s
|
||||
if sinks:
|
||||
return sinks[0]
|
||||
return graph["nodes"][-1]["id"]
|
||||
|
||||
|
||||
def _assemble_node_input(query: str, upstream: str, is_entry: bool) -> str:
|
||||
"""入口节点直接吃用户原问题;下游节点综合原问题 + 上游输出。"""
|
||||
if is_entry:
|
||||
return query or ""
|
||||
if upstream:
|
||||
return f"原始用户问题:\n{query}\n\n上游节点输出:\n{upstream}"
|
||||
return query or ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 模型解析 + LLM 调用原语
|
||||
# ---------------------------------------------------------------------------
|
||||
async def _resolve_model_and_endpoint(model_id, db):
|
||||
from app.models import AIModel, ModelProvider
|
||||
from ..capability_registry import resolve_auto_model
|
||||
|
||||
ai_model = None
|
||||
if model_id and model_id != "auto":
|
||||
ai_model = db.query(AIModel).filter(AIModel.id == model_id).first()
|
||||
if not ai_model:
|
||||
ai_model = resolve_auto_model("text", db)
|
||||
if not ai_model:
|
||||
raise OrchestrationError("无法解析到可用模型")
|
||||
provider = db.query(ModelProvider).filter(ModelProvider.id == ai_model.provider_id).first()
|
||||
api_key = ai_model.api_key or (provider.api_key if provider else None)
|
||||
api_url = ai_model.endpoint
|
||||
if not api_url and provider and provider.api_base_url:
|
||||
api_url = f"{provider.api_base_url.rstrip('/')}/v1/chat/completions"
|
||||
if not api_url or not api_key:
|
||||
raise OrchestrationError(f"节点模型 {ai_model.name} 缺少 endpoint 或 api_key")
|
||||
return ai_model, api_url, api_key
|
||||
|
||||
|
||||
async def _llm_complete(system_prompt: str, user_content: str, model_id, db, user_id) -> str:
|
||||
from app.api.chat_utils import _proxy_for_url
|
||||
from app.core.observability import span
|
||||
|
||||
ai_model, api_url, api_key = await _resolve_model_and_endpoint(model_id, db)
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_content},
|
||||
]
|
||||
request_body = {"model": ai_model.name, "messages": messages, "max_tokens": 2048}
|
||||
async with span("orchestrator.llm", model=ai_model.id):
|
||||
async with httpx.AsyncClient(
|
||||
timeout=ai_model.timeout or 120,
|
||||
proxy=_proxy_for_url(api_url),
|
||||
trust_env=False,
|
||||
) as client:
|
||||
resp = await client.post(
|
||||
api_url,
|
||||
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
||||
json=request_body,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
raise OrchestrationError(f"LLM 调用失败 HTTP {resp.status_code}: {resp.text[:200]}")
|
||||
data = resp.json()
|
||||
return data["choices"][0]["message"].get("content") or ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# local 节点:复用单 Agent system prompt + 受限 tool round + 3.1 沙箱
|
||||
# ---------------------------------------------------------------------------
|
||||
async def _run_local_agent_turn(agent: dict, node_input: str, db, user_id, conv_id) -> dict:
|
||||
from app.api.chat_skill import _resolve_market_skills, _execute_skill_tool
|
||||
from app.api.chat_context import _build_system_prompt, _retrieve_kb_context
|
||||
from app.api.chat_service import _resolve_agent_tool_policy, _check_tool_policy
|
||||
from app.api.chat_utils import _proxy_for_url
|
||||
|
||||
skills = _resolve_market_skills(db, agent.get("skill_ids") or [], user_id)
|
||||
kb_context = _retrieve_kb_context(db, agent.get("knowledge_base_ids") or [], node_input)
|
||||
system_prompt, function_definitions = _build_system_prompt(
|
||||
skills=skills,
|
||||
agents=[agent],
|
||||
kb_context=kb_context,
|
||||
user_content=node_input,
|
||||
auto_discovered=False,
|
||||
)
|
||||
ai_model, api_url, api_key = await _resolve_model_and_endpoint(agent.get("model_id"), db)
|
||||
tool_policy = _resolve_agent_tool_policy([agent])
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": node_input},
|
||||
]
|
||||
called_skills: list = []
|
||||
tool_call_count = 0
|
||||
max_tool_rounds = 5
|
||||
last_content = ""
|
||||
async with httpx.AsyncClient(
|
||||
timeout=ai_model.timeout or 120,
|
||||
proxy=_proxy_for_url(api_url),
|
||||
trust_env=False,
|
||||
) as client:
|
||||
for round_num in range(max_tool_rounds):
|
||||
request_body = {"model": ai_model.name, "messages": messages, "max_tokens": 2048}
|
||||
if function_definitions:
|
||||
request_body["tools"] = [
|
||||
{"type": "function", "function": fd} for fd in function_definitions
|
||||
]
|
||||
request_body["tool_choice"] = "auto"
|
||||
resp = await client.post(
|
||||
api_url,
|
||||
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
||||
json=request_body,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
raise OrchestrationError(f"local 节点 LLM 失败 HTTP {resp.status_code}: {resp.text[:200]}")
|
||||
data = resp.json()
|
||||
message = data["choices"][0]["message"]
|
||||
tool_calls = message.get("tool_calls") or []
|
||||
content = message.get("content") or ""
|
||||
last_content = content
|
||||
if not tool_calls or not function_definitions:
|
||||
return {"text": content, "used_tools": called_skills}
|
||||
messages.append({"role": "assistant", "content": content, "tool_calls": tool_calls})
|
||||
round_tool_count = 0
|
||||
for tc in tool_calls:
|
||||
func_name = tc.get("function", {}).get("name", "")
|
||||
check = _check_tool_policy(
|
||||
func_name, None, skills, tool_policy, tool_call_count, round_tool_count
|
||||
)
|
||||
if not check["allowed"]:
|
||||
messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tc.get("id", f"call_{round_num}"),
|
||||
"content": f"[工具调用已被策略拦截] {check.get('message', check.get('reason'))}",
|
||||
})
|
||||
tool_call_count += 1
|
||||
round_tool_count += 1
|
||||
continue
|
||||
tool_call_count += 1
|
||||
round_tool_count += 1
|
||||
try:
|
||||
func_args = json.loads(tc.get("function", {}).get("arguments", "{}") or "{}")
|
||||
except Exception:
|
||||
func_args = {}
|
||||
res = _execute_skill_tool(db, func_name, func_args, skills, node_input)
|
||||
called_skills.append(func_name)
|
||||
messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tc.get("id", f"call_{round_num}"),
|
||||
"content": res,
|
||||
})
|
||||
return {"text": last_content, "used_tools": called_skills}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 单节点执行(含独立 DB 会话,保证并发安全)
|
||||
# ---------------------------------------------------------------------------
|
||||
async def _run_node(node, graph, blackboard, query, user_id, conv_id, trace_id) -> dict:
|
||||
from app.api.chat_skill import _resolve_market_agents
|
||||
from app.database import SessionLocal
|
||||
|
||||
nid = node["id"]
|
||||
ntype = node["type"]
|
||||
preds = _predecessors(graph, nid)
|
||||
pred_results = [blackboard[p] for p in preds]
|
||||
is_entry = not preds
|
||||
|
||||
if pred_results:
|
||||
parts = []
|
||||
for b in pred_results:
|
||||
label = f"节点 {b.get('node_id')}/{b.get('node_type')}"
|
||||
if b.get("ok"):
|
||||
parts.append(f"【{label}】\n{b.get('output', '')}")
|
||||
else:
|
||||
parts.append(f"【{label} 执行失败】\n错误: {b.get('error', '')}")
|
||||
upstream = "\n\n".join(parts)
|
||||
else:
|
||||
upstream = ""
|
||||
|
||||
with SessionLocal() as db:
|
||||
if ntype in ("llm", "reduce"):
|
||||
if ntype == "llm":
|
||||
sys_prompt = node.get("system_prompt") or (
|
||||
"你是一个处理节点,请根据上游输入完成指定的文本变换、摘要或路由判断,"
|
||||
"只输出处理后的结果,不要多余解释。"
|
||||
)
|
||||
else:
|
||||
sys_prompt = node.get("system_prompt") or (
|
||||
"你是一个结果汇总器,请将多个子任务的输出综合成连贯、完整、面向用户的最终回答。"
|
||||
)
|
||||
uc = _assemble_node_input(query, upstream, is_entry)
|
||||
text = await _llm_complete(sys_prompt, uc, node.get("model_id"), db, user_id)
|
||||
return {"ok": True, "node_id": nid, "node_type": ntype, "output": text}
|
||||
|
||||
if ntype == "local":
|
||||
agents = _resolve_market_agents(db, [node["agent_id"]], user_id)
|
||||
if not agents:
|
||||
raise OrchestrationError(f"local 节点 {nid} 引用的 agent_id={node['agent_id']} 不存在")
|
||||
ni = _assemble_node_input(query, upstream, is_entry)
|
||||
res = await _run_local_agent_turn(agents[0], ni, db, user_id, conv_id)
|
||||
return {
|
||||
"ok": True,
|
||||
"node_id": nid,
|
||||
"node_type": "local",
|
||||
"output": res["text"],
|
||||
"used_tools": res.get("used_tools"),
|
||||
}
|
||||
|
||||
if ntype == "external":
|
||||
from app.api.chat_agents import _call_external_agent
|
||||
|
||||
agents = _resolve_market_agents(db, [node["agent_id"]], user_id)
|
||||
if not agents:
|
||||
raise OrchestrationError(f"external 节点 {nid} 引用的 agent_id={node['agent_id']} 不存在")
|
||||
ni = _assemble_node_input(query, upstream, is_entry)
|
||||
answer, _ = await _call_external_agent(
|
||||
db, agents[0], query=ni, user_id=user_id, ext_conversation_id=None, history=[]
|
||||
)
|
||||
return {"ok": True, "node_id": nid, "node_type": "external", "output": answer or ""}
|
||||
|
||||
raise OrchestrationError(f"未知节点类型 {ntype}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 入口
|
||||
# ---------------------------------------------------------------------------
|
||||
async def run_graph(graph, query, db, user_id, conv_id, trace_id=None):
|
||||
"""执行编排图,返回 (answer: str, trace: dict)。
|
||||
|
||||
- 拓扑分层后逐层 asyncio.gather 并行执行。
|
||||
- 节点异常按 on_error=halt|continue 处理(默认 continue,失败时下游拿到错误作为输入)。
|
||||
"""
|
||||
validate_graph(graph)
|
||||
layers = topological_layers(graph)
|
||||
blackboard: dict = {}
|
||||
trace = {"nodes": [], "trace_id": trace_id, "layers": layers}
|
||||
|
||||
for layer in layers:
|
||||
tasks = [
|
||||
_run_node(_node_by_id(graph, nid), graph, blackboard, query, user_id, conv_id, trace_id)
|
||||
for nid in layer
|
||||
]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
for nid, res in zip(layer, results):
|
||||
node = _node_by_id(graph, nid)
|
||||
if isinstance(res, Exception):
|
||||
on_error = node.get("on_error", "continue")
|
||||
err_msg = str(res)
|
||||
blackboard[nid] = {
|
||||
"ok": False,
|
||||
"node_id": nid,
|
||||
"node_type": node["type"],
|
||||
"error": err_msg,
|
||||
}
|
||||
trace["nodes"].append({"id": nid, "type": node["type"], "ok": False, "error": err_msg})
|
||||
logger.warning(f"[orchestrator] 节点 {nid} 失败 on_error={on_error}: {err_msg}")
|
||||
if on_error == "halt":
|
||||
raise OrchestrationError(f"节点 {nid} 执行失败(halt): {err_msg}")
|
||||
else:
|
||||
blackboard[nid] = res
|
||||
trace["nodes"].append({
|
||||
"id": nid,
|
||||
"type": res.get("node_type"),
|
||||
"ok": True,
|
||||
"used_tools": res.get("used_tools"),
|
||||
})
|
||||
|
||||
output_node = _find_output_node(graph)
|
||||
final = blackboard.get(output_node, {})
|
||||
if final.get("ok"):
|
||||
answer = final.get("output", "")
|
||||
else:
|
||||
answer = f"⚠️ 编排输出节点 {output_node} 执行失败: {final.get('error', '未知错误')}"
|
||||
trace["output_node"] = output_node
|
||||
return answer, trace
|
||||
"""app/core/orchestrator.py — Phase 3.3 多 Agent 编排图引擎。
|
||||
|
||||
设计目标:
|
||||
- 以轻量 DAG 描述多 Agent 协作:节点(llm/local/external/reduce) + 有向边(数据流)。
|
||||
- 拓扑分层后按层并行执行(asyncio.gather),互相独立的节点真正并发。
|
||||
- 失败隔离:每个节点异常不影响其它分支;per-node on_error=halt|continue。
|
||||
- 节点执行复用既有原语(_build_system_prompt / _execute_skill_tool / _call_external_agent /
|
||||
3.1 工具沙箱),完全向后兼容——agent.orchestration 为 null 时 chat_service 走原逻辑。
|
||||
- 并发安全:每个节点用独立 SessionLocal 会话做 DB 读,避免共享请求会话跨协程竞态。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger("aiflow.orchestrator")
|
||||
|
||||
VALID_NODE_TYPES = ("llm", "local", "external", "reduce")
|
||||
|
||||
|
||||
class OrchestrationError(ValueError):
|
||||
"""图结构非法或执行被 halt。"""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 校验 + 拓扑
|
||||
# ---------------------------------------------------------------------------
|
||||
def validate_graph(graph: Any) -> None:
|
||||
if not isinstance(graph, dict):
|
||||
raise OrchestrationError("orchestration 必须为对象")
|
||||
nodes = graph.get("nodes")
|
||||
edges = graph.get("edges")
|
||||
if not isinstance(nodes, list) or not nodes:
|
||||
raise OrchestrationError("nodes 必须为非空数组")
|
||||
if not isinstance(edges, list):
|
||||
raise OrchestrationError("edges 必须为数组")
|
||||
|
||||
ids = set()
|
||||
for i, n in enumerate(nodes):
|
||||
if not isinstance(n, dict):
|
||||
raise OrchestrationError(f"节点 #{i} 必须为对象")
|
||||
nid = n.get("id")
|
||||
if not nid or not isinstance(nid, str):
|
||||
raise OrchestrationError(f"节点 #{i} 缺少字符串 id")
|
||||
if nid in ids:
|
||||
raise OrchestrationError(f"节点 id 重复: {nid}")
|
||||
ids.add(nid)
|
||||
ntype = n.get("type")
|
||||
if ntype not in VALID_NODE_TYPES:
|
||||
raise OrchestrationError(
|
||||
f"节点 {nid} 类型非法: {ntype} (允许 {', '.join(VALID_NODE_TYPES)})"
|
||||
)
|
||||
if ntype in ("local", "external") and not n.get("agent_id"):
|
||||
raise OrchestrationError(f"节点 {nid} 类型 {ntype} 必须指定 agent_id")
|
||||
|
||||
for i, e in enumerate(edges):
|
||||
if not isinstance(e, dict):
|
||||
raise OrchestrationError(f"边 #{i} 必须为对象")
|
||||
f = e.get("from")
|
||||
t = e.get("to")
|
||||
if f not in ids:
|
||||
raise OrchestrationError(f"边 #{i} 的 from={f} 指向不存在的节点")
|
||||
if t not in ids:
|
||||
raise OrchestrationError(f"边 #{i} 的 to={t} 指向不存在的节点")
|
||||
if f == t:
|
||||
raise OrchestrationError(f"边 #{i} 存在自环: {f}")
|
||||
|
||||
|
||||
def topological_layers(graph: dict) -> list:
|
||||
"""Kahn 算法返回分层(每层内节点可并行);含环则抛错。"""
|
||||
nodes = graph["nodes"]
|
||||
edges = graph.get("edges", [])
|
||||
indeg = {n["id"]: 0 for n in nodes}
|
||||
adj = {n["id"]: [] for n in nodes}
|
||||
for e in edges:
|
||||
f, t = e["from"], e["to"]
|
||||
adj[f].append(t)
|
||||
indeg[t] += 1
|
||||
order = {n["id"]: i for i, n in enumerate(nodes)}
|
||||
layers: list = []
|
||||
remaining = set(indeg)
|
||||
while remaining:
|
||||
layer = sorted([nid for nid in remaining if indeg[nid] == 0], key=lambda x: order[x])
|
||||
if not layer:
|
||||
cycle = sorted(remaining, key=lambda x: order[x])
|
||||
raise OrchestrationError(f"检测到环路,无法拓扑排序: {cycle}")
|
||||
layers.append(layer)
|
||||
for nid in layer:
|
||||
remaining.discard(nid)
|
||||
for nxt in adj[nid]:
|
||||
indeg[nxt] -= 1
|
||||
return layers
|
||||
|
||||
|
||||
def _node_by_id(graph: dict, nid: str) -> dict:
|
||||
for n in graph["nodes"]:
|
||||
if n["id"] == nid:
|
||||
return n
|
||||
return {}
|
||||
|
||||
|
||||
def _predecessors(graph: dict, nid: str) -> list:
|
||||
return [e["from"] for e in graph.get("edges", []) if e["to"] == nid]
|
||||
|
||||
|
||||
def _successors(graph: dict, nid: str) -> list:
|
||||
return [e["to"] for e in graph.get("edges", []) if e["from"] == nid]
|
||||
|
||||
|
||||
def _find_output_node(graph: dict) -> str:
|
||||
if graph.get("output_node"):
|
||||
return graph["output_node"]
|
||||
sinks = [n["id"] for n in graph["nodes"] if not _successors(graph, n["id"])]
|
||||
if len(sinks) == 1:
|
||||
return sinks[0]
|
||||
for s in sinks:
|
||||
if _node_by_id(graph, s).get("type") == "reduce":
|
||||
return s
|
||||
if sinks:
|
||||
return sinks[0]
|
||||
return graph["nodes"][-1]["id"]
|
||||
|
||||
|
||||
def _assemble_node_input(query: str, upstream: str, is_entry: bool) -> str:
|
||||
"""入口节点直接吃用户原问题;下游节点综合原问题 + 上游输出。"""
|
||||
if is_entry:
|
||||
return query or ""
|
||||
if upstream:
|
||||
return f"原始用户问题:\n{query}\n\n上游节点输出:\n{upstream}"
|
||||
return query or ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 模型解析 + LLM 调用原语
|
||||
# ---------------------------------------------------------------------------
|
||||
async def _resolve_model_and_endpoint(model_id, db):
|
||||
from app.models import AIModel, ModelProvider
|
||||
from ..capability_registry import resolve_auto_model
|
||||
|
||||
ai_model = None
|
||||
if model_id and model_id != "auto":
|
||||
ai_model = db.query(AIModel).filter(AIModel.id == model_id).first()
|
||||
if not ai_model:
|
||||
ai_model = resolve_auto_model("text", db)
|
||||
if not ai_model:
|
||||
raise OrchestrationError("无法解析到可用模型")
|
||||
provider = db.query(ModelProvider).filter(ModelProvider.id == ai_model.provider_id).first()
|
||||
api_key = ai_model.api_key or (provider.api_key if provider else None)
|
||||
api_url = ai_model.endpoint
|
||||
if not api_url and provider and provider.api_base_url:
|
||||
api_url = f"{provider.api_base_url.rstrip('/')}/v1/chat/completions"
|
||||
if not api_url or not api_key:
|
||||
raise OrchestrationError(f"节点模型 {ai_model.name} 缺少 endpoint 或 api_key")
|
||||
return ai_model, api_url, api_key
|
||||
|
||||
|
||||
async def _llm_complete(system_prompt: str, user_content: str, model_id, db, user_id) -> str:
|
||||
from app.api.chat_utils import _proxy_for_url
|
||||
from app.core.observability import span
|
||||
|
||||
ai_model, api_url, api_key = await _resolve_model_and_endpoint(model_id, db)
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_content},
|
||||
]
|
||||
request_body = {"model": ai_model.name, "messages": messages, "max_tokens": 2048}
|
||||
async with span("orchestrator.llm", model=ai_model.id):
|
||||
async with httpx.AsyncClient(
|
||||
timeout=ai_model.timeout or 120,
|
||||
proxy=_proxy_for_url(api_url),
|
||||
trust_env=False,
|
||||
) as client:
|
||||
resp = await client.post(
|
||||
api_url,
|
||||
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
||||
json=request_body,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
raise OrchestrationError(f"LLM 调用失败 HTTP {resp.status_code}: {resp.text[:200]}")
|
||||
data = resp.json()
|
||||
return data["choices"][0]["message"].get("content") or ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# local 节点:复用单 Agent system prompt + 受限 tool round + 3.1 沙箱
|
||||
# ---------------------------------------------------------------------------
|
||||
async def _run_local_agent_turn(agent: dict, node_input: str, db, user_id, conv_id) -> dict:
|
||||
from app.api.chat_skill import _resolve_market_skills, _execute_skill_tool
|
||||
from app.api.chat_context import _build_system_prompt, _retrieve_kb_context
|
||||
from app.api.chat_service import _resolve_agent_tool_policy, _check_tool_policy
|
||||
from app.api.chat_utils import _proxy_for_url
|
||||
|
||||
skills = _resolve_market_skills(db, agent.get("skill_ids") or [], user_id)
|
||||
kb_context = await _retrieve_kb_context(db, agent.get("knowledge_base_ids") or [], node_input)
|
||||
system_prompt, function_definitions = _build_system_prompt(
|
||||
skills=skills,
|
||||
agents=[agent],
|
||||
kb_context=kb_context,
|
||||
user_content=node_input,
|
||||
auto_discovered=False,
|
||||
)
|
||||
ai_model, api_url, api_key = await _resolve_model_and_endpoint(agent.get("model_id"), db)
|
||||
tool_policy = _resolve_agent_tool_policy([agent])
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": node_input},
|
||||
]
|
||||
called_skills: list = []
|
||||
tool_call_count = 0
|
||||
max_tool_rounds = 5
|
||||
last_content = ""
|
||||
async with httpx.AsyncClient(
|
||||
timeout=ai_model.timeout or 120,
|
||||
proxy=_proxy_for_url(api_url),
|
||||
trust_env=False,
|
||||
) as client:
|
||||
for round_num in range(max_tool_rounds):
|
||||
request_body = {"model": ai_model.name, "messages": messages, "max_tokens": 2048}
|
||||
if function_definitions:
|
||||
request_body["tools"] = [
|
||||
{"type": "function", "function": fd} for fd in function_definitions
|
||||
]
|
||||
request_body["tool_choice"] = "auto"
|
||||
resp = await client.post(
|
||||
api_url,
|
||||
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
||||
json=request_body,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
raise OrchestrationError(f"local 节点 LLM 失败 HTTP {resp.status_code}: {resp.text[:200]}")
|
||||
data = resp.json()
|
||||
message = data["choices"][0]["message"]
|
||||
tool_calls = message.get("tool_calls") or []
|
||||
content = message.get("content") or ""
|
||||
last_content = content
|
||||
if not tool_calls or not function_definitions:
|
||||
return {"text": content, "used_tools": called_skills}
|
||||
messages.append({"role": "assistant", "content": content, "tool_calls": tool_calls})
|
||||
round_tool_count = 0
|
||||
for tc in tool_calls:
|
||||
func_name = tc.get("function", {}).get("name", "")
|
||||
check = _check_tool_policy(
|
||||
func_name, None, skills, tool_policy, tool_call_count, round_tool_count
|
||||
)
|
||||
if not check["allowed"]:
|
||||
messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tc.get("id", f"call_{round_num}"),
|
||||
"content": f"[工具调用已被策略拦截] {check.get('message', check.get('reason'))}",
|
||||
})
|
||||
tool_call_count += 1
|
||||
round_tool_count += 1
|
||||
continue
|
||||
tool_call_count += 1
|
||||
round_tool_count += 1
|
||||
try:
|
||||
func_args = json.loads(tc.get("function", {}).get("arguments", "{}") or "{}")
|
||||
except Exception:
|
||||
func_args = {}
|
||||
res = _execute_skill_tool(db, func_name, func_args, skills, node_input)
|
||||
called_skills.append(func_name)
|
||||
messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tc.get("id", f"call_{round_num}"),
|
||||
"content": res,
|
||||
})
|
||||
return {"text": last_content, "used_tools": called_skills}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 单节点执行(含独立 DB 会话,保证并发安全)
|
||||
# ---------------------------------------------------------------------------
|
||||
async def _run_node(node, graph, blackboard, query, user_id, conv_id, trace_id) -> dict:
|
||||
from app.api.chat_skill import _resolve_market_agents
|
||||
from app.database import SessionLocal
|
||||
|
||||
nid = node["id"]
|
||||
ntype = node["type"]
|
||||
preds = _predecessors(graph, nid)
|
||||
pred_results = [blackboard[p] for p in preds]
|
||||
is_entry = not preds
|
||||
|
||||
if pred_results:
|
||||
parts = []
|
||||
for b in pred_results:
|
||||
label = f"节点 {b.get('node_id')}/{b.get('node_type')}"
|
||||
if b.get("ok"):
|
||||
parts.append(f"【{label}】\n{b.get('output', '')}")
|
||||
else:
|
||||
parts.append(f"【{label} 执行失败】\n错误: {b.get('error', '')}")
|
||||
upstream = "\n\n".join(parts)
|
||||
else:
|
||||
upstream = ""
|
||||
|
||||
with SessionLocal() as db:
|
||||
if ntype in ("llm", "reduce"):
|
||||
if ntype == "llm":
|
||||
sys_prompt = node.get("system_prompt") or (
|
||||
"你是一个处理节点,请根据上游输入完成指定的文本变换、摘要或路由判断,"
|
||||
"只输出处理后的结果,不要多余解释。"
|
||||
)
|
||||
else:
|
||||
sys_prompt = node.get("system_prompt") or (
|
||||
"你是一个结果汇总器,请将多个子任务的输出综合成连贯、完整、面向用户的最终回答。"
|
||||
)
|
||||
uc = _assemble_node_input(query, upstream, is_entry)
|
||||
text = await _llm_complete(sys_prompt, uc, node.get("model_id"), db, user_id)
|
||||
return {"ok": True, "node_id": nid, "node_type": ntype, "output": text}
|
||||
|
||||
if ntype == "local":
|
||||
agents = _resolve_market_agents(db, [node["agent_id"]], user_id)
|
||||
if not agents:
|
||||
raise OrchestrationError(f"local 节点 {nid} 引用的 agent_id={node['agent_id']} 不存在")
|
||||
ni = _assemble_node_input(query, upstream, is_entry)
|
||||
res = await _run_local_agent_turn(agents[0], ni, db, user_id, conv_id)
|
||||
return {
|
||||
"ok": True,
|
||||
"node_id": nid,
|
||||
"node_type": "local",
|
||||
"output": res["text"],
|
||||
"used_tools": res.get("used_tools"),
|
||||
}
|
||||
|
||||
if ntype == "external":
|
||||
from app.api.chat_agents import _call_external_agent
|
||||
|
||||
agents = _resolve_market_agents(db, [node["agent_id"]], user_id)
|
||||
if not agents:
|
||||
raise OrchestrationError(f"external 节点 {nid} 引用的 agent_id={node['agent_id']} 不存在")
|
||||
ni = _assemble_node_input(query, upstream, is_entry)
|
||||
answer, _ = await _call_external_agent(
|
||||
db, agents[0], query=ni, user_id=user_id, ext_conversation_id=None, history=[]
|
||||
)
|
||||
return {"ok": True, "node_id": nid, "node_type": "external", "output": answer or ""}
|
||||
|
||||
raise OrchestrationError(f"未知节点类型 {ntype}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 入口
|
||||
# ---------------------------------------------------------------------------
|
||||
async def run_graph(graph, query, db, user_id, conv_id, trace_id=None):
|
||||
"""执行编排图,返回 (answer: str, trace: dict)。
|
||||
|
||||
- 拓扑分层后逐层 asyncio.gather 并行执行。
|
||||
- 节点异常按 on_error=halt|continue 处理(默认 continue,失败时下游拿到错误作为输入)。
|
||||
"""
|
||||
validate_graph(graph)
|
||||
layers = topological_layers(graph)
|
||||
blackboard: dict = {}
|
||||
trace = {"nodes": [], "trace_id": trace_id, "layers": layers}
|
||||
|
||||
for layer in layers:
|
||||
tasks = [
|
||||
_run_node(_node_by_id(graph, nid), graph, blackboard, query, user_id, conv_id, trace_id)
|
||||
for nid in layer
|
||||
]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
for nid, res in zip(layer, results):
|
||||
node = _node_by_id(graph, nid)
|
||||
if isinstance(res, Exception):
|
||||
on_error = node.get("on_error", "continue")
|
||||
err_msg = str(res)
|
||||
blackboard[nid] = {
|
||||
"ok": False,
|
||||
"node_id": nid,
|
||||
"node_type": node["type"],
|
||||
"error": err_msg,
|
||||
}
|
||||
trace["nodes"].append({"id": nid, "type": node["type"], "ok": False, "error": err_msg})
|
||||
logger.warning(f"[orchestrator] 节点 {nid} 失败 on_error={on_error}: {err_msg}")
|
||||
if on_error == "halt":
|
||||
raise OrchestrationError(f"节点 {nid} 执行失败(halt): {err_msg}")
|
||||
else:
|
||||
blackboard[nid] = res
|
||||
trace["nodes"].append({
|
||||
"id": nid,
|
||||
"type": res.get("node_type"),
|
||||
"ok": True,
|
||||
"used_tools": res.get("used_tools"),
|
||||
})
|
||||
|
||||
output_node = _find_output_node(graph)
|
||||
final = blackboard.get(output_node, {})
|
||||
if final.get("ok"):
|
||||
answer = final.get("output", "")
|
||||
else:
|
||||
answer = f"⚠️ 编排输出节点 {output_node} 执行失败: {final.get('error', '未知错误')}"
|
||||
trace["output_node"] = output_node
|
||||
return answer, trace
|
||||
@@ -1,360 +1,390 @@
|
||||
"""
|
||||
PGVector 向量检索模块
|
||||
====================
|
||||
使用 PostgreSQL + pgvector 实现语义向量检索,
|
||||
与现有 MySQL BM25 检索配合,实现混合检索。
|
||||
|
||||
使用方式:
|
||||
from app.pgvector_search import vector_search
|
||||
results = await vector_search(query, kb_id, top_k=5)
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import List, Dict, Optional
|
||||
from datetime import datetime
|
||||
|
||||
logger = logging.getLogger("aiflow.pgvector")
|
||||
|
||||
# PostgreSQL 连接配置
|
||||
PG_HOST = os.environ.get("POSTGRES_HOST", "pgvector")
|
||||
PG_PORT = int(os.environ.get("POSTGRES_PORT", "5432"))
|
||||
PG_USER = os.environ.get("POSTGRES_USER", "aiflow")
|
||||
PG_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "pwd2026")
|
||||
PG_DB = os.environ.get("POSTGRES_DB", "aiflow_vectors")
|
||||
|
||||
# 嵌入模型配置
|
||||
EMBEDDING_API_URL = os.environ.get("SILICONFLOW_BASE_URL", "https://api.siliconflow.cn")
|
||||
EMBEDDING_API_KEY = os.environ.get("SILICONFLOW_API_KEY", "")
|
||||
EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-large-zh-v1.5")
|
||||
|
||||
# 连接池
|
||||
_pg_pool = None
|
||||
_pool_lock = asyncio.Lock()
|
||||
|
||||
|
||||
async def get_pg_pool():
|
||||
"""获取 PostgreSQL 异步连接池"""
|
||||
global _pg_pool
|
||||
if _pg_pool is not None:
|
||||
return _pg_pool
|
||||
|
||||
async with _pool_lock:
|
||||
if _pg_pool is not None:
|
||||
return _pg_pool
|
||||
|
||||
try:
|
||||
import asyncpg
|
||||
except ImportError:
|
||||
logger.warning("asyncpg not installed, vector search disabled")
|
||||
return None
|
||||
|
||||
try:
|
||||
_pg_pool = await asyncpg.create_pool(
|
||||
host=PG_HOST,
|
||||
port=PG_PORT,
|
||||
user=PG_USER,
|
||||
password=PG_PASSWORD,
|
||||
database=PG_DB,
|
||||
min_size=2,
|
||||
max_size=10,
|
||||
command_timeout=30,
|
||||
)
|
||||
logger.info(f"PostgreSQL connection pool created: {PG_HOST}:{PG_PORT}/{PG_DB}")
|
||||
|
||||
# 初始化表结构
|
||||
async with _pg_pool.acquire() as conn:
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS document_embeddings (
|
||||
id SERIAL PRIMARY KEY,
|
||||
doc_id VARCHAR(64) NOT NULL,
|
||||
kb_id VARCHAR(64) NOT NULL,
|
||||
chunk_id VARCHAR(64),
|
||||
content TEXT NOT NULL,
|
||||
embedding vector(1024),
|
||||
metadata JSONB DEFAULT '{}',
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
""")
|
||||
await conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_embeddings_kb
|
||||
ON document_embeddings(kb_id);
|
||||
""")
|
||||
await conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_embeddings_doc
|
||||
ON document_embeddings(doc_id);
|
||||
""")
|
||||
await conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_embeddings_vector
|
||||
ON document_embeddings USING ivfflat (embedding vector_cosine_ops)
|
||||
WITH (lists = 100);
|
||||
""")
|
||||
logger.info("PGVector tables initialized")
|
||||
|
||||
return _pg_pool
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create PostgreSQL pool: {e}")
|
||||
_pg_pool = None
|
||||
return None
|
||||
|
||||
|
||||
async def get_embedding(text: str) -> Optional[List[float]]:
|
||||
"""调用嵌入API获取文本向量"""
|
||||
if not EMBEDDING_API_KEY:
|
||||
logger.warning("No embedding API key configured")
|
||||
return None
|
||||
|
||||
try:
|
||||
import httpx
|
||||
async with httpx.AsyncClient(timeout=30) as client:
|
||||
resp = await client.post(
|
||||
f"{EMBEDDING_API_URL}/v1/embeddings",
|
||||
headers={
|
||||
"Authorization": f"Bearer {EMBEDDING_API_KEY}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json={
|
||||
"model": EMBEDDING_MODEL,
|
||||
"input": text[:8000], # 截断防止超长
|
||||
},
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
embedding = data["data"][0]["embedding"]
|
||||
return embedding
|
||||
else:
|
||||
logger.error(f"Embedding API error: {resp.status_code} {resp.text[:200]}")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Embedding generation failed: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def vector_search(
|
||||
query: str,
|
||||
kb_id: Optional[str] = None,
|
||||
top_k: int = 5,
|
||||
threshold: float = 0.5,
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
向量语义检索
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
kb_id: 知识库ID (可选,不指定则搜索全部)
|
||||
top_k: 返回结果数
|
||||
threshold: 相似度阈值 (0-1)
|
||||
|
||||
Returns:
|
||||
[{"doc_id": ..., "content": ..., "score": ..., "metadata": ...}, ...]
|
||||
"""
|
||||
pool = await get_pg_pool()
|
||||
if pool is None:
|
||||
return []
|
||||
|
||||
# 获取查询向量
|
||||
query_embedding = await get_embedding(query)
|
||||
if query_embedding is None:
|
||||
return []
|
||||
|
||||
try:
|
||||
async with pool.acquire() as conn:
|
||||
# 将向量转换为 PostgreSQL 格式
|
||||
embedding_str = f"[{','.join(str(x) for x in query_embedding)}]"
|
||||
|
||||
if kb_id:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT doc_id, chunk_id, content, metadata,
|
||||
1 - (embedding <=> $1::vector) as score
|
||||
FROM document_embeddings
|
||||
WHERE kb_id = $2
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $3;
|
||||
""",
|
||||
embedding_str, kb_id, top_k,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT doc_id, chunk_id, content, metadata,
|
||||
1 - (embedding <=> $1::vector) as score
|
||||
FROM document_embeddings
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $2;
|
||||
""",
|
||||
embedding_str, top_k,
|
||||
)
|
||||
|
||||
results = []
|
||||
for row in rows:
|
||||
score = float(row["score"])
|
||||
if score >= threshold:
|
||||
results.append({
|
||||
"doc_id": row["doc_id"],
|
||||
"chunk_id": row["chunk_id"],
|
||||
"content": row["content"],
|
||||
"score": score,
|
||||
"metadata": row["metadata"] if isinstance(row["metadata"], dict) else {},
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Vector search failed: {e}")
|
||||
return []
|
||||
|
||||
|
||||
async def index_document(
|
||||
doc_id: str,
|
||||
kb_id: str,
|
||||
content: str,
|
||||
chunk_id: Optional[str] = None,
|
||||
metadata: Optional[Dict] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
将文档内容索引到 PGVector
|
||||
|
||||
Args:
|
||||
doc_id: 文档ID
|
||||
kb_id: 知识库ID
|
||||
content: 文档内容
|
||||
chunk_id: 分块ID (可选)
|
||||
metadata: 元数据
|
||||
|
||||
Returns:
|
||||
是否成功
|
||||
"""
|
||||
pool = await get_pg_pool()
|
||||
if pool is None:
|
||||
return False
|
||||
|
||||
# 分块处理(简单按长度分块)
|
||||
chunks = []
|
||||
chunk_size = 500
|
||||
overlap = 50
|
||||
if len(content) <= chunk_size:
|
||||
chunks = [(content, chunk_id or f"{doc_id}_0")]
|
||||
else:
|
||||
for i in range(0, len(content), chunk_size - overlap):
|
||||
chunk = content[i:i + chunk_size]
|
||||
if len(chunk.strip()) > 10:
|
||||
chunks.append((chunk, f"{doc_id}_{i // (chunk_size - overlap)}"))
|
||||
|
||||
success_count = 0
|
||||
for chunk_text, cid in chunks:
|
||||
embedding = await get_embedding(chunk_text)
|
||||
if embedding is None:
|
||||
continue
|
||||
|
||||
try:
|
||||
async with pool.acquire() as conn:
|
||||
embedding_str = f"[{','.join(str(x) for x in embedding)}]"
|
||||
import json
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO document_embeddings (doc_id, kb_id, chunk_id, content, embedding, metadata)
|
||||
VALUES ($1, $2, $3, $4, $5::vector, $6)
|
||||
ON CONFLICT DO NOTHING;
|
||||
""",
|
||||
doc_id, kb_id, cid, chunk_text, embedding_str,
|
||||
json.dumps(metadata or {}),
|
||||
)
|
||||
success_count += 1
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to index chunk {cid}: {e}")
|
||||
|
||||
logger.info(f"Indexed document {doc_id}: {success_count}/{len(chunks)} chunks")
|
||||
return success_count > 0
|
||||
|
||||
|
||||
async def hybrid_search(
|
||||
query: str,
|
||||
kb_id: Optional[str] = None,
|
||||
top_k: int = 5,
|
||||
bm25_results: Optional[List[Dict]] = None,
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
混合检索:BM25 + 向量检索融合
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
kb_id: 知识库ID
|
||||
top_k: 返回结果数
|
||||
bm25_results: BM25 检索结果(来自MySQL)
|
||||
|
||||
Returns:
|
||||
融合后的检索结果
|
||||
"""
|
||||
# 获取向量检索结果
|
||||
vec_results = await vector_search(query, kb_id, top_k)
|
||||
|
||||
if not vec_results and not bm25_results:
|
||||
return []
|
||||
|
||||
if not vec_results:
|
||||
return bm25_results or []
|
||||
|
||||
if not bm25_results:
|
||||
return vec_results
|
||||
|
||||
# 融合:Reciprocal Rank Fusion (RRF)
|
||||
rrf_k = 60
|
||||
score_map = {}
|
||||
|
||||
for rank, item in enumerate(bm25_results):
|
||||
doc_id = item.get("doc_id", item.get("id", ""))
|
||||
score_map[doc_id] = {
|
||||
"content": item.get("content", item.get("text", "")),
|
||||
"score": 1.0 / (rrf_k + rank + 1),
|
||||
"bm25_rank": rank + 1,
|
||||
"metadata": item.get("metadata", {}),
|
||||
}
|
||||
|
||||
for rank, item in enumerate(vec_results):
|
||||
doc_id = item["doc_id"]
|
||||
if doc_id in score_map:
|
||||
score_map[doc_id]["score"] += 1.0 / (rrf_k + rank + 1)
|
||||
score_map[doc_id]["vec_score"] = item["score"]
|
||||
score_map[doc_id]["vec_rank"] = rank + 1
|
||||
else:
|
||||
score_map[doc_id] = {
|
||||
"content": item["content"],
|
||||
"score": 1.0 / (rrf_k + rank + 1),
|
||||
"vec_score": item["score"],
|
||||
"vec_rank": rank + 1,
|
||||
"metadata": item.get("metadata", {}),
|
||||
}
|
||||
|
||||
# 按融合分数排序
|
||||
merged = sorted(score_map.values(), key=lambda x: x["score"], reverse=True)
|
||||
return merged[:top_k]
|
||||
|
||||
|
||||
async def delete_document_embeddings(doc_id: str) -> bool:
|
||||
"""删除文档的所有向量索引"""
|
||||
pool = await get_pg_pool()
|
||||
if pool is None:
|
||||
return False
|
||||
|
||||
try:
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"DELETE FROM document_embeddings WHERE doc_id = $1;",
|
||||
doc_id,
|
||||
)
|
||||
logger.info(f"Deleted embeddings for document {doc_id}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete embeddings: {e}")
|
||||
return False
|
||||
|
||||
|
||||
async def close_pool():
|
||||
"""关闭连接池"""
|
||||
global _pg_pool
|
||||
if _pg_pool is not None:
|
||||
await _pg_pool.close()
|
||||
_pg_pool = None
|
||||
logger.info("PostgreSQL connection pool closed")
|
||||
"""
|
||||
PGVector 向量检索模块
|
||||
====================
|
||||
使用 PostgreSQL + pgvector 实现语义向量检索,
|
||||
与现有 MySQL BM25 检索配合,实现混合检索。
|
||||
|
||||
使用方式:
|
||||
from app.pgvector_search import vector_search
|
||||
results = await vector_search(query, kb_id, top_k=5)
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import List, Dict, Optional
|
||||
from datetime import datetime
|
||||
|
||||
logger = logging.getLogger("aiflow.pgvector")
|
||||
|
||||
# PostgreSQL 连接配置
|
||||
PG_HOST = os.environ.get("POSTGRES_HOST", "qihu-aiflow-pgvector")
|
||||
PG_PORT = int(os.environ.get("POSTGRES_PORT", "5432"))
|
||||
PG_USER = os.environ.get("POSTGRES_USER", "aiflow")
|
||||
PG_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "pwd2026")
|
||||
PG_DB = os.environ.get("POSTGRES_DB", "aiflow_vectors")
|
||||
|
||||
# 嵌入模型配置
|
||||
EMBEDDING_API_URL = os.environ.get("SILICONFLOW_BASE_URL", "https://api.siliconflow.cn")
|
||||
EMBEDDING_API_KEY = os.environ.get("SILICONFLOW_API_KEY", "")
|
||||
EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-large-zh-v1.5")
|
||||
|
||||
# ---- 本地嵌入缓存(SQLite 持久化):破解"每次嵌入消耗外部 API 额度"的隐性成本障碍 ----
|
||||
import sqlite3, hashlib, json as _json
|
||||
_EMBED_CACHE_DB = os.environ.get("EMBED_CACHE_DB", "/app/backend/python/.embed_cache.db")
|
||||
def _embed_cache_get(text, model):
|
||||
try:
|
||||
key = hashlib.sha256(f"{model}::{text}".encode("utf-8")).hexdigest()
|
||||
conn = sqlite3.connect(_EMBED_CACHE_DB, timeout=5)
|
||||
row = conn.execute("SELECT vec FROM embed_cache WHERE k=?", (key,)).fetchone()
|
||||
conn.close()
|
||||
if row:
|
||||
return _json.loads(row[0])
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
def _embed_cache_put(text, model, vec):
|
||||
try:
|
||||
key = hashlib.sha256(f"{model}::{text}".encode("utf-8")).hexdigest()
|
||||
conn = sqlite3.connect(_EMBED_CACHE_DB, timeout=5)
|
||||
conn.execute("CREATE TABLE IF NOT EXISTS embed_cache (k TEXT PRIMARY KEY, vec TEXT, ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP)")
|
||||
conn.execute("INSERT OR REPLACE INTO embed_cache (k, vec) VALUES (?, ?)", (key, _json.dumps(vec)))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# 连接池
|
||||
_pg_pool = None
|
||||
_pool_lock = asyncio.Lock()
|
||||
|
||||
|
||||
async def get_pg_pool():
|
||||
"""获取 PostgreSQL 异步连接池"""
|
||||
global _pg_pool
|
||||
if _pg_pool is not None:
|
||||
return _pg_pool
|
||||
|
||||
async with _pool_lock:
|
||||
if _pg_pool is not None:
|
||||
return _pg_pool
|
||||
|
||||
try:
|
||||
import asyncpg
|
||||
except ImportError:
|
||||
logger.warning("asyncpg not installed, vector search disabled")
|
||||
return None
|
||||
|
||||
try:
|
||||
_pg_pool = await asyncpg.create_pool(
|
||||
host=PG_HOST,
|
||||
port=PG_PORT,
|
||||
user=PG_USER,
|
||||
password=PG_PASSWORD,
|
||||
database=PG_DB,
|
||||
min_size=2,
|
||||
max_size=10,
|
||||
command_timeout=30,
|
||||
)
|
||||
logger.info(f"PostgreSQL connection pool created: {PG_HOST}:{PG_PORT}/{PG_DB}")
|
||||
|
||||
# 初始化表结构
|
||||
async with _pg_pool.acquire() as conn:
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS document_embeddings (
|
||||
id SERIAL PRIMARY KEY,
|
||||
doc_id VARCHAR(64) NOT NULL,
|
||||
kb_id VARCHAR(64) NOT NULL,
|
||||
chunk_id VARCHAR(64),
|
||||
content TEXT NOT NULL,
|
||||
embedding vector(1024),
|
||||
metadata JSONB DEFAULT '{}',
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
""")
|
||||
await conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_embeddings_kb
|
||||
ON document_embeddings(kb_id);
|
||||
""")
|
||||
await conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_embeddings_doc
|
||||
ON document_embeddings(doc_id);
|
||||
""")
|
||||
await conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_embeddings_vector
|
||||
ON document_embeddings USING ivfflat (embedding vector_cosine_ops)
|
||||
WITH (lists = 100);
|
||||
""")
|
||||
logger.info("PGVector tables initialized")
|
||||
|
||||
return _pg_pool
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create PostgreSQL pool: {e}")
|
||||
_pg_pool = None
|
||||
return None
|
||||
|
||||
|
||||
async def get_embedding(text: str) -> Optional[List[float]]:
|
||||
"""调用嵌入API获取文本向量"""
|
||||
if not EMBEDDING_API_KEY:
|
||||
logger.warning("No embedding API key configured")
|
||||
return None
|
||||
|
||||
cached = _embed_cache_get(text, EMBEDDING_MODEL)
|
||||
if cached:
|
||||
return cached
|
||||
try:
|
||||
import httpx
|
||||
async with httpx.AsyncClient(timeout=30, trust_env=False) as client:
|
||||
resp = await client.post(
|
||||
f"{EMBEDDING_API_URL}/v1/embeddings",
|
||||
headers={
|
||||
"Authorization": f"Bearer {EMBEDDING_API_KEY}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json={
|
||||
"model": EMBEDDING_MODEL,
|
||||
"input": text[:8000], # 截断防止超长
|
||||
},
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
embedding = data["data"][0]["embedding"]
|
||||
_embed_cache_put(text, EMBEDDING_MODEL, embedding)
|
||||
return embedding
|
||||
else:
|
||||
logger.error(f"Embedding API error: {resp.status_code} {resp.text[:200]}")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Embedding generation failed: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def vector_search(
|
||||
query: str,
|
||||
kb_id: Optional[str] = None,
|
||||
top_k: int = 5,
|
||||
threshold: float = 0.5,
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
向量语义检索
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
kb_id: 知识库ID (可选,不指定则搜索全部)
|
||||
top_k: 返回结果数
|
||||
threshold: 相似度阈值 (0-1)
|
||||
|
||||
Returns:
|
||||
[{"doc_id": ..., "content": ..., "score": ..., "metadata": ...}, ...]
|
||||
"""
|
||||
pool = await get_pg_pool()
|
||||
if pool is None:
|
||||
return []
|
||||
|
||||
# 获取查询向量
|
||||
query_embedding = await get_embedding(query)
|
||||
if query_embedding is None:
|
||||
return []
|
||||
|
||||
try:
|
||||
async with pool.acquire() as conn:
|
||||
# 将向量转换为 PostgreSQL 格式
|
||||
embedding_str = f"[{','.join(str(x) for x in query_embedding)}]"
|
||||
|
||||
if kb_id:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT doc_id, chunk_id, content, metadata,
|
||||
1 - (embedding <=> $1::vector) as score
|
||||
FROM document_embeddings
|
||||
WHERE kb_id = $2
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $3;
|
||||
""",
|
||||
embedding_str, kb_id, top_k,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT doc_id, chunk_id, content, metadata,
|
||||
1 - (embedding <=> $1::vector) as score
|
||||
FROM document_embeddings
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $2;
|
||||
""",
|
||||
embedding_str, top_k,
|
||||
)
|
||||
|
||||
results = []
|
||||
for row in rows:
|
||||
score = float(row["score"])
|
||||
if score >= threshold:
|
||||
results.append({
|
||||
"doc_id": row["doc_id"],
|
||||
"chunk_id": row["chunk_id"],
|
||||
"content": row["content"],
|
||||
"score": score,
|
||||
"metadata": row["metadata"] if isinstance(row["metadata"], dict) else {},
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Vector search failed: {e}")
|
||||
return []
|
||||
|
||||
|
||||
async def index_document(
|
||||
doc_id: str,
|
||||
kb_id: str,
|
||||
content: str,
|
||||
chunk_id: Optional[str] = None,
|
||||
metadata: Optional[Dict] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
将文档内容索引到 PGVector
|
||||
|
||||
Args:
|
||||
doc_id: 文档ID
|
||||
kb_id: 知识库ID
|
||||
content: 文档内容
|
||||
chunk_id: 分块ID (可选)
|
||||
metadata: 元数据
|
||||
|
||||
Returns:
|
||||
是否成功
|
||||
"""
|
||||
pool = await get_pg_pool()
|
||||
if pool is None:
|
||||
return False
|
||||
|
||||
# 分块处理(简单按长度分块)
|
||||
chunks = []
|
||||
chunk_size = 500
|
||||
overlap = 50
|
||||
if len(content) <= chunk_size:
|
||||
chunks = [(content, chunk_id or f"{doc_id}_0")]
|
||||
else:
|
||||
for i in range(0, len(content), chunk_size - overlap):
|
||||
chunk = content[i:i + chunk_size]
|
||||
if len(chunk.strip()) > 10:
|
||||
chunks.append((chunk, f"{doc_id}_{i // (chunk_size - overlap)}"))
|
||||
|
||||
success_count = 0
|
||||
for chunk_text, cid in chunks:
|
||||
embedding = await get_embedding(chunk_text)
|
||||
if embedding is None:
|
||||
continue
|
||||
|
||||
try:
|
||||
async with pool.acquire() as conn:
|
||||
embedding_str = f"[{','.join(str(x) for x in embedding)}]"
|
||||
import json
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO document_embeddings (doc_id, kb_id, chunk_id, content, embedding, metadata)
|
||||
VALUES ($1, $2, $3, $4, $5::vector, $6)
|
||||
ON CONFLICT DO NOTHING;
|
||||
""",
|
||||
doc_id, kb_id, cid, chunk_text, embedding_str,
|
||||
json.dumps(metadata or {}),
|
||||
)
|
||||
success_count += 1
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to index chunk {cid}: {e}")
|
||||
|
||||
logger.info(f"Indexed document {doc_id}: {success_count}/{len(chunks)} chunks")
|
||||
return success_count > 0
|
||||
|
||||
|
||||
async def hybrid_search(
|
||||
query: str,
|
||||
kb_id: Optional[str] = None,
|
||||
top_k: int = 5,
|
||||
bm25_results: Optional[List[Dict]] = None,
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
混合检索:BM25 + 向量检索融合
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
kb_id: 知识库ID
|
||||
top_k: 返回结果数
|
||||
bm25_results: BM25 检索结果(来自MySQL)
|
||||
|
||||
Returns:
|
||||
融合后的检索结果
|
||||
"""
|
||||
# 获取向量检索结果
|
||||
vec_results = await vector_search(query, kb_id, top_k)
|
||||
|
||||
if not vec_results and not bm25_results:
|
||||
return []
|
||||
|
||||
if not vec_results:
|
||||
return bm25_results or []
|
||||
|
||||
if not bm25_results:
|
||||
return vec_results
|
||||
|
||||
# 融合:Reciprocal Rank Fusion (RRF)
|
||||
rrf_k = 60
|
||||
score_map = {}
|
||||
|
||||
for rank, item in enumerate(bm25_results):
|
||||
doc_id = item.get("doc_id", item.get("id", ""))
|
||||
score_map[doc_id] = {
|
||||
"content": item.get("content", item.get("text", "")),
|
||||
"score": 1.0 / (rrf_k + rank + 1),
|
||||
"bm25_rank": rank + 1,
|
||||
"metadata": item.get("metadata", {}),
|
||||
}
|
||||
|
||||
for rank, item in enumerate(vec_results):
|
||||
doc_id = item["doc_id"]
|
||||
if doc_id in score_map:
|
||||
score_map[doc_id]["score"] += 1.0 / (rrf_k + rank + 1)
|
||||
score_map[doc_id]["vec_score"] = item["score"]
|
||||
score_map[doc_id]["vec_rank"] = rank + 1
|
||||
else:
|
||||
score_map[doc_id] = {
|
||||
"content": item["content"],
|
||||
"score": 1.0 / (rrf_k + rank + 1),
|
||||
"vec_score": item["score"],
|
||||
"vec_rank": rank + 1,
|
||||
"metadata": item.get("metadata", {}),
|
||||
}
|
||||
|
||||
# 按融合分数排序
|
||||
merged = sorted(score_map.values(), key=lambda x: x["score"], reverse=True)
|
||||
return merged[:top_k]
|
||||
|
||||
|
||||
async def delete_document_embeddings(doc_id: str) -> bool:
|
||||
"""删除文档的所有向量索引"""
|
||||
pool = await get_pg_pool()
|
||||
if pool is None:
|
||||
return False
|
||||
|
||||
try:
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"DELETE FROM document_embeddings WHERE doc_id = $1;",
|
||||
doc_id,
|
||||
)
|
||||
logger.info(f"Deleted embeddings for document {doc_id}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete embeddings: {e}")
|
||||
return False
|
||||
|
||||
|
||||
async def close_pool():
|
||||
"""关闭连接池"""
|
||||
global _pg_pool
|
||||
if _pg_pool is not None:
|
||||
await _pg_pool.close()
|
||||
_pg_pool = None
|
||||
logger.info("PostgreSQL connection pool closed")
|
||||
@@ -81,6 +81,41 @@ PLATFORM_KNOWLEDGE = """# 七虎AI平台能力全览
|
||||
"""
|
||||
|
||||
|
||||
def get_platform_knowledge() -> str:
|
||||
"""获取平台全功能知识库文本,注入到AI对话系统提示中"""
|
||||
return PLATFORM_KNOWLEDGE
|
||||
def get_platform_knowledge(db=None, user_id=None) -> str:
|
||||
"""获取平台全功能知识库文本,注入到AI对话系统提示中。
|
||||
|
||||
若传入 db,则动态追加当前平台已上线的官方技能目录
|
||||
(source=official 且 status=published),确保管理员/前台导入或下架技能后,
|
||||
AI 能及时感知并使用「平台自身已有内容」,而非停留在陈旧硬编码清单(见 #712)。
|
||||
任何异常均安全回退到静态 PLATFORM_KNOWLEDGE。
|
||||
"""
|
||||
if db is None:
|
||||
return PLATFORM_KNOWLEDGE
|
||||
try:
|
||||
from .models import MarketApp
|
||||
|
||||
apps = (
|
||||
db.query(MarketApp)
|
||||
.filter(MarketApp.source == "official", MarketApp.status == "published")
|
||||
.order_by(MarketApp.capability, MarketApp.name)
|
||||
.all()
|
||||
)
|
||||
if not apps:
|
||||
return PLATFORM_KNOWLEDGE
|
||||
|
||||
lines = [
|
||||
"",
|
||||
"## 平台已上线技能(官方,优先推荐)",
|
||||
"以下为平台当前已发布上线的技能/应用,回答用户相关问题时**优先推荐使用**这些平台能力:",
|
||||
]
|
||||
for a in apps:
|
||||
name = getattr(a, "name", None) or "未命名"
|
||||
cap = getattr(a, "capability", None) or getattr(a, "category", None) or "通用"
|
||||
desc = (getattr(a, "description", "") or "").strip().replace("\n", " ")
|
||||
if len(desc) > 60:
|
||||
desc = desc[:60] + "…"
|
||||
lines.append(f"- **{name}**({cap}): {desc}")
|
||||
return PLATFORM_KNOWLEDGE + "\n" + "\n".join(lines)
|
||||
except Exception as e:
|
||||
logger.warning(f"[platform_knowledge] 动态构建失败,回退静态知识库: {e}")
|
||||
return PLATFORM_KNOWLEDGE
|
||||
Reference in new issue
Block a user