model_config.py 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  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.wraps import get_app_model
  7. from controllers.console.setup import setup_required
  8. from controllers.console.wraps import account_initialization_required
  9. from core.agent.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 AppMode, 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. @get_app_model(mode=[AppMode.AGENT_CHAT, AppMode.CHAT, AppMode.COMPLETION])
  22. def post(self, app_model):
  23. """Modify app model config"""
  24. # validate config
  25. model_configuration = AppModelConfigService.validate_configuration(
  26. tenant_id=current_user.current_tenant_id, config=request.json, app_mode=AppMode.value_of(app_model.mode)
  27. )
  28. new_app_model_config = AppModelConfig(
  29. app_id=app_model.id,
  30. )
  31. new_app_model_config = new_app_model_config.from_model_config_dict(model_configuration)
  32. if app_model.mode == AppMode.AGENT_CHAT.value or app_model.is_agent:
  33. # get original app model config
  34. original_app_model_config: AppModelConfig = (
  35. db.session.query(AppModelConfig).filter(AppModelConfig.id == app_model.app_model_config_id).first()
  36. )
  37. agent_mode = original_app_model_config.agent_mode_dict
  38. # decrypt agent tool parameters if it's secret-input
  39. parameter_map = {}
  40. masked_parameter_map = {}
  41. tool_map = {}
  42. for tool in agent_mode.get("tools") or []:
  43. if not isinstance(tool, dict) or len(tool.keys()) <= 3:
  44. continue
  45. agent_tool_entity = AgentToolEntity(**tool)
  46. # get tool
  47. try:
  48. tool_runtime = ToolManager.get_agent_tool_runtime(
  49. tenant_id=current_user.current_tenant_id,
  50. app_id=app_model.id,
  51. agent_tool=agent_tool_entity,
  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. identity_id=f"AGENT.{app_model.id}",
  59. )
  60. except Exception as e:
  61. continue
  62. # get decrypted parameters
  63. if agent_tool_entity.tool_parameters:
  64. parameters = manager.decrypt_tool_parameters(agent_tool_entity.tool_parameters or {})
  65. masked_parameter = manager.mask_tool_parameters(parameters or {})
  66. else:
  67. parameters = {}
  68. masked_parameter = {}
  69. key = f"{agent_tool_entity.provider_id}.{agent_tool_entity.provider_type}.{agent_tool_entity.tool_name}"
  70. masked_parameter_map[key] = masked_parameter
  71. parameter_map[key] = parameters
  72. tool_map[key] = tool_runtime
  73. # encrypt agent tool parameters if it's secret-input
  74. agent_mode = new_app_model_config.agent_mode_dict
  75. for tool in agent_mode.get("tools") or []:
  76. agent_tool_entity = AgentToolEntity(**tool)
  77. # get tool
  78. key = f"{agent_tool_entity.provider_id}.{agent_tool_entity.provider_type}.{agent_tool_entity.tool_name}"
  79. if key in tool_map:
  80. tool_runtime = tool_map[key]
  81. else:
  82. try:
  83. tool_runtime = ToolManager.get_agent_tool_runtime(
  84. tenant_id=current_user.current_tenant_id,
  85. app_id=app_model.id,
  86. agent_tool=agent_tool_entity,
  87. )
  88. except Exception as e:
  89. continue
  90. manager = ToolParameterConfigurationManager(
  91. tenant_id=current_user.current_tenant_id,
  92. tool_runtime=tool_runtime,
  93. provider_name=agent_tool_entity.provider_id,
  94. provider_type=agent_tool_entity.provider_type,
  95. identity_id=f"AGENT.{app_model.id}",
  96. )
  97. manager.delete_tool_parameters_cache()
  98. # override parameters if it equals to masked parameters
  99. if agent_tool_entity.tool_parameters:
  100. if key not in masked_parameter_map:
  101. continue
  102. for masked_key, masked_value in masked_parameter_map[key].items():
  103. if (
  104. masked_key in agent_tool_entity.tool_parameters
  105. and agent_tool_entity.tool_parameters[masked_key] == masked_value
  106. ):
  107. agent_tool_entity.tool_parameters[masked_key] = parameter_map[key].get(masked_key)
  108. # encrypt parameters
  109. if agent_tool_entity.tool_parameters:
  110. tool["tool_parameters"] = manager.encrypt_tool_parameters(agent_tool_entity.tool_parameters or {})
  111. # update app model config
  112. new_app_model_config.agent_mode = json.dumps(agent_mode)
  113. db.session.add(new_app_model_config)
  114. db.session.flush()
  115. app_model.app_model_config_id = new_app_model_config.id
  116. db.session.commit()
  117. app_model_config_was_updated.send(app_model, app_model_config=new_app_model_config)
  118. return {"result": "success"}
  119. api.add_resource(ModelConfigResource, "/apps/<uuid:app_id>/model-config")