218 lines
6.8 KiB
Python
218 lines
6.8 KiB
Python
#
|
|
# Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License
|
|
#
|
|
import logging
|
|
from datetime import datetime
|
|
import json
|
|
|
|
from fastapi import APIRouter, Depends
|
|
from fastapi.responses import JSONResponse
|
|
|
|
from api.db.db_models import APIToken
|
|
from api.db.services.api_service import APITokenService
|
|
from api.db.services.knowledgebase_service import KnowledgebaseService
|
|
from api.db.services.user_service import UserTenantService
|
|
from api import settings
|
|
from api.utils import current_timestamp, datetime_format
|
|
from api.utils.api_utils import (
|
|
get_json_result,
|
|
get_data_error_result,
|
|
server_error_response,
|
|
generate_confirmation_token,
|
|
)
|
|
from api.versions import get_ragflow_version
|
|
from rag.utils.storage_factory import STORAGE_IMPL, STORAGE_IMPL_TYPE
|
|
from timeit import default_timer as timer
|
|
|
|
from rag.utils.redis_conn import REDIS_CONN
|
|
from api.utils.health_utils import run_health_checks
|
|
from api.apps.models.auth_dependencies import get_current_user
|
|
|
|
# 创建路由器
|
|
router = APIRouter()
|
|
|
|
|
|
@router.get("/version")
|
|
async def version(
|
|
current_user = Depends(get_current_user)
|
|
):
|
|
"""获取应用程序当前版本"""
|
|
return get_json_result(data=get_ragflow_version())
|
|
|
|
|
|
@router.get("/status")
|
|
async def status(
|
|
current_user = Depends(get_current_user)
|
|
):
|
|
"""获取系统状态"""
|
|
res = {}
|
|
st = timer()
|
|
try:
|
|
res["doc_engine"] = settings.docStoreConn.health()
|
|
res["doc_engine"]["elapsed"] = "{:.1f}".format((timer() - st) * 1000.0)
|
|
except Exception as e:
|
|
res["doc_engine"] = {
|
|
"type": "unknown",
|
|
"status": "red",
|
|
"elapsed": "{:.1f}".format((timer() - st) * 1000.0),
|
|
"error": str(e),
|
|
}
|
|
|
|
st = timer()
|
|
try:
|
|
STORAGE_IMPL.health()
|
|
res["storage"] = {
|
|
"storage": STORAGE_IMPL_TYPE.lower(),
|
|
"status": "green",
|
|
"elapsed": "{:.1f}".format((timer() - st) * 1000.0),
|
|
}
|
|
except Exception as e:
|
|
res["storage"] = {
|
|
"storage": STORAGE_IMPL_TYPE.lower(),
|
|
"status": "red",
|
|
"elapsed": "{:.1f}".format((timer() - st) * 1000.0),
|
|
"error": str(e),
|
|
}
|
|
|
|
st = timer()
|
|
try:
|
|
KnowledgebaseService.get_by_id("x")
|
|
res["database"] = {
|
|
"database": settings.DATABASE_TYPE.lower(),
|
|
"status": "green",
|
|
"elapsed": "{:.1f}".format((timer() - st) * 1000.0),
|
|
}
|
|
except Exception as e:
|
|
res["database"] = {
|
|
"database": settings.DATABASE_TYPE.lower(),
|
|
"status": "red",
|
|
"elapsed": "{:.1f}".format((timer() - st) * 1000.0),
|
|
"error": str(e),
|
|
}
|
|
|
|
st = timer()
|
|
try:
|
|
if not REDIS_CONN.health():
|
|
raise Exception("Lost connection!")
|
|
res["redis"] = {
|
|
"status": "green",
|
|
"elapsed": "{:.1f}".format((timer() - st) * 1000.0),
|
|
}
|
|
except Exception as e:
|
|
res["redis"] = {
|
|
"status": "red",
|
|
"elapsed": "{:.1f}".format((timer() - st) * 1000.0),
|
|
"error": str(e),
|
|
}
|
|
|
|
task_executor_heartbeats = {}
|
|
try:
|
|
task_executors = REDIS_CONN.smembers("TASKEXE")
|
|
now = datetime.now().timestamp()
|
|
for task_executor_id in task_executors:
|
|
heartbeats = REDIS_CONN.zrangebyscore(task_executor_id, now - 60 * 30, now)
|
|
heartbeats = [json.loads(heartbeat) for heartbeat in heartbeats]
|
|
task_executor_heartbeats[task_executor_id] = heartbeats
|
|
except Exception:
|
|
logging.exception("get task executor heartbeats failed!")
|
|
res["task_executor_heartbeats"] = task_executor_heartbeats
|
|
|
|
return get_json_result(data=res)
|
|
|
|
|
|
@router.get("/healthz")
|
|
async def healthz():
|
|
"""健康检查"""
|
|
result, all_ok = run_health_checks()
|
|
return JSONResponse(content=result, status_code=200 if all_ok else 500)
|
|
|
|
|
|
@router.get("/ping")
|
|
async def ping():
|
|
"""心跳检测"""
|
|
return "pong"
|
|
|
|
|
|
@router.post("/new_token")
|
|
async def new_token(
|
|
current_user = Depends(get_current_user)
|
|
):
|
|
"""生成新的 API 令牌"""
|
|
try:
|
|
tenants = UserTenantService.query(user_id=current_user.id)
|
|
if not tenants:
|
|
return get_data_error_result(message="Tenant not found!")
|
|
|
|
tenant_id = [tenant for tenant in tenants if tenant.role == 'owner'][0].tenant_id
|
|
obj = {
|
|
"tenant_id": tenant_id,
|
|
"token": generate_confirmation_token(tenant_id),
|
|
"beta": generate_confirmation_token(generate_confirmation_token(tenant_id)).replace("ragflow-", "")[:32],
|
|
"create_time": current_timestamp(),
|
|
"create_date": datetime_format(datetime.now()),
|
|
"update_time": None,
|
|
"update_date": None,
|
|
}
|
|
|
|
if not APITokenService.save(**obj):
|
|
return get_data_error_result(message="Fail to new a dialog!")
|
|
|
|
return get_json_result(data=obj)
|
|
except Exception as e:
|
|
return server_error_response(e)
|
|
|
|
|
|
@router.get("/token_list")
|
|
async def token_list(
|
|
current_user = Depends(get_current_user)
|
|
):
|
|
"""列出当前用户的所有 API 令牌"""
|
|
try:
|
|
tenants = UserTenantService.query(user_id=current_user.id)
|
|
if not tenants:
|
|
return get_data_error_result(message="Tenant not found!")
|
|
|
|
tenant_id = [tenant for tenant in tenants if tenant.role == 'owner'][0].tenant_id
|
|
objs = APITokenService.query(tenant_id=tenant_id)
|
|
objs = [o.to_dict() for o in objs]
|
|
for o in objs:
|
|
if not o["beta"]:
|
|
o["beta"] = generate_confirmation_token(generate_confirmation_token(tenants[0].tenant_id)).replace(
|
|
"ragflow-", "")[:32]
|
|
APITokenService.filter_update([APIToken.tenant_id == tenant_id, APIToken.token == o["token"]], o)
|
|
return get_json_result(data=objs)
|
|
except Exception as e:
|
|
return server_error_response(e)
|
|
|
|
|
|
@router.delete("/token/{token}")
|
|
async def rm(
|
|
token: str,
|
|
current_user = Depends(get_current_user)
|
|
):
|
|
"""删除 API 令牌"""
|
|
APITokenService.filter_delete(
|
|
[APIToken.tenant_id == current_user.id, APIToken.token == token]
|
|
)
|
|
return get_json_result(data=True)
|
|
|
|
|
|
@router.get('/config')
|
|
async def get_config():
|
|
"""获取系统配置"""
|
|
return get_json_result(data={
|
|
"registerEnabled": settings.REGISTER_ENABLED
|
|
})
|