adai/MemoryBear
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'release/v0.3.1' into develop

Browse Source
This commit is contained in:
Ke Sun
2026-04-23 12:16:57 +08:00
parent 27f6d18a05 09393b2326
commit b8009074d5
54 changed files with 1240 additions and 707 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
api/app/core/api_key_auth.py
View File

@@ -96,40 +96,8 @@ def require_api_key(
resource_id=api_key_obj.resource_id,
)
# ── Tenant 级别限速（来自套餐配额 api_ops_rate_limit）──────────
try:
from app.models.workspace_model import Workspace
from premium.platform_admin.package_plan_service import TenantSubscriptionService
workspace = db.query(Workspace).filter(
Workspace.id == api_key_obj.workspace_id
).first()
if workspace:
quota = TenantSubscriptionService(db).get_effective_quota(workspace.tenant_id)
tenant_qps_limit = quota.get("api_ops_rate_limit") if quota else None
if tenant_qps_limit:
rate_limiter = RateLimiterService()
tenant_ok, tenant_info = await rate_limiter.check_tenant_rate_limit(
workspace.tenant_id, tenant_qps_limit
)
if not tenant_ok:
raise RateLimitException(
"租户 API 调用速率超限",
BizCode.API_KEY_QPS_LIMIT_EXCEEDED,
rate_headers={
"X-RateLimit-Tenant-Limit": str(tenant_info["limit"]),
"X-RateLimit-Tenant-Remaining": str(tenant_info["remaining"]),
"X-RateLimit-Tenant-Reset": str(tenant_info["reset"]),
}
)
except RateLimitException:
raise
except Exception as e:
logger.warning(f"Tenant 限速检查异常，跳过: {e}")
# ─────────────────────────────────────────────────────────────
rate_limiter = RateLimiterService()
is_allowed, error_msg, rate_headers = await rate_limiter.check_all_limits(api_key_obj)
is_allowed, error_msg, rate_headers = await rate_limiter.check_all_limits(api_key_obj, db=db)
if not is_allowed:
logger.warning("API Key 限流触发", extra={
"api_key_id": str(api_key_obj.id),
@@ -138,10 +106,12 @@ def require_api_key(
"error_msg": error_msg
})
# 根据错误消息判断限流类型
if "QPS" in error_msg:
code = BizCode.API_KEY_QPS_LIMIT_EXCEEDED
elif "Daily" in error_msg:
if "Daily" in error_msg:
code = BizCode.API_KEY_DAILY_LIMIT_EXCEEDED
elif "Tenant" in error_msg:
code = BizCode.API_KEY_QPS_LIMIT_EXCEEDED # 租户套餐速率超限，同属 QPS 类
elif "QPS" in error_msg:
code = BizCode.API_KEY_QPS_LIMIT_EXCEEDED
else:
code = BizCode.API_KEY_QUOTA_EXCEEDED

23
api/app/core/error_codes.py
View File

@@ -31,6 +31,9 @@ class BizCode(IntEnum):
API_KEY_QPS_LIMIT_EXCEEDED = 3014
API_KEY_DAILY_LIMIT_EXCEEDED = 3015
API_KEY_QUOTA_EXCEEDED = 3016
API_KEY_RATE_LIMIT_EXCEEDED = 3017
QUOTA_EXCEEDED = 3018
RATE_LIMIT_EXCEEDED = 3019
# 资源（4xxx）
NOT_FOUND = 4000
USER_NOT_FOUND = 4001
@@ -155,7 +158,8 @@ HTTP_MAPPING = {
BizCode.API_KEY_QPS_LIMIT_EXCEEDED: 429,
BizCode.API_KEY_DAILY_LIMIT_EXCEEDED: 429,
BizCode.API_KEY_QUOTA_EXCEEDED: 429,
BizCode.QUOTA_EXCEEDED: 402,
BizCode.MODEL_CONFIG_INVALID: 400,
BizCode.API_KEY_MISSING: 400,
BizCode.PROVIDER_NOT_SUPPORTED: 400,
@@ -184,4 +188,21 @@ HTTP_MAPPING = {
BizCode.DB_ERROR: 500,
BizCode.SERVICE_UNAVAILABLE: 503,
BizCode.RATE_LIMITED: 429,
BizCode.RATE_LIMIT_EXCEEDED: 429,
}
ERROR_CODE_TO_BIZ_CODE = {
"QUOTA_EXCEEDED": BizCode.QUOTA_EXCEEDED,
"RATE_LIMIT_EXCEEDED": BizCode.RATE_LIMIT_EXCEEDED,
"API_KEY_NOT_FOUND": BizCode.API_KEY_NOT_FOUND,
"API_KEY_INVALID": BizCode.API_KEY_INVALID,
"API_KEY_EXPIRED": BizCode.API_KEY_EXPIRED,
"WORKSPACE_NOT_FOUND": BizCode.WORKSPACE_NOT_FOUND,
"WORKSPACE_NO_ACCESS": BizCode.WORKSPACE_NO_ACCESS,
"PERMISSION_DENIED": BizCode.PERMISSION_DENIED,
"TOKEN_EXPIRED": BizCode.TOKEN_EXPIRED,
"TOKEN_INVALID": BizCode.TOKEN_INVALID,
"VALIDATION_FAILED": BizCode.VALIDATION_FAILED,
"INVALID_PARAMETER": BizCode.INVALID_PARAMETER,
"MISSING_PARAMETER": BizCode.MISSING_PARAMETER,
}

542
api/app/core/quota_manager.py
View File

@@ -6,7 +6,6 @@
2. 降级到 default_free_plan.py 配置文件（社区版兜底）
"""
import asyncio
import time
from functools import wraps
from typing import Optional, Callable, Dict, Any
from uuid import UUID
@@ -15,10 +14,13 @@ from sqlalchemy import func
from sqlalchemy.orm import Session
from app.core.logging_config import get_auth_logger
from app.i18n.exceptions import QuotaExceededError
from app.i18n.exceptions import QuotaExceededError, InternalServerError
logger = get_auth_logger()
# Redis key 格式常量，与 RateLimiterService.check_qps 保持一致（per api_key 独立计数）
API_KEY_QPS_REDIS_KEY = "rate_limit:qps:{api_key_id}"
def _get_user_from_kwargs(kwargs: dict):
"""从 kwargs 中获取 user 对象"""
@@ -28,6 +30,29 @@ def _get_user_from_kwargs(kwargs: dict):
return None
def _get_workspace_id_from_kwargs(kwargs: dict):
"""从 kwargs 中获取 workspace_id"""
# 优先从 kwargs['workspace_id'] 获取
workspace_id = kwargs.get("workspace_id")
if workspace_id:
return workspace_id
# 从 api_key_auth.workspace_id 获取（API Key 认证场景）
api_key_auth = kwargs.get("api_key_auth")
if api_key_auth and hasattr(api_key_auth, 'workspace_id'):
return api_key_auth.workspace_id
# 从 user.current_workspace_id 获取
user = _get_user_from_kwargs(kwargs)
if user:
ws_id = getattr(user, 'current_workspace_id', None)
if ws_id:
return ws_id
logger.warning(f"无法获取 workspace_id, kwargs keys: {list(kwargs.keys())}")
return None
def _get_tenant_id_from_kwargs(db: Session, kwargs: dict):
"""从 kwargs 中获取 tenant_id"""
user = _get_user_from_kwargs(kwargs)
@@ -65,7 +90,9 @@ def _get_tenant_id_from_kwargs(db: Session, kwargs: dict):
if share_record:
app = db.query(App).filter(App.id == share_record.app_id, App.is_active.is_(True)).first()
if app:
return app.workspace.tenant_id
workspace = db.query(Workspace).filter(Workspace.id == app.workspace_id).first()
if workspace:
return workspace.tenant_id
return None
@@ -73,31 +100,52 @@ def _get_tenant_id_from_kwargs(db: Session, kwargs: dict):
def _get_quota_config(db: Session, tenant_id: UUID) -> Optional[Dict[str, Any]]:
"""
获取租户的配额配置
优先级：
1. premium 模块的 tenant_subscriptions（SaaS 版）
2. default_free_plan.py 配置文件（社区版兜底）
"""
# 尝试从 premium 模块获取
# 尝试从 premium 模块获取（SaaS 版）
try:
from premium.platform_admin.package_plan_service import TenantSubscriptionService
# premium 模块存在，运行时错误不应被静默降级，直接抛出
quota_config = TenantSubscriptionService(db).get_effective_quota(tenant_id)
if quota_config:
logger.debug(f"从 premium 模块获取租户 {tenant_id} 配额配置")
return quota_config
except (ModuleNotFoundError, ImportError, Exception) as e:
logger.debug(f"无法从 premium 模块获取配额配置: {e}")
# premium 存在但该租户无订阅记录，降级到免费套餐
logger.debug(f"租户 {tenant_id} 无 premium 订阅，降级到免费套餐")
except (ModuleNotFoundError, ImportError):
# 社区版：premium 包不存在，正常降级
logger.debug("premium 模块不存在，使用社区版免费套餐配额")
# 降级到配置文件
# 降级到社区版配置文件
try:
from app.config.default_free_plan import DEFAULT_FREE_PLAN
logger.info(f"使用配置文件中的免费套餐配额: tenant={tenant_id}")
logger.debug(f"使用社区版免费套餐配额: tenant={tenant_id}")
return DEFAULT_FREE_PLAN.get("quotas")
except Exception as e:
logger.error(f"无法从配置文件获取配额: {e}")
return None
def get_api_ops_rate_limit(db: Session, tenant_id: UUID) -> Optional[int]:
"""
获取租户套餐的 API 操作速率限制（QPS 上限）
该函数兼容社区版和 SaaS 版：
- SaaS 版：从 premium 模块的套餐配额读取
- 社区版：从 default_free_plan.py 配置文件读取
Returns:
int: api_ops_rate_limit 值，如果未配置则返回 None
"""
quota_config = _get_quota_config(db, tenant_id)
if quota_config:
return quota_config.get("api_ops_rate_limit")
return None
class QuotaUsageRepository:
"""配额使用量数据访问层"""
@@ -111,15 +159,19 @@ class QuotaUsageRepository:
Workspace.is_active.is_(True)
).count()
def count_apps(self, tenant_id: UUID) -> int:
def count_apps(self, tenant_id: UUID, workspace_id: Optional[UUID] = None) -> int:
from app.models.app_model import App
from app.models.workspace_model import Workspace
return self.db.query(App).join(
query = self.db.query(App).join(
Workspace, App.workspace_id == Workspace.id
).filter(
Workspace.tenant_id == tenant_id,
App.is_active.is_(True)
).count()
)
if workspace_id:
query = query.filter(App.workspace_id == workspace_id)
else:
query = query.filter(Workspace.tenant_id == tenant_id)
return query.count()
def count_skills(self, tenant_id: UUID) -> int:
from app.models.skill_model import Skill
@@ -128,55 +180,76 @@ class QuotaUsageRepository:
Skill.is_active.is_(True)
).count()
def sum_knowledge_capacity_gb(self, tenant_id: UUID) -> float:
def sum_knowledge_capacity_gb(self, tenant_id: UUID, workspace_id: Optional[UUID] = None) -> float:
from app.models.document_model import Document
from app.models.knowledge_model import Knowledge
from app.models.workspace_model import Workspace
result = self.db.query(func.coalesce(func.sum(Document.file_size), 0)).join(
query = self.db.query(func.coalesce(func.sum(Document.file_size), 0)).join(
Knowledge, Document.kb_id == Knowledge.id
).join(
Workspace, Knowledge.workspace_id == Workspace.id
).filter(
Workspace.tenant_id == tenant_id,
Document.status == 1,
).scalar()
)
if workspace_id:
query = query.filter(Knowledge.workspace_id == workspace_id)
else:
query = query.filter(Workspace.tenant_id == tenant_id)
result = query.scalar()
return float(result) / (1024 ** 3) if result else 0.0
def count_memory_engines(self, tenant_id: UUID) -> int:
def count_memory_engines(self, tenant_id: UUID, workspace_id: Optional[UUID] = None) -> int:
from app.models.memory_config_model import MemoryConfig
from app.models.workspace_model import Workspace
return self.db.query(MemoryConfig).join(
query = self.db.query(MemoryConfig).join(
Workspace, MemoryConfig.workspace_id == Workspace.id
).filter(
Workspace.tenant_id == tenant_id
).count()
)
if workspace_id:
query = query.filter(MemoryConfig.workspace_id == workspace_id)
else:
query = query.filter(Workspace.tenant_id == tenant_id)
return query.count()
def count_end_users(self, tenant_id: UUID) -> int:
def count_end_users(self, tenant_id: UUID, workspace_id: Optional[UUID] = None) -> int:
from app.models.end_user_model import EndUser
from app.models.workspace_model import Workspace
return self.db.query(EndUser).join(
from app.models.user_model import User
query = self.db.query(EndUser).join(
Workspace, EndUser.workspace_id == Workspace.id
).filter(
Workspace.tenant_id == tenant_id
).count()
)
if workspace_id:
query = query.filter(EndUser.workspace_id == workspace_id)
else:
query = query.filter(Workspace.tenant_id == tenant_id)
trial_user_ids = [
str(u.id) for u in self.db.query(User.id).filter(User.tenant_id == tenant_id).all()
]
if trial_user_ids:
query = query.filter(~EndUser.other_id.in_(trial_user_ids))
return query.count()
def count_models(self, tenant_id: UUID) -> int:
from app.models.models_model import ModelConfig
return self.db.query(ModelConfig).filter(
ModelConfig.tenant_id == tenant_id,
ModelConfig.is_active == True
ModelConfig.is_active == True,
ModelConfig.is_composite == True
).count()
def count_ontology_projects(self, tenant_id: UUID) -> int:
def count_ontology_projects(self, tenant_id: UUID, workspace_id: Optional[UUID] = None) -> int:
from app.models.ontology_scene import OntologyScene
from app.models.workspace_model import Workspace
if workspace_id:
return self.db.query(OntologyScene).filter(
OntologyScene.workspace_id == workspace_id
).count()
return self.db.query(OntologyScene).join(
Workspace, OntologyScene.workspace_id == Workspace.id
).filter(
Workspace.tenant_id == tenant_id
).count()
def get_usage_by_quota_type(self, tenant_id: UUID, quota_type: str):
def get_usage_by_quota_type(self, tenant_id: UUID, quota_type: str, workspace_id: Optional[UUID] = None):
"""按配额类型分发，返回当前使用量"""
dispatch = {
"workspace_quota": self.count_workspaces,
@@ -189,6 +262,8 @@ class QuotaUsageRepository:
"ontology_project_quota": self.count_ontology_projects,
}
fn = dispatch.get(quota_type)
if workspace_id:
return fn(tenant_id, workspace_id) if fn else 0
return fn(tenant_id) if fn else 0
@@ -198,6 +273,7 @@ def _check_quota(
quota_type: str,
resource_name: str,
usage_func: Optional[Callable] = None,
workspace_id: Optional[UUID] = None,
) -> None:
"""核心配额检查逻辑：对比使用量和配额限制"""
try:
@@ -212,13 +288,13 @@ def _check_quota(
return
if usage_func:
current_usage = usage_func(db, tenant_id)
current_usage = usage_func(db, tenant_id, workspace_id) if workspace_id else usage_func(db, tenant_id)
else:
current_usage = QuotaUsageRepository(db).get_usage_by_quota_type(tenant_id, quota_type)
current_usage = QuotaUsageRepository(db).get_usage_by_quota_type(tenant_id, quota_type, workspace_id)
if current_usage >= quota_limit:
logger.warning(
f"配额不足: tenant={tenant_id}, type={quota_type}, "
f"配额不足: tenant={tenant_id}, workspace={workspace_id}, type={quota_type}, "
f"usage={current_usage}, limit={quota_limit}"
)
raise QuotaExceededError(
@@ -228,7 +304,7 @@ def _check_quota(
)
logger.debug(
f"配额检查通过: tenant={tenant_id}, type={quota_type}, "
f"配额检查通过: tenant={tenant_id}, workspace={workspace_id}, type={quota_type}, "
f"usage={current_usage}, limit={quota_limit}"
)
@@ -236,7 +312,7 @@ def _check_quota(
raise
except Exception as e:
logger.error(
f"配额检查异常: tenant={tenant_id}, type={quota_type}, "
f"配额检查异常: tenant={tenant_id}, workspace={workspace_id}, type={quota_type}, "
f"error_type={type(e).__name__}, error={str(e)}",
exc_info=True,
)
@@ -247,41 +323,82 @@ def _check_quota(
def check_workspace_quota(func: Callable) -> Callable:
@wraps(func)
def wrapper(*args, **kwargs):
async def async_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.warning("配额检查失败：缺少 db 或 user 参数")
return func(*args, **kwargs)
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, "workspace_quota", "workspace")
return await func(*args, **kwargs)
@wraps(func)
def sync_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, "workspace_quota", "workspace")
return func(*args, **kwargs)
return wrapper
return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
def check_skill_quota(func: Callable) -> Callable:
@wraps(func)
def wrapper(*args, **kwargs):
async def async_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.warning("配额检查失败：缺少 db 或 user 参数")
return func(*args, **kwargs)
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, "skill_quota", "skill")
return await func(*args, **kwargs)
@wraps(func)
def sync_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, "skill_quota", "skill")
return func(*args, **kwargs)
return wrapper
return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
def check_app_quota(func: Callable) -> Callable:
@wraps(func)
def wrapper(*args, **kwargs):
async def async_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.warning("配额检查失败：缺少 db 或 user 参数")
return func(*args, **kwargs)
_check_quota(db, user.tenant_id, "app_quota", "app")
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
workspace_id = _get_workspace_id_from_kwargs(kwargs)
if not workspace_id:
logger.error(f"配额检查失败：{func.__name__} 无法获取 workspace_id，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, "app_quota", "app", workspace_id=workspace_id)
return await func(*args, **kwargs)
@wraps(func)
def sync_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
workspace_id = _get_workspace_id_from_kwargs(kwargs)
if not workspace_id:
logger.error(f"配额检查失败：{func.__name__} 无法获取 workspace_id，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, "app_quota", "app", workspace_id=workspace_id)
return func(*args, **kwargs)
return wrapper
return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
def check_knowledge_capacity_quota(func: Callable) -> Callable:
@@ -289,13 +406,17 @@ def check_knowledge_capacity_quota(func: Callable) -> Callable:
async def async_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
if not db:
logger.warning("配额检查失败：缺少 db 参数")
return await func(*args, **kwargs)
logger.error(f"配额检查失败：{func.__name__} 缺少 db 参数，拒绝请求")
raise InternalServerError()
tenant_id = _get_tenant_id_from_kwargs(db, kwargs)
if not tenant_id:
logger.warning("配额检查失败：无法获取 tenant_id")
return await func(*args, **kwargs)
_check_quota(db, tenant_id, "knowledge_capacity_quota", "knowledge_capacity")
logger.error(f"配额检查失败：{func.__name__} 无法获取 tenant_id，拒绝请求")
raise InternalServerError()
workspace_id = _get_workspace_id_from_kwargs(kwargs)
if not workspace_id:
logger.error(f"配额检查失败：{func.__name__} 无法获取 workspace_id，拒绝请求")
raise InternalServerError()
_check_quota(db, tenant_id, "knowledge_capacity_quota", "knowledge_capacity", workspace_id=workspace_id)
return await func(*args, **kwargs)
@wraps(func)
@@ -303,9 +424,13 @@ def check_knowledge_capacity_quota(func: Callable) -> Callable:
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.warning("配额检查失败：缺少 db 或 user 参数")
return func(*args, **kwargs)
_check_quota(db, user.tenant_id, "knowledge_capacity_quota", "knowledge_capacity")
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
workspace_id = _get_workspace_id_from_kwargs(kwargs)
if not workspace_id:
logger.error(f"配额检查失败：{func.__name__} 无法获取 workspace_id，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, "knowledge_capacity_quota", "knowledge_capacity", workspace_id=workspace_id)
return func(*args, **kwargs)
return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
@@ -313,15 +438,36 @@ def check_knowledge_capacity_quota(func: Callable) -> Callable:
def check_memory_engine_quota(func: Callable) -> Callable:
@wraps(func)
def wrapper(*args, **kwargs):
async def async_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
logger.debug(f"check_memory_engine_quota async_wrapper: db={db is not None}, user={user}, kwargs_keys={list(kwargs.keys())}")
if not db or not user:
logger.warning("配额检查失败：缺少 db 或 user 参数")
return func(*args, **kwargs)
_check_quota(db, user.tenant_id, "memory_engine_quota", "memory_engine")
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
workspace_id = _get_workspace_id_from_kwargs(kwargs)
if not workspace_id:
logger.error(f"配额检查失败：{func.__name__} 无法获取 workspace_id，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, "memory_engine_quota", "memory_engine", workspace_id=workspace_id)
return await func(*args, **kwargs)
@wraps(func)
def sync_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
logger.debug(f"check_memory_engine_quota sync_wrapper: db={db is not None}, user={user}, kwargs_keys={list(kwargs.keys())}")
if not db or not user:
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
workspace_id = _get_workspace_id_from_kwargs(kwargs)
if not workspace_id:
logger.error(f"配额检查失败：{func.__name__} 无法获取 workspace_id，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, "memory_engine_quota", "memory_engine", workspace_id=workspace_id)
return func(*args, **kwargs)
return wrapper
return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
def check_end_user_quota(func: Callable) -> Callable:
@@ -329,26 +475,34 @@ def check_end_user_quota(func: Callable) -> Callable:
async def async_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
if not db:
logger.warning("配额检查失败：缺少 db 参数")
return await func(*args, **kwargs)
logger.error(f"配额检查失败：{func.__name__} 缺少 db 参数，拒绝请求")
raise InternalServerError()
tenant_id = _get_tenant_id_from_kwargs(db, kwargs)
if not tenant_id:
logger.warning("配额检查失败：无法获取 tenant_id")
return await func(*args, **kwargs)
_check_quota(db, tenant_id, "end_user_quota", "end_user")
logger.error(f"配额检查失败：{func.__name__} 无法获取 tenant_id，拒绝请求")
raise InternalServerError()
workspace_id = _get_workspace_id_from_kwargs(kwargs)
if not workspace_id:
logger.error(f"配额检查失败：{func.__name__} 无法获取 workspace_id，拒绝请求")
raise InternalServerError()
_check_quota(db, tenant_id, "end_user_quota", "end_user", workspace_id=workspace_id)
return await func(*args, **kwargs)
@wraps(func)
def sync_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
if not db:
logger.warning("配额检查失败：缺少 db 参数")
return func(*args, **kwargs)
logger.error(f"配额检查失败：{func.__name__} 缺少 db 参数，拒绝请求")
raise InternalServerError()
tenant_id = _get_tenant_id_from_kwargs(db, kwargs)
if not tenant_id:
logger.warning("配额检查失败：无法获取 tenant_id")
return func(*args, **kwargs)
_check_quota(db, tenant_id, "end_user_quota", "end_user")
logger.error(f"配额检查失败：{func.__name__} 无法获取 tenant_id，拒绝请求")
raise InternalServerError()
workspace_id = _get_workspace_id_from_kwargs(kwargs)
if not workspace_id:
logger.error(f"配额检查失败：{func.__name__} 无法获取 workspace_id，拒绝请求")
raise InternalServerError()
_check_quota(db, tenant_id, "end_user_quota", "end_user", workspace_id=workspace_id)
return func(*args, **kwargs)
return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
@@ -356,89 +510,171 @@ def check_end_user_quota(func: Callable) -> Callable:
def check_ontology_project_quota(func: Callable) -> Callable:
@wraps(func)
def wrapper(*args, **kwargs):
async def async_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.warning("配额检查失败：缺少 db 或 user 参数")
return func(*args, **kwargs)
_check_quota(db, user.tenant_id, "ontology_project_quota", "ontology_project")
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
workspace_id = _get_workspace_id_from_kwargs(kwargs)
if not workspace_id:
logger.error(f"配额检查失败：{func.__name__} 无法获取 workspace_id，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, "ontology_project_quota", "ontology_project", workspace_id=workspace_id)
return await func(*args, **kwargs)
@wraps(func)
def sync_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
workspace_id = _get_workspace_id_from_kwargs(kwargs)
if not workspace_id:
logger.error(f"配额检查失败：{func.__name__} 无法获取 workspace_id，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, "ontology_project_quota", "ontology_project", workspace_id=workspace_id)
return func(*args, **kwargs)
return wrapper
return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
def check_model_quota(func: Callable) -> Callable:
@wraps(func)
def wrapper(*args, **kwargs):
async def async_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.warning("配额检查失败：缺少 db 或 user 参数")
return func(*args, **kwargs)
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, "model_quota", "model")
return await func(*args, **kwargs)
@wraps(func)
def sync_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, "model_quota", "model")
return func(*args, **kwargs)
return wrapper
return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
def check_model_activation_quota(func: Callable) -> Callable:
"""模型激活时的配额检查装饰器"""
@wraps(func)
def wrapper(*args, **kwargs):
async def async_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.warning("配额检查失败：缺少 db 或 user 参数")
return func(*args, **kwargs)
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
model_id = kwargs.get("model_id") or (args[1] if len(args) > 1 else None)
model_data = kwargs.get("model_data")
if not model_id or not model_data:
logger.warning("模型激活配额检查失败：缺少 model_id 或 model_data 参数")
return func(*args, **kwargs)
if model_data.is_active is True:
return await func(*args, **kwargs)
if model_data.is_active:
try:
from app.models.models_model import ModelConfig
from app.services.model_service import ModelConfigService
existing_model = ModelConfigService.get_model_by_id(
db=db,
model_id=model_id,
db=db,
model_id=model_id,
tenant_id=user.tenant_id
)
if not existing_model.is_active:
logger.info(f"模型激活操作，检查配额: model_id={model_id}, tenant_id={user.tenant_id}")
_check_quota(db, user.tenant_id, "model_quota", "model")
except Exception as e:
logger.error(f"模型激活配额检查异常: model_id={model_id}, error={str(e)}")
raise
return await func(*args, **kwargs)
@wraps(func)
def sync_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
model_id = kwargs.get("model_id") or (args[1] if len(args) > 1 else None)
model_data = kwargs.get("model_data")
if not model_id or not model_data:
logger.warning("模型激活配额检查失败：缺少 model_id 或 model_data 参数")
return func(*args, **kwargs)
if model_data.is_active:
try:
from app.services.model_service import ModelConfigService
existing_model = ModelConfigService.get_model_by_id(
db=db,
model_id=model_id,
tenant_id=user.tenant_id
)
if not existing_model.is_active:
logger.info(f"模型激活操作，检查配额: model_id={model_id}, tenant_id={user.tenant_id}")
_check_quota(db, user.tenant_id, "model_quota", "model")
except Exception as e:
logger.error(f"模型激活配额检查异常: model_id={model_id}, error={str(e)}")
raise
return func(*args, **kwargs)
return wrapper
return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
def check_quota(quota_type: str, resource_name: str, usage_func: Optional[Callable] = None):
"""通用配额检查装饰器，支持自定义使用量获取函数"""
def decorator(func: Callable) -> Callable:
@wraps(func)
def wrapper(*args, **kwargs):
async def async_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.warning("配额检查失败：缺少 db 或 user 参数")
return func(*args, **kwargs)
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, quota_type, resource_name, usage_func)
return await func(*args, **kwargs)
@wraps(func)
def sync_wrapper(*args, **kwargs):
db: Session = kwargs.get("db")
user = _get_user_from_kwargs(kwargs)
if not db or not user:
logger.error(f"配额检查失败：{func.__name__} 缺少 db 或 user 参数，拒绝请求")
raise InternalServerError()
_check_quota(db, user.tenant_id, quota_type, resource_name, usage_func)
return func(*args, **kwargs)
return wrapper
return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
return decorator
# ─── 配额使用统计 ────────────────────────────────────────────────────────────
def get_quota_usage(db: Session, tenant_id: UUID) -> dict:
"""获取租户所有配额的使用情况"""
async def get_quota_usage(db: Session, tenant_id: UUID) -> dict:
"""获取租户所有配额的使用情况
对于 workspace 级别的配额（app/knowledge_capacity/memory_engine/end_user）：
- used: 租户汇总（所有空间加总）
- limit: quota × 活跃工作区数（有效总限额，使汇总数据自洽）
- per_workspace: 各空间明细，包含 workspace_id、workspace_name、used、limit、percentage
- 配额检查逻辑不变：仍按单个空间独立检查
"""
quota_config = _get_quota_config(db, tenant_id)
if not quota_config:
return {}
@@ -457,29 +693,99 @@ def get_quota_usage(db: Session, tenant_id: UUID) -> dict:
model_count = repo.count_models(tenant_id)
ontology_count = repo.count_ontology_projects(tenant_id)
# 获取租户下所有活跃工作区，用于按空间拆分明细
from app.models.workspace_model import Workspace
active_workspaces = db.query(Workspace).filter(
Workspace.tenant_id == tenant_id,
Workspace.is_active.is_(True)
).all()
# 构建各空间的 workspace 级配额明细
def _build_per_workspace_detail(count_func, per_unit_limit):
"""为 workspace 级配额构建 per_workspace 明细列表"""
if not per_unit_limit or not active_workspaces:
return []
details = []
for ws in active_workspaces:
ws_used = count_func(tenant_id, ws.id)
details.append({
"workspace_id": str(ws.id),
"workspace_name": ws.name,
"used": ws_used,
"limit": per_unit_limit,
"percentage": pct(ws_used, per_unit_limit),
})
return details
# workspace 级配额的每空间限额
app_quota_per_ws = quota_config.get("app_quota")
knowledge_quota_per_ws = quota_config.get("knowledge_capacity_quota")
memory_quota_per_ws = quota_config.get("memory_engine_quota")
end_user_quota_per_ws = quota_config.get("end_user_quota")
ontology_quota_per_ws = quota_config.get("ontology_project_quota")
# workspace 级配额的有效总限额 = 每空间限额 × 活跃工作区数
app_effective_limit = app_quota_per_ws * workspace_count if app_quota_per_ws is not None and workspace_count > 0 else app_quota_per_ws
knowledge_effective_limit = knowledge_quota_per_ws * workspace_count if knowledge_quota_per_ws is not None and workspace_count > 0 else knowledge_quota_per_ws
memory_effective_limit = memory_quota_per_ws * workspace_count if memory_quota_per_ws is not None and workspace_count > 0 else memory_quota_per_ws
end_user_effective_limit = end_user_quota_per_ws * workspace_count if end_user_quota_per_ws is not None and workspace_count > 0 else end_user_quota_per_ws
ontology_effective_limit = ontology_quota_per_ws * workspace_count if ontology_quota_per_ws is not None and workspace_count > 0 else ontology_quota_per_ws
api_ops_current = 0
try:
from app.core.config import settings
import redis
_now = time.time()
_rk = f"rate_limit:tenant_qps:{tenant_id}"
_r = redis.StrictRedis(
host=settings.REDIS_HOST, port=settings.REDIS_PORT,
db=settings.REDIS_DB, password=settings.REDIS_PASSWORD,
decode_responses=True
)
api_ops_current = int(_r.zcount(_rk, _now - 1, "+inf"))
except Exception:
pass
from app.aioRedis import aio_redis as _aio_redis
from app.models.api_key_model import ApiKey
# api_ops_rate_limit 限的是每个 api_key 每秒最高限额
# 展示当前最接近触发限流的 key 的 QPS（取最大值）
api_key_ids = db.query(ApiKey.id).join(
Workspace, ApiKey.workspace_id == Workspace.id
).filter(
Workspace.tenant_id == tenant_id,
ApiKey.is_active.is_(True)
).all()
for (key_id,) in api_key_ids:
_rk = API_KEY_QPS_REDIS_KEY.format(api_key_id=key_id)
val = await _aio_redis.get(_rk)
count = int(val) if val else 0
if count > api_ops_current:
api_ops_current = count
except Exception as e:
logger.warning(f"获取 api_ops_current 失败，返回 0: {type(e).__name__}: {e}")
return {
"workspace": {"used": workspace_count, "limit": quota_config.get("workspace_quota"), "percentage": pct(workspace_count, quota_config.get("workspace_quota"))},
"skill": {"used": skill_count, "limit": quota_config.get("skill_quota"), "percentage": pct(skill_count, quota_config.get("skill_quota"))},
"app": {"used": app_count, "limit": quota_config.get("app_quota"), "percentage": pct(app_count, quota_config.get("app_quota"))},
"knowledge_capacity": {"used": round(knowledge_gb, 2), "limit": quota_config.get("knowledge_capacity_quota"), "percentage": pct(knowledge_gb, quota_config.get("knowledge_capacity_quota")), "unit": "GB"},
"memory_engine": {"used": memory_count, "limit": quota_config.get("memory_engine_quota"), "percentage": pct(memory_count, quota_config.get("memory_engine_quota"))},
"end_user": {"used": end_user_count, "limit": quota_config.get("end_user_quota"), "percentage": pct(end_user_count, quota_config.get("end_user_quota"))},
"ontology_project": {"used": ontology_count, "limit": quota_config.get("ontology_project_quota"), "percentage": pct(ontology_count, quota_config.get("ontology_project_quota"))},
"app": {
"used": app_count,
"limit": app_effective_limit,
"percentage": pct(app_count, app_effective_limit),
"per_workspace": _build_per_workspace_detail(repo.count_apps, app_quota_per_ws),
},
"knowledge_capacity": {
"used": round(knowledge_gb, 2),
"limit": knowledge_effective_limit,
"percentage": pct(knowledge_gb, knowledge_effective_limit),
"unit": "GB",
"per_workspace": _build_per_workspace_detail(repo.sum_knowledge_capacity_gb, knowledge_quota_per_ws),
},
"memory_engine": {
"used": memory_count,
"limit": memory_effective_limit,
"percentage": pct(memory_count, memory_effective_limit),
"per_workspace": _build_per_workspace_detail(repo.count_memory_engines, memory_quota_per_ws),
},
"end_user": {
"used": end_user_count,
"limit": end_user_effective_limit,
"percentage": pct(end_user_count, end_user_effective_limit),
"per_workspace": _build_per_workspace_detail(repo.count_end_users, end_user_quota_per_ws),
},
"ontology_project": {
"used": ontology_count,
"limit": ontology_effective_limit,
"percentage": pct(ontology_count, ontology_effective_limit),
"per_workspace": _build_per_workspace_detail(repo.count_ontology_projects, ontology_quota_per_ws),
},
"model": {"used": model_count, "limit": quota_config.get("model_quota"), "percentage": pct(model_count, quota_config.get("model_quota"))},
"api_ops_rate_limit": {"current": api_ops_current, "limit": quota_config.get("api_ops_rate_limit"), "percentage": None, "unit": "次/秒"},
}

2
api/app/core/quota_stub.py
View File

@@ -18,6 +18,7 @@ from app.core.quota_manager import (
get_quota_usage,
_check_quota,
QuotaUsageRepository,
API_KEY_QPS_REDIS_KEY,
)
__all__ = [
@@ -33,4 +34,5 @@ __all__ = [
"get_quota_usage",
"_check_quota",
"QuotaUsageRepository",
"API_KEY_QPS_REDIS_KEY",
]

8
api/app/core/rag/common/connection_utils.py
View File

@@ -33,18 +33,16 @@ def timeout(seconds: float | int | str = None, attempts: int = 2, *, exception:
thread.daemon = True
thread.start()
effective_timeout = seconds if seconds else 120 # 默认 120 秒超时
for a in range(attempts):
try:
if os.environ.get("ENABLE_TIMEOUT_ASSERTION"):
result = result_queue.get(timeout=seconds)
else:
result = result_queue.get()
result = result_queue.get(timeout=effective_timeout)
if isinstance(result, Exception):
raise result
return result
except queue.Empty:
pass
raise TimeoutError(f"Function '{func.__name__}' timed out after {seconds} seconds and {attempts} attempts.")
raise TimeoutError(f"Function '{func.__name__}' timed out after {effective_timeout} seconds and {attempts} attempts.")
@wraps(func)
async def async_wrapper(*args, **kwargs) -> Any:

2
api/app/core/rag/nlp/search.py
View File

@@ -113,7 +113,7 @@ def knowledge_retrieval(
continue
# Use the specified reranker for re-ranking
if reranker_id:
if reranker_id and all_results:
try:
all_results = rerank(db=db, reranker_id=reranker_id, query=query, docs=all_results, top_k=reranker_top_k)
except Exception as rerank_error:

4
api/app/core/rag/utils/es_conn.py
View File

@@ -68,9 +68,9 @@ class ESConnection(DocStoreConnection):
client_config = {
"hosts": [hosts],
"basic_auth": (os.getenv("ELASTICSEARCH_USERNAME", "elastic"), os.getenv("ELASTICSEARCH_PASSWORD", "elastic")),
"request_timeout": int(os.getenv("ELASTICSEARCH_REQUEST_TIMEOUT", 100000)),
"request_timeout": int(os.getenv("ELASTICSEARCH_REQUEST_TIMEOUT", 30)),
"retry_on_timeout": os.getenv("ELASTICSEARCH_RETRY_ON_TIMEOUT", True) == "true",
"max_retries": int(os.getenv("ELASTICSEARCH_MAX_RETRIES", 10000)),
"max_retries": int(os.getenv("ELASTICSEARCH_MAX_RETRIES", 3)),
}
# Only add SSL settings if using HTTPS

193
api/app/core/rag/vdb/elasticsearch/elasticsearch_vector.py
View File

@@ -1,25 +1,22 @@
import os
import logging
from typing import Any, cast
import threading
from typing import Any
from urllib.parse import urlparse
import uuid
import requests
from elasticsearch import Elasticsearch, helpers
from elasticsearch.helpers import BulkIndexError
from packaging.version import parse as parse_version
from pydantic import BaseModel, model_validator
from abc import ABC
# langchain-community
# langchain-xinference
# from langchain_community.embeddings import XinferenceEmbeddings
# from langchain_xinference import XinferenceRerank
from langchain_core.documents import Document
from app.core.models.base import RedBearModelConfig
from app.core.models import RedBearLLM, RedBearRerank
from app.core.models import RedBearRerank
from app.core.models.embedding import RedBearEmbeddings
from app.models.models_model import ModelConfig, ModelApiKey
from app.services.model_service import ModelConfigService
from app.models.models_model import ModelApiKey
from app.models.knowledge_model import Knowledge
from app.core.rag.vdb.field import Field
@@ -29,37 +26,9 @@ from app.core.rag.models.chunk import DocumentChunk
logger = logging.getLogger(__name__)
class ElasticSearchConfig(BaseModel):
# Regular Elasticsearch config
host: str | None = None
port: int | None = None
username: str | None = None
password: str | None = None
# Common config
ca_certs: str | None = None
verify_certs: bool = False
request_timeout: int = 100000
retry_on_timeout: bool = True
max_retries: int = 10000
@model_validator(mode="before")
@classmethod
def validate_config(cls, values: dict):
# Regular Elasticsearch validation
if not values.get("host"):
raise ValueError("config HOST is required for regular Elasticsearch")
if not values.get("port"):
raise ValueError("config PORT is required for regular Elasticsearch")
if not values.get("username"):
raise ValueError("config USERNAME is required for regular Elasticsearch")
if not values.get("password"):
raise ValueError("config PASSWORD is required for regular Elasticsearch")
return values
class ElasticSearchVector(BaseVector):
def __init__(self, index_name: str, config: ElasticSearchConfig, embedding_config: ModelApiKey, reranker_config: ModelApiKey):
def __init__(self, index_name: str, client: Elasticsearch,
embedding_config: ModelApiKey, reranker_config: ModelApiKey):
super().__init__(index_name.lower())
# 初始化 Embedding 模型（自动支持火山引擎多模态）
@@ -77,58 +46,8 @@ class ElasticSearchVector(BaseVector):
api_key=reranker_config.api_key,
base_url=reranker_config.api_base
))
self._client = self._init_client(config)
self._version = self._get_version()
self._check_version()
def _init_client(self, config: ElasticSearchConfig) -> Elasticsearch:
"""
Initialize Elasticsearch client for regular Elasticsearch.
"""
try:
# Regular Elasticsearch configuration
parsed_url = urlparse(config.host or "")
if parsed_url.scheme in {"http", "https"}:
hosts = f"{config.host}:{config.port}"
use_https = parsed_url.scheme == "https"
else:
hosts = f"https://{config.host}:{config.port}"
use_https = False
client_config = {
"hosts": [hosts],
"basic_auth": (config.username, config.password),
"request_timeout": config.request_timeout,
"retry_on_timeout": config.retry_on_timeout,
"max_retries": config.max_retries,
}
# Only add SSL settings if using HTTPS
if use_https:
client_config["verify_certs"] = config.verify_certs
if config.ca_certs:
client_config["ca_certs"] = config.ca_certs
client = Elasticsearch(**client_config)
# Test connection
if not client.ping():
raise ConnectionError("Failed to connect to Elasticsearch")
except requests.ConnectionError as e:
raise ConnectionError(f"Vector database connection error: {str(e)}")
except Exception as e:
raise ConnectionError(f"Elasticsearch client initialization failed: {str(e)}")
return client
def _get_version(self) -> str:
info = self._client.info()
return cast(str, info["version"]["number"])
def _check_version(self):
if parse_version(self._version) < parse_version("8.0.0"):
raise ValueError("Elasticsearch vector database version must be greater than 8.0.0")
# 使用外部传入的共享客户端
self._client = client
def get_type(self) -> str:
return "elasticsearch"
@@ -745,29 +664,79 @@ class ElasticSearchVector(BaseVector):
class ElasticSearchVectorFactory:
@staticmethod
def init_vector(knowledge: Knowledge) -> ElasticSearchVector:
"""ES 向量服务工厂 - 单例共享连接"""
_client: Elasticsearch | None = None
_lock = threading.Lock()
_version_checked = False
@classmethod
def _get_shared_client(cls) -> Elasticsearch:
"""获取共享的 ES 客户端（线程安全的懒加载单例）"""
if cls._client is not None:
return cls._client
with cls._lock:
# 双重检查，防止并发时重复创建
if cls._client is not None:
return cls._client
try:
parsed_url = urlparse(os.getenv("ELASTICSEARCH_HOST", "127.0.0.1") or "")
if parsed_url.scheme in {"http", "https"}:
hosts = f'{os.getenv("ELASTICSEARCH_HOST")}:{os.getenv("ELASTICSEARCH_PORT", 9200)}'
use_https = parsed_url.scheme == "https"
else:
hosts = f'https://{os.getenv("ELASTICSEARCH_HOST", "127.0.0.1")}:{os.getenv("ELASTICSEARCH_PORT", 9200)}'
use_https = False
client_config = {
"hosts": [hosts],
"basic_auth": (
os.getenv("ELASTICSEARCH_USERNAME", "elastic"),
os.getenv("ELASTICSEARCH_PASSWORD", "elastic"),
),
"request_timeout": int(os.getenv("ELASTICSEARCH_REQUEST_TIMEOUT", 30)),
"retry_on_timeout": True,
"max_retries": int(os.getenv("ELASTICSEARCH_MAX_RETRIES", 3)),
"connections_per_node": int(os.getenv("ELASTICSEARCH_CONNECTIONS_PER_NODE", 10)),
}
if use_https:
client_config["verify_certs"] = os.getenv("ELASTICSEARCH_VERIFY_CERTS", "false") == "true"
ca_certs = os.getenv("ELASTICSEARCH_CA_CERTS")
if ca_certs:
client_config["ca_certs"] = str(ca_certs)
client = Elasticsearch(**client_config)
if not client.ping():
raise ConnectionError("Failed to connect to Elasticsearch")
# 版本检查只做一次
if not cls._version_checked:
info = client.info()
version = info["version"]["number"]
if parse_version(version) < parse_version("8.0.0"):
raise ValueError(f"Elasticsearch version must be >= 8.0.0, got {version}")
cls._version_checked = True
logger.info(f"Elasticsearch shared client initialized, version: {version}")
cls._client = client
except requests.ConnectionError as e:
raise ConnectionError(f"Vector database connection error: {str(e)}")
except Exception as e:
raise ConnectionError(f"Elasticsearch client initialization failed: {str(e)}")
return cls._client
@classmethod
def init_vector(cls, knowledge: Knowledge) -> ElasticSearchVector:
"""创建向量服务实例（共享 ES 连接）"""
client = cls._get_shared_client()
collection_name = f"Vector_index_{knowledge.id}_Node"
# Use regular Elasticsearch with config values
config_dict = {
"host": os.getenv("ELASTICSEARCH_HOST", "127.0.0.1"),
"port": os.getenv("ELASTICSEARCH_PORT", 9200),
"username": os.getenv("ELASTICSEARCH_USERNAME", "elastic"),
"password": os.getenv("ELASTICSEARCH_PASSWORD", "elastic"),
}
# Common configuration
config_dict.update(
{
"ca_certs": str(os.getenv("ELASTICSEARCH_CA_CERTS")) if os.getenv("ELASTICSEARCH_CA_CERTS") else None,
"verify_certs": os.getenv("ELASTICSEARCH_VERIFY_CERTS", False) == "true",
"request_timeout": int(os.getenv("ELASTICSEARCH_REQUEST_TIMEOUT", 100000)),
"retry_on_timeout": os.getenv("ELASTICSEARCH_RETRY_ON_TIMEOUT", True) == "true",
"max_retries": int(os.getenv("ELASTICSEARCH_MAX_RETRIES", 10000)),
}
)
if knowledge.embedding is None:
raise ValueError(f"embedding_id config error: {str(knowledge.embedding_id)}")
if knowledge.reranker is None:
@@ -775,9 +744,9 @@ class ElasticSearchVectorFactory:
return ElasticSearchVector(
index_name=collection_name,
config=ElasticSearchConfig(**config_dict),
client=client,
embedding_config=knowledge.embedding.api_keys[0],
reranker_config=knowledge.reranker.api_keys[0]
reranker_config=knowledge.reranker.api_keys[0],
)

13
api/app/core/workflow/nodes/tool/node.py
View File

@@ -11,6 +11,7 @@ from app.core.workflow.nodes.tool.config import ToolNodeConfig
from app.core.workflow.variable.base_variable import VariableType
from app.db import get_db_read
from app.services.tool_service import ToolService
from app.models.tool_model import ToolType
logger = logging.getLogger(__name__)
@@ -76,6 +77,18 @@ class ToolNode(BaseNode):
# 执行工具
with get_db_read() as db:
tool_service = ToolService(db)
# MCP 工具：将 operation 映射为 tool_name，其余参数包装进 arguments
tool_instance = tool_service.get_tool_instance(self.typed_config.tool_id, tenant_id)
if tool_instance and tool_instance.tool_type == ToolType.MCP:
operation = rendered_parameters.pop("operation", None)
if operation:
old_params = rendered_parameters
rendered_parameters = {
"tool_name": operation,
"arguments": old_params
}
result = await tool_service.execute_tool(
tool_id=self.typed_config.tool_id,
parameters=rendered_parameters,