# # 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 })