app.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277
  1. import os
  2. import sys
  3. from logging.handlers import RotatingFileHandler
  4. if not os.environ.get("DEBUG") or os.environ.get("DEBUG").lower() != 'true':
  5. from gevent import monkey
  6. monkey.patch_all()
  7. # if os.environ.get("VECTOR_STORE") == 'milvus':
  8. import grpc.experimental.gevent
  9. grpc.experimental.gevent.init_gevent()
  10. import json
  11. import logging
  12. import threading
  13. import time
  14. import warnings
  15. from flask import Flask, Response, request
  16. from flask_cors import CORS
  17. from werkzeug.exceptions import Unauthorized
  18. from commands import register_commands
  19. from config import CloudEditionConfig, Config
  20. # DO NOT REMOVE BELOW
  21. from events import event_handlers
  22. from extensions import (
  23. ext_celery,
  24. ext_code_based_extension,
  25. ext_compress,
  26. ext_database,
  27. ext_hosting_provider,
  28. ext_login,
  29. ext_mail,
  30. ext_migrate,
  31. ext_redis,
  32. ext_sentry,
  33. ext_storage,
  34. )
  35. from extensions.ext_database import db
  36. from extensions.ext_login import login_manager
  37. from libs.passport import PassportService
  38. from models import account, dataset, model, source, task, tool, tools, web
  39. from services.account_service import AccountService
  40. # DO NOT REMOVE ABOVE
  41. warnings.simplefilter("ignore", ResourceWarning)
  42. # fix windows platform
  43. if os.name == "nt":
  44. os.system('tzutil /s "UTC"')
  45. else:
  46. os.environ['TZ'] = 'UTC'
  47. time.tzset()
  48. class DifyApp(Flask):
  49. pass
  50. # -------------
  51. # Configuration
  52. # -------------
  53. config_type = os.getenv('EDITION', default='SELF_HOSTED') # ce edition first
  54. # ----------------------------
  55. # Application Factory Function
  56. # ----------------------------
  57. def create_app(test_config=None) -> Flask:
  58. app = DifyApp(__name__)
  59. if test_config:
  60. app.config.from_object(test_config)
  61. else:
  62. if config_type == "CLOUD":
  63. app.config.from_object(CloudEditionConfig())
  64. else:
  65. app.config.from_object(Config())
  66. app.secret_key = app.config['SECRET_KEY']
  67. log_handlers = None
  68. log_file = app.config.get('LOG_FILE')
  69. if log_file:
  70. log_dir = os.path.dirname(log_file)
  71. os.makedirs(log_dir, exist_ok=True)
  72. log_handlers = [
  73. RotatingFileHandler(
  74. filename=log_file,
  75. maxBytes=1024 * 1024 * 1024,
  76. backupCount=5
  77. ),
  78. logging.StreamHandler(sys.stdout)
  79. ]
  80. logging.basicConfig(
  81. level=app.config.get('LOG_LEVEL'),
  82. format=app.config.get('LOG_FORMAT'),
  83. datefmt=app.config.get('LOG_DATEFORMAT'),
  84. handlers=log_handlers
  85. )
  86. initialize_extensions(app)
  87. register_blueprints(app)
  88. register_commands(app)
  89. return app
  90. def initialize_extensions(app):
  91. # Since the application instance is now created, pass it to each Flask
  92. # extension instance to bind it to the Flask application instance (app)
  93. ext_compress.init_app(app)
  94. ext_code_based_extension.init()
  95. ext_database.init_app(app)
  96. ext_migrate.init(app, db)
  97. ext_redis.init_app(app)
  98. ext_storage.init_app(app)
  99. ext_celery.init_app(app)
  100. ext_login.init_app(app)
  101. ext_mail.init_app(app)
  102. ext_hosting_provider.init_app(app)
  103. ext_sentry.init_app(app)
  104. # Flask-Login configuration
  105. @login_manager.request_loader
  106. def load_user_from_request(request_from_flask_login):
  107. """Load user based on the request."""
  108. if request.blueprint in ['console', 'inner_api']:
  109. # Check if the user_id contains a dot, indicating the old format
  110. auth_header = request.headers.get('Authorization', '')
  111. if not auth_header:
  112. auth_token = request.args.get('_token')
  113. if not auth_token:
  114. raise Unauthorized('Invalid Authorization token.')
  115. else:
  116. if ' ' not in auth_header:
  117. raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')
  118. auth_scheme, auth_token = auth_header.split(None, 1)
  119. auth_scheme = auth_scheme.lower()
  120. if auth_scheme != 'bearer':
  121. raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')
  122. decoded = PassportService().verify(auth_token)
  123. user_id = decoded.get('user_id')
  124. return AccountService.load_user(user_id)
  125. else:
  126. return None
  127. @login_manager.unauthorized_handler
  128. def unauthorized_handler():
  129. """Handle unauthorized requests."""
  130. return Response(json.dumps({
  131. 'code': 'unauthorized',
  132. 'message': "Unauthorized."
  133. }), status=401, content_type="application/json")
  134. # register blueprint routers
  135. def register_blueprints(app):
  136. from controllers.console import bp as console_app_bp
  137. from controllers.files import bp as files_bp
  138. from controllers.inner_api import bp as inner_api_bp
  139. from controllers.service_api import bp as service_api_bp
  140. from controllers.web import bp as web_bp
  141. CORS(service_api_bp,
  142. allow_headers=['Content-Type', 'Authorization', 'X-App-Code'],
  143. methods=['GET', 'PUT', 'POST', 'DELETE', 'OPTIONS', 'PATCH']
  144. )
  145. app.register_blueprint(service_api_bp)
  146. CORS(web_bp,
  147. resources={
  148. r"/*": {"origins": app.config['WEB_API_CORS_ALLOW_ORIGINS']}},
  149. supports_credentials=True,
  150. allow_headers=['Content-Type', 'Authorization', 'X-App-Code'],
  151. methods=['GET', 'PUT', 'POST', 'DELETE', 'OPTIONS', 'PATCH'],
  152. expose_headers=['X-Version', 'X-Env']
  153. )
  154. app.register_blueprint(web_bp)
  155. CORS(console_app_bp,
  156. resources={
  157. r"/*": {"origins": app.config['CONSOLE_CORS_ALLOW_ORIGINS']}},
  158. supports_credentials=True,
  159. allow_headers=['Content-Type', 'Authorization'],
  160. methods=['GET', 'PUT', 'POST', 'DELETE', 'OPTIONS', 'PATCH'],
  161. expose_headers=['X-Version', 'X-Env']
  162. )
  163. app.register_blueprint(console_app_bp)
  164. CORS(files_bp,
  165. allow_headers=['Content-Type'],
  166. methods=['GET', 'PUT', 'POST', 'DELETE', 'OPTIONS', 'PATCH']
  167. )
  168. app.register_blueprint(files_bp)
  169. app.register_blueprint(inner_api_bp)
  170. # create app
  171. app = create_app()
  172. celery = app.extensions["celery"]
  173. if app.config['TESTING']:
  174. print("App is running in TESTING mode")
  175. @app.after_request
  176. def after_request(response):
  177. """Add Version headers to the response."""
  178. response.set_cookie('remember_token', '', expires=0)
  179. response.headers.add('X-Version', app.config['CURRENT_VERSION'])
  180. response.headers.add('X-Env', app.config['DEPLOY_ENV'])
  181. return response
  182. @app.route('/health')
  183. def health():
  184. return Response(json.dumps({
  185. 'status': 'ok',
  186. 'version': app.config['CURRENT_VERSION']
  187. }), status=200, content_type="application/json")
  188. @app.route('/threads')
  189. def threads():
  190. num_threads = threading.active_count()
  191. threads = threading.enumerate()
  192. thread_list = []
  193. for thread in threads:
  194. thread_name = thread.name
  195. thread_id = thread.ident
  196. is_alive = thread.is_alive()
  197. thread_list.append({
  198. 'name': thread_name,
  199. 'id': thread_id,
  200. 'is_alive': is_alive
  201. })
  202. return {
  203. 'thread_num': num_threads,
  204. 'threads': thread_list
  205. }
  206. @app.route('/db-pool-stat')
  207. def pool_stat():
  208. engine = db.engine
  209. return {
  210. 'pool_size': engine.pool.size(),
  211. 'checked_in_connections': engine.pool.checkedin(),
  212. 'checked_out_connections': engine.pool.checkedout(),
  213. 'overflow_connections': engine.pool.overflow(),
  214. 'connection_timeout': engine.pool.timeout(),
  215. 'recycle_time': db.engine.pool._recycle
  216. }
  217. if __name__ == '__main__':
  218. app.run(host='0.0.0.0', port=5001)