model_config.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. import json
  2. from flask import request
  3. from flask_login import current_user
  4. from flask_restful import Resource
  5. from controllers.console import api
  6. from controllers.console.app import _get_app
  7. from controllers.console.setup import setup_required
  8. from controllers.console.wraps import account_initialization_required
  9. from core.entities.application_entities import AgentToolEntity
  10. from core.tools.tool_manager import ToolManager
  11. from core.tools.utils.configuration import ToolParameterConfigurationManager
  12. from events.app_event import app_model_config_was_updated
  13. from extensions.ext_database import db
  14. from libs.login import login_required
  15. from models.model import AppModelConfig
  16. from services.app_model_config_service import AppModelConfigService
  17. class ModelConfigResource(Resource):
  18. @setup_required
  19. @login_required
  20. @account_initialization_required
  21. def post(self, app_id):
  22. """Modify app model config"""
  23. app_id = str(app_id)
  24. app = _get_app(app_id)
  25. # validate config
  26. model_configuration = AppModelConfigService.validate_configuration(
  27. tenant_id=current_user.current_tenant_id,
  28. account=current_user,
  29. config=request.json,
  30. app_mode=app.mode
  31. )
  32. new_app_model_config = AppModelConfig(
  33. app_id=app.id,
  34. )
  35. new_app_model_config = new_app_model_config.from_model_config_dict(model_configuration)
  36. # get original app model config
  37. original_app_model_config: AppModelConfig = db.session.query(AppModelConfig).filter(
  38. AppModelConfig.id == app.app_model_config_id
  39. ).first()
  40. agent_mode = original_app_model_config.agent_mode_dict
  41. # decrypt agent tool parameters if it's secret-input
  42. parameter_map = {}
  43. masked_parameter_map = {}
  44. tool_map = {}
  45. for tool in agent_mode.get('tools') or []:
  46. agent_tool_entity = AgentToolEntity(**tool)
  47. # get tool
  48. tool_runtime = ToolManager.get_agent_tool_runtime(
  49. tenant_id=current_user.current_tenant_id,
  50. agent_tool=agent_tool_entity,
  51. agent_callback=None
  52. )
  53. manager = ToolParameterConfigurationManager(
  54. tenant_id=current_user.current_tenant_id,
  55. tool_runtime=tool_runtime,
  56. provider_name=agent_tool_entity.provider_id,
  57. provider_type=agent_tool_entity.provider_type,
  58. )
  59. # get decrypted parameters
  60. if agent_tool_entity.tool_parameters:
  61. parameters = manager.decrypt_tool_parameters(agent_tool_entity.tool_parameters or {})
  62. masked_parameter = manager.mask_tool_parameters(parameters or {})
  63. else:
  64. parameters = {}
  65. masked_parameter = {}
  66. key = f'{agent_tool_entity.provider_id}.{agent_tool_entity.provider_type}.{agent_tool_entity.tool_name}'
  67. masked_parameter_map[key] = masked_parameter
  68. parameter_map[key] = parameters
  69. tool_map[key] = tool_runtime
  70. # encrypt agent tool parameters if it's secret-input
  71. agent_mode = new_app_model_config.agent_mode_dict
  72. for tool in agent_mode.get('tools') or []:
  73. agent_tool_entity = AgentToolEntity(**tool)
  74. # get tool
  75. key = f'{agent_tool_entity.provider_id}.{agent_tool_entity.provider_type}.{agent_tool_entity.tool_name}'
  76. if key in tool_map:
  77. tool_runtime = tool_map[key]
  78. else:
  79. tool_runtime = ToolManager.get_agent_tool_runtime(
  80. tenant_id=current_user.current_tenant_id,
  81. agent_tool=agent_tool_entity,
  82. agent_callback=None
  83. )
  84. manager = ToolParameterConfigurationManager(
  85. tenant_id=current_user.current_tenant_id,
  86. tool_runtime=tool_runtime,
  87. provider_name=agent_tool_entity.provider_id,
  88. provider_type=agent_tool_entity.provider_type,
  89. )
  90. manager.delete_tool_parameters_cache()
  91. # override parameters if it equals to masked parameters
  92. if agent_tool_entity.tool_parameters:
  93. if key not in masked_parameter_map:
  94. continue
  95. if agent_tool_entity.tool_parameters == masked_parameter_map[key]:
  96. agent_tool_entity.tool_parameters = parameter_map[key]
  97. # encrypt parameters
  98. if agent_tool_entity.tool_parameters:
  99. tool['tool_parameters'] = manager.encrypt_tool_parameters(agent_tool_entity.tool_parameters or {})
  100. # update app model config
  101. new_app_model_config.agent_mode = json.dumps(agent_mode)
  102. db.session.add(new_app_model_config)
  103. db.session.flush()
  104. app.app_model_config_id = new_app_model_config.id
  105. db.session.commit()
  106. app_model_config_was_updated.send(
  107. app,
  108. app_model_config=new_app_model_config
  109. )
  110. return {'result': 'success'}
  111. api.add_resource(ModelConfigResource, '/apps/<uuid:app_id>/model-config')