Преглед изворни кода

refactor: tool parameter cache (#3703)

tags/0.6.5
Yeuoly пре 1 година
родитељ
комит
3480f1c59e
No account linked to committer's email address

+ 3
- 44
api/controllers/console/app/app.py Прегледај датотеку

import json

from flask_login import current_user from flask_login import current_user
from flask_restful import Resource, inputs, marshal_with, reqparse from flask_restful import Resource, inputs, marshal_with, reqparse
from werkzeug.exceptions import BadRequest, Forbidden from werkzeug.exceptions import BadRequest, Forbidden
from controllers.console.app.wraps import get_app_model from controllers.console.app.wraps import get_app_model
from controllers.console.setup import setup_required from controllers.console.setup import setup_required
from controllers.console.wraps import account_initialization_required, cloud_edition_billing_resource_check from controllers.console.wraps import account_initialization_required, cloud_edition_billing_resource_check
from core.agent.entities import AgentToolEntity
from core.tools.tool_manager import ToolManager
from core.tools.utils.configuration import ToolParameterConfigurationManager
from extensions.ext_database import db
from fields.app_fields import ( from fields.app_fields import (
app_detail_fields, app_detail_fields,
app_detail_fields_with_site, app_detail_fields_with_site,
app_pagination_fields, app_pagination_fields,
) )
from libs.login import login_required from libs.login import login_required
from models.model import App, AppMode, AppModelConfig
from services.app_service import AppService from services.app_service import AppService


ALLOW_CREATE_APP_MODES = ['chat', 'agent-chat', 'advanced-chat', 'workflow', 'completion'] ALLOW_CREATE_APP_MODES = ['chat', 'agent-chat', 'advanced-chat', 'workflow', 'completion']
@marshal_with(app_detail_fields_with_site) @marshal_with(app_detail_fields_with_site)
def get(self, app_model): def get(self, app_model):
"""Get app detail""" """Get app detail"""
# get original app model config
if app_model.mode == AppMode.AGENT_CHAT.value or app_model.is_agent:
model_config: AppModelConfig = app_model.app_model_config
agent_mode = model_config.agent_mode_dict
# decrypt agent tool parameters if it's secret-input
for tool in agent_mode.get('tools') or []:
if not isinstance(tool, dict) or len(tool.keys()) <= 3:
continue
agent_tool_entity = AgentToolEntity(**tool)
# get tool
try:
tool_runtime = ToolManager.get_agent_tool_runtime(
tenant_id=current_user.current_tenant_id,
agent_tool=agent_tool_entity,
)
manager = ToolParameterConfigurationManager(
tenant_id=current_user.current_tenant_id,
tool_runtime=tool_runtime,
provider_name=agent_tool_entity.provider_id,
provider_type=agent_tool_entity.provider_type,
)

# get decrypted parameters
if agent_tool_entity.tool_parameters:
parameters = manager.decrypt_tool_parameters(agent_tool_entity.tool_parameters or {})
masked_parameter = manager.mask_tool_parameters(parameters or {})
else:
masked_parameter = {}

# override tool parameters
tool['tool_parameters'] = masked_parameter
except Exception as e:
pass

# override agent mode
model_config.agent_mode = json.dumps(agent_mode)
db.session.commit()
app_service = AppService()

app_model = app_service.get_app(app_model)


return app_model return app_model



+ 9
- 3
api/controllers/console/app/model_config.py Прегледај датотеку

try: try:
tool_runtime = ToolManager.get_agent_tool_runtime( tool_runtime = ToolManager.get_agent_tool_runtime(
tenant_id=current_user.current_tenant_id, tenant_id=current_user.current_tenant_id,
app_id=app_model.id,
agent_tool=agent_tool_entity, agent_tool=agent_tool_entity,
) )
manager = ToolParameterConfigurationManager( manager = ToolParameterConfigurationManager(
tool_runtime=tool_runtime, tool_runtime=tool_runtime,
provider_name=agent_tool_entity.provider_id, provider_name=agent_tool_entity.provider_id,
provider_type=agent_tool_entity.provider_type, provider_type=agent_tool_entity.provider_type,
identity_id=f'AGENT.{app_model.id}'
) )
except Exception as e: except Exception as e:
continue continue
try: try:
tool_runtime = ToolManager.get_agent_tool_runtime( tool_runtime = ToolManager.get_agent_tool_runtime(
tenant_id=current_user.current_tenant_id, tenant_id=current_user.current_tenant_id,
app_id=app_model.id,
agent_tool=agent_tool_entity, agent_tool=agent_tool_entity,
) )
except Exception as e: except Exception as e:
tool_runtime=tool_runtime, tool_runtime=tool_runtime,
provider_name=agent_tool_entity.provider_id, provider_name=agent_tool_entity.provider_id,
provider_type=agent_tool_entity.provider_type, provider_type=agent_tool_entity.provider_type,
identity_id=f'AGENT.{app_model.id}'
) )
manager.delete_tool_parameters_cache() manager.delete_tool_parameters_cache()


if agent_tool_entity.tool_parameters: if agent_tool_entity.tool_parameters:
if key not in masked_parameter_map: if key not in masked_parameter_map:
continue continue

if agent_tool_entity.tool_parameters == masked_parameter_map[key]:
agent_tool_entity.tool_parameters = parameter_map[key]
for masked_key, masked_value in masked_parameter_map[key].items():
if masked_key in agent_tool_entity.tool_parameters and \
agent_tool_entity.tool_parameters[masked_key] == masked_value:
agent_tool_entity.tool_parameters[masked_key] = parameter_map[key].get(masked_key)


# encrypt parameters # encrypt parameters
if agent_tool_entity.tool_parameters: if agent_tool_entity.tool_parameters:

+ 1
- 0
api/core/agent/base_agent_runner.py Прегледај датотеку

""" """
tool_entity = ToolManager.get_agent_tool_runtime( tool_entity = ToolManager.get_agent_tool_runtime(
tenant_id=self.tenant_id, tenant_id=self.tenant_id,
app_id=self.app_config.app_id,
agent_tool=tool, agent_tool=tool,
) )
tool_entity.load_variables(self.variables_pool) tool_entity.load_variables(self.variables_pool)

+ 6
- 5
api/core/helper/tool_parameter_cache.py Прегледај датотеку



class ToolParameterCache: class ToolParameterCache:
def __init__(self, def __init__(self,
tenant_id: str,
provider: str,
tool_name: str,
cache_type: ToolParameterCacheType
tenant_id: str,
provider: str,
tool_name: str,
cache_type: ToolParameterCacheType,
identity_id: str
): ):
self.cache_key = f"{cache_type.value}_secret:tenant_id:{tenant_id}:provider:{provider}:tool_name:{tool_name}"
self.cache_key = f"{cache_type.value}_secret:tenant_id:{tenant_id}:provider:{provider}:tool_name:{tool_name}:identity_id:{identity_id}"


def get(self) -> Optional[dict]: def get(self) -> Optional[dict]:
""" """

+ 4
- 2
api/core/tools/tool_manager.py Прегледај датотеку

return parameter_value return parameter_value


@classmethod @classmethod
def get_agent_tool_runtime(cls, tenant_id: str, agent_tool: AgentToolEntity) -> Tool:
def get_agent_tool_runtime(cls, tenant_id: str, app_id: str, agent_tool: AgentToolEntity) -> Tool:
""" """
get the agent tool runtime get the agent tool runtime
""" """
tool_runtime=tool_entity, tool_runtime=tool_entity,
provider_name=agent_tool.provider_id, provider_name=agent_tool.provider_id,
provider_type=agent_tool.provider_type, provider_type=agent_tool.provider_type,
identity_id=f'AGENT.{app_id}'
) )
runtime_parameters = encryption_manager.decrypt_tool_parameters(runtime_parameters) runtime_parameters = encryption_manager.decrypt_tool_parameters(runtime_parameters)


return tool_entity return tool_entity


@classmethod @classmethod
def get_workflow_tool_runtime(cls, tenant_id: str, workflow_tool: ToolEntity):
def get_workflow_tool_runtime(cls, tenant_id: str, app_id: str, node_id: str, workflow_tool: ToolEntity):
""" """
get the workflow tool runtime get the workflow tool runtime
""" """
tool_runtime=tool_entity, tool_runtime=tool_entity,
provider_name=workflow_tool.provider_id, provider_name=workflow_tool.provider_id,
provider_type=workflow_tool.provider_type, provider_type=workflow_tool.provider_type,
identity_id=f'WORKFLOW.{app_id}.{node_id}'
) )


if runtime_parameters: if runtime_parameters:

+ 8
- 3
api/core/tools/utils/configuration.py Прегледај датотеку

tool_runtime: Tool tool_runtime: Tool
provider_name: str provider_name: str
provider_type: str provider_type: str
identity_id: str


def _deep_copy(self, parameters: dict[str, Any]) -> dict[str, Any]: def _deep_copy(self, parameters: dict[str, Any]) -> dict[str, Any]:
""" """
deep copy parameters deep copy parameters
""" """
return {key: value for key, value in parameters.items()}
return deepcopy(parameters)
def _merge_parameters(self) -> list[ToolParameter]: def _merge_parameters(self) -> list[ToolParameter]:
""" """
# override parameters # override parameters
current_parameters = self._merge_parameters() current_parameters = self._merge_parameters()


parameters = self._deep_copy(parameters)

for parameter in current_parameters: for parameter in current_parameters:
if parameter.form == ToolParameter.ToolParameterForm.FORM and parameter.type == ToolParameter.ToolParameterType.SECRET_INPUT: if parameter.form == ToolParameter.ToolParameterForm.FORM and parameter.type == ToolParameter.ToolParameterType.SECRET_INPUT:
if parameter.name in parameters: if parameter.name in parameters:
tenant_id=self.tenant_id, tenant_id=self.tenant_id,
provider=f'{self.provider_type}.{self.provider_name}', provider=f'{self.provider_type}.{self.provider_name}',
tool_name=self.tool_runtime.identity.name, tool_name=self.tool_runtime.identity.name,
cache_type=ToolParameterCacheType.PARAMETER
cache_type=ToolParameterCacheType.PARAMETER,
identity_id=self.identity_id
) )
cached_parameters = cache.get() cached_parameters = cache.get()
if cached_parameters: if cached_parameters:
tenant_id=self.tenant_id, tenant_id=self.tenant_id,
provider=f'{self.provider_type}.{self.provider_name}', provider=f'{self.provider_type}.{self.provider_name}',
tool_name=self.tool_runtime.identity.name, tool_name=self.tool_runtime.identity.name,
cache_type=ToolParameterCacheType.PARAMETER
cache_type=ToolParameterCacheType.PARAMETER,
identity_id=self.identity_id
) )
cache.delete() cache.delete()



+ 2
- 1
api/core/workflow/nodes/tool/tool_node.py Прегледај датотеку

parameters = self._generate_parameters(variable_pool, node_data) parameters = self._generate_parameters(variable_pool, node_data)
# get tool runtime # get tool runtime
try: try:
tool_runtime = ToolManager.get_workflow_tool_runtime(self.tenant_id, node_data)
self.app_id
tool_runtime = ToolManager.get_workflow_tool_runtime(self.tenant_id, self.app_id, self.node_id, node_data)
except Exception as e: except Exception as e:
return NodeRunResult( return NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED, status=WorkflowNodeExecutionStatus.FAILED,

+ 1
- 0
api/events/event_handlers/delete_tool_parameters_cache_when_sync_draft_workflow.py Прегледај датотеку

tool_runtime=tool_runtime, tool_runtime=tool_runtime,
provider_name=tool_entity.provider_name, provider_name=tool_entity.provider_name,
provider_type=tool_entity.provider_type, provider_type=tool_entity.provider_type,
identity_id=f'WORKFLOW.{app.id}.{node_data.get("id")}'
) )
manager.delete_tool_parameters_cache() manager.delete_tool_parameters_cache()

+ 62
- 0
api/services/app_service.py Прегледај датотеку



import yaml import yaml
from flask import current_app from flask import current_app
from flask_login import current_user
from flask_sqlalchemy.pagination import Pagination from flask_sqlalchemy.pagination import Pagination


from constants.model_template import default_app_templates from constants.model_template import default_app_templates
from core.agent.entities import AgentToolEntity
from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
from core.model_manager import ModelManager from core.model_manager import ModelManager
from core.model_runtime.entities.model_entities import ModelPropertyKey, ModelType from core.model_runtime.entities.model_entities import ModelPropertyKey, ModelType
from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
from core.tools.tool_manager import ToolManager
from core.tools.utils.configuration import ToolParameterConfigurationManager
from events.app_event import app_model_config_was_updated, app_was_created, app_was_deleted from events.app_event import app_model_config_was_updated, app_was_created, app_was_deleted
from extensions.ext_database import db from extensions.ext_database import db
from models.account import Account from models.account import Account


return yaml.dump(export_data) return yaml.dump(export_data)


def get_app(self, app: App) -> App:
"""
Get App
"""
# get original app model config
if app.mode == AppMode.AGENT_CHAT.value or app.is_agent:
model_config: AppModelConfig = app.app_model_config
agent_mode = model_config.agent_mode_dict
# decrypt agent tool parameters if it's secret-input
for tool in agent_mode.get('tools') or []:
if not isinstance(tool, dict) or len(tool.keys()) <= 3:
continue
agent_tool_entity = AgentToolEntity(**tool)
# get tool
try:
tool_runtime = ToolManager.get_agent_tool_runtime(
tenant_id=current_user.current_tenant_id,
app_id=app.id,
agent_tool=agent_tool_entity,
)
manager = ToolParameterConfigurationManager(
tenant_id=current_user.current_tenant_id,
tool_runtime=tool_runtime,
provider_name=agent_tool_entity.provider_id,
provider_type=agent_tool_entity.provider_type,
identity_id=f'AGENT.{app.id}'
)

# get decrypted parameters
if agent_tool_entity.tool_parameters:
parameters = manager.decrypt_tool_parameters(agent_tool_entity.tool_parameters or {})
masked_parameter = manager.mask_tool_parameters(parameters or {})
else:
masked_parameter = {}

# override tool parameters
tool['tool_parameters'] = masked_parameter
except Exception as e:
pass

# override agent mode
model_config.agent_mode = json.dumps(agent_mode)

class ModifiedApp(App):
"""
Modified App class
"""
def __init__(self, app):
self.__dict__.update(app.__dict__)

@property
def app_model_config(self):
return model_config
app = ModifiedApp(app)

return app

def update_app(self, app: App, args: dict) -> App: def update_app(self, app: App, args: dict) -> App:
""" """
Update app Update app

Loading…
Откажи
Сачувај